/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * channelmix.cu — RWKV-7 channel-mix (squared-ReLU FFN) fwd+bwd on the K80 (sm_37).
 *
 * GPU counterpart of the CPU oracle rx_channelmix_forward_seq /
 * rx_channelmix_backward (rwkvx/rx_pretrain.c).  EXCLUDES the pre-block ln2 and
 * the residual add (those live in the block).  Per token:
 *     sh  = xx + (prev - xx) * ffn_x_k          (token shift / lerp)
 *     fk  = sh @ ffn_key_w                       ([C] -> [F])
 *     k   = relu(fk)^2                           (SQUARED ReLU activation)
 *     out = k @ ffn_value_w                      ([F] -> [C])
 * prev at t=0 is x_prev_init (or zero); otherwise the previous token's input.
 *
 * Both matmuls go through cuBLAS SGEMM over the whole [B*T,*] token batch via the
 * row-major->column-major rm_gemm() helper (swap A/B, pass N,M,K; same convention
 * as fc_cuda.c / timemix.cu).  Token-shift + squared-ReLU are elementwise kernels.
 * The token-shift backward folds each token's own (1-lerp) term and the NEXT
 * token's lerp term in one pass (no dx race); param grads use atomicAdd; weight
 * grads fall out of single batched GEMMs (beta=0 == full token-sum).
 */
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include "rwkvx/rx_pretrain.h"

#define TCK(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 TBK(call) do{ cublasStatus_t s_=(call); if(s_!=CUBLAS_STATUS_SUCCESS){ \
  fprintf(stderr,"cuBLAS %s:%d status %d\n",__FILE__,__LINE__,(int)s_); return -1; } }while(0)

#define TPB 256
static inline int nblk(size_t n){ return (int)((n+TPB-1)/TPB); }

/* ---- row-major GEMM via cuBLAS: C[M,N] = op(A) @ op(B) ---- */
static cublasStatus_t rm_gemm(cublasHandle_t H,int tA,int tB,int M,int N,int K,
    float alpha,const float*A,const float*B,float beta,float*C){
    int lda = tA? M : K;
    int ldb = tB? K : N;
    return cublasSgemm(H, tB?CUBLAS_OP_T:CUBLAS_OP_N, tA?CUBLAS_OP_T:CUBLAS_OP_N,
                       N, M, K, &alpha, B, ldb, A, lda, &beta, C, N);
}

/* ---- elementwise kernels ---- */
__global__ void k_cm_shift(const float*x,const float*xprev,const float*x_k,
    float*sh,int B,int T,int C){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)B*T*C; if(g>=tot)return;
    int i=g%C; size_t token=g/C; int t=(int)(token%T); size_t b=token/T;
    float xx=x[g];
    float prev=(t==0)?(xprev?xprev[b*C+i]:0.0f):x[g-C];
    sh[g]=xx+(prev-xx)*x_k[i];
}
__global__ void k_sqrelu(const float*fk,float*ka,size_t n){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=n)return;
    float v=fk[g]>0.0f?fk[g]:0.0f; ka[g]=v*v;
}
/* squared-relu backward: k=relu(fk)^2 => dfk = dk * 2*relu(fk) */
__global__ void k_sqrelu_bwd(const float*fk,const float*dka,float*dfk,size_t n){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=n)return;
    float r=fk[g]>0.0f?fk[g]:0.0f; dfk[g]=dka[g]*2.0f*r;
}
/* token-shift backward: sh = xx + (prev-xx)*x_k.
 * dx[t] = dsh[t]*(1-x_k) + dsh[t+1]*x_k ; d_xprev_init = dsh[t=0]*x_k. */
