// AdamW optimizer kernel -- f32 native.
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, (), sqrtf(), powf(),
// 0.0f, 1.0f, etc.

extern "C" __global__
void adamw_update(float* __restrict__ param,
                  const float* __restrict__ grad,
                  float* __restrict__ m,
                  float* __restrict__ v,
                  float lr,
                  float beta1,
                  float beta2,
                  float epsilon,
                  float weight_decay,
                  int t,
                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float g = grad[i];
        float p = param[i];
        float b1 = (beta1);
        float b2 = (beta2);
        float one = 1.0f;
        float lr_bf = (lr);
        float eps = (epsilon);

        // Decoupled weight decay: p = p * (1 - lr * wd)
        p = p * (one - lr_bf * (weight_decay));

        // Moment updates
        float mi = b1 * m[i] + (one - b1) * g;
        float vi = b2 * v[i] + (one - b2) * g * g;
        m[i] = mi;
        v[i] = vi;

        // Bias correction
        float bc1 = one - powf(b1, (float)t);
        float bc2 = one - powf(b2, (float)t);
        float m_hat = mi / bc1;
        float v_hat = vi / bc2;

        // Parameter update
        p = p - lr_bf * m_hat / (sqrtf(v_hat) + eps);
        param[i] = p;
    }
}
