/**
 * Backward pass utility kernels for the cuBLAS backward pass.
 *
 * All kernels operate on f32.
 *
 * 1. relu_mask_kernel — zero gradient where activation <= 0 (ReLU backward)
 * 2. bias_grad_reduce_f32_p1/p2 — 2-phase shared-memory sum of f32 dY across batch into dB
 *
 * Launch config: grid=(ceil(n/256), 1, 1) or grid=(ceil(out_dim/256), 1, 1),
 *                block=(256, 1, 1).
 */

/* ReLU mask for f32 backward dX.
 *
 * Both dX and activation are f32 (the entire training pipeline is f32).
 * Zeros elements where the saved activation <= 0. */

extern "C" __global__ void relu_mask_kernel(
    float* __restrict__ dy,
    const float* __restrict__ activation,
    int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    if (activation[i] <= 0.0f) {
        dy[i] = 0.0f;
    }
}

/* Phase 1: block-level shared-memory reduce over batch dimension.
 * Grid: (num_blocks, out_dim), Block: 256, shared: 256*sizeof(float).
 */
extern "C" __global__ void bias_grad_reduce_f32_p1(
    const float* __restrict__ dy,
    float* __restrict__ partials,
    int out_dim, int batch_size
) {
    extern __shared__ float sdata[];
    int j   = blockIdx.y;
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;
    float val = (gid < batch_size) ? dy[gid * out_dim + j] : 0.0f;
    sdata[tid] = val;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }
    if (tid == 0) partials[blockIdx.x * out_dim + j] = sdata[0];
}

/* Phase 2: final reduce across block partials + accumulate into db.
 * Grid: ceil(out_dim/256), Block: 256.
 * IMPORTANT: db[j] += sum (accumulate, not overwrite) to match existing behavior.
 */
extern "C" __global__ void bias_grad_reduce_f32_p2(
    const float* __restrict__ partials,
    float* __restrict__ db,
    int out_dim, int num_blocks
) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= out_dim) return;
    float sum = 0.0f;
    for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + j];
    db[j] += sum;
}
