proofman_exps_codegen/
lib.rs

1//! Per-AIR expression -> straight-line CUDA kernel codegen.
2//!
3//! Two kernel families are emitted into each AIR's self-contained
4//! `<base>.exps.so` (placed next to that AIR's `.bin`):
5//! * the Q (cExp) kernel -- register-bounded, chunk size autotuned to zero
6//!   register spill;
7//! * one small trace-domain kernel per other covered expression (hint fields,
8//!   im columns, ...), dispatched by expId via `exps_expr_covered` /
9//!   `exps_launch_expr` (see `emit_exprs_tu`).
10//!
11//! The prover `dlopen`s the library by convention and falls back to the
12//! bytecode interpreter for anything absent: missing `.so`, missing symbol,
13//! or an uncovered expression.
14//!
15//! Two entry points:
16//! * [`generate_air`]  — one AIR dir -> its `.exps.so`.
17//! * [`generate_all`]  — a provingKey dir -> every AIR's `.exps.so`.
18
19mod autotune;
20mod emit;
21mod ir;
22mod model;
23mod toolchain;
24
25use anyhow::{Context, Result};
26use ir::{plan_chunks, UnhandledOperand};
27use model::{ExpressionsInfo, StarkInfo};
28use rayon::prelude::*;
29use std::path::{Path, PathBuf};
30use toolchain::Toolchain;
31
32const DEFAULT_CAP: usize = 40000; // skip an AIR whose Q has more ops than this
33const DEFAULT_CHUNK: usize = 512; // fixed ops/chunk when autotuning is off
34const SLOTS_CAP: u64 = 1000; // skip an AIR whose cross-chunk cut exceeds this
35
36/// Codegen configuration. Defaults: CAP=40000, autotune on, arch=auto.
37#[derive(Debug, Clone)]
38pub struct GenConfig {
39    /// Skip an AIR whose Q has more than this many ops (-> interpreter).
40    pub cap: usize,
41    /// Fixed ops/chunk for every AIR; `None` turns the no-spill autotuner ON.
42    pub chunk: Option<usize>,
43    /// CUDA arch spec: `auto` (default), `major`, or a list like `89,120`.
44    pub archspec: String,
45    /// pil2-stark source root; `None` resolves it relative to this crate.
46    pub stark_src: Option<PathBuf>,
47    /// Retain the generated `.cu`/`.o` here; `None` uses a temp dir removed on exit.
48    pub keep_dir: Option<PathBuf>,
49    /// Emit the `.cu` sources only — skip compiling/linking the `.so` (so the
50    /// provingKey is untouched). Requires `keep_dir`. Used for inspecting the
51    /// generated sources.
52    pub dry_run: bool,
53}
54
55impl Default for GenConfig {
56    fn default() -> Self {
57        GenConfig {
58            cap: DEFAULT_CAP,
59            chunk: None,
60            archspec: "auto".into(),
61            stark_src: None,
62            keep_dir: None,
63            dry_run: false,
64        }
65    }
66}
67
68/// One AIR whose kernel was generated.
69#[derive(Debug, Clone)]
70pub struct GeneratedAir {
71    pub name: String,
72    pub base: String,
73    pub sym: String,
74    pub nbits: u64,
75    pub cexp: i64,
76    pub n_ops: usize,
77    pub slots: u64,
78}
79
80/// Outcome of a codegen run.
81#[derive(Debug, Default)]
82pub struct GenSummary {
83    pub generated: Vec<GeneratedAir>,
84    pub skipped: Vec<(String, String)>,
85    pub placed: usize,
86    pub max_scratch_bytes: u64,
87}
88
89/// A discovered, codegen-eligible AIR (unique per `sym`).
90struct Candidate {
91    stark_info: StarkInfo,
92    expr_info: ExpressionsInfo,
93    sym: String,
94    nbits: u64,
95    cexp: i64,
96    name: String,
97    n_ops: usize,
98    base: String,
99}
100
101/// One `.so` destination (every eligible AIR, not deduped by sym).
102struct Placement {
103    name: String,
104    base: String,
105    sym: String,
106    air_dir: PathBuf,
107}
108
109/// Recursively collect `*.starkinfo.json` paths under `root`, sorted.
110fn find_starkinfos(root: &Path) -> Vec<PathBuf> {
111    let mut out = Vec::new();
112    let mut stack = vec![root.to_path_buf()];
113    while let Some(dir) = stack.pop() {
114        let Ok(entries) = std::fs::read_dir(&dir) else { continue };
115        for e in entries.flatten() {
116            let p = e.path();
117            if p.is_dir() {
118                stack.push(p);
119            } else if p.file_name().and_then(|s| s.to_str()).is_some_and(|s| s.ends_with(".starkinfo.json")) {
120                out.push(p);
121            }
122        }
123    }
124    out.sort();
125    out
126}
127
128/// The proof phase a circuit belongs to, derived from the AIR dir's leaf
129/// component. Part of the kernel identity so circuits from different phases
130/// never collide on `(airgroupId, airId, nBits, cExpId)`.
131fn proof_phase(air_dir: &Path) -> String {
132    let comp = air_dir.file_name().and_then(|s| s.to_str()).unwrap_or("");
133    match comp {
134        "air" => "basic",
135        "compressor" => "compressor",
136        "recursive1" | "recursive2" => "recursive",
137        c if c.starts_with("vadcop_final") => "final",
138        other => other, // unknown layout: keep the raw dir name so it still disambiguates
139    }
140    .to_string()
141}
142
143/// `sym` — an AIR's kernel identity string `<phase>_a<airgroupId>_<airId>_b<nBits>_e<cExpId>`
144/// (phase ∈ basic|compressor|recursive|final). Dedup key, file stem, and C/CUDA
145/// symbol stem all at once.
146fn make_sym(si: &StarkInfo, phase: &str) -> String {
147    format!("{phase}_a{}_{}_b{}_e{}", si.airgroup_id, si.air_id, si.stark_struct.n_bits, si.c_exp_id)
148}
149
150/// Parse one starkinfo + sibling expressionsinfo into a Candidate, or return a
151/// skip reason (string), or `None` if the file pair isn't a codegen target.
152fn load_candidate(
153    stark_info_path: &Path,
154    root: &Path,
155    cap: usize,
156) -> Result<Option<Candidate>, Option<(String, String)>> {
157    let air_dir = stark_info_path.parent().unwrap();
158    let fname = stark_info_path.file_name().unwrap().to_string_lossy();
159    let base = fname.strip_suffix(".starkinfo.json").unwrap().to_string();
160    let expr_info_path = air_dir.join(format!("{base}.expressionsinfo.json"));
161    if !expr_info_path.exists() {
162        return Err(None);
163    }
164    let name = air_dir.strip_prefix(root).unwrap_or(air_dir).to_string_lossy().to_string();
165
166    let stark_info: StarkInfo = match std::fs::read(stark_info_path).ok().and_then(|b| serde_json::from_slice(&b).ok())
167    {
168        Some(si) => si,
169        None => return Err(None), // unparseable / not a full AIR starkinfo
170    };
171    let expr_info: ExpressionsInfo =
172        match std::fs::read(&expr_info_path).ok().and_then(|b| serde_json::from_slice(&b).ok()) {
173            Some(ei) => ei,
174            None => return Err(None),
175        };
176
177    let cexp = stark_info.c_exp_id;
178    let Some(code) = expr_info.expressions_code.iter().find(|e| e.exp_id == cexp) else {
179        return Err(None); // cExpId not in expressionsCode
180    };
181    let n_ops = code.code.len();
182    let nbits = stark_info.stark_struct.n_bits;
183    if n_ops > cap {
184        return Err(Some((name, format!("{n_ops} ops > CAP"))));
185    }
186    let sym = make_sym(&stark_info, &proof_phase(air_dir));
187    Ok(Some(Candidate { stark_info, expr_info, sym, nbits, cexp, name, n_ops, base }))
188}
189
190/// Generate every AIR's `.exps.so` under `proving_key`. Returns a summary of
191/// what was generated, skipped, and placed.
192pub fn generate_all(proving_key: &Path, cfg: &GenConfig) -> Result<GenSummary> {
193    if cfg.dry_run && cfg.keep_dir.is_none() {
194        anyhow::bail!("dry_run requires keep_dir (nowhere to write the .cu otherwise)");
195    }
196    let work = WorkDir::new(cfg.keep_dir.clone())?;
197    std::fs::write(work.path().join("gen_common.cuh"), emit::COMMON_CUH)?;
198    let tc = Toolchain::new(cfg.stark_src.clone(), &cfg.archspec, work.path())?;
199    eprintln!("[exps-codegen] generating kernels for {} (archs: {})", proving_key.display(), tc.arch_summary());
200
201    // Phase 1: discovery — unique candidates (by sym) + every placement.
202    let mut candidates: Vec<Candidate> = Vec::new();
203    let mut placements: Vec<Placement> = Vec::new();
204    let mut skipped: Vec<(String, String)> = Vec::new();
205    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
206    for si_path in find_starkinfos(proving_key) {
207        match load_candidate(&si_path, proving_key, cfg.cap) {
208            Ok(Some(c)) => {
209                placements.push(Placement {
210                    name: c.name.clone(),
211                    base: c.base.clone(),
212                    sym: c.sym.clone(),
213                    air_dir: si_path.parent().unwrap().to_path_buf(),
214                });
215                if seen.insert(c.sym.clone()) {
216                    candidates.push(c);
217                }
218            }
219            Ok(None) => {}
220            Err(Some(skip)) => skipped.push(skip),
221            Err(None) => {}
222        }
223    }
224
225    let mut summary = run_pipeline(&tc, work.path(), &candidates, &placements, cfg)?;
226    summary.skipped.extend(skipped);
227    summary.skipped.sort();
228    print_summary(&summary, cfg);
229    Ok(summary)
230}
231
232/// Generate the `.exps.so` for a single AIR directory (the dir containing its
233/// `*.starkinfo.json` + `*.expressionsinfo.json`). Returns the `.so` path, or
234/// an error if the AIR was skipped (unhandled operand, over CAP, or spills).
235pub fn generate_air(air_dir: &Path, cfg: &GenConfig) -> Result<PathBuf> {
236    let si_path = find_starkinfos(air_dir)
237        .into_iter()
238        .find(|p| p.parent() == Some(air_dir))
239        .with_context(|| format!("no *.starkinfo.json in {}", air_dir.display()))?;
240
241    let work = WorkDir::new(cfg.keep_dir.clone())?;
242    std::fs::write(work.path().join("gen_common.cuh"), emit::COMMON_CUH)?;
243    let tc = Toolchain::new(cfg.stark_src.clone(), &cfg.archspec, work.path())?;
244
245    let root = air_dir;
246    let candidate = match load_candidate(&si_path, root, cfg.cap) {
247        Ok(Some(c)) => c,
248        Ok(None) => anyhow::bail!("{} is not a codegen target", air_dir.display()),
249        Err(Some((_, why))) => anyhow::bail!("skipped: {why}"),
250        Err(None) => anyhow::bail!("{} missing/invalid expressionsinfo", air_dir.display()),
251    };
252    let placement = Placement {
253        name: candidate.name.clone(),
254        base: candidate.base.clone(),
255        sym: candidate.sym.clone(),
256        air_dir: air_dir.to_path_buf(),
257    };
258    let dest = air_dir.join(format!("{}.exps.so", candidate.base));
259
260    let summary =
261        run_pipeline(&tc, work.path(), std::slice::from_ref(&candidate), std::slice::from_ref(&placement), cfg)?;
262    if summary.placed == 1 {
263        Ok(dest)
264    } else {
265        let why = summary.skipped.first().map(|(_, w)| w.clone()).unwrap_or_else(|| "unknown".into());
266        anyhow::bail!("{}: not generated ({why})", candidate.name)
267    }
268}
269
270fn run_pipeline(
271    tc: &Toolchain,
272    work: &Path,
273    candidates: &[Candidate],
274    placements: &[Placement],
275    cfg: &GenConfig,
276) -> Result<GenSummary> {
277    let autotune = cfg.chunk.is_none();
278
279    // Build IR for every candidate (catches unhandled operands here). Each entry
280    // is (candidate, Ok(ir) | Err(skip-reason)).
281    let built: Vec<(&Candidate, std::result::Result<ir::Ir, String>)> = candidates
282        .iter()
283        .map(|c| {
284            let r = match ir::build_ir(&c.stark_info, &c.expr_info) {
285                Ok(ir) => Ok(ir),
286                Err(e) if e.downcast_ref::<UnhandledOperand>().is_some() => Err("unhandled operand".to_string()),
287                Err(e) => Err(format!("build_ir error: {e}")),
288            };
289            (c, r)
290        })
291        .collect();
292
293    // Phase 2: autotune the no-spill chunk size per AIR (parallel). Maps sym -> Some(chunk) | None(spills).
294    let chunk_map: std::collections::HashMap<String, Option<usize>> = if autotune {
295        built
296            .par_iter()
297            .filter_map(|(c, r)| {
298                r.as_ref()
299                    .ok()
300                    .map(|ir| autotune::tune_chunk(tc, ir, &c.sym, c.n_ops, work).map(|ck| (c.sym.clone(), ck)))
301            })
302            .collect::<Result<std::collections::HashMap<_, _>>>()?
303    } else {
304        Default::default()
305    };
306
307    // Phase 3: emit the .cu sources; record per-sym slot counts.
308    let mut slots_by_sym: std::collections::HashMap<String, u64> = std::collections::HashMap::new();
309    let mut exprs_by_sym: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
310    let mut generated: Vec<GeneratedAir> = Vec::new();
311    let mut skipped: Vec<(String, String)> = Vec::new();
312    let mut max_scratch: u64 = 0;
313    for (c, r) in &built {
314        let ir = match r {
315            Ok(ir) => ir,
316            Err(why) => {
317                skipped.push((c.name.clone(), why.clone()));
318                continue;
319            }
320        };
321        let chunk = if autotune {
322            match chunk_map.get(&c.sym).copied().flatten() {
323                Some(ck) => ck,
324                None => {
325                    skipped.push((c.name.clone(), format!("{} ops: still spills at CHUNK_MIN", c.n_ops)));
326                    continue;
327                }
328            }
329        } else {
330            cfg.chunk.unwrap_or(DEFAULT_CHUNK)
331        };
332        let plan = plan_chunks(ir, chunk, &c.sym)?;
333        if plan.total_slots > SLOTS_CAP {
334            skipped.push((c.name.clone(), format!("slots {} > SLOTS_CAP (wide cut)", plan.total_slots)));
335            continue;
336        }
337        for (fname, text) in emit::emit_air(ir, &plan, &c.sym) {
338            std::fs::write(work.join(&fname), text)?;
339        }
340        // Generic (non-Q) expression kernels: cover every trace-domain
341        // expression the interpreter might be asked for (hint fields, im
342        // columns, ...). Small straight-line kernels only; anything odd
343        // (zerofier use, non-canonical output shape, oversized) is skipped
344        // and stays on the interpreter.
345        {
346            const EXPR_CAP: usize = 512;
347            let mut items: Vec<(i64, ir::Ir, u64)> = Vec::new();
348            for ec in &c.expr_info.expressions_code {
349                if ec.exp_id == c.cexp || ec.code.is_empty() || ec.code.len() > EXPR_CAP {
350                    continue;
351                }
352                let Ok(eir) = ir::build_ir_expr(&c.stark_info, &c.expr_info, ec.exp_id, false) else {
353                    continue;
354                };
355                if eir.uses_zi() {
356                    continue;
357                }
358                let Some(od) = eir.out_dim() else { continue };
359                if od != 1 && od != 3 {
360                    continue;
361                }
362                items.push((ec.exp_id, eir, od));
363            }
364            if !items.is_empty() {
365                let n_exprs = items.len();
366                std::fs::write(work.join(format!("gen_{}_cexprs.cu", c.sym)), emit::emit_exprs_tu(&c.sym, &items))?;
367                exprs_by_sym.insert(c.sym.clone(), n_exprs);
368            }
369        }
370        let n_ext = 1u64 << c.stark_info.stark_struct.n_bits_ext;
371        max_scratch = max_scratch.max(plan.total_slots * n_ext);
372        slots_by_sym.insert(c.sym.clone(), plan.total_slots);
373        generated.push(GeneratedAir {
374            name: c.name.clone(),
375            base: c.base.clone(),
376            sym: c.sym.clone(),
377            nbits: c.nbits,
378            cexp: c.cexp,
379            n_ops: c.n_ops,
380            slots: plan.total_slots,
381        });
382    }
383
384    // Phase 3b: compile every emitted .cu that does not already have its .o
385    // (parallel). The autotuner leaves the winning Q objects in `work`; the
386    // generic-expression TUs (and, without autotune, the Q TUs) are compiled
387    // here so the link step below is uniformly object-based.
388    {
389        let mut jobs: Vec<(PathBuf, PathBuf)> = Vec::new();
390        for sym in slots_by_sym.keys() {
391            for cu in collect_artifacts(work, sym, "cu") {
392                let obj = cu.with_extension("o");
393                if !obj.exists() {
394                    jobs.push((cu, obj));
395                }
396            }
397        }
398        // Plain scoped-thread fan-out (NOT rayon: a worker blocked on a child
399        // nvcc inside the global pool can deadlock against the pipeline's
400        // outer parallel bridges).
401        let par = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(8);
402        for batch in jobs.chunks(par) {
403            let errs: Vec<String> = std::thread::scope(|scope| {
404                let handles: Vec<_> = batch
405                    .iter()
406                    .map(|(cu, obj)| {
407                        scope.spawn(move || -> Option<String> {
408                            match tc.compile_tu(cu, obj, Some(work)) {
409                                Ok((true, _)) => None,
410                                Ok((false, log)) => Some(format!("nvcc failed for {}: {log}", cu.display())),
411                                Err(e) => Some(format!("nvcc spawn failed for {}: {e}", cu.display())),
412                            }
413                        })
414                    })
415                    .collect();
416                handles.into_iter().filter_map(|h| h.join().ok().flatten()).collect()
417            });
418            if let Some(e) = errs.into_iter().next() {
419                anyhow::bail!(e);
420            }
421        }
422    }
423
424    // Phase 4: link one self-contained .so per placement whose sym was generated (parallel).
425    let placed: Vec<&Placement> = placements.iter().filter(|p| slots_by_sym.contains_key(&p.sym)).collect();
426    write_gen_log(work, &placed, &slots_by_sym)?;
427    if !cfg.dry_run {
428        placed.par_iter().try_for_each(|p| -> Result<()> {
429            let dest = p.air_dir.join(format!("{}.exps.so", p.base));
430            link_one(tc, work, &p.sym, &dest)
431        })?;
432    }
433
434    Ok(GenSummary { placed: placed.len(), generated, skipped, max_scratch_bytes: max_scratch * 8 })
435}
436
437/// Link (or compile+link) one AIR's objects/sources into `dest`.
438fn link_one(tc: &Toolchain, work: &Path, sym: &str, dest: &Path) -> Result<()> {
439    let objs = collect_artifacts(work, sym, "o");
440    if !objs.is_empty() {
441        tc.link_objs(&objs, dest)
442    } else {
443        let cus = collect_artifacts(work, sym, "cu");
444        tc.compile_link_cus(&cus, dest)
445    }
446}
447
448/// `gen_<sym>.<ext>` + `gen_<sym>_c*.<ext>` present in `work`.
449fn collect_artifacts(work: &Path, sym: &str, ext: &str) -> Vec<PathBuf> {
450    let mut out = Vec::new();
451    let main = work.join(format!("gen_{sym}.{ext}"));
452    if main.exists() {
453        out.push(main);
454    }
455    let prefix = format!("gen_{sym}_c");
456    if let Ok(entries) = std::fs::read_dir(work) {
457        let mut chunks: Vec<PathBuf> = entries
458            .flatten()
459            .map(|e| e.path())
460            .filter(|p| {
461                p.file_name()
462                    .and_then(|s| s.to_str())
463                    .is_some_and(|s| s.starts_with(&prefix) && s.ends_with(&format!(".{ext}")))
464            })
465            .collect();
466        chunks.sort();
467        out.extend(chunks);
468    }
469    out
470}
471
472/// gen.log: one TAB-separated line per placement, written into the work dir as
473/// an inspection aid (useful with `--keep-dir`); nothing reads it back.
474fn write_gen_log(
475    work: &Path,
476    placed: &[&Placement],
477    slots_by_sym: &std::collections::HashMap<String, u64>,
478) -> Result<()> {
479    let mut log = String::new();
480    for p in placed {
481        log.push_str(&format!("{}\t{}\t{}\t{}\n", p.name, p.base, p.sym, slots_by_sym[&p.sym]));
482    }
483    std::fs::write(work.join("gen.log"), log)?;
484    Ok(())
485}
486
487fn print_summary(s: &GenSummary, cfg: &GenConfig) {
488    let chunk_info = if let Some(chunk) = cfg.chunk { format!(", chunk={}", chunk) } else { String::new() };
489    eprintln!(
490        "generated {} kernels -> {} per-AIR .exps.so (CAP={}{}, max scratch {:.0}MB):",
491        s.generated.len(),
492        s.placed,
493        cfg.cap,
494        chunk_info,
495        s.max_scratch_bytes as f64 / 1e6
496    );
497    for g in &s.generated {
498        let chunked = if g.slots > 0 { format!("CHUNKED slots={}", g.slots) } else { "single".into() };
499        eprintln!("  {:40} {}.exps.so  nBits={} cExp={} ops={} {}", g.name, g.base, g.nbits, g.cexp, g.n_ops, chunked);
500    }
501    for (name, why) in &s.skipped {
502        eprintln!("  SKIP {name:38} {why}");
503    }
504}
505
506/// A work dir that is either user-provided (kept) or a unique temp dir removed on drop.
507struct WorkDir {
508    path: PathBuf,
509    temp: bool,
510}
511
512impl WorkDir {
513    fn new(keep_dir: Option<PathBuf>) -> Result<Self> {
514        match keep_dir {
515            Some(p) => {
516                std::fs::create_dir_all(&p)?;
517                eprintln!("[exps-codegen] keeping generated code in {}", p.display());
518                Ok(WorkDir { path: p, temp: false })
519            }
520            None => {
521                use std::sync::atomic::{AtomicU64, Ordering};
522                static SEQ: AtomicU64 = AtomicU64::new(0);
523                let seq = SEQ.fetch_add(1, Ordering::Relaxed);
524                let p = std::env::temp_dir().join(format!("genexps_{}_{}", std::process::id(), seq));
525                let _ = std::fs::remove_dir_all(&p);
526                std::fs::create_dir_all(&p)?;
527                Ok(WorkDir { path: p, temp: true })
528            }
529        }
530    }
531    fn path(&self) -> &Path {
532        &self.path
533    }
534}
535
536impl Drop for WorkDir {
537    fn drop(&mut self) {
538        if self.temp {
539            let _ = std::fs::remove_dir_all(&self.path);
540        }
541    }
542}
proofman_exps_codegen/lib.rs

proofman_exps_codegen/
lib.rs
