logicaffeine_compile/codegen_cuda/

fused_row.rs

//! The **fused row-op** CUDA emitter — the real one-block-per-row shared-memory reduction kernel.
//!
//! This generalizes the hand-written `cuda/qwen_forward.cu::rmsnorm` (one block per row, a strided partial
//! into shared memory, a tree reduction, an `rsqrt`, a rescale) into a *generated* kernel parameterized by a
//! recognized [`FusedRowSpec`]. It replaces the toy serial `CUDA_RMSNORM_TEMPLATE` / `CUDA_SOFTMAX_TEMPLATE`:
//! those reduced on a single thread over one flat vector; this reduces per row in parallel over `rows × cols`,
//! which is the transformer-normalization shape SOL-ExecBench scores. The flat sugar is the `rows == 1` case.
//!
//! It is **chip-parametric**: the provenance header records the target part (from [`CudaKnobs::chip`]), so the
//! same recognized op emits a kernel stamped for the RTX 3060 or the B200. (Precision/bf16 loads and an
//! `extern "C"` harness entrypoint layer on in the following steps; today the datapath is f32, matching the
//! `Real` sugar it lowers, and the host prints `Σ Y` — the checksum contract the reduction lock reads.)

use crate::codegen_cuda::knobs::CudaKnobs;
use crate::tensor_recognize::{Epsilon, FusedRowSpec, RowReduceKind, SwiGluSpec};

/// Threads per block for the row reduction — a power of two so the shared-memory tree reduction is exact;
/// the strided load loop covers any `cols`. 256 is the oracle's width and fits every part's block budget.
const ROW_THREADS: usize = 256;

/// Emit a self-contained CUDA program for a recognized fused row op: one `__global__` block per row doing a
/// shared-memory tree reduction over `cols`, plus a host that fills `X` deterministically and prints `Σ Y`.
/// `knobs.chip` stamps the provenance header (and drives the arch the caller compiles for).
pub fn emit_fused_row(spec: &FusedRowSpec, knobs: &CudaKnobs) -> String {
    let rows = spec.rows.max(1);
    let cols = spec.cols.max(1);
    let chip = knobs.chip;
    let is_rms = matches!(spec.reduce, RowReduceKind::RmsNorm);
    let layernorm = matches!(spec.reduce, RowReduceKind::LayerNorm);
    let residual = is_rms && spec.residual.is_some();
    let weighted = is_rms && spec.weight.is_some();
    let has_w = weighted || layernorm; // LayerNorm always scales by a per-column weight
    let has_b = layernorm; //           …and shifts by a per-column bias
    let kernel_body = match spec.reduce {
        RowReduceKind::RmsNorm => rmsnorm_kernel(residual, weighted, spec.eps),
        RowReduceKind::Softmax => SOFTMAX_KERNEL.to_string(),
        RowReduceKind::LayerNorm => layernorm_kernel(spec.eps),
    };
    // Host input plumbing, in kernel-parameter order `x, [res,] [w,] [b,] y`. The self-contained host fills
    // each input deterministically (`res[i]=((i+1)%5)-2`, `w[c]=1+0.5·(c%3)`, `b[c]=0.1·(c+1)`) so the emitted
    // program stays runnable and its checksum is reproducible against a CPU reference over the same fills.
    let mut allocs = String::new();
    let mut fills = String::new();
    let mut args = String::from("x,");
    if residual {
        allocs.push_str("    float* res; if(cudaMallocManaged(&res,n*sizeof(float))){ fprintf(stderr,\"malloc\\n\"); return 1; }\n");
        fills.push_str("    for(size_t i=0;i<n;++i) res[i]=(float)((int)((i+1)%5)-2);\n");
        args.push_str("res,");
    }
    if has_w {
        allocs.push_str("    float* w; if(cudaMallocManaged(&w,(size_t)COLS*sizeof(float))){ fprintf(stderr,\"malloc\\n\"); return 1; }\n");
        fills.push_str("    for(int c=0;c<COLS;++c) w[c]=1.0f+0.5f*(float)(c%3);\n");
        args.push_str("w,");
    }
    if has_b {
        allocs.push_str("    float* b; if(cudaMallocManaged(&b,(size_t)COLS*sizeof(float))){ fprintf(stderr,\"malloc\\n\"); return 1; }\n");
        fills.push_str("    for(int c=0;c<COLS;++c) b[c]=0.1f*(float)(c+1);\n");
        args.push_str("b,");
    }
    args.push('y');
    let tag = format!(
        "{}{}{}",
        if residual { " (fused-add)" } else { "" },
        if weighted { " (weighted)" } else { "" },
        if layernorm { " (weight+bias)" } else { "" },
    );
    format!(
        r#"// Auto-generated by LOGOS codegen_cuda::fused_row — a recognized LOGOS {op}{tag} lowered to CUDA.
// Per-row reduction over rows={rows} x cols={cols}; one block per row, shared-memory tree reduction.
// Target: {chip_name} (sm_{arch}). Not a template — the launch bounds are a function of the recognized shape.
#include <cstdio>
#include <cmath>
#define ROWS {rows}
#define COLS {cols}
#define THREADS {threads}
{kernel}
int main(void){{
    float *x,*y;
    size_t n = (size_t)ROWS*COLS;
    if(cudaMallocManaged(&x,n*sizeof(float))||cudaMallocManaged(&y,n*sizeof(float))){{ fprintf(stderr,"malloc\n"); return 1; }}
{allocs}    for(size_t i=0;i<n;++i) x[i]=(float)((int)(i%7)-3);
{fills}    fused_row<<<ROWS,THREADS>>>({args});
    if(cudaDeviceSynchronize()!=cudaSuccess){{ fprintf(stderr,"sync\n"); return 1; }}
    double sum=0.0; for(size_t i=0;i<n;++i) sum+=y[i];
    printf("%.6f\n", sum);
    return 0;
}}
"#,
        op = spec.reduce.name(),
        tag = tag,
        rows = rows,
        cols = cols,
        threads = ROW_THREADS,
        chip_name = chip.name,
        arch = chip.sm_arch,
        kernel = kernel_body,
        allocs = allocs,
        fills = fills,
        args = args,
    )
}

