/* SPDX-License-Identifier: AGPL-3.0-or-later AND Apache-2.0 */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * Portions derived from wkv7_cuda_fp32.cu in RWKV-LM
 * (https://github.com/BlinkDL/RWKV-LM), by Bo Peng (BlinkDL),
 * licensed under the Apache License, Version 2.0. Modified by
 * David (Corey) Herrera, 2026. See NOTICE and LICENSES/Apache-2.0.txt.
 */
/*
 * wkv7.cu — RWKV-7 WKV7 forward + backward-through-time CUDA kernel (sm_37/K80).
 *
 * Per-head sequential scan (NOT a matmul): one CUDA block per (batch,head),
 * N threads per block (thread i owns state row i), sequential over time.
 * fp32 throughout.  Reproduces the fp32 numerics and the exact grad contracts
 * of the CPU oracle rx_wkv7_forward_seq / rx_wkv7_backward (rwkvx/rx_pretrain.c),
 * which is a scalar-C transcription of BlinkDL Apache-2.0 wkv7_cuda_fp32.cu.
 *
 * Layouts (identical to the CPU oracle):
 *   input vectors r,w,k,v,a,b,dy,y  : [B][T][H][N], idx = (((b*T+t)*H+h)*N+n)
 *   state S / S_cache[t] / dS        : [B][H][N][N], idx = (((b*H+h)*N+i)*N+j)
 *   sa_cache[t]                      : [B][H][N],    idx = ((b*H+h)*N+i)
 *   a = nkk = -kk ;  b = kka = kk*iclr ; w = EFFECTIVE decay (dw is bare, §5).
 *
 * Recurrence (per head, per t):
 *   sa[i]   = sum_j a[j]*S_prev[i][j]
 *   S[i][j] = S_prev[i][j]*w[j] + k[j]*v[i] + sa[i]*b[j]
 *   y[i]    = sum_j S[i][j]*r[j]
 *
 * Reverse scan (carry dS = dL/dS_t, init 0):
 *   dr[j]  = sum_i S_t[i][j]*dy[i];  dS[i][j] += dy[i]*r[j]
 *   dk[j]  = sum_i dS[i][j]*v[i];  db[j] = sum_i dS[i][j]*sa[i];  dw[j] = sum_i dS[i][j]*S_prev[i][j]
 *   dv[i]  = sum_j dS[i][j]*k[j];  dsa[i] = sum_j dS[i][j]*b[j]
 *   da[j]  = sum_i dsa[i]*S_prev[i][j]
 *   dS[i][j] = dS[i][j]*w[j] + dsa[i]*a[j]
 *
 * Shared-mem budget (per block): fwd = (N*N + 7*N) floats.  For the backward we
 * keep ONLY two N*N tiles in shared — the carried dS and the S_prev tile read
 * from the global forward trajectory — plus 9*N vectors, i.e. (2*N*N + 9*N)
 * floats.  S_t is never materialised: dr[j] = sum_i S_t[i][j]*dy[i] is expanded
 * as w[j]*sum_i S_prev[i][j]*dy[i] + k[j]*sum_i v[i]*dy[i] + b[j]*sum_i sa[i]*dy[i]
 * (division-free; all operands are S_prev + two scalars).  At N=64 this is
 * 2*4096 + 576 = 8768 floats = 34.25 KiB, well under the sm_37 48 KiB/block cap
 * (the old 3*N*N layout was 51456 B and overflowed at N=64).
 */
#include <cuda_runtime.h>
#include <stdio.h>

/* ---------------- forward + trajectory cache ---------------- */
__global__ void wkv7_fwd_kernel(
    const float *r, const float *w, const float *k, const float *v,
    const float *a, const float *b, const float *s_init,
    float *y, float *s_final, float *S_cache, float *sa_cache,
    int B, int T, int H, int N)
{
    int bh  = blockIdx.x;          /* 0..B*H-1 */
    int bch = bh / H;
    int h   = bh % H;
    int i   = threadIdx.x;         /* row */
    if (i >= N) return;

    extern __shared__ float sh[];
    float *sS = sh;                /* N*N */
    float *sr = sS + N * N;        /* N */
    float *sw = sr + N;
    float *sk = sw + N;
    float *sv = sk + N;
    float *sa = sv + N;
    float *sb = sa + N;

    for (int j = 0; j < N; j++)
        sS[i * N + j] = s_init ? s_init[(((size_t)(bch * H + h) * N + i) * N) + j] : 0.0f;
    __syncthreads();

    size_t state_sz = (size_t)B * H * N * N;
    for (int t = 0; t < T; t++) {
        int idx = (((bch * T + t) * H + h) * N + i);
        sr[i] = r[idx]; sw[i] = w[idx]; sk[i] = k[idx];
        sv[i] = v[idx]; sa[i] = a[idx]; sb[i] = b[idx];
        __syncthreads();

        /* cache S_prev, compute sa[i] */
        float saij = 0.0f;
        float *Sc = S_cache + (size_t)t * state_sz;
        size_t sbase = ((size_t)(bch * H + h) * N + i) * N;
        for (int j = 0; j < N; j++) {
            float sval = sS[i * N + j];
            Sc[sbase + j] = sval;
            saij += sa[j] * sval;
        }
        sa_cache[(size_t)t * B * H * N + (bch * H + h) * N + i] = saij;

        /* update row i, read out y[i] */
        float vi = sv[i];
        float yy = 0.0f;
        for (int j = 0; j < N; j++) {
            float sval = sS[i * N + j] * sw[j] + sk[j] * vi + saij * sb[j];
            sS[i * N + j] = sval;
            yy += sval * sr[j];
        }
        if (y) y[idx] = yy;
        __syncthreads();
    }
    if (s_final) {
        size_t sbase = ((size_t)(bch * H + h) * N + i) * N;
        for (int j = 0; j < N; j++) s_final[sbase + j] = sS[i * N + j];
    }
}

/* ---------------- backward-through-time (reverse scan) ----------------
 * Shared tiles: sdS (carried dS) and sSp (S_prev read from the global forward
 * trajectory).  S_t is NOT held in shared; dr is computed analytically from
 * S_prev + two scalars (see file header).  Budget: 2*N*N + 9*N floats. */
__global__ void wkv7_bwd_kernel(
    const float *r, const float *w, const float *k, const float *v,
    const float *a, const float *b, const float *dy,
    const float *S_cache, const float *sa_cache,
    float *dr, float *dw, float *dk, float *dv, float *da, float *db,
    float *dS_out,
    int B, int T, int H, int N)
{
    int bh  = blockIdx.x;
    int bch = bh / H;
    int h   = bh % H;
    int tid = threadIdx.x;
    if (tid >= N) return;

    extern __shared__ float sh[];
    float *sdS  = sh;              /* N*N carried dS */
    float *sSp  = sdS + N * N;     /* N*N S_prev (from global trajectory) */
    float *sr   = sSp + N * N;     /* N each */
    float *sw   = sr + N;
    float *sk   = sw + N;
    float *sv   = sk + N;
    float *saV  = sv + N;
    float *sb   = saV + N;
    float *sdy  = sb + N;
    float *ssa  = sdy + N;
    float *sdsa = ssa + N;

    for (int j = 0; j < N; j++) sdS[tid * N + j] = 0.0f;
    __syncthreads();

    size_t state_sz = (size_t)B * H * N * N;
    for (int t = T - 1; t >= 0; t--) {
        int idx = (((bch * T + t) * H + h) * N + tid);
        sr[tid] = r[idx]; sw[tid] = w[idx]; sk[tid] = k[idx]; sv[tid] = v[idx];
        saV[tid] = a[idx]; sb[tid] = b[idx]; sdy[tid] = dy[idx];
        ssa[tid] = sa_cache[(size_t)t * B * H * N + (bch * H + h) * N + tid];
        const float *Sc = S_cache + (size_t)t * state_sz;
        size_t sbase = ((size_t)(bch * H + h) * N + tid) * N;
        for (int j = 0; j < N; j++) sSp[tid * N + j] = Sc[sbase + j];
        __syncthreads();

        /* (1) read-out grad dr[j] WITHOUT materialising S_t:
         *   dr[j] = sum_i S_t[i][j]*dy[i]
         *         = w[j]*sum_i S_prev[i][j]*dy[i]
         *         + k[j]*sum_i v[i]*dy[i]
         *         + b[j]*sum_i sa[i]*dy[i]
         * (column tid of S_prev + two thread-invariant scalars; division-free).
         * Then fold dy*r into the dS carry. */
        float colSpDY = 0.0f, vDY = 0.0f, saDY = 0.0f;
        for (int i = 0; i < N; i++) {
            float dyi = sdy[i];
            colSpDY += sSp[i * N + tid] * dyi;   /* column tid of S_prev */
            vDY     += sv[i]  * dyi;
            saDY    += ssa[i] * dyi;
        }
        dr[idx] = sw[tid] * colSpDY + sk[tid] * vDY + sb[tid] * saDY;
        float dyt = sdy[tid];
        for (int j = 0; j < N; j++) sdS[tid * N + j] += dyt * sr[j];
        __syncthreads();

        /* (2) column grads (dk,db,dw) sum over i ; row grads (dv,dsa) sum over j */
        float dkj = 0.0f, dbj = 0.0f, dwj = 0.0f;
        for (int i = 0; i < N; i++) {
            float dSij = sdS[i * N + tid];
            dkj += dSij * sv[i];
            dbj += dSij * ssa[i];
            dwj += dSij * sSp[i * N + tid];
        }
        dk[idx] = dkj; db[idx] = dbj; dw[idx] = dwj;
        float dvi = 0.0f, dsai = 0.0f;
        for (int j = 0; j < N; j++) {
            float dSij = sdS[tid * N + j];
            dvi  += dSij * sk[j];
            dsai += dSij * sb[j];
        }
        dv[idx] = dvi; sdsa[tid] = dsai;
        __syncthreads();

        /* (3) da[j] = sum_i dsa[i]*S_prev[i][j] ; dS carry -> t-1 */
        float daj = 0.0f;
        for (int i = 0; i < N; i++) daj += sdsa[i] * sSp[i * N + tid];
        da[idx] = daj;
        float dsai2 = sdsa[tid];
        for (int j = 0; j < N; j++)
            sdS[tid * N + j] = sdS[tid * N + j] * sw[j] + dsai2 * saV[j];
        __syncthreads();
    }
    if (dS_out) {
        size_t sbase = ((size_t)(bch * H + h) * N + tid) * N;
        for (int j = 0; j < N; j++) dS_out[sbase + j] = sdS[tid * N + j];
    }
}

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

/* forward: fills y (may be NULL), and S_cache/sa_cache for the backward. */
extern "C" int wkv7_forward_cuda(
    const float *r, const float *w, const float *k, const float *v,
    const float *a, const float *b, const float *s_init,
    float *y, float *s_final, float *S_cache_out, float *sa_cache_out,
    int B, int T, int H, int N)
{
    size_t nvec = (size_t)B * T * H * N;
    size_t stsz = (size_t)B * H * N * N;
    size_t sasz = (size_t)T * B * H * N;
    size_t scsz = (size_t)T * stsz;
    float *dr_,*dw_,*dk_,*dv_,*da_,*db_,*dsi=NULL,*dy_,*dsf=NULL,*dSc,*dsa;
    CK(cudaMalloc(&dr_,nvec*sizeof(float))); CK(cudaMemcpy(dr_,r,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&dw_,nvec*sizeof(float))); CK(cudaMemcpy(dw_,w,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&dk_,nvec*sizeof(float))); CK(cudaMemcpy(dk_,k,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&dv_,nvec*sizeof(float))); CK(cudaMemcpy(dv_,v,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&da_,nvec*sizeof(float))); CK(cudaMemcpy(da_,a,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&db_,nvec*sizeof(float))); CK(cudaMemcpy(db_,b,nvec*sizeof(float),cudaMemcpyHostToDevice));
    if (s_init){ CK(cudaMalloc(&dsi,stsz*sizeof(float))); CK(cudaMemcpy(dsi,s_init,stsz*sizeof(float),cudaMemcpyHostToDevice)); }
    CK(cudaMalloc(&dy_,nvec*sizeof(float)));
    if (s_final) CK(cudaMalloc(&dsf,stsz*sizeof(float)));
    CK(cudaMalloc(&dSc,scsz*sizeof(float)));
    CK(cudaMalloc(&dsa,sasz*sizeof(float)));

    size_t shmem = ((size_t)N*N + 7*N)*sizeof(float);
    wkv7_fwd_kernel<<<B*H, N, shmem>>>(dr_,dw_,dk_,dv_,da_,db_,dsi,
        y?dy_:NULL, dsf, dSc, dsa, B,T,H,N);
    CK(cudaGetLastError()); CK(cudaDeviceSynchronize());

    if (y) CK(cudaMemcpy(y,dy_,nvec*sizeof(float),cudaMemcpyDeviceToHost));
    if (s_final) CK(cudaMemcpy(s_final,dsf,stsz*sizeof(float),cudaMemcpyDeviceToHost));
    if (S_cache_out) CK(cudaMemcpy(S_cache_out,dSc,scsz*sizeof(float),cudaMemcpyDeviceToHost));
    if (sa_cache_out) CK(cudaMemcpy(sa_cache_out,dsa,sasz*sizeof(float),cudaMemcpyDeviceToHost));

    cudaFree(dr_);cudaFree(dw_);cudaFree(dk_);cudaFree(dv_);cudaFree(da_);cudaFree(db_);
    if(dsi)cudaFree(dsi); cudaFree(dy_); if(dsf)cudaFree(dsf); cudaFree(dSc); cudaFree(dsa);
    return 0;
}

extern "C" int wkv7_backward_cuda(
    const float *r, const float *w, const float *k, const float *v,
    const float *a, const float *b, const float *dy, const float *s_init,
    float *dr, float *dw, float *dk, float *dv, float *da, float *db, float *dS_out,
    int B, int T, int H, int N)
{
    size_t nvec = (size_t)B * T * H * N;
    size_t stsz = (size_t)B * H * N * N;
    size_t sasz = (size_t)T * B * H * N;
    size_t scsz = (size_t)T * stsz;
    /* forward pass on device to fill caches */
    float *dSc, *dsa;
    CK(cudaMalloc(&dSc,scsz*sizeof(float)));
    CK(cudaMalloc(&dsa,sasz*sizeof(float)));
    /* device input copies */
    float *R,*W,*K,*V,*A,*Bb,*DY,*Si=NULL;
    CK(cudaMalloc(&R,nvec*sizeof(float))); CK(cudaMemcpy(R,r,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&W,nvec*sizeof(float))); CK(cudaMemcpy(W,w,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&K,nvec*sizeof(float))); CK(cudaMemcpy(K,k,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&V,nvec*sizeof(float))); CK(cudaMemcpy(V,v,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&A,nvec*sizeof(float))); CK(cudaMemcpy(A,a,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&Bb,nvec*sizeof(float))); CK(cudaMemcpy(Bb,b,nvec*sizeof(float),cudaMemcpyHostToDevice));
    CK(cudaMalloc(&DY,nvec*sizeof(float))); CK(cudaMemcpy(DY,dy,nvec*sizeof(float),cudaMemcpyHostToDevice));
    if (s_init){ CK(cudaMalloc(&Si,stsz*sizeof(float))); CK(cudaMemcpy(Si,s_init,stsz*sizeof(float),cudaMemcpyHostToDevice)); }

    size_t shf = ((size_t)N*N + 7*N)*sizeof(float);
    wkv7_fwd_kernel<<<B*H, N, shf>>>(R,W,K,V,A,Bb,Si,NULL,NULL,dSc,dsa,B,T,H,N);
    CK(cudaGetLastError()); CK(cudaDeviceSynchronize());

    /* backward grads */
    float *dR,*dW,*dK,*dV,*dA,*dB,*dSo;
    CK(cudaMalloc(&dR,nvec*sizeof(float))); CK(cudaMalloc(&dW,nvec*sizeof(float)));
    CK(cudaMalloc(&dK,nvec*sizeof(float))); CK(cudaMalloc(&dV,nvec*sizeof(float)));
    CK(cudaMalloc(&dA,nvec*sizeof(float))); CK(cudaMalloc(&dB,nvec*sizeof(float)));
    CK(cudaMalloc(&dSo,stsz*sizeof(float)));

    size_t shb = (2*(size_t)N*N + 9*N)*sizeof(float);
    wkv7_bwd_kernel<<<B*H, N, shb>>>(R,W,K,V,A,Bb,DY,dSc,dsa,
        dR,dW,dK,dV,dA,dB,dSo,B,T,H,N);
    CK(cudaGetLastError()); CK(cudaDeviceSynchronize());

    if(dr)CK(cudaMemcpy(dr,dR,nvec*sizeof(float),cudaMemcpyDeviceToHost));
    if(dw)CK(cudaMemcpy(dw,dW,nvec*sizeof(float),cudaMemcpyDeviceToHost));
    if(dk)CK(cudaMemcpy(dk,dK,nvec*sizeof(float),cudaMemcpyDeviceToHost));
    if(dv)CK(cudaMemcpy(dv,dV,nvec*sizeof(float),cudaMemcpyDeviceToHost));
    if(da)CK(cudaMemcpy(da,dA,nvec*sizeof(float),cudaMemcpyDeviceToHost));
    if(db)CK(cudaMemcpy(db,dB,nvec*sizeof(float),cudaMemcpyDeviceToHost));
    if(dS_out)CK(cudaMemcpy(dS_out,dSo,stsz*sizeof(float),cudaMemcpyDeviceToHost));

    cudaFree(dSc);cudaFree(dsa);cudaFree(R);cudaFree(W);cudaFree(K);cudaFree(V);
    cudaFree(A);cudaFree(Bb);cudaFree(DY); if(Si)cudaFree(Si);
    cudaFree(dR);cudaFree(dW);cudaFree(dK);cudaFree(dV);cudaFree(dA);cudaFree(dB);cudaFree(dSo);
    return 0;
}