// ===================================================================
// Linear layer helper kernels -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, (), 0.0f, 1.0f, etc.
//
// - add_bias_2d: output[i] += bias[i % out_dim]
// - reduce_sum_axis0: per-column sum of [rows, cols] -> [cols]
// ===================================================================

// Add bias to a 2D tensor: Y[b, j] += bias[j]
// total = rows * cols
extern "C" __global__
void add_bias_2d_kernel(
    float* __restrict__ output,
    const float* __restrict__ bias,
    int rows,
    int cols
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int total = rows * cols;
    if (i < total) {
        int col = i % cols;
        output[i] = output[i] + bias[col];
    }
}

// Reduce a [rows, cols] tensor to [cols] by summing over axis 0.
// One thread per column — simple for small col counts typical in DQN.
extern "C" __global__
void reduce_sum_axis0_kernel(
    const float* __restrict__ data,
    float* __restrict__ out,
    int rows,
    int cols
) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col < cols) {
        float acc = 0.0f;
        for (int r = 0; r < rows; r++) {
            acc = acc + data[r * cols + col];
        }
        out[col] = acc;
    }
}