__global__ void k_cm_shift_bwd(const float*x,const float*xprev,const float*dsh,
    const float*x_k,float*G_x_k,float*dx,float*dxprev,int B,int T,int C){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)B*T*C; if(g>=tot)return;
    int i=g%C; size_t token=g/C; int t=(int)(token%T); size_t b=token/T;
    float xx=x[g];
    float prev=(t==0)?(xprev?xprev[b*C+i]:0.0f):x[g-C];
    atomicAdd(&G_x_k[i], dsh[g]*(prev-xx));
    float own=dsh[g]*(1.0f-x_k[i]);
    float nxt=0.0f;
    if(t<T-1) nxt=dsh[g+C]*x_k[i];
    dx[g]=own+nxt;
    if(t==0 && dxprev) dxprev[b*C+i]=dsh[g]*x_k[i];
}

/* ---- upload helper ---- */
static float* cup(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 channelmix_forward_cuda(const rx_cm_params*P,const float*x,
    const float*x_prev_init,float*out,int B,int T){
    int C=P->c, F=P->ffn_dim; size_t BT=(size_t)B*T, tot=BT*(size_t)C, totF=BT*(size_t)F;
    cublasHandle_t H; if(cublasCreate(&H)!=CUBLAS_STATUS_SUCCESS) return -1;
    float*dx=cup(x,tot),*dxk=cup(P->ffn_x_k,(size_t)C);
    float*dkw=cup(P->ffn_key_w,(size_t)C*F),*dvw=cup(P->ffn_value_w,(size_t)F*C);
    float*dxprev = x_prev_init? cup(x_prev_init,(size_t)B*C) : NULL;
    float*sh=NULL,*fk=NULL,*ka=NULL,*dout=NULL;
    TCK(cudaMalloc(&sh,tot*sizeof(float))); TCK(cudaMalloc(&fk,totF*sizeof(float)));
    TCK(cudaMalloc(&ka,totF*sizeof(float))); TCK(cudaMalloc(&dout,tot*sizeof(float)));
    k_cm_shift<<<nblk(tot),TPB>>>(dx,dxprev,dxk,sh,B,T,C);
    TBK(rm_gemm(H,0,0,BT,F,C,1.0f,sh,dkw,0.0f,fk));
    k_sqrelu<<<nblk(totF),TPB>>>(fk,ka,totF);
    TBK(rm_gemm(H,0,0,BT,C,F,1.0f,ka,dvw,0.0f,dout));
    TCK(cudaDeviceSynchronize());
    TCK(cudaMemcpy(out,dout,tot*sizeof(float),cudaMemcpyDeviceToHost));
    cudaFree(dx);cudaFree(dxk);cudaFree(dkw);cudaFree(dvw);if(dxprev)cudaFree(dxprev);
    cudaFree(sh);cudaFree(fk);cudaFree(ka);cudaFree(dout);
    cublasDestroy(H);
    return 0;
}

/* ---- host-facing backward parity entry ---- */
extern "C" int channelmix_backward_cuda(const rx_cm_params*P,const float*x,
    const float*x_prev_init,const float*dy,rx_cm_grads*G,float*dx,float*d_xprev_init,
    int B,int T){
    int C=P->c, F=P->ffn_dim; size_t BT=(size_t)B*T, tot=BT*(size_t)C, totF=BT*(size_t)F;
    cublasHandle_t H; if(cublasCreate(&H)!=CUBLAS_STATUS_SUCCESS) return -1;
    float*dxin=cup(x,tot),*dxk=cup(P->ffn_x_k,(size_t)C);
    float*dkw=cup(P->ffn_key_w,(size_t)C*F),*dvw=cup(P->ffn_value_w,(size_t)F*C);
    float*ddy=cup(dy,tot);
    float*dxprev = x_prev_init? cup(x_prev_init,(size_t)B*C) : NULL;

    /* device grads (zeroed) */
    float*G_xk=NULL,*G_kw=NULL,*G_vw=NULL;
    TCK(cudaMalloc(&G_xk,(size_t)C*sizeof(float)));   TCK(cudaMemset(G_xk,0,(size_t)C*sizeof(float)));
    TCK(cudaMalloc(&G_kw,(size_t)C*F*sizeof(float))); TCK(cudaMemset(G_kw,0,(size_t)C*F*sizeof(float)));
    TCK(cudaMalloc(&G_vw,(size_t)F*C*sizeof(float))); TCK(cudaMemset(G_vw,0,(size_t)F*C*sizeof(float)));

    /* scratch + recompute forward (sh, fk, ka) */
    float*sh=NULL,*fk=NULL,*ka=NULL,*dka=NULL,*dfk=NULL,*dsh=NULL,*ddx=NULL,*ddxprev=NULL;
    TCK(cudaMalloc(&sh,tot*sizeof(float))); TCK(cudaMalloc(&fk,totF*sizeof(float)));
    TCK(cudaMalloc(&ka,totF*sizeof(float))); TCK(cudaMalloc(&dka,totF*sizeof(float)));
    TCK(cudaMalloc(&dfk,totF*sizeof(float))); TCK(cudaMalloc(&dsh,tot*sizeof(float)));
    TCK(cudaMalloc(&ddx,tot*sizeof(float)));
    if(d_xprev_init){ TCK(cudaMalloc(&ddxprev,(size_t)B*C*sizeof(float)));
        TCK(cudaMemset(ddxprev,0,(size_t)B*C*sizeof(float))); }

    k_cm_shift<<<nblk(tot),TPB>>>(dxin,dxprev,dxk,sh,B,T,C);
    TBK(rm_gemm(H,0,0,BT,F,C,1.0f,sh,dkw,0.0f,fk));
    k_sqrelu<<<nblk(totF),TPB>>>(fk,ka,totF);

    /* value proj backward: out = ka @ Wv */
    TBK(rm_gemm(H,0,1,BT,F,C,1.0f,ddy,dvw,0.0f,dka));   /* dka = dy @ Wv^T */
    TBK(rm_gemm(H,1,0,F,C,BT,1.0f,ka,ddy,0.0f,G_vw));   /* G_Wv = ka^T @ dy */
    /* squared-relu backward */
    k_sqrelu_bwd<<<nblk(totF),TPB>>>(fk,dka,dfk,totF);
    /* key proj backward: fk = sh @ Wk */
    TBK(rm_gemm(H,0,1,BT,C,F,1.0f,dfk,dkw,0.0f,dsh));   /* dsh = dfk @ Wk^T */
    TBK(rm_gemm(H,1,0,C,F,BT,1.0f,sh,dfk,0.0f,G_kw));   /* G_Wk = sh^T @ dfk */
    /* token-shift backward -> dx, d_xprev_init, G_x_k */
    k_cm_shift_bwd<<<nblk(tot),TPB>>>(dxin,dxprev,dsh,dxk,G_xk,ddx,ddxprev,B,T,C);
    TCK(cudaDeviceSynchronize());

    /* download */
    TCK(cudaMemcpy(dx,ddx,tot*sizeof(float),cudaMemcpyDeviceToHost));
    if(d_xprev_init) TCK(cudaMemcpy(d_xprev_init,ddxprev,(size_t)B*C*sizeof(float),cudaMemcpyDeviceToHost));
    if(G->ffn_x_k)    TCK(cudaMemcpy(G->ffn_x_k,G_xk,(size_t)C*sizeof(float),cudaMemcpyDeviceToHost));
    if(G->ffn_key_w)  TCK(cudaMemcpy(G->ffn_key_w,G_kw,(size_t)C*F*sizeof(float),cudaMemcpyDeviceToHost));
    if(G->ffn_value_w)TCK(cudaMemcpy(G->ffn_value_w,G_vw,(size_t)F*C*sizeof(float),cudaMemcpyDeviceToHost));

    cudaFree(dxin);cudaFree(dxk);cudaFree(dkw);cudaFree(dvw);cudaFree(ddy);if(dxprev)cudaFree(dxprev);
    cudaFree(G_xk);cudaFree(G_kw);cudaFree(G_vw);
    cudaFree(sh);cudaFree(fk);cudaFree(ka);cudaFree(dka);cudaFree(dfk);cudaFree(dsh);
    cudaFree(ddx);if(ddxprev)cudaFree(ddxprev);
    cublasDestroy(H);
    return 0;
}