logicaffeine_compile/codegen_cuda/
solbench.rs

1//! SOL-ExecBench harness bridge — turn a recognized LOGOS fused-row op into a submittable **`solution.json`**
2//! plus a torch-binding CUDA kernel that conforms to the benchmark's ABI.
3//!
4//! The benchmark evaluates `sol-execbench <problem_dir> --solution solution.json`. A CUDA/C++ solution is a
5//! JSON manifest referencing source files, with an `entry_point` `"file::fn"`. The entry point is a
6//! **Destination-Passing-Style** binding — `void run(torch::Tensor inputs…, torch::Tensor output)` — that
7//! takes the inputs (in Definition order) by value and writes the pre-allocated output in place. The harness
8//! materializes the reference `run()` (PyTorch) and compares within `(atol, rtol, matched_ratio)` on a B200.
9//!
10//! Our emitters produce a *self-contained checksum* program for the ≡-oracle GPU locks; this bridge wraps the
11//! SAME kernel math in the harness ABI instead: a runtime-shaped kernel (rows/cols are kernel args, not baked)
12//! and a `run` binding that up-casts the (bf16/…) inputs to f32, launches, and writes the output back.
13
14use crate::codegen_cuda::knobs::CudaKnobs;
15use crate::tensor_recognize::{FusedRowSpec, RowReduceKind};
16
17/// A generated SOL-ExecBench solution: the manifest and the CUDA source it references.
18#[derive(Debug, Clone)]
19pub struct SolSolution {
20    /// The `solution.json` manifest (the `--solution` artifact).
21    pub solution_json: String,
22    /// The CUDA source referenced by `sources[0].content` (kernel + `run` binding).
23    pub kernel_cu: String,
24    /// The `entry_point` — `"kernel.cu::run"`.
25    pub entry_point: String,
26}
27
28/// Emit a SOL-ExecBench solution for a recognized fused-row op (RMSNorm plain/weighted/fused-add, or LayerNorm).
29/// `definition` is the problem's Definition name; `inputs` the input tensor names in Definition order; `output`
30/// the output name. The emitted `run` binding matches that order (DPS).
31pub fn emit_fused_row_solution(
32    spec: &FusedRowSpec,
33    knobs: &CudaKnobs,
34    definition: &str,
35    inputs: &[&str],
36    output: &str,
37) -> SolSolution {
38    let chip = knobs.chip;
39    let layernorm = matches!(spec.reduce, RowReduceKind::LayerNorm);
40    let residual = !layernorm && spec.residual.is_some();
41    let weighted = layernorm || spec.weight.is_some();
42    let bias = layernorm;
43    let eps = format!("{:e}f", spec.eps.get() as f32);
44
45    // Kernel: one block per row, runtime `cols`, shared-memory tree reduction over a fixed 256-wide block.
46    let kernel_math = if layernorm {
47        // two-pass mean/variance, then (x-mean)*scale*w + b
48        r#"    float s = 0.0f;
49    for(int i=threadIdx.x;i<cols;i+=blockDim.x) s += xr[i];
50    sh[threadIdx.x]=s; __syncthreads();
51    for(int k=blockDim.x/2;k>0;k>>=1){ if(threadIdx.x<k) sh[threadIdx.x]+=sh[threadIdx.x+k]; __syncthreads(); }
52    float mean = sh[0]/(float)cols; __syncthreads();
53    float v = 0.0f;
54    for(int i=threadIdx.x;i<cols;i+=blockDim.x){ float d=xr[i]-mean; v+=d*d; }
55    sh[threadIdx.x]=v; __syncthreads();
56    for(int k=blockDim.x/2;k>0;k>>=1){ if(threadIdx.x<k) sh[threadIdx.x]+=sh[threadIdx.x+k]; __syncthreads(); }
57    float scale = rsqrtf(sh[0]/(float)cols + EPS);
58    for(int i=threadIdx.x;i<cols;i+=blockDim.x) yr[i] = (xr[i]-mean)*scale*w[i]+b[i];"#
59            .replace("EPS", &eps)
60    } else {
61        let elem = if residual { "(xr[i]+rr[i])" } else { "xr[i]" };
62        let wmul = if weighted { "*w[i]" } else { "" };
63        format!(
64            r#"    float s = 0.0f;
65    for(int i=threadIdx.x;i<cols;i+=blockDim.x){{ float e={elem}; s+=e*e; }}
66    sh[threadIdx.x]=s; __syncthreads();
67    for(int k=blockDim.x/2;k>0;k>>=1){{ if(threadIdx.x<k) sh[threadIdx.x]+=sh[threadIdx.x+k]; __syncthreads(); }}
68    float scale = rsqrtf(sh[0]/(float)cols + {eps});
69    for(int i=threadIdx.x;i<cols;i+=blockDim.x){{ float e={elem}; yr[i] = e*scale{wmul}; }}"#,
70            elem = elem,
71            wmul = wmul,
72            eps = eps,
73        )
74    };
75
76    // Kernel signature + row-pointer prologue for whichever inputs the op consumes.
77    let kparams = {
78        let mut p = String::from("const float* x");
79        if residual {
80            p.push_str(", const float* res");
81        }
82        if weighted {
83            p.push_str(", const float* w");
84        }
85        if bias {
86            p.push_str(", const float* b");
87        }
88        p.push_str(", float* y, int rows, int cols");
89        p
90    };
91    let res_row = if residual {
92        "    const float* rr = res + (long)row*cols;\n"
93    } else {
94        ""
95    };
96
97    // `run` binding: convert each Definition input to f32/contiguous, launch, write output back in its dtype.
98    // Inputs are (in Definition order) x, [residual], [weight], [bias].
99    let mut run_params = Vec::new();
100    let mut conversions = String::new();
101    let mut kernel_args = Vec::new();
102    let mut it = inputs.iter();
103    let xn = it.next().copied().unwrap_or("x");
104    run_params.push(format!("torch::Tensor {xn}"));
105    conversions.push_str(&format!(
106        "    auto x_ = {xn}.to(torch::kFloat32).contiguous();\n"
107    ));
108    kernel_args.push("x_.data_ptr<float>()".to_string());
109    if residual {
110        let rn = it.next().copied().unwrap_or("residual");
111        run_params.push(format!("torch::Tensor {rn}"));
112        conversions.push_str(&format!(
113            "    auto res_ = {rn}.to(torch::kFloat32).contiguous();\n"
114        ));
115        kernel_args.push("res_.data_ptr<float>()".to_string());
116    }
117    if weighted {
118        let wn = it.next().copied().unwrap_or("weight");
119        run_params.push(format!("torch::Tensor {wn}"));
120        conversions.push_str(&format!(
121            "    auto w_ = {wn}.to(torch::kFloat32).contiguous();\n"
122        ));
123        kernel_args.push("w_.data_ptr<float>()".to_string());
124    }
125    if bias {
126        let bn = it.next().copied().unwrap_or("bias");
127        run_params.push(format!("torch::Tensor {bn}"));
128        conversions.push_str(&format!(
129            "    auto b_ = {bn}.to(torch::kFloat32).contiguous();\n"
130        ));
131        kernel_args.push("b_.data_ptr<float>()".to_string());
132    }
133    run_params.push(format!("torch::Tensor {output}"));
134    kernel_args.push("y_.data_ptr<float>()".to_string());
135    kernel_args.push("rows".to_string());
136    kernel_args.push("cols".to_string());
137
138    let kernel_cu = format!(
139        r#"// Auto-generated by LOGOS codegen_cuda::solbench — a recognized LOGOS {op} as a SOL-ExecBench solution.
140// Target: {chip_name} (sm_{arch}). Destination-passing-style torch binding; runtime rows/cols.
141#include <torch/extension.h>
142#include <cuda_runtime.h>
143#include <cmath>
144
145__global__ void fused_row({kparams}){{
146    int row = blockIdx.x; if(row>=rows) return;
147    const float* xr = x + (long)row*cols;
148{res_row}    float* yr = y + (long)row*cols;
149    __shared__ float sh[256];
150{kernel_math}
151}}
152
153void run({run_params}){{
154{conversions}    int rows = x_.size(0);
155    int cols = x_.size(x_.dim()-1);
156    auto y_ = torch::empty_like(x_);
157    fused_row<<<rows, 256>>>({kernel_args});
158    cudaDeviceSynchronize();
159    {output}.copy_(y_.to({output}.dtype()).view_as({output}));
160}}
161"#,
162        op = spec.reduce.name(),
163        chip_name = chip.name,
164        arch = chip.sm_arch,
165        kparams = kparams,
166        res_row = res_row,
167        kernel_math = kernel_math,
168        run_params = run_params.join(", "),
169        conversions = conversions,
170        kernel_args = kernel_args.join(", "),
171        output = output,
172    );
173
174    let entry_point = "kernel.cu::run".to_string();
175    // JSON: hand-built (small, fixed shape) so we don't pull a serde dep into the emitter.
176    let escaped = json_escape(&kernel_cu);
177    let arch_flag = knobs.chip.nvcc_arch_flag();
178    let solution_json = format!(
179        r#"{{
180  "name": "logos_{op}_{def_slug}",
181  "definition": "{definition}",
182  "author": "logicaffeine",
183  "description": "Compiler-generated (LOGOS -> codegen_cuda) fused-row kernel.",
184  "spec": {{
185    "languages": ["cuda_cpp"],
186    "target_hardware": ["B200", "LOCAL"],
187    "entry_point": "{entry_point}",
188    "dependencies": ["torch"],
189    "destination_passing_style": true,
190    "compile_options": {{ "cuda_cflags": ["-O3", "--use_fast_math", "{arch_flag}"] }}
191  }},
192  "sources": [
193    {{ "path": "kernel.cu", "content": "{escaped}" }}
194  ]
195}}
196"#,
197        op = spec.reduce.name(),
198        def_slug = definition.replace(|c: char| !c.is_ascii_alphanumeric(), "_"),
199        definition = definition,
200        entry_point = entry_point,
201        arch_flag = arch_flag,
202        escaped = escaped,
203    );
204
205    SolSolution {
206        solution_json,
207        kernel_cu,
208        entry_point,
209    }
210}
211
212/// Minimal JSON string escaping for embedding the CUDA source in `sources[].content`.
213fn json_escape(s: &str) -> String {
214    let mut out = String::with_capacity(s.len() + 16);
215    for c in s.chars() {
216        match c {
217            '"' => out.push_str("\\\""),
218            '\\' => out.push_str("\\\\"),
219            '\n' => out.push_str("\\n"),
220            '\r' => out.push_str("\\r"),
221            '\t' => out.push_str("\\t"),
222            c => out.push(c),
223        }
224    }
225    out
226}
logicaffeine_compile/codegen_cuda/solbench.rs

logicaffeine_compile/codegen_cuda/
solbench.rs
