/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * layernorm.cu — RWKV-x LayerNorm forward+backward on the K80 (sm_37).
 *
 * GPU counterpart of the CPU oracle rx_layernorm_forward_seq /
 * rx_layernorm_backward (rwkvx/rx_pretrain.c).  LayerNorm over the full channel
 * dim C for each of the B*T tokens, eps=1e-5:
 *     mean = (1/C) sum_i x_i ; var = (1/C) sum_i (x_i-mean)^2
 *     rstd = 1/sqrt(var+eps) ; out_i = (x_i-mean)*rstd*w_i + b_i
 *
 * Numerics: the mean/variance AND both backward reductions (mean(dnorm),
 * mean(dnorm*norm)) are cancellation-prone.  One thread owns one token's entire
 * C-length reduction, accumulating in DOUBLE — matching the CPU oracle's double
 * discipline exactly.  The backward RE-derives mean/var/rstd in double from x
 * (the oracle does the same; it does not reuse a truncated fp32 cache), so the
 * two paths are bit-for-bit in the reduction that matters.  dw/db are summed
 * across all tokens with atomicAdd (order-independent, well under the 1e-3 bar).
 *
 * Token layout [B][T][C] row-major == BT rows of length C (no reshape).
 */
#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include "rwkvx/rx_pretrain.h"

#define LCK(call) do{ cudaError_t e_=(call); if(e_!=cudaSuccess){ \
  fprintf(stderr,"CUDA %s:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e_)); return -1; } }while(0)

#define RX_LN_EPS 1e-5
#define TPB 128
static inline int nblk(size_t n){ return (int)((n+TPB-1)/TPB); }

/* forward: one thread per token, double-accumulated mean/var.
 * mean_o/rstd_o are optional fp32 stat caches (pass NULL to skip). */
__global__ void k_ln_fwd(const float*x,const float*w,const float*b,float*out,
    float*mean_o,float*rstd_o,int C,size_t BT){
    size_t tk=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(tk>=BT)return;
    const float*xx=x+tk*C; float*o=out+tk*C;
    double mean=0.0; for(int i=0;i<C;i++) mean+=xx[i]; mean/=C;
    double var=0.0; for(int i=0;i<C;i++){ double d=(double)xx[i]-mean; var+=d*d; } var/=C;
    double rstd=1.0/sqrt(var+RX_LN_EPS);
    if(mean_o) mean_o[tk]=(float)mean;
    if(rstd_o) rstd_o[tk]=(float)rstd;
    for(int i=0;i<C;i++) o[i]=(float)(((double)xx[i]-mean)*rstd)*w[i]+b[i];
}

/* backward: one thread per token.  Re-derives mean/var/rstd in double (matches
 * oracle).  dnorm=dy*w; mdn=mean(dnorm); mdnn=mean(dnorm*norm);
 * dx = rstd*(dnorm - mdn - norm*mdnn).  dw+=dy*norm; db+=dy (atomic). */
__global__ void k_ln_bwd(const float*x,const float*w,const float*dy,
    float*dx,float*dw,float*db,int C,size_t BT){
    size_t tk=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(tk>=BT)return;
    const float*xx=x+tk*C; const float*dyt=dy+tk*C; float*dxt=dx+tk*C;
    double mean=0.0; for(int i=0;i<C;i++) mean+=xx[i]; mean/=C;
    double var=0.0; for(int i=0;i<C;i++){ double d=(double)xx[i]-mean; var+=d*d; } var/=C;
    double rstd=1.0/sqrt(var+RX_LN_EPS);
    double mdn=0.0, mdnn=0.0;
    for(int i=0;i<C;i++){
        double nv=((double)xx[i]-mean)*rstd;
        double dnorm=(double)dyt[i]*(double)w[i];
        mdn+=dnorm; mdnn+=dnorm*nv;
        if(dw) atomicAdd(&dw[i], dyt[i]*(float)nv);
        if(db) atomicAdd(&db[i], dyt[i]);
    }
    mdn/=C; mdnn/=C;
    for(int i=0;i<C;i++){
        double nv=((double)xx[i]-mean)*rstd;
        double dnorm=(double)dyt[i]*(double)w[i];
        dxt[i]=(float)(rstd*(dnorm - mdn - nv*mdnn));
    }
}

/* ---- host->device upload helper (NULL passthrough) ---- */
static float* lup(const float*h,size_t n){
    if(!h) return NULL; float*d=NULL; if(cudaMalloc(&d,n*sizeof(float))!=cudaSuccess) return NULL;
    if(cudaMemcpy(d,h,n*sizeof(float),cudaMemcpyHostToDevice)!=cudaSuccess){cudaFree(d);return NULL;} return d; }

/* ---- host-facing forward parity entry ---- */
extern "C" int layernorm_forward_cuda(const float*x,const float*w,const float*b,
    float*out,int B,int T,int C){
    size_t BT=(size_t)B*T, tot=BT*(size_t)C;
    float*dx=lup(x,tot),*dw=lup(w,(size_t)C),*db=lup(b,(size_t)C),*dout=NULL;
    if(!dx||!dw||!db) return -1;
    LCK(cudaMalloc(&dout,tot*sizeof(float)));
    k_ln_fwd<<<nblk(BT),TPB>>>(dx,dw,db,dout,NULL,NULL,C,BT);
    LCK(cudaDeviceSynchronize());
    LCK(cudaMemcpy(out,dout,tot*sizeof(float),cudaMemcpyDeviceToHost));
    cudaFree(dx);cudaFree(dw);cudaFree(db);cudaFree(dout);
    return 0;
}

/* ---- host-facing backward parity entry ---- */
extern "C" int layernorm_backward_cuda(const float*x,const float*w,const float*dy,
    float*dx,float*dw,float*db,int B,int T,int C){
    size_t BT=(size_t)B*T, tot=BT*(size_t)C;
    float* dxin=lup(x,tot),*dwt=lup(w,(size_t)C),*ddy=lup(dy,tot);
    float*ddx=NULL,*ddw=NULL,*ddb=NULL;
    if(!dxin||!dwt||!ddy) return -1;
    LCK(cudaMalloc(&ddx,tot*sizeof(float)));
    LCK(cudaMalloc(&ddw,(size_t)C*sizeof(float))); LCK(cudaMemset(ddw,0,(size_t)C*sizeof(float)));
    LCK(cudaMalloc(&ddb,(size_t)C*sizeof(float))); LCK(cudaMemset(ddb,0,(size_t)C*sizeof(float)));
    k_ln_bwd<<<nblk(BT),TPB>>>(dxin,dwt,ddy,ddx, dw?ddw:NULL, db?ddb:NULL, C,BT);
    LCK(cudaDeviceSynchronize());
    LCK(cudaMemcpy(dx,ddx,tot*sizeof(float),cudaMemcpyDeviceToHost));
    if(dw) LCK(cudaMemcpy(dw,ddw,(size_t)C*sizeof(float),cudaMemcpyDeviceToHost));
    if(db) LCK(cudaMemcpy(db,ddb,(size_t)C*sizeof(float),cudaMemcpyDeviceToHost));
    cudaFree(dxin);cudaFree(dwt);cudaFree(ddy);cudaFree(ddx);cudaFree(ddw);cudaFree(ddb);
    return 0;
}