// Softmax and log-softmax CUDA kernels -- f32 native.
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, (), expf(), logf(), fmaxf(), etc.
//
// Each warp/block handles one row of [dim] elements.

extern "C" __global__ void softmax_forward(
    float* __restrict__ output,
    const float* __restrict__ input,
    int batch_size,
    int dim
) {
    int row = blockIdx.x;
    if (row >= batch_size) return;

    const float* in_row = input + row * dim;
    float* out_row = output + row * dim;

    // Find max for numerical stability
    float max_val = (-1e4f);
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float v = in_row[i];
        max_val = fmaxf(max_val, v);
    }
    // Warp reduction for max
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, max_val, offset);
        max_val = fmaxf(max_val, other);
    }
    // Broadcast max from lane 0
    max_val = (__shfl_sync(0xffffffff, (max_val), 0));

    // Compute exp(x - max) and sum
    float sum = 0.0f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float e = expf(in_row[i] - max_val);
        out_row[i] = e;
        sum = sum + e;
    }
    // Warp reduction for sum
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        sum = sum + __shfl_down_sync(0xffffffff, sum, offset);
    }
    sum = (__shfl_sync(0xffffffff, (sum), 0));

    // Normalize
    float inv_sum = 1.0f / sum;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        out_row[i] = out_row[i] * inv_sum;
    }
}

extern "C" __global__ void log_softmax_forward(
    float* __restrict__ output,
    const float* __restrict__ input,
    int batch_size,
    int dim
) {
    int row = blockIdx.x;
    if (row >= batch_size) return;

    const float* in_row = input + row * dim;
    float* out_row = output + row * dim;

    // Find max for numerical stability
    float max_val = (-1e4f);
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float v = in_row[i];
        max_val = fmaxf(max_val, v);
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, max_val, offset);
        max_val = fmaxf(max_val, other);
    }
    max_val = (__shfl_sync(0xffffffff, (max_val), 0));

    // Compute sum of exp(x - max)
    float sum = 0.0f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        sum = sum + expf(in_row[i] - max_val);
    }
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        sum = sum + __shfl_down_sync(0xffffffff, sum, offset);
    }
    sum = (__shfl_sync(0xffffffff, (sum), 0));

    float log_sum = logf(sum);

    // log_softmax(x_i) = x_i - max - log(sum(exp(x - max)))
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        out_row[i] = in_row[i] - max_val - log_sum;
    }
}
