/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_channelmix_cuda.cu — parity harness for GPU RWKV-7 channel-mix
 * (cuda_kernels/channelmix.cu) vs the CPU oracle rx_channelmix_forward_seq /
 * rx_channelmix_backward (rwkvx/rx_pretrain.c), on GPU 0.  Fixed-seed fixtures
 * incl. the real config C=768, F=4*C=3072.  Parity metric: max-abs-error /
 * tensor-scale (inf-norm), bar < 1e-3.  Covers out, dx, d_xprev_init, and every
 * param grad (ffn_x_k, ffn_key_w, ffn_value_w).
 */
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#include <time.h>

#include "cuda_kernels/channelmix.cu"

extern "C" int rx_channelmix_forward_seq(const rx_cm_params*,const float*,const float*,float*,int,int);
extern "C" int rx_channelmix_backward(const rx_cm_params*,const float*,const float*,const float*,
    rx_cm_grads*,float*,float*,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* 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("    %-13s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",tname,rel,maxabs,scale,st);
}

static void mkP(rx_cm_params*P,int C,int F){
    P->c=C; P->ffn_dim=F;
    P->ffn_x_k=rnd((size_t)C,0.5f);
    P->ffn_key_w=rnd((size_t)C*F,0.3f);
    P->ffn_value_w=rnd((size_t)F*C,0.3f);
}
static void mkG(rx_cm_grads*G,int C,int F){
    G->ffn_x_k=zeros((size_t)C); G->ffn_key_w=zeros((size_t)C*F); G->ffn_value_w=zeros((size_t)F*C);
}

static void run_case(const char*name,int B,int T,int C,int F,uint64_t seed){
    sm_state=seed; size_t BT=(size_t)B*T, tot=BT*(size_t)C;
    rx_cm_params P; memset(&P,0,sizeof(P)); mkP(&P,C,F);
    float*x=rnd(tot,1.0f),*xprev=rnd((size_t)B*C,1.0f),*dy=rnd(tot,1.0f);
    float*out_c=(float*)malloc(tot*4),*out_g=(float*)malloc(tot*4);
    rx_channelmix_forward_seq(&P,x,xprev,out_c,B,T);
    if(channelmix_forward_cuda(&P,x,xprev,out_g,B,T)){printf("  %s GPU fwd ERR\n",name);g_fail=1;return;}
    rx_cm_grads Gc,Gg; mkG(&Gc,C,F); mkG(&Gg,C,F);
    float*dx_c=zeros(tot),*dx_g=zeros(tot),*dxp_c=zeros((size_t)B*C),*dxp_g=zeros((size_t)B*C);
    rx_channelmix_backward(&P,x,xprev,dy,&Gc,dx_c,dxp_c,B,T);
    if(channelmix_backward_cuda(&P,x,xprev,dy,&Gg,dx_g,dxp_g,B,T)){printf("  %s GPU bwd ERR\n",name);g_fail=1;return;}
    printf("  %s (B=%d T=%d C=%d F=%d)\n",name,B,T,C,F);
    chk("out",out_g,out_c,tot); chk("dx",dx_g,dx_c,tot); chk("d_xprev",dxp_g,dxp_c,(size_t)B*C);
    chk("ffn_x_k",Gg.ffn_x_k,Gc.ffn_x_k,(size_t)C);
    chk("ffn_key_w",Gg.ffn_key_w,Gc.ffn_key_w,(size_t)C*F);
    chk("ffn_value_w",Gg.ffn_value_w,Gc.ffn_value_w,(size_t)F*C);
}

static void bench(int B,int T,int C,int F,int iters){
    sm_state=99; size_t BT=(size_t)B*T, tot=BT*(size_t)C;
    rx_cm_params P; memset(&P,0,sizeof(P)); mkP(&P,C,F);
    float*x=rnd(tot,1.0f),*xprev=rnd((size_t)B*C,1.0f),*out=(float*)malloc(tot*4);
    channelmix_forward_cuda(&P,x,xprev,out,B,T); /* warmup */
    struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
    for(int i=0;i<iters;i++) channelmix_forward_cuda(&P,x,xprev,out,B,T);
    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("  channelmix fwd B=%d T=%d C=%d F=%d: %d it in %.4fs -> %.0f tokens/sec (one K80 die; incl H2D+alloc)\n",
        B,T,C,F,iters,sec,toks/sec);
}

int main(void){
    int dev=0; cudaSetDevice(dev); cudaDeviceProp p; cudaGetDeviceProperties(&p,dev);
    printf("=== GPU channel-mix 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,32,   11);
    run_case("mid",   2,16,128,512,37);
    run_case("C768",  4,32,768,3072,41);
    printf("\n=== throughput (channel-mix forward, one die) ===\n");
    bench(8,512,768,3072,20);
    printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
    return g_fail;
}