// LSTM gate fusion + bias-add kernels -- f32 native.
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, (), expf(), tanhf(),
// 0.0f, 1.0f, etc.

extern "C" __global__ void lstm_gates(
    float* __restrict__ h_new,
    float* __restrict__ c_new,
    const float* __restrict__ gates,
    const float* __restrict__ c_old,
    int hidden_dim,
    int batch_size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * hidden_dim;
    if (idx >= total) return;

    int b = idx / hidden_dim;
    int h = idx % hidden_dim;
    int gates_offset = b * 4 * hidden_dim;

    float one = 1.0f;

    float i_gate = one / (one + expf(0.0f - gates[gates_offset + h]));
    float f_gate = one / (one + expf(0.0f - gates[gates_offset + hidden_dim + h]));
    float g_gate = tanhf(gates[gates_offset + 2 * hidden_dim + h]);
    float o_gate = one / (one + expf(0.0f - gates[gates_offset + 3 * hidden_dim + h]));

    float c_val = f_gate * c_old[idx] + i_gate * g_gate;
    float h_val = o_gate * tanhf(c_val);

    c_new[idx] = c_val;
    h_new[idx] = h_val;
}

extern "C" __global__ void bias_add_lstm(
    float* __restrict__ output,
    const float* __restrict__ bias,
    int out_features,
    int total_elements
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total_elements) {
        int col = idx % out_features;
        output[idx] = output[idx] + bias[col];
    }
}
