// adamw_step.cu — element-wise AdamW update.
//
//   m_t = beta1 * m_{t-1} + (1 - beta1) * g
//   v_t = beta2 * v_{t-1} + (1 - beta2) * g^2
//   m_hat = m_t / (1 - beta1^step)
//   v_hat = v_t / (1 - beta2^step)
//   theta -= lr * (m_hat / (sqrt(v_hat) + eps) + wd * theta)
//
// One thread per parameter. Block tree-reduce not needed; this is pure
// element-wise. No atomicAdd.
//
// `step` is a host-supplied i32 holding the current 1-indexed training
// step. The Rust launcher increments a host-side counter and passes it
// as a kernel argument each launch, eliminating the separate
// `adamw_increment_counter` kernel (saves 3407 launches per training
// step).

extern "C" __global__ void adamw_step(
    float* __restrict__ theta,
    const float* __restrict__ grad,
    float* __restrict__ m,
    float* __restrict__ v,
    int n_params,
    float lr,
    float beta1,
    float beta2,
    float eps,
    float wd,
    int step
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n_params) return;

    const float g = grad[i];
    const float m_n = beta1 * m[i] + (1.0f - beta1) * g;
    const float v_n = beta2 * v[i] + (1.0f - beta2) * g * g;
    m[i] = m_n;
    v[i] = v_n;

    const float bc1 = 1.0f - powf(beta1, (float) step);
    const float bc2 = 1.0f - powf(beta2, (float) step);
    const float m_hat = m_n / fmaxf(bc1, 1e-12f);
    const float v_hat = v_n / fmaxf(bc2, 1e-12f);

    theta[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * theta[i]);
}