/// Emit a self-contained CUDA program for a recognized SwiGLU map — `y[i] = silu(gate[i]) * up[i]`, one thread
/// per element, self-contained host (deterministic fills + `Σ Y` checksum). Chip-stamped like the fused-row.
pub fn emit_swiglu(spec: &SwiGluSpec, knobs: &CudaKnobs) -> String {
    let n = spec.total.max(1);
    let chip = knobs.chip;
    format!(
        r#"// Auto-generated by LOGOS codegen_cuda::fused_row — a recognized LOGOS swiglu lowered to CUDA.
// Y[i] = silu(gate[i]) * up[i]; silu(g) = g / (1 + e^-g); N={n} elements, one thread each.
// Target: {chip_name} (sm_{arch}). Not a template — N is a function of the recognized shape.
#include <cstdio>
#include <cmath>
#define N {n}
#define THREADS 256
__global__ void swiglu(const float* gate, const float* up, float* y){{
    int i = blockIdx.x*blockDim.x + threadIdx.x; if(i>=N) return;
    float g = gate[i];
    y[i] = (g/(1.0f+expf(-g))) * up[i];
}}
int main(void){{
    float *gate,*up,*y;
    size_t n=(size_t)N;
    if(cudaMallocManaged(&gate,n*sizeof(float))||cudaMallocManaged(&up,n*sizeof(float))||cudaMallocManaged(&y,n*sizeof(float))){{ fprintf(stderr,"malloc\n"); return 1; }}
    for(size_t i=0;i<n;++i){{ gate[i]=(float)((int)(i%7)-3); up[i]=(float)((int)((i+1)%5)-2); }}
    swiglu<<<(N+THREADS-1)/THREADS,THREADS>>>(gate,up,y);
    if(cudaDeviceSynchronize()!=cudaSuccess){{ fprintf(stderr,"sync\n"); return 1; }}
    double sum=0.0; for(size_t i=0;i<n;++i) sum+=y[i];
    printf("%.6f\n", sum);
    return 0;
}}
"#,
        n = n,
        chip_name = chip.name,
        arch = chip.sm_arch,
    )
}

