/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_layernorm_cuda.cu — parity harness for GPU LayerNorm (cuda_kernels/
 * layernorm.cu) vs the CPU oracle rx_layernorm_forward_seq /
 * rx_layernorm_backward (rwkvx/rx_pretrain.c), on GPU 0.  Fixed-seed fixtures
 * incl. the real config C=768.  Parity metric: max-abs-error / tensor-scale
 * (inf-norm), bar < 1e-3 (immune to signed-reduction near-zero entries).
 * Covers out, dx, dw, db.
 */
#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#include <time.h>

#include "cuda_kernels/layernorm.cu"

extern "C" int rx_layernorm_forward_seq(const float*,const float*,const float*,float*,int,int,int);
extern "C" int rx_layernorm_backward(const float*,const float*,const float*,const float*,
    float*,float*,float*,int,int,int);

static uint64_t sm_state;
static uint64_t sm_next(void){ uint64_t z=(sm_state+=0x9E3779B97F4A7C15ULL);
    z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL; return z^(z>>31); }
static float sm_uf(void){ return (float)((sm_next()>>11)*(1.0/9007199254740992.0)); }
static float sm_sym(float s){ return (2.0f*sm_uf()-1.0f)*s; }
static float* rnd(size_t n,float s){ float*p=(float*)malloc(n*4); for(size_t i=0;i<n;i++)p[i]=sm_sym(s); return p; }
static float* rndln(size_t n){ float*p=(float*)malloc(n*4); for(size_t i=0;i<n;i++)p[i]=1.0f+0.2f*sm_sym(1.0f); return p; }
static float* zeros(size_t n){ return (float*)calloc(n,4); }

static int g_fail=0;
static void chk(const char*tname,const float*gpu,const float*cpu,size_t n){
    if(!gpu||!cpu) return;
    double maxabs=0.0,scale=0.0;
    for(size_t i=0;i<n;i++){ double d=fabs((double)gpu[i]-(double)cpu[i]); if(d>maxabs)maxabs=d;
        double c=fabs((double)cpu[i]); if(c>scale)scale=c; }
    double rel=maxabs/(scale+1e-30); const char*st=(rel<1e-3)?"PASS":"FAIL"; if(rel>=1e-3)g_fail=1;
    printf("    %-8s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",tname,rel,maxabs,scale,st);
}

static void run_case(const char*name,int B,int T,int C,uint64_t seed){
    sm_state=seed; size_t BT=(size_t)B*T, tot=BT*(size_t)C;
    float*x=rnd(tot,1.0f),*w=rndln((size_t)C),*b=rnd((size_t)C,0.2f),*dy=rnd(tot,1.0f);
    float*out_c=(float*)malloc(tot*4),*out_g=(float*)malloc(tot*4);
    rx_layernorm_forward_seq(x,w,b,out_c,B,T,C);
    if(layernorm_forward_cuda(x,w,b,out_g,B,T,C)){printf("  %s GPU fwd ERR\n",name);g_fail=1;return;}
    float*dx_c=zeros(tot),*dx_g=zeros(tot);
    float*dw_c=zeros((size_t)C),*dw_g=zeros((size_t)C),*db_c=zeros((size_t)C),*db_g=zeros((size_t)C);
    rx_layernorm_backward(x,w,b,dy,dx_c,dw_c,db_c,B,T,C);
    if(layernorm_backward_cuda(x,w,dy,dx_g,dw_g,db_g,B,T,C)){printf("  %s GPU bwd ERR\n",name);g_fail=1;return;}
    printf("  %s (B=%d T=%d C=%d)\n",name,B,T,C);
    chk("out",out_g,out_c,tot); chk("dx",dx_g,dx_c,tot);
    chk("dw",dw_g,dw_c,(size_t)C); chk("db",db_g,db_c,(size_t)C);
    free(x);free(w);free(b);free(dy);free(out_c);free(out_g);
    free(dx_c);free(dx_g);free(dw_c);free(dw_g);free(db_c);free(db_g);
}

static void bench(int B,int T,int C,int iters){
    sm_state=99; size_t BT=(size_t)B*T, tot=BT*(size_t)C;
    float*x=rnd(tot,1.0f),*w=rndln((size_t)C),*b=rnd((size_t)C,0.2f),*out=(float*)malloc(tot*4);
    layernorm_forward_cuda(x,w,b,out,B,T,C); /* warmup */
    struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
    for(int i=0;i<iters;i++) layernorm_forward_cuda(x,w,b,out,B,T,C);
    clock_gettime(CLOCK_MONOTONIC,&t1);
    double sec=(t1.tv_sec-t0.tv_sec)+(t1.tv_nsec-t0.tv_nsec)*1e-9, toks=(double)BT*iters;
    printf("  layernorm fwd B=%d T=%d C=%d: %d it in %.4fs -> %.0f tokens/sec (one K80 die; incl H2D+alloc)\n",
        B,T,C,iters,sec,toks/sec);
    free(x);free(w);free(b);free(out);
}

int main(void){
    int dev=0; cudaSetDevice(dev); cudaDeviceProp p; cudaGetDeviceProperties(&p,dev);
    printf("=== GPU LayerNorm FORWARD+BACKWARD parity vs CPU oracle ===\n");
    printf("device %d: %s cc %d.%d\n\n",dev,p.name,p.major,p.minor);
    run_case("small", 1,5,8,   11);
    run_case("mid",   2,16,128,37);
    run_case("C768",  4,64,768,41);
    printf("\n=== throughput (LayerNorm forward, one die) ===\n");
    bench(8,512,768,50);
    printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
    return g_fail;
}