// 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_ptr` is a device pointer to a single i32 holding the current
// 1-indexed training step. This eliminates the host scalar arg that
// blocks CUDA Graph capture (host scalars are baked into kernel args
// at capture time; replays would freeze the step counter). The trainer
// must launch `adamw_increment_counter` once per training step (inside
// the captured region) to advance the counter.

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,
    const int* __restrict__ step_ptr
) {
    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 int step = step_ptr[0];
    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]);
}

// Increment the device-resident step counter. Single-thread kernel;
// launched once per training step (inside the captured graph, AFTER
// all adamw_step launches that read the current step value).
extern "C" __global__ void adamw_increment_counter(int* __restrict__ step_ptr) {
    if (threadIdx.x == 0 && blockIdx.x == 0) {
        step_ptr[0] += 1;
    }
}