/// Build the per-row RMSNorm kernel for the four `(residual, weighted)` variants: the reduced value is
/// `X[r,i]` or `X[r,i] + res[r,i]` (fused-add), and the rescale is `·rsqrt(mean+eps)` optionally `·W[i]`, with the
/// program's `eps`.
fn rmsnorm_kernel(residual: bool, weighted: bool, eps: Epsilon) -> String {
    let params = format!(
        "const float* x, {}{}float* y",
        if residual { "const float* res, " } else { "" },
        if weighted { "const float* w, " } else { "" },
    );
    let res_row = if residual {
        "    const float* rr = res + (long)row*COLS;\n"
    } else {
        ""
    };
    let elem = if residual { "(xr[i]+rr[i])" } else { "xr[i]" };
    let wmul = if weighted { "*w[i]" } else { "" };
    format!(
        r#"__global__ void fused_row({params}){{
    int row = blockIdx.x;
    const float* xr = x + (long)row*COLS;
{res_row}    float* yr = y + (long)row*COLS;
    __shared__ float sh[THREADS];
    float s = 0.0f;
    for(int i = threadIdx.x; i < COLS; i += blockDim.x){{ float e = {elem}; s += e*e; }}
    sh[threadIdx.x] = s; __syncthreads();
    for(int k = blockDim.x/2; k > 0; k >>= 1){{ if(threadIdx.x < k) sh[threadIdx.x] += sh[threadIdx.x+k]; __syncthreads(); }}
    float scale = rsqrtf(sh[0]/(float)COLS + {eps});
    for(int i = threadIdx.x; i < COLS; i += blockDim.x){{ float e = {elem}; yr[i] = e*scale{wmul}; }}
}}"#,
        params = params,
        res_row = res_row,
        elem = elem,
        wmul = wmul,
        eps = eps_literal(eps),
    )
}

/// The `f32` literal of a norm's epsilon, as the kernels spell it (`1e-6f`).
fn eps_literal(eps: Epsilon) -> String {
    format!("{:e}f", eps.get() as f32)
}

/// Build the per-row LayerNorm kernel: two-pass mean/variance, then `(x-mean)·rsqrt(var+eps)·W[i] + B[i]`, with the
/// program's `eps`.
fn layernorm_kernel(eps: Epsilon) -> String {
    r#"__global__ void fused_row(const float* x, const float* w, const float* b, float* y){
    int row = blockIdx.x;
    const float* xr = x + (long)row*COLS;
    float* yr = y + (long)row*COLS;
    __shared__ float sh[THREADS];
    float s = 0.0f;
    for(int i = threadIdx.x; i < COLS; i += blockDim.x) s += xr[i];
    sh[threadIdx.x] = s; __syncthreads();
    for(int k = blockDim.x/2; k > 0; k >>= 1){ if(threadIdx.x < k) sh[threadIdx.x] += sh[threadIdx.x+k]; __syncthreads(); }
    float mean = sh[0]/(float)COLS; __syncthreads();
    float v = 0.0f;
    for(int i = threadIdx.x; i < COLS; i += blockDim.x){ float d = xr[i]-mean; v += d*d; }
    sh[threadIdx.x] = v; __syncthreads();
    for(int k = blockDim.x/2; k > 0; k >>= 1){ if(threadIdx.x < k) sh[threadIdx.x] += sh[threadIdx.x+k]; __syncthreads(); }
    float scale = rsqrtf(sh[0]/(float)COLS + EPS);
    for(int i = threadIdx.x; i < COLS; i += blockDim.x) yr[i] = (xr[i]-mean)*scale*w[i]+b[i];
}"#
    .replace("EPS", &eps_literal(eps))
}

/// Per-row softmax: `Y[r,i] = exp(X[r,i]-max)/Σ exp(X[r,·]-max)`. One block per row; two tree reductions.
const SOFTMAX_KERNEL: &str = r#"__global__ void fused_row(const float* x, float* y){
    int row = blockIdx.x;
    const float* xr = x + (long)row*COLS;
    float* yr = y + (long)row*COLS;
    __shared__ float sh[THREADS];
    float m = -1e30f;
    for(int i = threadIdx.x; i < COLS; i += blockDim.x) m = fmaxf(m, xr[i]);
    sh[threadIdx.x] = m; __syncthreads();
    for(int k = blockDim.x/2; k > 0; k >>= 1){ if(threadIdx.x < k) sh[threadIdx.x] = fmaxf(sh[threadIdx.x], sh[threadIdx.x+k]); __syncthreads(); }
    float mx = sh[0]; __syncthreads();
    float s = 0.0f;
    for(int i = threadIdx.x; i < COLS; i += blockDim.x) s += expf(xr[i]-mx);
    sh[threadIdx.x] = s; __syncthreads();
    for(int k = blockDim.x/2; k > 0; k >>= 1){ if(threadIdx.x < k) sh[threadIdx.x] += sh[threadIdx.x+k]; __syncthreads(); }
    float sm = sh[0];
    for(int i = threadIdx.x; i < COLS; i += blockDim.x) yr[i] = expf(xr[i]-mx)/sm;
}"#;