diff --git a/Cargo.lock b/Cargo.lock index 2868ac78d3..c62bdcb4b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7046,18 +7046,14 @@ dependencies = [ name = "sp1-gpu-air" version = "6.2.1" dependencies = [ - "itertools 0.14.0", "lazy_static", "slop-air", "slop-algebra", - "slop-koala-bear", "slop-matrix", "sp1-core-executor", "sp1-core-machine", - "sp1-curves", "sp1-hypercube", "sp1-primitives", - "sp1-recursion-machine", ] [[package]] @@ -7417,6 +7413,7 @@ version = "6.2.1" dependencies = [ "clap", "serde", + "slop-air", "slop-algebra", "slop-basefold", "slop-bn254", @@ -7472,7 +7469,9 @@ version = "6.2.1" dependencies = [ "criterion", "rand 0.8.6", + "serde_json", "serial_test", + "slop-air", "slop-algebra", "slop-alloc", "slop-basefold", diff --git a/sp1-gpu/crates/air/Cargo.toml b/sp1-gpu/crates/air/Cargo.toml index 77524666c5..4ec49c262b 100644 --- a/sp1-gpu/crates/air/Cargo.toml +++ b/sp1-gpu/crates/air/Cargo.toml @@ -11,18 +11,13 @@ categories.workspace = true [dependencies] slop-air = { workspace = true } slop-matrix = { workspace = true } -slop-koala-bear = { workspace = true } slop-algebra = { workspace = true } sp1-core-machine = { workspace = true } -sp1-recursion-machine = { workspace = true } sp1-hypercube = { workspace = true } -sp1-curves = { workspace = true } sp1-core-executor = { workspace = true } sp1-primitives = { workspace = true } -itertools = { workspace = true } - lazy_static = "1.4.0" [features] diff --git a/sp1-gpu/crates/air/examples/bytecode.rs b/sp1-gpu/crates/air/examples/bytecode.rs new file mode 100644 index 0000000000..5201b6c39c --- /dev/null +++ b/sp1-gpu/crates/air/examples/bytecode.rs @@ -0,0 +1,140 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! Smoke test for the v2 bytecode lowering. +//! +//! Builds DAG → analyzes → chunks → enumerates lowerings → lowers each +//! Sequential plan to bytecode. Prints per-chunk bytecode stats so we can +//! eyeball that the lowering's working correctly. + +use sp1_core_machine::riscv::RiscvAir; +use sp1_gpu_air::ir::{ + analyze_constraints, build_dag, chunk_dag, enumerate_lowerings, lower_sequential, BcOp, + ChunkBudget, Lowering, +}; +use sp1_gpu_air::F; +use sp1_hypercube::air::MachineAir; + +const FOCUS: &[&str] = &["Add", "Bitwise", "Mul", "Branch", "KeccakPermute"]; + +fn main() { + let machine = RiscvAir::::machine(); + let budget = ChunkBudget::recommended(); + + println!( + "{:<22} {:>5} {:>6} {:>5} {:>5} {:>5} {:>5} {:>5} {:>8}", + "chip", "nChnk", "instrs", "leaf", "const", "pub", "asrt", "maxR", "compat?" + ); + println!("{}", "-".repeat(78)); + + for chip in machine.chips() { + let name = chip.name(); + if !FOCUS.contains(&name) { + continue; + } + let dag = build_dag(chip.air.as_ref()); + let infos = analyze_constraints(&dag); + let chunks = chunk_dag(&infos, &budget); + + let mut total = LoweredStats::default(); + for chunk in &chunks { + let lowerings = enumerate_lowerings(chunk, &infos, &dag); + let plan = lowerings + .iter() + .find_map(|l| match l { + Lowering::Sequential(p) => Some(p), + _ => None, + }) + .unwrap(); + let bc = lower_sequential(chunk, &infos, &dag, plan); + total.add(&bc); + // Sanity assertions: + assert_eq!(bc.n_constraints, chunk.constraint_indices.len() as u32); + assert!(bc.max_reg as usize <= plan.topo_order.len()); + // Every assertion's reg index must be inside `max_reg`. + for &(reg, _) in &bc.asserts { + assert!(reg < bc.max_reg, "assert reg {} >= max_reg {}", reg, bc.max_reg); + } + } + println!( + "{:<22} {:>5} {:>6} {:>5} {:>5} {:>5} {:>5} {:>5} {:>8}", + name, + chunks.len(), + total.instrs, + total.leaves, + total.consts, + total.publics, + total.asserts, + total.max_reg_seen, + if total.all_compatible { "yes" } else { "MIXED" }, + ); + } + + println!("\nOpcode distribution across all chunks (KeccakPermute):"); + for chip in machine.chips() { + if chip.name() != "KeccakPermute" { + continue; + } + let dag = build_dag(chip.air.as_ref()); + let infos = analyze_constraints(&dag); + let chunks = chunk_dag(&infos, &budget); + let mut counts = [0u64; 8]; + for chunk in &chunks { + let lowerings = enumerate_lowerings(chunk, &infos, &dag); + let plan = lowerings + .iter() + .find_map(|l| match l { + Lowering::Sequential(p) => Some(p), + _ => None, + }) + .unwrap(); + let bc = lower_sequential(chunk, &infos, &dag, plan); + for instr in &bc.instrs { + counts[instr.opcode as usize] += 1; + } + } + let names = + ["LoadLeaf", "LoadConst", "LoadPublic", "AddF", "SubF", "MulF", "NegF", "AssertF"]; + for (op, &n) in names.iter().zip(counts.iter()) { + if n > 0 { + println!(" {:<12} {}", op, n); + } + } + break; + } +} + +#[derive(Default)] +struct LoweredStats { + instrs: usize, + leaves: usize, + consts: usize, + publics: usize, + asserts: usize, + max_reg_seen: u16, + all_compatible: bool, +} + +impl LoweredStats { + fn add(&mut self, bc: &sp1_gpu_air::ir::ChunkBytecode) { + self.instrs += bc.instrs.len(); + self.leaves += bc.leaves.len(); + self.consts += bc.consts.len(); + self.publics += bc.publics.len(); + self.asserts += bc.asserts.len(); + self.max_reg_seen = self.max_reg_seen.max(bc.max_reg); + self.all_compatible = bc.is_compatible_with_v1_kernel() || self.all_compatible; + } +} + +// Suppress unused-import warning for BcOp (used in opcode-array indexing). +const _: fn() = || { + let _ = BcOp::LoadLeaf; +}; diff --git a/sp1-gpu/crates/air/examples/chunk_efficiency.rs b/sp1-gpu/crates/air/examples/chunk_efficiency.rs new file mode 100644 index 0000000000..57b94d5b1b --- /dev/null +++ b/sp1-gpu/crates/air/examples/chunk_efficiency.rs @@ -0,0 +1,140 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::doc_lazy_continuation +)] +//! Arithmetization / chunking efficiency report. +//! +//! For every `RiscvAir` chip, measure how many column loads the chunker emits +//! relative to how many distinct columns the chip's constraints actually +//! reference. This is the direct measure of *chunking quality*. +//! +//! Definitions (all at `ColumnLeaf` granularity — a `(source, col)` pair, so a +//! column read both as `MainLocal` and `MainNext` counts as two): +//! +//! * `refd` — distinct column leaves referenced anywhere in the chip's +//! constraints (the union over all chunks). Each must be loaded at least +//! once; this is the irreducible floor. +//! * `loaded` — `Σ over chunks of |chunk.leafset|`. The fused sequential kernel +//! materialises one register slot per leaf *per chunk* and reuses it within +//! the chunk, so a column shared by K chunks is fetched from memory K times. +//! * `reload` — `loaded / refd`. 1.00 = perfect (every column loaded once); +//! >1.00 = the chunk split forces redundant column re-fetches. +//! +//! Separately, against the chip's physical width (`main + preprocessed` +//! columns), `touch = refd / (2 × width)` shows what fraction of the chip's +//! leaf space (each column × {local, next}) the arithmetization actually +//! reads — a constraint-density signal, not a chunking one. +//! +//! Run with: cargo run --release --example chunk_efficiency -p sp1-gpu-air + +use slop_air::BaseAir; +use sp1_core_machine::riscv::RiscvAir; +use sp1_gpu_air::ir::{analyze_constraints, build_dag, chunk_dag, Chunk, ChunkBudget}; +use sp1_gpu_air::F; +use sp1_hypercube::air::MachineAir; + +struct Row { + name: String, + width: usize, // physical columns: main + preprocessed + refd: usize, // distinct leaves referenced + loaded: usize, // Σ |chunk leafset| + nchunk: usize, + ovr: usize, // chunks that are a single over-budget constraint +} + +fn measure(name: &str, width: usize, chunks: &[Chunk]) -> Row { + let loaded: usize = chunks.iter().map(|c| c.leafset.len()).sum(); + let ovr = chunks.iter().filter(|c| c.oversize_singleton).count(); + let mut refd_set = std::collections::HashSet::new(); + for chunk in chunks { + for &leaf in &chunk.leafset { + refd_set.insert(leaf); + } + } + Row { name: name.to_string(), width, refd: refd_set.len(), loaded, nchunk: chunks.len(), ovr } +} + +fn main() { + let machine = RiscvAir::::machine(); + + let budgets = [ + ChunkBudget { max_leafset: 64, max_constraints_per_chunk: 1024 }, + ChunkBudget { max_leafset: 128, max_constraints_per_chunk: 1024 }, + ChunkBudget { max_leafset: 256, max_constraints_per_chunk: 1024 }, + ]; + + for budget in &budgets { + let mut rows: Vec = Vec::new(); + for chip in machine.chips() { + let width = chip.width() + chip.preprocessed_width(); + let dag = build_dag(chip.air.as_ref()); + let infos = analyze_constraints(&dag); + let chunks = chunk_dag(&infos, budget); + if chunks.is_empty() { + continue; // chip with no column-reading constraints + } + rows.push(measure(chip.name(), width, &chunks)); + } + // Heaviest column-load chips first. + rows.sort_by_key(|r| std::cmp::Reverse(r.loaded)); + + println!("\n=== chunk efficiency — max_leafset={} ===", budget.max_leafset); + println!( + "{:<28} {:>5} {:>6} {:>6} {:>4} {:>6} {:>8} {:>7}", + "chip", "width", "refd", "nChnk", "ovr", "loaded", "reload x", "touch %" + ); + let (mut tot_refd, mut tot_loaded, mut tot_chunks, mut tot_ovr) = (0usize, 0, 0, 0); + let mut tot_leafspace = 0usize; + for r in &rows { + let reload = r.loaded as f64 / r.refd.max(1) as f64; + let touch = 100.0 * r.refd as f64 / (2 * r.width).max(1) as f64; + println!( + "{:<28} {:>5} {:>6} {:>6} {:>4} {:>6} {:>8.3} {:>6.1}%", + truncate(&r.name, 28), + r.width, + r.refd, + r.nchunk, + r.ovr, + r.loaded, + reload, + touch, + ); + tot_refd += r.refd; + tot_loaded += r.loaded; + tot_chunks += r.nchunk; + tot_ovr += r.ovr; + tot_leafspace += 2 * r.width; + } + println!("{:-<28}", ""); + let machine_reload = tot_loaded as f64 / tot_refd.max(1) as f64; + let machine_touch = 100.0 * tot_refd as f64 / tot_leafspace.max(1) as f64; + println!( + "{:<28} {:>5} {:>6} {:>6} {:>4} {:>6} {:>8.3} {:>6.1}%", + format!("TOTAL ({} chips)", rows.len()), + "", + tot_refd, + tot_chunks, + tot_ovr, + tot_loaded, + machine_reload, + machine_touch, + ); + let waste = tot_loaded.saturating_sub(tot_refd); + println!( + " machine-wide: {tot_loaded} column loads for {tot_refd} distinct columns \ + → {waste} redundant re-fetches ({:.1}% overhead)", + 100.0 * waste as f64 / tot_refd.max(1) as f64, + ); + } +} + +fn truncate(s: &str, n: usize) -> String { + if s.len() <= n { + s.to_string() + } else { + format!("{}...", &s[..n.saturating_sub(3)]) + } +} diff --git a/sp1-gpu/crates/air/examples/chunker.rs b/sp1-gpu/crates/air/examples/chunker.rs new file mode 100644 index 0000000000..e67e7d63d3 --- /dev/null +++ b/sp1-gpu/crates/air/examples/chunker.rs @@ -0,0 +1,143 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! Phase 1 validation: run the v2 builder + analysis + chunker on every chip +//! in `RiscvAir` and reproduce the chunk-count and union-leafset distributions +//! we saw in the earlier `analyze_chips` study. +//! +//! The numbers here should match v1 to within rounding; differences are a +//! warning sign the v2 path lost something. + +use sp1_core_machine::riscv::RiscvAir; +use sp1_gpu_air::ir::{ + analyze_constraints, build_dag, chunk_dag, Chunk, ChunkBudget, ConstraintShape, +}; +use sp1_gpu_air::F; +use sp1_hypercube::air::MachineAir; + +const FOCUS: &[&str] = &[ + "KeccakPermute", + "Secp256k1AddAssign", + "Secp256k1DoubleAssign", + "Bn254AddAssign", + "Bls12381AddAssign", + "Bls12381FpOpAssign", + "EdAddAssign", + "Add", + "Bitwise", + "Mul", + "Branch", + "Global", +]; + +fn main() { + let machine = RiscvAir::::machine(); + + let budgets = [ + ChunkBudget { max_leafset: 64, max_constraints_per_chunk: 1024 }, + ChunkBudget { max_leafset: 128, max_constraints_per_chunk: 1024 }, + ChunkBudget { max_leafset: 256, max_constraints_per_chunk: 1024 }, + ]; + + for budget in &budgets { + println!("\n=== budget: max_leafset={} ===", budget.max_leafset); + println!( + "{:<28} {:>5} {:>6} {:>5} {:>6} {:>5} {:>5} {:>5} {:>7} {:>3}", + "chip", "cons", "nodes", "depth", "nChnk", "maxU", "p50U", "meanC", "linear%", "ovr" + ); + for chip in machine.chips() { + let name = chip.name(); + if !FOCUS.contains(&name) { + continue; + } + let dag = build_dag(chip.air.as_ref()); + let infos = analyze_constraints(&dag); + let chunks = chunk_dag(&infos, budget); + print_row(name, &infos, &chunks); + } + } + + // Show that GKR-shape detection works: synthesize one chip's worth of + // LinearWeightedSum constraints and confirm the chunker tags them. + println!("\n=== shape-detection report ==="); + let machine = RiscvAir::::machine(); + for chip in machine.chips() { + let dag = build_dag(chip.air.as_ref()); + let infos = analyze_constraints(&dag); + let n_linear = + infos.iter().filter(|c| matches!(c.shape, ConstraintShape::LinearWeightedSum)).count(); + if n_linear > 0 { + println!( + " {:<32} {:>4} / {:<4} constraints are LinearWeightedSum", + chip.name(), + n_linear, + infos.len() + ); + } + } +} + +fn print_row(name: &str, infos: &[sp1_gpu_air::ir::ConstraintInfo], chunks: &[Chunk]) { + let cons = infos.len(); + let nodes: usize = infos.iter().map(|c| c.total_nodes as usize).sum(); + let depth_max = infos.iter().map(|c| c.depth).max().unwrap_or(0); + + let unions: Vec = chunks.iter().map(|c| c.leafset.len()).collect(); + let max_u = unions.iter().copied().max().unwrap_or(0); + let p50_u = pct(&unions, 0.5); + let mean_c = if chunks.is_empty() { + 0.0 + } else { + chunks.iter().map(|c| c.constraint_indices.len()).sum::() as f64 + / chunks.len() as f64 + }; + let linear_pct = if chunks.is_empty() { + 0 + } else { + 100 * chunks + .iter() + .filter(|c| matches!(c.shape, ConstraintShape::LinearWeightedSum)) + .count() + / chunks.len() + }; + let oversize = chunks.iter().filter(|c| c.oversize_singleton).count(); + + println!( + "{:<28} {:>5} {:>6} {:>5} {:>6} {:>5} {:>5} {:>5.1} {:>6}% {:>3}", + truncate(name, 28), + cons, + nodes, + depth_max, + chunks.len(), + max_u, + p50_u, + mean_c, + linear_pct, + oversize, + ); +} + +fn pct(xs: &[T], p: f64) -> T { + if xs.is_empty() { + return T::default(); + } + let mut s = xs.to_vec(); + s.sort(); + let idx = ((s.len() as f64) * p).floor() as usize; + s[idx.min(s.len() - 1)] +} + +fn truncate(s: &str, n: usize) -> String { + if s.len() <= n { + s.to_string() + } else { + format!("{}...", &s[..n.saturating_sub(3)]) + } +} diff --git a/sp1-gpu/crates/air/examples/dump_chip_layout.rs b/sp1-gpu/crates/air/examples/dump_chip_layout.rs new file mode 100644 index 0000000000..8f75651bdd --- /dev/null +++ b/sp1-gpu/crates/air/examples/dump_chip_layout.rs @@ -0,0 +1,37 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! Dump every `RiscvAir` chip as a JSON layout entry (height 0). Feed the +//! output into the zerocheck bench's JSON layout source, with heights edited +//! to whatever distribution you want to test. +//! +//! Run with: cargo run --release --example dump_chip_layout -p sp1-gpu-air + +use slop_air::BaseAir; +use sp1_core_machine::riscv::RiscvAir; +use sp1_gpu_air::F; +use sp1_hypercube::air::MachineAir; + +fn main() { + let machine = RiscvAir::::machine(); + let chips: Vec<_> = machine.chips().iter().collect(); + println!("["); + for (i, chip) in chips.iter().enumerate() { + let comma = if i + 1 < chips.len() { "," } else { "" }; + println!( + " {{\"name\": \"{}\", \"preprocessed_width\": {}, \"main_width\": {}, \"height\": 0}}{}", + chip.name(), + chip.preprocessed_width(), + chip.width(), + comma, + ); + } + println!("]"); +} diff --git a/sp1-gpu/crates/air/examples/scheduler.rs b/sp1-gpu/crates/air/examples/scheduler.rs new file mode 100644 index 0000000000..9a7b7279e5 --- /dev/null +++ b/sp1-gpu/crates/air/examples/scheduler.rs @@ -0,0 +1,90 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! Phase 1 capstone: build a chip's DAG, chunk it, enumerate lowerings, and +//! show what the heuristic cost model picks per round across the sumcheck. +//! +//! Validates that: +//! - column-tile is elected for linear-weighted-sum chunks +//! - sequential is elected when row × eval ≥ warp_size +//! - lane-starved late rounds gracefully fall through to sequential +//! (the placeholder until Phase 5) + +use sp1_core_machine::riscv::RiscvAir; +use sp1_gpu_air::ir::{ + analyze_constraints, build_dag, chunk_dag, enumerate_lowerings, pick_lowering, ChunkBudget, + GpuCaps, LaneBudget, Lowering, +}; +use sp1_gpu_air::F; +use sp1_hypercube::air::MachineAir; + +fn main() { + let machine = RiscvAir::::machine(); + let caps = GpuCaps::ampere_default(); + let budget = ChunkBudget { max_leafset: 64, max_constraints_per_chunk: 1024 }; + + // Simulate per-round lane budgets for a few chip log-row-counts and report + // per-(chunk, round) decisions aggregated by lowering kind. + let scenarios: &[(&str, u32)] = &[ + ("KeccakPermute", 14), // ~16k initial rows + ("Bls12381FpOpAssign", 8), + ("Add", 18), + ("Bitwise", 18), + ("Branch", 16), + ]; + + for &(chip_name, log_rows) in scenarios { + let chip = match machine.chips().iter().find(|c| c.name() == chip_name) { + Some(c) => c, + None => continue, + }; + let dag = build_dag(chip.air.as_ref()); + let infos = analyze_constraints(&dag); + let chunks = chunk_dag(&infos, &budget); + let plans: Vec> = + chunks.iter().map(|c| enumerate_lowerings(c, &infos, &dag)).collect(); + + println!( + "\n=== {} (initial log_rows = {}, chunks = {}) ===", + chip_name, + log_rows, + chunks.len() + ); + println!( + "{:>4} {:>10} {:>10} {:>10} {:>10}", + "rnd", "lanes", "sequential", "columntile", "starved" + ); + for r in 0..=log_rows { + let cur_log_rows = log_rows.saturating_sub(r); + let lane_budget = LaneBudget { log_rows: cur_log_rows, eval_points: 3 }; + let lanes = lane_budget.total_lanes(); + + let mut counts = [0u32; 3]; // [sequential_warpfull, columntile, starved] + for (chunk, lowerings) in chunks.iter().zip(plans.iter()) { + let pick = pick_lowering(chunk, lowerings, lane_budget, caps); + match &lowerings[pick.lowering_idx] { + Lowering::ColumnTile(_) => counts[1] += 1, + Lowering::Sequential(_) => { + if lanes >= caps.warp_size { + counts[0] += 1; + } else { + counts[2] += 1; + } + } + _ => {} + } + } + println!( + "{:>4} {:>10} {:>10} {:>10} {:>10}", + r, lanes, counts[0], counts[1], counts[2] + ); + } + } +} diff --git a/sp1-gpu/crates/air/examples/smoke.rs b/sp1-gpu/crates/air/examples/smoke.rs new file mode 100644 index 0000000000..bff5570063 --- /dev/null +++ b/sp1-gpu/crates/air/examples/smoke.rs @@ -0,0 +1,163 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! Smoke test for the v2 DAG-native builder. +//! +//! Runs `build_dag` on a few chips, prints node/constraint counts and shows +//! cross-constraint sharing. Sanity check before further phases. + +use std::collections::HashSet; + +use sp1_core_machine::riscv::RiscvAir; +use sp1_gpu_air::ir::{build_dag, ConstraintDag, DagNode, NodeId, TraceSource}; +use sp1_gpu_air::F; +use sp1_hypercube::air::MachineAir; + +fn main() { + let machine = RiscvAir::::machine(); + + let focus: &[&str] = &["Add", "Bitwise", "Mul", "Branch", "KeccakPermute"]; + + println!( + "{:<22} {:>5} {:>5} {:>6} {:>6} {:>7} {:>7} {:>9}", + "chip", "main", "prep", "nodes", "cons", "leaves", "shared", "leafShar" + ); + println!("{}", "-".repeat(78)); + + for chip in machine.chips() { + let name = chip.name(); + if !focus.contains(&name) { + continue; + } + let dag = build_dag(chip.air.as_ref()); + print_summary(name, &dag); + } + + println!("\nDetailed leaf-share check (KeccakPermute):"); + for chip in machine.chips() { + if chip.name() == "KeccakPermute" { + let dag = build_dag(chip.air.as_ref()); + print_sharing_details(&dag); + break; + } + } +} + +fn print_summary(name: &str, dag: &ConstraintDag) { + let leaves: HashSet = dag + .nodes + .iter() + .enumerate() + .filter_map(|(i, n)| match n { + DagNode::InputLeaf { .. } => Some(i as NodeId), + _ => None, + }) + .collect(); + // For each leaf, count how many parent edges point at it. + let mut leaf_uses: std::collections::HashMap = + leaves.iter().map(|&l| (l, 0u32)).collect(); + for n in &dag.nodes { + for c in children(n) { + if let Some(count) = leaf_uses.get_mut(&c) { + *count += 1; + } + } + } + // Constraints also reference leaves through their roots; root → leaf walks happen below. + let shared_leaves = leaf_uses.values().filter(|&&c| c > 1).count(); + let mean_uses = if leaf_uses.is_empty() { + 0.0 + } else { + leaf_uses.values().sum::() as f64 / leaf_uses.len() as f64 + }; + + println!( + "{:<22} {:>5} {:>5} {:>6} {:>6} {:>7} {:>7} {:>9.2}", + name, + dag.main_width, + dag.preprocessed_width, + dag.nodes.len(), + dag.constraints.len(), + leaves.len(), + shared_leaves, + mean_uses, + ); +} + +fn children(node: &DagNode) -> Vec { + match *node { + DagNode::InputLeaf { .. } + | DagNode::PublicValue { .. } + | DagNode::GlobalCumulativeSum { .. } + | DagNode::ConstF { .. } + | DagNode::ConstEF { .. } + | DagNode::IsFirstRow + | DagNode::IsLastRow + | DagNode::IsTransition => vec![], + DagNode::AddF { a, b } + | DagNode::SubF { a, b } + | DagNode::MulF { a, b } + | DagNode::AddEF { a, b } + | DagNode::SubEF { a, b } + | DagNode::MulEF { a, b } + | DagNode::EFAddF { a, b } + | DagNode::EFSubF { a, b } + | DagNode::EFMulF { a, b } => vec![a, b], + DagNode::NegF { a } | DagNode::NegEF { a } | DagNode::EFFromF { a } => vec![a], + } +} + +fn print_sharing_details(dag: &ConstraintDag) { + // Walk back from each constraint root and collect transitive column leaves. + // Compare per-constraint leaf set sizes with the SSA-tape result from v1. + let mut per_constraint_main: Vec> = Vec::new(); + for c in &dag.constraints { + let mut visited = HashSet::new(); + let mut main_cols = HashSet::new(); + collect_main_columns(&dag.nodes, c.root, &mut visited, &mut main_cols); + per_constraint_main.push(main_cols); + } + let max_main = per_constraint_main.iter().map(|s| s.len()).max().unwrap_or(0); + let mean_main = if per_constraint_main.is_empty() { + 0.0 + } else { + per_constraint_main.iter().map(|s| s.len()).sum::() as f64 + / per_constraint_main.len() as f64 + }; + let union_main: HashSet = per_constraint_main.iter().flatten().copied().collect(); + let sum_main: usize = per_constraint_main.iter().map(|s| s.len()).sum(); + + println!(" per-constraint MainLocal leaf counts: max={} mean={:.1}", max_main, mean_main); + println!(" union(MainLocal) across all constraints: {}", union_main.len()); + println!( + " overlap factor (sum/union): {:.2}x (compare with v1 analyze_chips ≈ 6.0x)", + sum_main as f64 / union_main.len().max(1) as f64 + ); +} + +fn collect_main_columns( + nodes: &[DagNode], + root: NodeId, + visited: &mut HashSet, + main_cols: &mut HashSet, +) { + if !visited.insert(root) { + return; + } + let node = &nodes[root as usize]; + if let DagNode::InputLeaf { source: TraceSource::MainLocal, col } + | DagNode::InputLeaf { source: TraceSource::MainNext, col } = node + { + main_cols.insert(*col); + } + for c in children(node) { + collect_main_columns(nodes, c, visited, main_cols); + } +} diff --git a/sp1-gpu/crates/air/src/air_block.rs b/sp1-gpu/crates/air/src/air_block.rs deleted file mode 100644 index 4933caef14..0000000000 --- a/sp1-gpu/crates/air/src/air_block.rs +++ /dev/null @@ -1,903 +0,0 @@ -use crate::{ - symbolic_expr_f::SymbolicExprF, symbolic_var_f::SymbolicVarF, SymbolicProverFolder, F, -}; -use itertools::Itertools; -use slop_air::{Air, AirBuilder, PairBuilder}; -use slop_algebra::AbstractField; -use slop_matrix::Matrix; -use sp1_core_executor::events::FieldOperation; -use sp1_core_executor::SyscallCode; -use sp1_core_machine::air::{MemoryAirBuilder, SP1CoreAirBuilder}; -use sp1_core_machine::operations::{ - AddrAddOperation, AddressSlicePageProtOperation, SyscallAddrOperation, -}; -use sp1_core_machine::riscv::{WeierstrassAddAssignChip, WeierstrassDoubleAssignChip}; -use sp1_core_machine::syscall::precompiles::weierstrass::{ - WeierstrassAddAssignCols, WeierstrassDoubleAssignCols, -}; -use sp1_core_machine::utils::limbs_to_words; -use sp1_core_machine::{ - riscv::{KeccakPermuteChip, RiscvAir}, - syscall::precompiles::keccak256::{columns::KeccakMemCols, constants::rc_value_bit}, -}; -use sp1_core_machine::{TrustMode, UserMode}; -use sp1_curves::k256::elliptic_curve::generic_array::typenum::Unsigned; -use sp1_curves::params::FieldParameters; -use sp1_curves::params::{Limbs, NumLimbs}; -use sp1_curves::weierstrass::WeierstrassParameters; -use sp1_curves::{BigUint, CurveType, EllipticCurve}; -use sp1_hypercube::air::InstructionAirBuilder; -#[cfg(feature = "mprotect")] -use sp1_hypercube::air::MachineAirBuilder; -use sp1_hypercube::operations::poseidon2::air::{eval_external_round, eval_internal_rounds}; -use sp1_hypercube::operations::poseidon2::WIDTH; -use sp1_hypercube::Word; -use sp1_hypercube::{ - air::{AirInteraction, InteractionScope, MachineAir, MessageBuilder}, - InteractionKind, -}; -use sp1_primitives::consts::{PROT_READ, PROT_WRITE}; -use sp1_primitives::polynomial::Polynomial; -use sp1_primitives::SP1Field; -use sp1_recursion_machine::builder::RecursionAirBuilder; -use sp1_recursion_machine::chips::poseidon2_wide::columns::preprocessed::Poseidon2PreprocessedColsWide; -use sp1_recursion_machine::chips::poseidon2_wide::Poseidon2WideChip; -use sp1_recursion_machine::RecursionAir; -use std::borrow::Borrow; -use std::iter::once; -pub trait BlockAir: Air + MachineAir + 'static + Send + Sync { - fn num_blocks(&self) -> usize { - 1 - } - - fn eval_block(&self, builder: &mut AB, index: usize) { - assert!(index == 0); - self.eval(builder); - } -} - -impl<'a> BlockAir> for RiscvAir { - fn num_blocks(&self) -> usize { - match self { - RiscvAir::KeccakP(keccak) => keccak.num_blocks(), - RiscvAir::Secp256k1Add(secp256k1_add) => secp256k1_add.num_blocks(), - RiscvAir::Secp256k1AddUser(secp256k1_add) => secp256k1_add.num_blocks(), - RiscvAir::Secp256k1Double(secp256k1_double) => secp256k1_double.num_blocks(), - RiscvAir::Secp256k1DoubleUser(secp256k1_double) => secp256k1_double.num_blocks(), - _ => 1, - } - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder<'a>, index: usize) { - match self { - RiscvAir::KeccakP(keccak) => keccak.eval_block(builder, index), - RiscvAir::Secp256k1Add(secp256k1_add) => secp256k1_add.eval_block(builder, index), - RiscvAir::Secp256k1AddUser(secp256k1_add) => secp256k1_add.eval_block(builder, index), - RiscvAir::Secp256k1Double(secp256k1_double) => { - secp256k1_double.eval_block(builder, index) - } - RiscvAir::Secp256k1DoubleUser(secp256k1_double) => { - secp256k1_double.eval_block(builder, index) - } - _ => { - assert!(index == 0); - self.eval(builder); - } - } - } -} - -impl<'a, const DEGREE: usize, const VAR_EVENTS_PER_ROW: usize> BlockAir> - for RecursionAir -{ - fn num_blocks(&self) -> usize { - match self { - RecursionAir::Poseidon2Wide(poseidon2_wide) => poseidon2_wide.num_blocks(), - _ => 1, - } - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder<'a>, index: usize) { - match self { - RecursionAir::Poseidon2Wide(poseidon2_wide) => { - poseidon2_wide.eval_block(builder, index) - } - _ => { - assert!(index == 0); - self.eval(builder); - } - } - } -} - -impl<'a> BlockAir> for KeccakPermuteChip { - fn num_blocks(&self) -> usize { - 11 - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder<'a>, index: usize) { - const NUM_ROUNDS: usize = 24; - const BITS_PER_LIMB: usize = 16; - const U64_LIMBS: usize = 4; - - let main = builder.main(); - let local = main.row_slice(0); - let local: &KeccakMemCols = (*local).borrow(); - - // Keccak AIRs from Plonky3. - let andn_gen = |a: SymbolicExprF, b: SymbolicExprF| b - a * b; - let xor_gen = |a: SymbolicExprF, b: SymbolicExprF| a + b - a * b.double(); - let xor3_gen = - |a: SymbolicExprF, b: SymbolicExprF, c: SymbolicExprF| xor_gen(a, xor_gen(b, c)); - - match index { - 0 => { - builder.assert_bool(local.is_real); - // Flag constraints. - let mut sum_flags = SymbolicExprF::zero(); - let mut computed_index = SymbolicExprF::zero(); - for i in 0..NUM_ROUNDS { - builder.assert_bool(local.keccak.step_flags[i]); - sum_flags = sum_flags + local.keccak.step_flags[i]; - computed_index = computed_index - + SymbolicExprF::from_canonical_u32(i as u32) * local.keccak.step_flags[i]; - } - builder.assert_one(sum_flags); - builder.when(local.is_real).assert_eq(computed_index, local.index); - - // C'[x, z] = xor(C[x, z], C[x - 1, z], C[x + 1, z - 1]). - for x in 0..5 { - for z in 0..64 { - builder.assert_bool(local.keccak.c[x][z]); - let xor = xor3_gen( - local.keccak.c[x][z].into(), - local.keccak.c[(x + 4) % 5][z].into(), - local.keccak.c[(x + 1) % 5][(z + 63) % 64].into(), - ); - let c_prime = local.keccak.c_prime[x][z]; - builder.assert_eq(c_prime, xor); - } - } - } - 1..=5 => { - // Check that the input limbs are consistent with A' and D. - // A[x, y, z] = xor(A'[x, y, z], D[x, y, z]) - // = xor(A'[x, y, z], C[x - 1, z], C[x + 1, z - 1]) - // = xor(A'[x, y, z], C[x, z], C'[x, z]). - // The last step is valid based on the identity we checked above. - // It isn't required, but makes this check a bit cleaner. - let y = index - 1; - for x in 0..5 { - let get_bit = |z| { - let a_prime: SymbolicVarF = local.keccak.a_prime[y][x][z]; - let c: SymbolicVarF = local.keccak.c[x][z]; - let c_prime: SymbolicVarF = local.keccak.c_prime[x][z]; - xor3_gen(a_prime.into(), c.into(), c_prime.into()) - }; - - for limb in 0..U64_LIMBS { - let a_limb = local.keccak.a[y][x][limb]; - let computed_limb = (limb * BITS_PER_LIMB..(limb + 1) * BITS_PER_LIMB) - .rev() - .fold(SymbolicExprF::zero(), |acc, z| { - builder.assert_bool(local.keccak.a_prime[y][x][z]); - acc.double() + get_bit(z) - }); - builder.assert_eq(computed_limb, a_limb); - } - } - } - 6 => { - for x in 0..5 { - for z in 0..64 { - let sum: SymbolicExprF = - (0..5).map(|y| local.keccak.a_prime[y][x][z].into()).sum(); - let diff = sum - local.keccak.c_prime[x][z]; - let four = SymbolicExprF::from_canonical_u8(4); - builder.assert_zero(diff * (diff - SymbolicExprF::two()) * (diff - four)); - } - } - } - 7..=9 => { - let y_range = match index { - 7 => 0..2, - 8 => 2..4, - 9 => 4..5, - _ => unreachable!(), - }; - // A''[x, y] = xor(B[x, y], andn(B[x + 1, y], B[x + 2, y])). - for y in y_range { - for x in 0..5 { - let get_bit = |z| { - let andn = andn_gen( - local.keccak.b((x + 1) % 5, y, z).into(), - local.keccak.b((x + 2) % 5, y, z).into(), - ); - xor_gen(local.keccak.b(x, y, z).into(), andn) - }; - - for limb in 0..U64_LIMBS { - let computed_limb = (limb * BITS_PER_LIMB..(limb + 1) * BITS_PER_LIMB) - .rev() - .fold(SymbolicExprF::zero(), |acc, z| acc.double() + get_bit(z)); - builder - .assert_eq(computed_limb, local.keccak.a_prime_prime[y][x][limb]); - } - } - } - } - 10 => { - // A'''[0, 0] = A''[0, 0] XOR RC - for limb in 0..U64_LIMBS { - let computed_a_prime_prime_0_0_limb = (limb * BITS_PER_LIMB - ..(limb + 1) * BITS_PER_LIMB) - .rev() - .fold(SymbolicExprF::zero(), |acc, z| { - builder.assert_bool(local.keccak.a_prime_prime_0_0_bits[z]); - acc.double() + local.keccak.a_prime_prime_0_0_bits[z] - }); - let a_prime_prime_0_0_limb = local.keccak.a_prime_prime[0][0][limb]; - builder.assert_eq(computed_a_prime_prime_0_0_limb, a_prime_prime_0_0_limb); - } - - let get_xored_bit = |i| { - let mut rc_bit_i = SymbolicExprF::zero(); - for r in 0..NUM_ROUNDS { - let this_round = local.keccak.step_flags[r]; - let this_round_constant = - SymbolicExprF::from_canonical_u8(rc_value_bit(r, i)); - rc_bit_i = rc_bit_i + this_round * this_round_constant; - } - - xor_gen(local.keccak.a_prime_prime_0_0_bits[i].into(), rc_bit_i) - }; - - for limb in 0..U64_LIMBS { - let a_prime_prime_prime_0_0_limb = - local.keccak.a_prime_prime_prime_0_0_limbs[limb]; - let computed_a_prime_prime_prime_0_0_limb = (limb * BITS_PER_LIMB - ..(limb + 1) * BITS_PER_LIMB) - .rev() - .fold(SymbolicExprF::zero(), |acc, z| acc.double() + get_xored_bit(z)); - builder.assert_eq( - computed_a_prime_prime_prime_0_0_limb, - a_prime_prime_prime_0_0_limb, - ); - } - // Receive state. - let receive_values = - once(local.clk_high) - .chain(once(local.clk_low)) - .chain(local.state_addr) - .chain(once(local.index)) - .chain(local.keccak.a.into_iter().flat_map(|two_d| { - two_d.into_iter().flat_map(|one_d| one_d.into_iter()) - })) - .collect::>(); - - builder.receive( - AirInteraction::new(receive_values, local.is_real, InteractionKind::Keccak), - InteractionScope::Local, - ); - - // Send state. - let send_values = once(local.clk_high.into()) - .chain(once(local.clk_low.into())) - .chain(local.state_addr.map(Into::into)) - .chain(once(local.index + SymbolicExprF::one())) - .chain((0..5).flat_map(|y| { - (0..5).flat_map(move |x| { - (0..4).map(move |limb| { - local.keccak.a_prime_prime_prime(y, x, limb).into() - }) - }) - })) - .collect::>(); - - builder.send( - AirInteraction::new(send_values, local.is_real.into(), InteractionKind::Keccak), - InteractionScope::Local, - ); - } - _ => unreachable!(), - }; - } -} - -impl<'a, E: EllipticCurve + WeierstrassParameters, M: TrustMode> BlockAir> - for WeierstrassAddAssignChip -where - Limbs::Limbs>: Copy, -{ - fn num_blocks(&self) -> usize { - 11 - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder<'a>, index: usize) { - let main = builder.main(); - let local = main.row_slice(0); - let local: &WeierstrassAddAssignCols = (*local).borrow(); - - // Fetch the syscall id for the curve type. - let syscall_id_felt = match E::CURVE_TYPE { - CurveType::Secp256k1 => F::from_canonical_u32(SyscallCode::SECP256K1_ADD.syscall_id()), - CurveType::Secp256r1 => F::from_canonical_u32(SyscallCode::SECP256R1_ADD.syscall_id()), - CurveType::Bn254 => F::from_canonical_u32(SyscallCode::BN254_ADD.syscall_id()), - CurveType::Bls12381 => F::from_canonical_u32(SyscallCode::BLS12381_ADD.syscall_id()), - _ => panic!("Unsupported curve"), - }; - let num_words_field_element = ::Limbs::USIZE / 8; - - // It's very important that the `generate_limbs` function do not call `assert_zero`. - let p_x_limbs = builder - .generate_limbs(&local.p_access[0..num_words_field_element], local.is_real.into()); - let p_x: Limbs::Limbs> = - Limbs(p_x_limbs.try_into().expect("failed to convert limbs")); - let p_y_limbs = builder - .generate_limbs(&local.p_access[num_words_field_element..], local.is_real.into()); - let p_y: Limbs::Limbs> = - Limbs(p_y_limbs.try_into().expect("failed to convert limbs")); - let q_x_limbs = builder - .generate_limbs(&local.q_access[0..num_words_field_element], local.is_real.into()); - let q_x: Limbs::Limbs> = - Limbs(q_x_limbs.try_into().expect("failed to convert limbs")); - let q_y_limbs = builder - .generate_limbs(&local.q_access[num_words_field_element..], local.is_real.into()); - let q_y: Limbs::Limbs> = - Limbs(q_y_limbs.try_into().expect("failed to convert limbs")); - - let is_not_trap: SymbolicExprF = local.is_real.into(); - let trap_code = SymbolicExprF::zero(); - - match index { - 0 => { - let mut is_not_trap = local.is_real.into(); - let mut trap_code = SymbolicExprF::zero(); - - #[cfg(feature = "mprotect")] - builder.assert_eq( - builder.extract_public_values().is_untrusted_programs_enabled, - SymbolicExprF::from_bool(!M::IS_TRUSTED), - ); - - if !M::IS_TRUSTED { - let local = main.row_slice(0); - let local: &WeierstrassAddAssignCols = - (*local).borrow(); - - #[cfg(not(feature = "mprotect"))] - builder.assert_zero(local.is_real); - - AddressSlicePageProtOperation::::eval( - builder, - local.clk_high.into(), - local.clk_low.into(), - &local.q_ptr.addr.map(Into::into), - &local.q_addrs[local.q_addrs.len() - 1].value.map(Into::into), - PROT_READ, - &local.read_slice_page_prot_access, - &mut is_not_trap, - &mut trap_code, - ); - - let clk_low: SymbolicExprF = local.clk_low.into(); - - AddressSlicePageProtOperation::::eval( - builder, - local.clk_high.into(), - clk_low + SymbolicExprF::one(), - &local.p_ptr.addr.map(Into::into), - &local.p_addrs[local.p_addrs.len() - 1].value.map(Into::into), - PROT_READ | PROT_WRITE, - &local.write_slice_page_prot_access, - &mut is_not_trap, - &mut trap_code, - ); - - let x3_result_words = - limbs_to_words::(local.x3_ins.result.0.to_vec()); - let y3_result_words = - limbs_to_words::(local.y3_ins.result.0.to_vec()); - let result_words = - x3_result_words.into_iter().chain(y3_result_words).collect_vec(); - - builder.eval_memory_access_slice_read( - local.clk_high, - local.clk_low, - &local.q_addrs.iter().map(|addr| addr.value.map(Into::into)).collect_vec(), - &local.q_access.iter().map(|access| access.memory_access).collect_vec(), - is_not_trap, - ); - // We read p at +1 since p, q could be the same. - builder.eval_memory_access_slice_write( - local.clk_high, - local.clk_low + F::from_canonical_u32(1), - &local.p_addrs.iter().map(|addr| addr.value.map(Into::into)).collect_vec(), - &local.p_access.iter().map(|access| access.memory_access).collect_vec(), - result_words, - is_not_trap, - ); - builder.receive_syscall( - local.clk_high, - local.clk_low, - syscall_id_felt, - trap_code, - local.p_ptr.addr.map(Into::into), - local.q_ptr.addr.map(Into::into), - local.is_real, - InteractionScope::Local, - ); - } - - local.slope_numerator.eval(builder, &q_y, &p_y, FieldOperation::Sub, local.is_real); - } - 1 => { - local.slope_denominator.eval( - builder, - &q_x, - &p_x, - FieldOperation::Sub, - local.is_real, - ); - } - 2 => { - // We check that (q.x - p.x) is non-zero in the base field, by computing 1 / (q.x - p.x). - let mut coeff_1 = Vec::new(); - coeff_1.resize(::Limbs::USIZE, SymbolicExprF::zero()); - coeff_1[0] = SymbolicExprF::one(); - let one_polynomial = Polynomial::from_coefficients(&coeff_1); - - local.inverse_check.eval( - builder, - &one_polynomial, - &local.slope_denominator.result, - FieldOperation::Div, - local.is_real, - ); - } - 3 => { - local.slope.eval( - builder, - &local.slope_numerator.result, - &local.slope_denominator.result, - FieldOperation::Div, - local.is_real, - ); - } - 4 => { - local.slope_squared.eval( - builder, - &local.slope.result, - &local.slope.result, - FieldOperation::Mul, - local.is_real, - ); - } - 5 => { - local.p_x_plus_q_x.eval(builder, &p_x, &q_x, FieldOperation::Add, local.is_real); - } - 6 => { - local.x3_ins.eval( - builder, - &local.slope_squared.result, - &local.p_x_plus_q_x.result, - FieldOperation::Sub, - local.is_real, - ); - } - 7 => { - local.p_x_minus_x.eval( - builder, - &p_x, - &local.x3_ins.result, - FieldOperation::Sub, - local.is_real, - ); - } - 8 => { - local.slope_times_p_x_minus_x.eval( - builder, - &local.slope.result, - &local.p_x_minus_x.result, - FieldOperation::Mul, - local.is_real, - ); - } - 9 => { - local.y3_ins.eval( - builder, - &local.slope_times_p_x_minus_x.result, - &p_y, - FieldOperation::Sub, - local.is_real, - ); - } - 10 => { - let modulus = - E::BaseField::to_limbs_field::(&E::BaseField::modulus()); - local.x3_range.eval(builder, &local.x3_ins.result, &modulus, local.is_real); - local.y3_range.eval(builder, &local.y3_ins.result, &modulus, local.is_real); - let x3_result_words = - limbs_to_words::(local.x3_ins.result.0.to_vec()); - let y3_result_words = - limbs_to_words::(local.y3_ins.result.0.to_vec()); - let result_words = x3_result_words.into_iter().chain(y3_result_words).collect_vec(); - - let p_ptr = SyscallAddrOperation::::eval( - builder, - E::NB_LIMBS as u32 * 2, - local.p_ptr, - local.is_real.into(), - ); - let q_ptr = SyscallAddrOperation::::eval( - builder, - E::NB_LIMBS as u32 * 2, - local.q_ptr, - local.is_real.into(), - ); - - // p_addrs[i] = p_ptr + 8 * i - for i in 0..local.p_addrs.len() { - AddrAddOperation::::eval( - builder, - Word([ - p_ptr[0].into(), - p_ptr[1].into(), - p_ptr[2].into(), - SymbolicExprF::zero(), - ]), - Word::from(8 * i as u64), - local.p_addrs[i], - local.is_real.into(), - ); - } - - // q_addrs[i] = q_ptr + 8 * i - for i in 0..local.q_addrs.len() { - AddrAddOperation::::eval( - builder, - Word([ - q_ptr[0].into(), - q_ptr[1].into(), - q_ptr[2].into(), - SymbolicExprF::zero(), - ]), - Word::from(8 * i as u64), - local.q_addrs[i], - local.is_real.into(), - ); - } - - if M::IS_TRUSTED { - builder.eval_memory_access_slice_read( - local.clk_high, - local.clk_low, - &local.q_addrs.iter().map(|addr| addr.value.map(Into::into)).collect_vec(), - &local.q_access.iter().map(|access| access.memory_access).collect_vec(), - is_not_trap, - ); - // We read p at +1 since p, q could be the same. - builder.eval_memory_access_slice_write( - local.clk_high, - local.clk_low + F::from_canonical_u32(1), - &local.p_addrs.iter().map(|addr| addr.value.map(Into::into)).collect_vec(), - &local.p_access.iter().map(|access| access.memory_access).collect_vec(), - result_words, - is_not_trap, - ); - - builder.receive_syscall( - local.clk_high, - local.clk_low, - syscall_id_felt, - trap_code, - p_ptr.map(Into::into), - q_ptr.map(Into::into), - local.is_real, - InteractionScope::Local, - ); - } - } - _ => unreachable!(), - }; - } -} - -impl<'a, E: EllipticCurve + WeierstrassParameters, M: TrustMode> BlockAir> - for WeierstrassDoubleAssignChip -where - Limbs::Limbs>: Copy, -{ - fn num_blocks(&self) -> usize { - 12 - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder<'a>, index: usize) { - let main = builder.main(); - let local = main.row_slice(0); - let local: &WeierstrassDoubleAssignCols = (*local).borrow(); - - let num_words_field_element = ::Limbs::USIZE / 8; - - // Fetch the syscall id for the curve type. - let syscall_id_felt = match E::CURVE_TYPE { - CurveType::Secp256k1 => { - F::from_canonical_u32(SyscallCode::SECP256K1_DOUBLE.syscall_id()) - } - CurveType::Secp256r1 => { - F::from_canonical_u32(SyscallCode::SECP256R1_DOUBLE.syscall_id()) - } - CurveType::Bn254 => F::from_canonical_u32(SyscallCode::BN254_DOUBLE.syscall_id()), - CurveType::Bls12381 => F::from_canonical_u32(SyscallCode::BLS12381_DOUBLE.syscall_id()), - _ => panic!("Unsupported curve"), - }; - - // It's very important that the `generate_limbs` function do not call `assert_zero`. - let p_x_limbs = builder - .generate_limbs(&local.p_access[0..num_words_field_element], local.is_real.into()); - let p_x: Limbs::Limbs> = - Limbs(p_x_limbs.try_into().expect("failed to convert limbs")); - let p_y_limbs = builder - .generate_limbs(&local.p_access[num_words_field_element..], local.is_real.into()); - let p_y: Limbs::Limbs> = - Limbs(p_y_limbs.try_into().expect("failed to convert limbs")); - // `a` in the Weierstrass form: y^2 = x^3 + a * x + b. - let a = E::BaseField::to_limbs_field::(&E::a_int()); - - let is_not_trap: SymbolicExprF = local.is_real.into(); - let trap_code = SymbolicExprF::zero(); - - match index { - 0 => { - let mut is_not_trap = local.is_real.into(); - let mut trap_code = SymbolicExprF::zero(); - - #[cfg(feature = "mprotect")] - builder.assert_eq( - builder.extract_public_values().is_untrusted_programs_enabled, - SymbolicExprF::from_bool(!M::IS_TRUSTED), - ); - - if !M::IS_TRUSTED { - let local = main.row_slice(0); - let local: &WeierstrassDoubleAssignCols = - (*local).borrow(); - - #[cfg(not(feature = "mprotect"))] - builder.assert_zero(local.is_real); - AddressSlicePageProtOperation::::eval( - builder, - local.clk_high.into(), - local.clk_low.into(), - &local.p_ptr.addr.map(Into::into), - &local.p_addrs[local.p_addrs.len() - 1].value.map(Into::into), - PROT_READ | PROT_WRITE, - &local.write_slice_page_prot_access, - &mut is_not_trap, - &mut trap_code, - ); - - let x3_result_words = - limbs_to_words::(local.x3_ins.result.0.to_vec()); - let y3_result_words = - limbs_to_words::(local.y3_ins.result.0.to_vec()); - let result_words = - x3_result_words.into_iter().chain(y3_result_words).collect_vec(); - - builder.eval_memory_access_slice_write( - local.clk_high, - local.clk_low, - &local.p_addrs.iter().map(|addr| addr.value.map(Into::into)).collect_vec(), - &local.p_access.iter().map(|access| access.memory_access).collect_vec(), - result_words, - is_not_trap, - ); - - builder.receive_syscall( - local.clk_high, - local.clk_low, - syscall_id_felt, - trap_code, - local.p_ptr.addr.map(Into::into), - [SymbolicExprF::zero(), SymbolicExprF::zero(), SymbolicExprF::zero()], - local.is_real, - InteractionScope::Local, - ); - } - local.p_x_squared.eval(builder, &p_x, &p_x, FieldOperation::Mul, local.is_real); - } - 1 => { - local.p_x_squared_times_3.eval( - builder, - &local.p_x_squared.result, - &E::BaseField::to_limbs_field::(&BigUint::from(3u32)), - FieldOperation::Mul, - local.is_real, - ); - } - 2 => { - local.slope_numerator.eval( - builder, - &a, - &local.p_x_squared_times_3.result, - FieldOperation::Add, - local.is_real, - ); - } - 3 => { - local.slope_denominator.eval( - builder, - &E::BaseField::to_limbs_field::(&BigUint::from(2u32)), - &p_y, - FieldOperation::Mul, - local.is_real, - ); - } - 4 => { - local.slope.eval( - builder, - &local.slope_numerator.result, - &local.slope_denominator.result, - FieldOperation::Div, - local.is_real, - ); - } - 5 => { - local.slope_squared.eval( - builder, - &local.slope.result, - &local.slope.result, - FieldOperation::Mul, - local.is_real, - ); - } - 6 => { - local.p_x_plus_p_x.eval(builder, &p_x, &p_x, FieldOperation::Add, local.is_real); - } - 7 => { - local.x3_ins.eval( - builder, - &local.slope_squared.result, - &local.p_x_plus_p_x.result, - FieldOperation::Sub, - local.is_real, - ); - } - 8 => { - local.p_x_minus_x.eval( - builder, - &p_x, - &local.x3_ins.result, - FieldOperation::Sub, - local.is_real, - ); - } - 9 => { - local.slope_times_p_x_minus_x.eval( - builder, - &local.slope.result, - &local.p_x_minus_x.result, - FieldOperation::Mul, - local.is_real, - ); - } - 10 => { - local.y3_ins.eval( - builder, - &local.slope_times_p_x_minus_x.result, - &p_y, - FieldOperation::Sub, - local.is_real, - ); - } - 11 => { - let modulus = - E::BaseField::to_limbs_field::(&E::BaseField::modulus()); - local.x3_range.eval(builder, &local.x3_ins.result, &modulus, local.is_real); - local.y3_range.eval(builder, &local.y3_ins.result, &modulus, local.is_real); - - let x3_result_words = - limbs_to_words::(local.x3_ins.result.0.to_vec()); - let y3_result_words = - limbs_to_words::(local.y3_ins.result.0.to_vec()); - let result_words = x3_result_words.into_iter().chain(y3_result_words).collect_vec(); - - let p_ptr = SyscallAddrOperation::::eval( - builder, - E::NB_LIMBS as u32 * 2, - local.p_ptr, - local.is_real.into(), - ); - - // p_addrs[i] = p_ptr + 8 * i - for i in 0..local.p_addrs.len() { - AddrAddOperation::::eval( - builder, - Word([ - p_ptr[0].into(), - p_ptr[1].into(), - p_ptr[2].into(), - SymbolicExprF::zero(), - ]), - Word::from(8 * i as u64), - local.p_addrs[i], - local.is_real.into(), - ); - } - - if M::IS_TRUSTED { - builder.eval_memory_access_slice_write( - local.clk_high, - local.clk_low, - &local.p_addrs.iter().map(|addr| addr.value.map(Into::into)).collect_vec(), - &local.p_access.iter().map(|access| access.memory_access).collect_vec(), - result_words, - is_not_trap, - ); - - builder.receive_syscall( - local.clk_high, - local.clk_low, - syscall_id_felt, - trap_code, - p_ptr.map(Into::into), - [SymbolicExprF::zero(), SymbolicExprF::zero(), SymbolicExprF::zero()], - local.is_real, - InteractionScope::Local, - ); - } - } - - _ => unreachable!(), - }; - } -} - -impl<'a, const DEGREE: usize> BlockAir> for Poseidon2WideChip { - fn num_blocks(&self) -> usize { - 9 - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder<'a>, index: usize) { - let main = builder.main(); - let prepr = builder.preprocessed(); - let local_row = Self::convert::(main.row_slice(0)); - let prep_local = prepr.row_slice(0); - let prep_local: &Poseidon2PreprocessedColsWide<_> = (*prep_local).borrow(); - - match index { - 0 => { - // Dummy constraints to normalize to DEGREE. - let lhs = (0..DEGREE) - .map(|_| local_row.external_rounds_state()[0][0].into()) - .product::(); - let rhs = (0..DEGREE) - .map(|_| local_row.external_rounds_state()[0][0].into()) - .product::(); - builder.assert_eq(lhs, rhs); - - (0..WIDTH).for_each(|i| { - builder.receive_single( - prep_local.input[i], - local_row.external_rounds_state()[0][i], - prep_local.is_real, - ) - }); - - (0..WIDTH).for_each(|i| { - builder.send_single( - prep_local.output[i].addr, - local_row.perm_output()[i], - prep_local.output[i].mult, - ) - }); - eval_external_round(builder, local_row.as_ref(), 0); - } - 1..8 => { - eval_external_round(builder, local_row.as_ref(), index); - } - 8 => eval_internal_rounds(builder, local_row.as_ref()), - _ => unreachable!(), - } - } -} diff --git a/sp1-gpu/crates/air/src/instruction.rs b/sp1-gpu/crates/air/src/instruction.rs deleted file mode 100644 index 557b072f1f..0000000000 --- a/sp1-gpu/crates/air/src/instruction.rs +++ /dev/null @@ -1,896 +0,0 @@ -use std::fmt::Debug; -use std::mem::size_of; - -use crate::{ - symbolic_expr_ef::SymbolicExprEF, symbolic_expr_f::SymbolicExprF, - symbolic_var_ef::SymbolicVarEF, symbolic_var_f::SymbolicVarF, CUDA_P3_EVAL_EF_CONSTANTS, - CUDA_P3_EVAL_F_CONSTANTS, EF, F, -}; - -pub const INSTRUCTION_32_SIZE: usize = size_of::(); -pub const INSTRUCTION_16_SIZE: usize = size_of::(); - -#[derive(Clone, Copy)] -pub struct Instruction32 { - pub opcode: u8, - pub b_variant: u8, - pub c_variant: u8, - pub a: u32, - pub b: u32, - pub c: u32, -} - -#[derive(Debug, Clone, Copy)] -#[repr(C)] -pub struct Instruction16 { - pub opcode: u8, - pub b_variant: u8, - pub c_variant: u8, - pub a: u16, - pub b: u16, - pub c: u16, -} - -#[derive(Debug, Clone, Copy)] -#[repr(u8)] -pub enum Opcode { - Empty = 0, - - FAssignC = 1, - FAssignV = 2, - FAssignE = 3, - - FAddVC = 4, - FAddVV = 5, - FAddVE = 6, - - FAddEC = 7, - FAddEV = 8, - FAddEE = 9, - FAddAssignE = 10, - - FSubVC = 11, - FSubVV = 12, - FSubVE = 13, - - FSubEC = 14, - FSubEV = 15, - FSubEE = 16, - FSubAssignE = 17, - - FMulVC = 18, - FMulVV = 19, - FMulVE = 20, - - FMulEC = 21, - FMulEV = 22, - FMulEE = 23, - FMulAssignE = 24, - - FNegE = 25, - - EAssignC = 26, - EAssignV = 27, - EAssignE = 28, - - EAddVC = 29, - EAddVV = 30, - EAddVE = 31, - - EAddEC = 32, - EAddEV = 33, - EAddEE = 34, - EAddAssignE = 35, - - ESubVC = 36, - ESubVV = 37, - ESubVE = 38, - - ESubEC = 39, - ESubEV = 40, - ESubEE = 41, - ESubAssignE = 42, - - EMulVC = 43, - EMulVV = 44, - EMulVE = 45, - - EMulEC = 46, - EMulEV = 47, - EMulEE = 48, - EMulAssignE = 49, - - ENegE = 50, - - EFFromE = 51, - EFAddEE = 52, - EFAddAssignE = 53, - EFSubEE = 54, - EFSubAssignE = 55, - EFMulEE = 56, - EFMulAssignE = 57, - EFAsBaseSlice = 58, - - FAssertZero = 59, - EAssertZero = 60, -} - -impl Opcode { - pub fn is_f_assign(&self) -> bool { - let value = *self as u8; - (1..26).contains(&value) || value == 59 - } - - pub fn is_e_assign(&self) -> bool { - let value = *self as u8; - (26..59).contains(&value) || value == 60 - } - - pub fn is_f_arg1(&self) -> bool { - matches!( - self, - Opcode::FAssignE - | Opcode::FAddEC - | Opcode::FAddEV - | Opcode::FAddEE - | Opcode::FAddAssignE - | Opcode::FSubEC - | Opcode::FSubEV - | Opcode::FSubEE - | Opcode::FSubAssignE - | Opcode::FMulEC - | Opcode::FMulEV - | Opcode::FMulEE - | Opcode::FMulAssignE - | Opcode::FNegE - | Opcode::EFFromE - | Opcode::EFAddAssignE - | Opcode::EFSubAssignE - | Opcode::EFMulAssignE - ) - } - - pub fn is_f_arg2(&self) -> bool { - matches!( - self, - Opcode::FAddVE - | Opcode::FAddEE - | Opcode::FSubVE - | Opcode::FSubEE - | Opcode::FMulVE - | Opcode::FMulEE - | Opcode::EFAddEE - | Opcode::EFSubEE - | Opcode::EFMulEE - ) - } - - pub fn is_e_arg1(&self) -> bool { - matches!( - self, - Opcode::EAssignE - | Opcode::EAddEC - | Opcode::EAddEV - | Opcode::EAddEE - | Opcode::EAddAssignE - | Opcode::ESubEC - | Opcode::ESubEV - | Opcode::ESubEE - | Opcode::ESubAssignE - | Opcode::EMulEC - | Opcode::EMulEV - | Opcode::EMulEE - | Opcode::EMulAssignE - | Opcode::ENegE - | Opcode::EFAddEE - | Opcode::EFSubEE - | Opcode::EFMulEE - ) - } - - pub fn is_e_arg2(&self) -> bool { - matches!( - self, - Opcode::EAddVE - | Opcode::EAddEE - | Opcode::ESubVE - | Opcode::ESubEE - | Opcode::EMulVE - | Opcode::EMulEE - ) - } -} - -impl From for Opcode { - fn from(value: u8) -> Self { - unsafe { std::mem::transmute(value) } - } -} - -impl Instruction32 { - pub fn f_assign_c(a: SymbolicExprF, b: F) -> Self { - let b = f_constant(b); - Self { opcode: Opcode::FAssignC as u8, a: a.data(), b_variant: 0, b, c_variant: 0, c: 0 } - } - - pub fn f_assign_v(a: SymbolicExprF, b: SymbolicVarF) -> Self { - Self { - opcode: Opcode::FAssignV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn f_assign_e(a: SymbolicExprF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn f_add_vc(a: SymbolicExprF, b: SymbolicVarF, c: F) -> Self { - let c = f_constant(c); - Self { - opcode: Opcode::FAddVC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn f_add_vv(a: SymbolicExprF, b: SymbolicVarF, c: SymbolicVarF) -> Self { - Self { - opcode: Opcode::FAddVV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_add_ve(a: SymbolicExprF, b: SymbolicVarF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FAddVE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_add_ec(a: SymbolicExprF, b: SymbolicExprF, c: F) -> Self { - let c = f_constant(c); - Self { - opcode: Opcode::FAddEC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn f_add_ev(a: SymbolicExprF, b: SymbolicExprF, c: SymbolicVarF) -> Self { - Self { - opcode: Opcode::FAddEV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_add_ee(a: SymbolicExprF, b: SymbolicExprF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FAddEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_add_assign_e(a: SymbolicExprF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FAddAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn f_sub_vc(a: SymbolicExprF, b: SymbolicVarF, c: F) -> Self { - let c = f_constant(c); - Self { - opcode: Opcode::FSubVC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn f_sub_vv(a: SymbolicExprF, b: SymbolicVarF, c: SymbolicVarF) -> Self { - Self { - opcode: Opcode::FSubVV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_sub_ve(a: SymbolicExprF, b: SymbolicVarF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FSubVE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_sub_ec(a: SymbolicExprF, b: SymbolicExprF, c: F) -> Self { - let c = f_constant(c); - Self { - opcode: Opcode::FSubEC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn f_sub_ev(a: SymbolicExprF, b: SymbolicExprF, c: SymbolicVarF) -> Self { - Self { - opcode: Opcode::FSubEV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_sub_ee(a: SymbolicExprF, b: SymbolicExprF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FSubEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_sub_assign_e(a: SymbolicExprF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FSubAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn f_mul_vc(a: SymbolicExprF, b: SymbolicVarF, c: F) -> Self { - let c = f_constant(c); - Self { - opcode: Opcode::FMulVC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn f_mul_vv(a: SymbolicExprF, b: SymbolicVarF, c: SymbolicVarF) -> Self { - Self { - opcode: Opcode::FMulVV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_mul_ve(a: SymbolicExprF, b: SymbolicVarF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FMulVE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: c.data(), - } - } - - pub fn f_mul_ec(a: SymbolicExprF, b: SymbolicExprF, c: F) -> Self { - let c = f_constant(c); - Self { - opcode: Opcode::FMulEC as u8, - a: a.data(), - b_variant: 0, - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn f_mul_ev(a: SymbolicExprF, b: SymbolicExprF, c: SymbolicVarF) -> Self { - Self { - opcode: Opcode::FMulEV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_mul_ee(a: SymbolicExprF, b: SymbolicExprF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FMulEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn f_mul_assign_e(a: SymbolicExprF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FMulAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn f_neg_e(a: SymbolicExprF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FNegE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn e_assign_c(a: SymbolicExprEF, b: EF) -> Self { - let b = ef_constant(b); - Self { opcode: Opcode::EAssignC as u8, a: a.data(), b_variant: 0, b, c_variant: 0, c: 0 } - } - - pub fn e_assign_v(a: SymbolicExprEF, b: SymbolicVarEF) -> Self { - Self { - opcode: Opcode::EAssignV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn e_assign_e(a: SymbolicExprEF, b: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::EAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn e_add_vc(a: SymbolicExprEF, b: SymbolicVarEF, c: EF) -> Self { - let c = ef_constant(c); - Self { - opcode: Opcode::EAddVC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn e_add_vv(a: SymbolicExprEF, b: SymbolicVarEF, c: SymbolicVarEF) -> Self { - Self { - opcode: Opcode::EAddVV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_add_ve(a: SymbolicExprEF, b: SymbolicVarEF, c: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::EAddVE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_add_ec(a: SymbolicExprEF, b: SymbolicExprEF, c: EF) -> Self { - let c = ef_constant(c); - Self { - opcode: Opcode::EAddEC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn e_add_ev(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicVarEF) -> Self { - Self { - opcode: Opcode::EAddEV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_add_ee(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::EAddEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_add_assign_e(a: SymbolicExprEF, b: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::EAddAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn e_sub_vc(a: SymbolicExprEF, b: SymbolicVarEF, c: EF) -> Self { - let c = ef_constant(c); - Self { - opcode: Opcode::ESubVC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn e_sub_vv(a: SymbolicExprEF, b: SymbolicVarEF, c: SymbolicVarEF) -> Self { - Self { - opcode: Opcode::ESubVV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_sub_ve(a: SymbolicExprEF, b: SymbolicVarEF, c: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::ESubVE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_sub_ec(a: SymbolicExprEF, b: SymbolicExprEF, c: EF) -> Self { - let c = ef_constant(c); - Self { - opcode: Opcode::ESubEC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn e_sub_ev(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicVarEF) -> Self { - Self { - opcode: Opcode::ESubEV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_sub_ee(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::ESubEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_sub_assign_e(a: SymbolicExprEF, b: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::ESubAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn e_mul_vc(a: SymbolicExprEF, b: SymbolicVarEF, c: EF) -> Self { - let c = ef_constant(c); - Self { - opcode: Opcode::EMulVC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn e_mul_vv(a: SymbolicExprEF, b: SymbolicVarEF, c: SymbolicVarEF) -> Self { - Self { - opcode: Opcode::EMulVV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_mul_ve(a: SymbolicExprEF, b: SymbolicVarEF, c: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::EMulVE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_mul_ec(a: SymbolicExprEF, b: SymbolicExprEF, c: EF) -> Self { - let c = ef_constant(c); - Self { - opcode: Opcode::EMulEC as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c, - } - } - - pub fn e_mul_ev(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicVarEF) -> Self { - Self { - opcode: Opcode::EMulEV as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_mul_ee(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::EMulEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn e_mul_assign_e(a: SymbolicExprEF, b: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::EMulAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn e_neg_e(a: SymbolicExprEF, b: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::ENegE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn ef_from_e(a: SymbolicExprEF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::EFFromE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn ef_add_ee(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::EFAddEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn ef_add_assign_e(a: SymbolicExprEF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::EFAddAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn ef_sub_ee(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::EFSubEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn ef_sub_assign_e(a: SymbolicExprEF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::EFSubAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn ef_mul_ee(a: SymbolicExprEF, b: SymbolicExprEF, c: SymbolicExprF) -> Self { - Self { - opcode: Opcode::EFMulEE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: c.variant(), - c: c.data(), - } - } - - pub fn ef_mul_assign_e(a: SymbolicExprEF, b: SymbolicExprF) -> Self { - Self { - opcode: Opcode::EFMulAssignE as u8, - a: a.data(), - b_variant: b.variant(), - b: b.data(), - c_variant: 0, - c: 0, - } - } - - pub fn f_assert_zero(a: SymbolicExprF) -> Self { - Self { - opcode: Opcode::FAssertZero as u8, - a: a.data(), - b_variant: 0, - b: 0, - c_variant: 0, - c: 0, - } - } - - pub fn e_assert_zero(a: SymbolicExprEF) -> Self { - Self { - opcode: Opcode::EAssertZero as u8, - a: a.data(), - b_variant: 0, - b: 0, - c_variant: 0, - c: 0, - } - } -} - -impl Default for Instruction32 { - fn default() -> Self { - Self { opcode: Opcode::Empty as u8, a: 0, b_variant: 0, b: 0, c_variant: 0, c: 0 } - } -} - -impl Debug for Instruction32 { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let opcode = Opcode::from(self.opcode); - write!( - f, - "Instruction {{ opcode: {:?}, a: {}, b_variant: {}, b: {}, c_variant: {}, c: {} }}", - opcode, self.a, self.b_variant, self.b, self.c_variant, self.c - ) - } -} - -pub fn f_constant(c: F) -> u32 { - let mut tmp = CUDA_P3_EVAL_F_CONSTANTS.lock().unwrap(); - if let Some(pos) = tmp.iter().position(|&x| x == c) { - pos as u32 - } else { - tmp.push(c); - (tmp.len() - 1) as u32 - } -} - -pub fn ef_constant(c: EF) -> u32 { - let mut tmp = CUDA_P3_EVAL_EF_CONSTANTS.lock().unwrap(); - if let Some(pos) = tmp.iter().position(|&x| x == c) { - pos as u32 - } else { - tmp.push(c); - (tmp.len() - 1) as u32 - } -} diff --git a/sp1-gpu/crates/air/src/ir/analysis.rs b/sp1-gpu/crates/air/src/ir/analysis.rs new file mode 100644 index 0000000000..5b0444d4bb --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/analysis.rs @@ -0,0 +1,247 @@ +//! Per-constraint DAG analysis. +//! +//! For each `ConstraintRef`, walks back from the root to compute: +//! - the transitive column-leaf set (used by the chunker) +//! - work (count of arithmetic ops; leaves and constants are free) +//! - depth (longest dependency chain to a leaf) +//! - a structural shape tag (e.g. `LinearWeightedSum` for GKR-eligible chunks) +//! +//! Output is purely derived from the `ConstraintDag` — no mutation. + +use std::collections::{HashMap, HashSet}; + +use crate::ir::dag::{ConstraintDag, ConstraintField, DagNode, NodeId, TraceSource}; + +/// Identifies a column-ref leaf for chunker bookkeeping. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct ColumnLeaf { + pub source: TraceSource, + pub col: u32, +} + +/// Structural tag derived from the DAG shape. Used by the scheduler to pick +/// non-default lowerings (e.g. column-tile for GKR-shape). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ConstraintShape { + /// `Σ_i (coeff_i · leaf_i)` form. Each leaf appears in exactly one + /// multiplication-by-coefficient before being summed into the root. + /// This is the shape of GKR corrections and any linear combination. + LinearWeightedSum, + /// Anything else. + General, +} + +#[derive(Debug, Clone)] +pub struct ConstraintInfo { + /// Index into `dag.constraints`. + pub constraint_idx: usize, + pub root: NodeId, + pub field: ConstraintField, + pub alpha_index: u32, + + /// Total nodes in the transitive closure (leaves + consts + arithmetic). + pub total_nodes: u32, + /// Arithmetic op count only (excludes leaves and constants). + pub work: u32, + /// Longest dependency chain from root to a leaf. + pub depth: u32, + /// Distinct column-ref leaves reachable from root. + pub column_leaves: HashSet, + pub shape: ConstraintShape, +} + +impl ConstraintInfo { + pub fn leafset_size(&self) -> u32 { + self.column_leaves.len() as u32 + } +} + +/// Analyze all constraints in a DAG. +pub fn analyze_constraints(dag: &ConstraintDag) -> Vec { + dag.constraints + .iter() + .enumerate() + .map(|(i, c)| analyze_one(dag, i, c.root, c.field, c.alpha_index)) + .collect() +} + +fn analyze_one( + dag: &ConstraintDag, + constraint_idx: usize, + root: NodeId, + field: ConstraintField, + alpha_index: u32, +) -> ConstraintInfo { + let mut depth_of: HashMap = HashMap::new(); + let mut column_leaves: HashSet = HashSet::new(); + let mut work: u32 = 0; + let mut total: u32 = 0; + + let depth = walk(dag, root, &mut depth_of, &mut column_leaves, &mut work, &mut total); + let shape = detect_shape(dag, root, &column_leaves); + + ConstraintInfo { + constraint_idx, + root, + field, + alpha_index, + total_nodes: total, + work, + depth, + column_leaves, + shape, + } +} + +fn walk( + dag: &ConstraintDag, + reg: NodeId, + depth_of: &mut HashMap, + column_leaves: &mut HashSet, + work: &mut u32, + total: &mut u32, +) -> u32 { + if let Some(&d) = depth_of.get(®) { + return d; + } + let node = &dag.nodes[reg as usize]; + *total += 1; + let d = match *node { + DagNode::InputLeaf { source, col } => { + if matches!( + source, + TraceSource::MainLocal + | TraceSource::MainNext + | TraceSource::PreprocessedLocal + | TraceSource::PreprocessedNext + ) { + column_leaves.insert(ColumnLeaf { source, col }); + } + 0 + } + DagNode::ConstF { .. } + | DagNode::ConstEF { .. } + | DagNode::PublicValue { .. } + | DagNode::GlobalCumulativeSum { .. } + | DagNode::IsFirstRow + | DagNode::IsLastRow + | DagNode::IsTransition => 0, + DagNode::AddF { a, b } + | DagNode::SubF { a, b } + | DagNode::MulF { a, b } + | DagNode::AddEF { a, b } + | DagNode::SubEF { a, b } + | DagNode::MulEF { a, b } + | DagNode::EFAddF { a, b } + | DagNode::EFSubF { a, b } + | DagNode::EFMulF { a, b } => { + *work += 1; + let da = walk(dag, a, depth_of, column_leaves, work, total); + let db = walk(dag, b, depth_of, column_leaves, work, total); + 1 + da.max(db) + } + DagNode::NegF { a } | DagNode::NegEF { a } | DagNode::EFFromF { a } => { + *work += 1; + let da = walk(dag, a, depth_of, column_leaves, work, total); + 1 + da + } + }; + depth_of.insert(reg, d); + d +} + +/// Detect `Σ_i (coeff_i · leaf_i)` shape. +/// +/// Walks the root's spine looking for: a chain of `Add` nodes whose leaves +/// are `Mul(const_or_public, leaf)`, terminating in a similar `Mul` or a +/// bare leaf. If the structure matches and every column leaf in the +/// constraint appears in exactly one such product, the constraint is +/// `LinearWeightedSum`. +/// +/// This is the structural test for "GKR-shape" — the column-tile lowering +/// applies when this returns true. +fn detect_shape( + dag: &ConstraintDag, + root: NodeId, + column_leaves: &HashSet, +) -> ConstraintShape { + if column_leaves.is_empty() { + return ConstraintShape::General; + } + let mut leaves_seen: HashSet = HashSet::new(); + if !walk_linear_sum(dag, root, &mut leaves_seen) { + return ConstraintShape::General; + } + if leaves_seen == *column_leaves { + ConstraintShape::LinearWeightedSum + } else { + ConstraintShape::General + } +} + +/// Recurses into an Add-chain; each leaf of the chain must be a product +/// of a coefficient (constant / public) and a single column leaf. +fn walk_linear_sum( + dag: &ConstraintDag, + node_id: NodeId, + leaves_seen: &mut HashSet, +) -> bool { + match dag.nodes[node_id as usize] { + DagNode::AddF { a, b } | DagNode::SubF { a, b } => { + walk_linear_sum(dag, a, leaves_seen) && walk_linear_sum(dag, b, leaves_seen) + } + DagNode::MulF { a, b } => match (coefficient(dag, a), coefficient(dag, b)) { + (Some(_), None) => { + extract_column_leaf(dag, b).map(|c| leaves_seen.insert(c)).unwrap_or(false) + } + (None, Some(_)) => { + extract_column_leaf(dag, a).map(|c| leaves_seen.insert(c)).unwrap_or(false) + } + _ => false, + }, + DagNode::InputLeaf { source, col } + if matches!( + source, + TraceSource::MainLocal + | TraceSource::MainNext + | TraceSource::PreprocessedLocal + | TraceSource::PreprocessedNext + ) => + { + leaves_seen.insert(ColumnLeaf { source, col }) + } + _ => false, + } +} + +/// True iff `node_id` references a constant / public / cumsum (i.e. not a +/// column read and not an arithmetic op). +fn coefficient(dag: &ConstraintDag, node_id: NodeId) -> Option<()> { + match dag.nodes[node_id as usize] { + DagNode::ConstF { .. } + | DagNode::ConstEF { .. } + | DagNode::PublicValue { .. } + | DagNode::GlobalCumulativeSum { .. } + | DagNode::IsFirstRow + | DagNode::IsLastRow + | DagNode::IsTransition => Some(()), + _ => None, + } +} + +fn extract_column_leaf(dag: &ConstraintDag, node_id: NodeId) -> Option { + match dag.nodes[node_id as usize] { + DagNode::InputLeaf { source, col } + if matches!( + source, + TraceSource::MainLocal + | TraceSource::MainNext + | TraceSource::PreprocessedLocal + | TraceSource::PreprocessedNext + ) => + { + Some(ColumnLeaf { source, col }) + } + _ => None, + } +} diff --git a/sp1-gpu/crates/air/src/ir/builder.rs b/sp1-gpu/crates/air/src/ir/builder.rs new file mode 100644 index 0000000000..6c82b11564 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/builder.rs @@ -0,0 +1,181 @@ +//! `DagBuilder` and the trait implementations needed to drive a chip's `eval`. +//! +//! `DagBuilder` holds matrix views and a constraint counter, and defers actual +//! node allocation to the global DAG state. The trait impls (`AirBuilder`, +//! `ExtensionBuilder`, etc.) carry `unimplemented!()` stubs for the builder +//! methods SP1 chips don't use (transition windows, permutation columns). + +use slop_air::{ + Air, AirBuilder, AirBuilderWithPublicValues, ExtensionBuilder, PairBuilder, + PermutationAirBuilder, +}; +use slop_matrix::dense::{DenseMatrix, RowMajorMatrixView}; + +use sp1_core_machine::air::TrivialOperationBuilder; +use sp1_hypercube::air::{EmptyMessageBuilder, MachineAir}; +use sp1_hypercube::{AirOpenedValues, PROOF_MAX_NUM_PVS}; + +use crate::ir::dag::{ConstraintDag, ConstraintField, ConstraintRef}; +use crate::ir::expr::{DagExprEF, DagExprF}; +use crate::ir::state::{with_state, DAG_BUILDER_LOCK, DAG_STATE}; +use crate::ir::var::{DagVarEF, DagVarF}; +use crate::{EF, F}; + +/// Drives a chip's `eval` for the DAG-native path. +/// +/// Holds the matrix views the chip's `eval` will read (constructed before the +/// folder is built) and a per-folder constraint counter. All node allocation +/// flows through the global `DAG_STATE`. +pub struct DagBuilder<'a> { + pub preprocessed: RowMajorMatrixView<'a, DagVarF>, + pub main: RowMajorMatrixView<'a, DagVarF>, + pub public_values: &'a [DagVarF], + pub num_constraints: u32, +} + +impl<'a> AirBuilder for DagBuilder<'a> { + type F = F; + type Expr = DagExprF; + type Var = DagVarF; + type M = RowMajorMatrixView<'a, DagVarF>; + + fn main(&self) -> Self::M { + self.main + } + + fn is_first_row(&self) -> Self::Expr { + unimplemented!(); + } + + fn is_last_row(&self) -> Self::Expr { + unimplemented!(); + } + + fn is_transition_window(&self, _: usize) -> Self::Expr { + unimplemented!(); + } + + fn assert_zero>(&mut self, x: I) { + let x: Self::Expr = x.into(); + let alpha_index = self.num_constraints; + with_state(|s| { + s.constraints.push(ConstraintRef { + root: x.0, + alpha_index, + field: ConstraintField::Base, + }); + }); + self.num_constraints += 1; + } +} + +impl ExtensionBuilder for DagBuilder<'_> { + type EF = EF; + type ExprEF = DagExprEF; + type VarEF = DagVarEF; + + fn assert_zero_ext(&mut self, x: I) + where + I: Into, + { + let x: Self::ExprEF = x.into(); + let alpha_index = self.num_constraints; + with_state(|s| { + s.constraints.push(ConstraintRef { + root: x.0, + alpha_index, + field: ConstraintField::Extension, + }); + }); + self.num_constraints += 1; + } +} + +impl<'a> PermutationAirBuilder for DagBuilder<'a> { + type MP = RowMajorMatrixView<'a, DagVarEF>; + type RandomVar = DagVarEF; + + fn permutation(&self) -> Self::MP { + unimplemented!(); + } + + fn permutation_randomness(&self) -> &[Self::RandomVar] { + unimplemented!(); + } +} + +impl PairBuilder for DagBuilder<'_> { + fn preprocessed(&self) -> Self::M { + self.preprocessed + } +} + +impl AirBuilderWithPublicValues for DagBuilder<'_> { + type PublicVar = DagVarF; + + fn public_values(&self) -> &[Self::PublicVar] { + self.public_values + } +} + +impl EmptyMessageBuilder for DagBuilder<'_> {} + +impl TrivialOperationBuilder for DagBuilder<'_> {} + +// ============================================================================ +// Entry point +// ============================================================================ + +/// Run a chip's `eval` against the DAG-native builder and return the resulting +/// `ConstraintDag`. +/// +/// Acquires the global guard, resets state, runs `eval` once (auto-chunking +/// partitions the DAG downstream), snapshots the state, and releases. +pub fn build_dag(air: &A) -> ConstraintDag +where + A: MachineAir + for<'a> Air>, +{ + let _guard = DAG_BUILDER_LOCK.lock().unwrap(); + + // Reset before we begin. + { + let mut state = DAG_STATE.lock().unwrap(); + state.reset(); + } + + let preprocessed_width = air.preprocessed_width() as u32; + let main_width = air.width() as u32; + + // Eagerly intern all preprocessed and main column refs. Chip `eval` + // reads them through the matrix view. + let prep_vars: Vec = + (0..preprocessed_width).map(DagVarF::preprocessed_local).collect(); + let main_vars: Vec = (0..main_width).map(DagVarF::main_local).collect(); + let public_values: Vec = + (0..PROOF_MAX_NUM_PVS as u32).map(DagVarF::public_value).collect(); + + let prep_matrix = DenseMatrix::new(prep_vars, preprocessed_width.max(1) as usize); + let main_matrix = DenseMatrix::new(main_vars, main_width.max(1) as usize); + let preprocessed_view = AirOpenedValues { local: prep_matrix.values.clone() }; + let main_view = AirOpenedValues { local: main_matrix.values.clone() }; + + let mut folder = DagBuilder { + preprocessed: preprocessed_view.view(), + main: main_view.view(), + public_values: &public_values, + num_constraints: 0, + }; + + air.eval(&mut folder); + + // Snapshot. + let (nodes, constraints) = { + let mut state = DAG_STATE.lock().unwrap(); + let nodes = std::mem::take(&mut state.nodes); + let constraints = std::mem::take(&mut state.constraints); + state.reset(); + (nodes, constraints) + }; + + ConstraintDag { nodes, constraints, preprocessed_width, main_width } +} diff --git a/sp1-gpu/crates/air/src/ir/bytecode.rs b/sp1-gpu/crates/air/src/ir/bytecode.rs new file mode 100644 index 0000000000..cd9441f142 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/bytecode.rs @@ -0,0 +1,304 @@ +//! Bytecode that the Sequential kernel interprets per chunk. +//! +//! Each `ChunkBytecode` is the lowered form of one `Chunk`+`SequentialPlan`: +//! - `leaves` : per-leaf trace reference (which column, which source). +//! Kernel loads `(zero, one)` pairs into shared memory at CTA preamble. +//! - `consts` : pool of base-field constants. Indexed by `OpLoadConstF` instrs. +//! - `publics` : indices into the global public-values buffer. +//! - `instrs` : flat bytecode in topological order. Each instr writes +//! to `out`; reads from `a` / `b` are either reg slots, +//! leaf-cache indices, or pool indices depending on opcode. +//! - `max_reg` : size of the per-lane register file (max-live count). +//! - `roots` : the assertion roots, paired with their alpha index. The +//! kernel reads each root reg, multiplies by `α^k`, and adds +//! to the accumulator. +//! +//! Uses explicit leaf/const/public pools per chunk (rather than per-chip +//! globals), enabling shared-memory staging. + +use crate::ir::analysis::ConstraintInfo; +use crate::ir::chunker::Chunk; +use crate::ir::dag::{ConstraintDag, ConstraintField, DagNode, NodeId, TraceSource}; +use crate::ir::lowering::SequentialPlan; +use crate::F; +use std::collections::HashMap; + +/// Bytecode opcodes. Must mirror the constants in `sequential.cuh`. +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BcOp { + /// out = lerp(leaf_cache[a].zero, leaf_cache[a].one, eval_pt[lane]) + LoadLeaf = 0, + /// out = const_pool[a] + LoadConst = 1, + /// out = public_values[a] + LoadPublic = 2, + /// out = regs[a] + regs[b] + AddF = 3, + /// out = regs[a] - regs[b] + SubF = 4, + /// out = regs[a] * regs[b] + MulF = 5, + /// out = -regs[a] + NegF = 6, + /// accumulator += powers_of_alpha[a] * regs[b] + /// (`out` is unused for asserts; we pack a sentinel.) + AssertF = 7, +} + +/// One bytecode instruction. 8 bytes. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct DagInstr { + pub opcode: u8, + pub _pad: u8, + pub out: u16, + pub a: u16, + pub b: u16, +} + +impl DagInstr { + pub fn new(op: BcOp, out: u16, a: u16, b: u16) -> Self { + Self { opcode: op as u8, _pad: 0, out, a, b } + } +} + +/// Trace reference for a leaf. The kernel uses this at CTA preamble to load +/// `(zero, one)` pairs into shared memory. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct LeafRef { + /// 2 = PreprocessedLocal, 3 = PreprocessedNext, 4 = MainLocal, 5 = MainNext. + /// This encoding matches the jagged-mle layout code's column-variant tags. + pub source: u8, + pub _pad: u8, + /// Column index within the chip's preprocessed or main trace. + pub col: u32, +} + +/// Lowered, ready-to-launch chunk. +#[derive(Debug, Default, Clone)] +pub struct ChunkBytecode { + pub leaves: Vec, + pub consts: Vec, + pub publics: Vec, + pub instrs: Vec, + /// (reg, alpha_index) per assertion in this chunk. The kernel applies + /// `accumulator += alpha^k * regs[reg]` at the end. Kept separate + /// from `instrs` so the schedule can drive the alpha-table read. + pub asserts: Vec<(u16, u32)>, + pub max_reg: u16, + pub n_constraints: u32, + /// If non-zero, the kernel appends a per-row GKR sweep after the bytecode + /// and asserts pass — accumulating `Σ_i gkr_powers[i] · col_i(row)` over + /// `gkr_main_width` main cols and `gkr_prep_width` prep cols. This fuses + /// what would otherwise be a separate ColumnTile launch into the Sequential + /// pass, sharing column loads with the constraint bytecode in L1. + pub gkr_main_width: u32, + pub gkr_prep_width: u32, +} + +/// Lower a `SequentialPlan` (topological order over the chunk's DAG subgraph) +/// to flat bytecode. +/// +/// Performs liveness-based register allocation: each physical register slot +/// is reused once its current occupant's last use has passed. The resulting +/// `max_reg` is the chunk's peak live count, not its node count — typically +/// a small constant for shallow constraints. +pub fn lower_sequential( + chunk: &Chunk, + constraints: &[ConstraintInfo], + dag: &ConstraintDag, + plan: &SequentialPlan, +) -> ChunkBytecode { + let mut bc = ChunkBytecode { + n_constraints: chunk.constraint_indices.len() as u32, + ..ChunkBytecode::default() + }; + + let phys_of = liveness_allocate(chunk, constraints, dag, plan); + + let mut leaf_of: HashMap<(u8, u32), u16> = HashMap::new(); + let mut const_of: HashMap = HashMap::new(); + let mut public_of: HashMap = HashMap::new(); + + // Helper: produces the operand register for a node that's already been + // emitted (must be in `phys_of`). + let reg = |n: NodeId| -> u16 { *phys_of.get(&n).expect("topo order broken") }; + + for &node_id in &plan.topo_order { + let node = &dag.nodes[node_id as usize]; + match *node { + DagNode::InputLeaf { source, col } => { + let src_byte = match source { + TraceSource::PreprocessedLocal => 2, + TraceSource::PreprocessedNext => 3, + TraceSource::MainLocal => 4, + TraceSource::MainNext => 5, + }; + let leaf_idx = *leaf_of.entry((src_byte, col)).or_insert_with(|| { + let i = bc.leaves.len() as u16; + bc.leaves.push(LeafRef { source: src_byte, _pad: 0, col }); + i + }); + bc.instrs.push(DagInstr::new(BcOp::LoadLeaf, reg(node_id), leaf_idx, 0)); + } + DagNode::ConstF { value } => { + use slop_algebra::PrimeField32; + let key = value.as_canonical_u32(); + let cidx = *const_of.entry(key).or_insert_with(|| { + let i = bc.consts.len() as u16; + bc.consts.push(value); + i + }); + bc.instrs.push(DagInstr::new(BcOp::LoadConst, reg(node_id), cidx, 0)); + } + DagNode::PublicValue { idx } => { + let pidx = *public_of.entry(idx).or_insert_with(|| { + let i = bc.publics.len() as u16; + bc.publics.push(idx); + i + }); + bc.instrs.push(DagInstr::new(BcOp::LoadPublic, reg(node_id), pidx, 0)); + } + DagNode::AddF { a, b } => { + bc.instrs.push(DagInstr::new(BcOp::AddF, reg(node_id), reg(a), reg(b))); + } + DagNode::SubF { a, b } => { + bc.instrs.push(DagInstr::new(BcOp::SubF, reg(node_id), reg(a), reg(b))); + } + DagNode::MulF { a, b } => { + bc.instrs.push(DagInstr::new(BcOp::MulF, reg(node_id), reg(a), reg(b))); + } + DagNode::NegF { a } => { + bc.instrs.push(DagInstr::new(BcOp::NegF, reg(node_id), reg(a), 0)); + } + // EF / mixed / singletons / cumsum are not handled by the + // Phase 2 kernel — same scope as v1 production. + _ => { + panic!( + "Phase 2 Sequential kernel does not handle node kind {:?}; node {}", + node, node_id + ); + } + } + } + + // Append per-constraint assertions. + for &ci in &chunk.constraint_indices { + let info = &constraints[ci]; + // Phase 2 base-field only: skip EF constraints. + if !matches!(info.field, ConstraintField::Base) { + continue; + } + bc.asserts.push((reg(info.root), info.alpha_index)); + } + + // `max_reg` = peak physical-reg index used + 1. + bc.max_reg = phys_of.values().copied().max().map(|m| m + 1).unwrap_or(0); + bc +} + +/// Compute a `NodeId -> physical-register-slot` mapping by linear-scan over +/// the topological order, reusing slots whose previous occupant's last use +/// has passed. +/// +/// Constraint roots are kept live to the very end so the post-topo assert +/// pass can still read them. +fn liveness_allocate( + chunk: &Chunk, + constraints: &[ConstraintInfo], + dag: &ConstraintDag, + plan: &SequentialPlan, +) -> HashMap { + let topo = &plan.topo_order; + let pos_of: HashMap = topo.iter().enumerate().map(|(i, &n)| (n, i)).collect(); + + // Last-use position per node. Self-use (at the def site) counts as i. + let mut last_use: HashMap = HashMap::new(); + for (i, &node_id) in topo.iter().enumerate() { + last_use.insert(node_id, i); + } + for (i, &node_id) in topo.iter().enumerate() { + let node = &dag.nodes[node_id as usize]; + for child in node_children(node).into_iter().flatten() { + if pos_of.contains_key(&child) { + let e = last_use.entry(child).or_insert(0); + if i > *e { + *e = i; + } + } + } + } + // Constraint roots must remain live through the assert pass. + let end = topo.len(); + for &ci in &chunk.constraint_indices { + let root = constraints[ci].root; + if pos_of.contains_key(&root) { + last_use.insert(root, end); + } + } + + // Linear-scan: at each position, free regs whose occupants died before + // this position; allocate from pool (else bump). + let mut active: Vec<(u16, NodeId)> = Vec::new(); // (phys, node) + let mut free_pool: Vec = Vec::new(); + let mut next_phys: u16 = 0; + let mut phys_of: HashMap = HashMap::new(); + + for (i, &node_id) in topo.iter().enumerate() { + // Free dead regs. + active.retain(|&(p, n)| { + if last_use[&n] < i { + free_pool.push(p); + false + } else { + true + } + }); + + // Allocate. + let phys = free_pool.pop().unwrap_or_else(|| { + let p = next_phys; + next_phys += 1; + p + }); + active.push((phys, node_id)); + phys_of.insert(node_id, phys); + } + + phys_of +} + +/// DAG-node child enumeration (returns the operand `NodeId`s). +fn node_children(node: &DagNode) -> [Option; 2] { + use crate::ir::dag::DagNode::*; + match *node { + InputLeaf { .. } + | PublicValue { .. } + | GlobalCumulativeSum { .. } + | ConstF { .. } + | ConstEF { .. } + | IsFirstRow + | IsLastRow + | IsTransition => [None, None], + AddF { a, b } + | SubF { a, b } + | MulF { a, b } + | AddEF { a, b } + | SubEF { a, b } + | MulEF { a, b } + | EFAddF { a, b } + | EFSubF { a, b } + | EFMulF { a, b } => [Some(a), Some(b)], + NegF { a } | NegEF { a } | EFFromF { a } => [Some(a), None], + } +} + +/// Convenience that detects bytecode incompatibility before launch. +impl ChunkBytecode { + pub fn is_compatible_with_v1_kernel(&self) -> bool { + // Phase 2 Sequential covers the base-field subset of v1. + self.instrs.iter().all(|i| i.opcode <= BcOp::NegF as u8) && !self.asserts.is_empty() + } +} diff --git a/sp1-gpu/crates/air/src/ir/chunker.rs b/sp1-gpu/crates/air/src/ir/chunker.rs new file mode 100644 index 0000000000..8e8d469e40 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/chunker.rs @@ -0,0 +1,227 @@ +//! Greedy first-fit-decreasing chunker. +//! +//! Given the per-constraint analysis, packs constraints into chunks bounded +//! by a register-pressure +//! budget (max leafset per chunk → downstream `max_reg` → MAX_REGS template +//! tier). Heuristic: sort constraints by descending leafset size, then for +//! each, pick the chunk that minimizes the number of NEW unique leaves +//! introduced. +//! +//! If a single constraint already exceeds the budget, it becomes its own +//! chunk — flagged for the scheduler to route to the escape-valve lowering. + +use std::collections::HashSet; + +use crate::ir::analysis::{ColumnLeaf, ConstraintInfo, ConstraintShape}; + +/// Hard caps the chunker won't exceed when adding a constraint to an existing chunk. +/// +/// Both fields correspond to real per-kernel resources: +/// +/// * `max_leafset` — the chunk's column-leaf footprint. The fused sequential +/// kernel materialises one register slot per leaf and reuses them across +/// every assertion in the chunk; the chunk's downstream `max_reg` (peak +/// live count in `K regs[MAX_REGS][3]`) is `leafset_size + small_overhead`, +/// so this directly controls which `MAX_REGS={32,64,128,256,512,1024}` +/// template tier the chunk falls into. Cross a tier and the per-thread +/// `regs[]` array doubles, spilling out of the L1-cached local-memory +/// window. +/// * `max_constraints_per_chunk` — caps the size of the chunk's `instrs[]` +/// and `assert_regs[]` arrays in global memory. Rarely the binding +/// constraint in practice; lives here so a pathological AIR can't generate +/// an unbounded instruction stream per chunk. +#[derive(Debug, Clone, Copy)] +pub struct ChunkBudget { + pub max_leafset: u32, + pub max_constraints_per_chunk: u32, +} + +impl ChunkBudget { + /// Defaults derived from sweep across real-RSP shards (n_chips ∈ {6,7,34}) + /// and synthetic core / all-chips clusters at 2^25..2^27. + /// + /// `max_leafset = 64` keeps each chunk's downstream `max_reg` ≲ 128, so + /// the fused sequential kernel runs on the `MAX_REGS=128` template — its + /// per-thread `regs[]` footprint stays small enough to remain L1-resident + /// across the active block set. Pushing the leafset higher cross-tiers + /// into `MAX_REGS={256,512}`, doubling the per-thread footprint and + /// turning the kernel memory-latency-bound; pushing lower over-fragments + /// big chips so shared columns get redundantly reloaded across chunks. + pub fn recommended() -> Self { + let env_u32 = |k: &str, default: u32| -> u32 { + std::env::var(k).ok().and_then(|s| s.parse().ok()).unwrap_or(default) + }; + Self { + max_leafset: env_u32("CHUNKER_MAX_LEAFSET", 64), + max_constraints_per_chunk: env_u32("CHUNKER_MAX_CONSTRAINTS", 512), + } + } +} + +/// A grouping of constraints sharing a leaf cache. +#[derive(Debug)] +pub struct Chunk { + /// Indices into the input `&[ConstraintInfo]`. + pub constraint_indices: Vec, + /// Union of column leaves across contained constraints. + pub leafset: HashSet, + /// Max depth among contained constraints. + pub depth_max: u32, + /// Aggregate shape — `LinearWeightedSum` iff every contained constraint + /// is itself `LinearWeightedSum`. The scheduler uses this to elect the + /// column-tile lowering. + pub shape: ConstraintShape, + /// True iff this chunk is a single constraint that exceeded the budget + /// on its own. The scheduler must route it through the escape-valve + /// path (global-scratch leaves). + pub oversize_singleton: bool, +} + +/// Shape-aware chunker. +/// +/// Segregates `LinearWeightedSum` constraints from `General` ones, then chunks +/// each subset independently. Reason: a chunk's `shape` is `LinearWeightedSum` +/// only if *every* member is `LinearWeightedSum`; mixing kills column-tile +/// dispatch eligibility. So we don't let the leaf-greedy heuristic mix shapes. +/// +/// Both subsets are packed with the same budget. Future work: linear-sum +/// chunks could use a larger leaf budget since ColumnTile doesn't pay for a +/// shared-mem cache, but Phase 3 keeps it uniform. +pub fn chunk_dag(constraints: &[ConstraintInfo], budget: &ChunkBudget) -> Vec { + let (linear, general): (Vec<_>, Vec<_>) = (0..constraints.len()) + .filter(|&i| { + let c = &constraints[i]; + !c.column_leaves.is_empty() || c.work > 0 + }) + .partition(|&i| matches!(constraints[i].shape, ConstraintShape::LinearWeightedSum)); + + let mut chunks = Vec::new(); + chunks.extend(chunk_subset(constraints, &linear, budget)); + chunks.extend(chunk_subset(constraints, &general, budget)); + chunks +} + +/// Inner workhorse: greedy first-fit-decreasing over a constraint subset. +fn chunk_subset( + constraints: &[ConstraintInfo], + indices: &[usize], + budget: &ChunkBudget, +) -> Vec { + // First-fit decreasing: sort indices by descending leafset size, then by + // descending work. + let mut order: Vec = indices.to_vec(); + order.sort_by(|&a, &b| { + constraints[b] + .column_leaves + .len() + .cmp(&constraints[a].column_leaves.len()) + .then_with(|| constraints[b].work.cmp(&constraints[a].work)) + }); + + let mut chunks: Vec = Vec::new(); + for ci in order { + let c = &constraints[ci]; + + // Find an existing chunk that can absorb this constraint with the + // fewest new unique leaves. + let mut best: Option<(usize, usize)> = None; // (chunk_idx, new_leaf_count) + for (i, chunk) in chunks.iter().enumerate() { + // Quick rejection: would exceed constraint cap. + if chunk.constraint_indices.len() as u32 + 1 > budget.max_constraints_per_chunk { + continue; + } + // Compute the prospective union size. + let new_leaves = c.column_leaves.difference(&chunk.leafset).count(); + let new_union = chunk.leafset.len() + new_leaves; + if new_union as u32 > budget.max_leafset { + continue; + } + if best.is_none_or(|(_, bn)| new_leaves < bn) { + best = Some((i, new_leaves)); + } + } + + match best { + Some((idx, _)) => { + let chunk = &mut chunks[idx]; + for &leaf in &c.column_leaves { + chunk.leafset.insert(leaf); + } + chunk.constraint_indices.push(ci); + chunk.depth_max = chunk.depth_max.max(c.depth); + if !matches!(c.shape, ConstraintShape::LinearWeightedSum) { + chunk.shape = ConstraintShape::General; + } + } + None => { + // Could not fit into any existing chunk → emit a fresh one. + // If the constraint alone exceeds the budget, flag it. + let oversize_singleton = c.column_leaves.len() as u32 > budget.max_leafset; + let mut leafset = HashSet::with_capacity(c.column_leaves.len()); + for &leaf in &c.column_leaves { + leafset.insert(leaf); + } + chunks.push(Chunk { + constraint_indices: vec![ci], + leafset, + depth_max: c.depth, + shape: c.shape, + oversize_singleton, + }); + } + } + } + chunks +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ir::dag::TraceSource; + + fn cinfo(idx: usize, cols: &[u32], work: u32, depth: u32) -> ConstraintInfo { + let column_leaves = cols + .iter() + .copied() + .map(|c| ColumnLeaf { source: TraceSource::MainLocal, col: c }) + .collect(); + ConstraintInfo { + constraint_idx: idx, + root: 0, + field: crate::ir::dag::ConstraintField::Base, + alpha_index: idx as u32, + total_nodes: 0, + work, + depth, + column_leaves, + shape: ConstraintShape::General, + } + } + + #[test] + fn small_constraints_pack_into_one_chunk() { + let c = vec![cinfo(0, &[0, 1, 2], 5, 3), cinfo(1, &[0, 1, 3], 5, 3)]; + let chunks = chunk_dag(&c, &ChunkBudget { max_leafset: 16, max_constraints_per_chunk: 16 }); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].leafset.len(), 4); + assert_eq!(chunks[0].constraint_indices.len(), 2); + } + + #[test] + fn budget_forces_split() { + let c = vec![ + cinfo(0, &(0..10).collect::>(), 5, 3), + cinfo(1, &(10..20).collect::>(), 5, 3), + ]; + let chunks = chunk_dag(&c, &ChunkBudget { max_leafset: 12, max_constraints_per_chunk: 16 }); + assert_eq!(chunks.len(), 2); + } + + #[test] + fn oversize_singleton_flagged() { + let c = vec![cinfo(0, &(0..100).collect::>(), 50, 5)]; + let chunks = chunk_dag(&c, &ChunkBudget { max_leafset: 16, max_constraints_per_chunk: 16 }); + assert_eq!(chunks.len(), 1); + assert!(chunks[0].oversize_singleton); + } +} diff --git a/sp1-gpu/crates/air/src/ir/column_tile_bytecode.rs b/sp1-gpu/crates/air/src/ir/column_tile_bytecode.rs new file mode 100644 index 0000000000..a31e354257 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/column_tile_bytecode.rs @@ -0,0 +1,169 @@ +//! Bytecode for the ColumnTile lowering. +//! +//! ColumnTile applies to `LinearWeightedSum` chunks: each chunk's program is +//! `Σ_k α^k · (Σ_i coeff_{k,i} · leaf_{k,i})`. We linearize that into a flat +//! list of `(leaf, coeff, alpha_idx)` term entries. The kernel runs one lane +//! per `(term, row, eval_point)` tuple, no shared cache, no per-thread tape +//! interpretation — lane variation IS the program. + +use crate::ir::analysis::ConstraintInfo; +use crate::ir::bytecode::LeafRef; +use crate::ir::chunker::Chunk; +use crate::ir::dag::{ConstraintDag, DagNode, TraceSource}; +use crate::ir::lowering::ColumnTilePlan; +use crate::F; + +/// Kind tag for `ColumnTermEntry.coeff_kind`. Must match the CUDA header. +pub const COEFF_KIND_CONST: u32 = 0; +pub const COEFF_KIND_PUBLIC: u32 = 1; +/// Coefficient is an extension-field value supplied at launch time via the +/// kernel's `runtime_coeffs` buffer. Used by GKR-correction chunks: each +/// `batching_power_i` is a per-proof random challenge, not a DAG constant. +pub const COEFF_KIND_RUNTIME: u32 = 2; + +/// One term: `α^k · coeff · leaf_i(row, eval)`. +/// +/// `coeff_kind` discriminates between the const pool and the public-values +/// pool. `coeff_idx` indexes into the respective table. Layout matches the +/// device-side struct in `column_tile.cuh`. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct ColumnTermEntry { + pub leaf_idx: u32, + pub coeff_kind: u32, + pub coeff_idx: u32, + pub alpha_idx: u32, +} + +#[derive(Debug, Default, Clone)] +pub struct ColumnTileBytecode { + pub leaves: Vec, + pub consts: Vec, + pub publics: Vec, + pub terms: Vec, + pub n_constraints: u32, + /// True iff this is a `synthesize_gkr_chunk` output — the GKR-sweep + /// carrier for a chip with no Sequential chunk. Its terms are weighted + /// by `EF::one()`, so at launch the kernel must read the `powers_of_alpha` + /// slot holding `1` (the last slot) rather than applying the per-chip + /// alpha shift. (A chip with 0 constraints has `chip_alpha_offset` equal + /// to the table length, so the per-chip shift would read out of bounds.) + pub is_gkr_carrier: bool, +} + +/// Lower a `ColumnTilePlan` to flat term-entry bytecode. +/// +/// Returns `None` if any coefficient is not a constant or public value — +/// the plan-detector accepts a wider set (IsFirstRow etc.) but those will +/// be handled in a future phase. Returns `None` rather than panicking so +/// the scheduler can fall back to Sequential. +pub fn lower_column_tile( + chunk: &Chunk, + _constraints: &[ConstraintInfo], + dag: &ConstraintDag, + plan: &ColumnTilePlan, +) -> Option { + let mut bc = ColumnTileBytecode { + n_constraints: chunk.constraint_indices.len() as u32, + ..ColumnTileBytecode::default() + }; + + // Intern leaves, consts, publics by (source, col) / value / idx respectively. + let mut leaf_lookup: std::collections::HashMap<(u8, u32), u32> = Default::default(); + let mut const_lookup: std::collections::HashMap = Default::default(); + let mut public_lookup: std::collections::HashMap = Default::default(); + + for t in &plan.terms { + // Leaf side. + let (src_byte, col) = match dag.nodes[t.leaf_node as usize] { + DagNode::InputLeaf { source, col } => { + let s = match source { + TraceSource::PreprocessedLocal => 2u8, + TraceSource::PreprocessedNext => 3, + TraceSource::MainLocal => 4, + TraceSource::MainNext => 5, + }; + (s, col) + } + _ => return None, + }; + let leaf_idx = *leaf_lookup.entry((src_byte, col)).or_insert_with(|| { + let i = bc.leaves.len() as u32; + bc.leaves.push(LeafRef { source: src_byte, _pad: 0, col }); + i + }); + + // Coefficient side. + let (coeff_kind, coeff_idx) = match dag.nodes[t.coeff_node as usize] { + DagNode::ConstF { value } => { + use slop_algebra::PrimeField32; + let key = value.as_canonical_u32(); + let idx = *const_lookup.entry(key).or_insert_with(|| { + let i = bc.consts.len() as u32; + bc.consts.push(value); + i + }); + (COEFF_KIND_CONST, idx) + } + DagNode::PublicValue { idx } => { + let pidx = *public_lookup.entry(idx).or_insert_with(|| { + let i = bc.publics.len() as u32; + bc.publics.push(idx); + i + }); + (COEFF_KIND_PUBLIC, pidx) + } + // Other coefficient kinds (IsFirstRow, EF consts, cumsum) — not yet supported. + _ => return None, + }; + + bc.terms.push(ColumnTermEntry { leaf_idx, coeff_kind, coeff_idx, alpha_idx: t.alpha_idx }); + } + + Some(bc) +} + +/// Synthesize a GKR-correction chunk for a chip. +/// +/// Produces a `ColumnTileBytecode` whose program is +/// `Σ_i batching_power_i · col_i`, summed across all main and preprocessed +/// columns. Coefficient kind is `RUNTIME` — the caller supplies a per-launch +/// `runtime_coeffs: &[EF]` of length `main_width + preprocessed_width`. +/// +/// Term ordering matches v1's `jaggedConstraintPolyEval` GKR loop: main +/// columns first (indices `0 .. main_width`), then preprocessed +/// (indices `main_width .. main_width + preprocessed_width`). The caller +/// must lay out `runtime_coeffs` in this order. +/// +/// The GKR correction is added without an α weighting (weight `EF::one()`). +/// Terms store `alpha_idx = 0` and the chunk is flagged `is_gkr_carrier`; at +/// launch the kernel reads the `powers_of_alpha` slot holding `1` (the last +/// slot) directly, bypassing the per-chip alpha shift — which would be +/// out-of-bounds for a chip with 0 constraints. +pub fn synthesize_gkr_chunk(main_width: u32, preprocessed_width: u32) -> ColumnTileBytecode { + let mut bc = + ColumnTileBytecode { is_gkr_carrier: true, n_constraints: 1, ..Default::default() }; + + let push_term = |bc: &mut ColumnTileBytecode, source: u8, col: u32| { + let leaf_idx = bc.leaves.len() as u32; + bc.leaves.push(LeafRef { source, _pad: 0, col }); + let coeff_idx = bc.terms.len() as u32; + bc.terms.push(ColumnTermEntry { + leaf_idx, + coeff_kind: COEFF_KIND_RUNTIME, + coeff_idx, + alpha_idx: 0, + }); + }; + + // Main columns first (source byte 4 = MainLocal). + for col in 0..main_width { + push_term(&mut bc, 4, col); + } + // Then preprocessed (source byte 2 = PreprocessedLocal). + for col in 0..preprocessed_width { + push_term(&mut bc, 2, col); + } + + bc +} diff --git a/sp1-gpu/crates/air/src/ir/dag.rs b/sp1-gpu/crates/air/src/ir/dag.rs new file mode 100644 index 0000000000..47a1ae66f0 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/dag.rs @@ -0,0 +1,140 @@ +//! DAG-native IR for AIR constraints. +//! +//! Each chip's eval produces a single `ConstraintDag` with explicit +//! cross-constraint sharing via shared `NodeId`s. Leaves and constants are +//! interned during construction; arithmetic nodes are not (Rust-variable +//! reuse handles the common case, full CSE deferred to a post-pass if +//! measurement justifies). + +use crate::{EF, F}; + +/// A handle into `ConstraintDag::nodes`. Allocated monotonically by the builder. +pub type NodeId = u32; + +/// Which slice of trace a leaf comes from. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum TraceSource { + PreprocessedLocal, + PreprocessedNext, + MainLocal, + MainNext, +} + +/// Whether a constraint's assertion is over the base field or the extension field. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ConstraintField { + Base, + Extension, +} + +/// One node in the DAG. +/// +/// Leaves carry their identity (column index, public-value index, constant value) +/// so the chunker can group by leaf overlap without an extra lookup. Internal +/// nodes reference their inputs by `NodeId`; sharing is whenever the same +/// `NodeId` appears as input to multiple parents. +#[derive(Debug, Clone, Copy)] +pub enum DagNode { + // ----- Leaves (interned) ----- + InputLeaf { + source: TraceSource, + col: u32, + }, + PublicValue { + idx: u32, + }, + GlobalCumulativeSum { + idx: u32, + }, + ConstF { + value: F, + }, + ConstEF { + value: EF, + }, + IsFirstRow, + IsLastRow, + IsTransition, + + // ----- Base-field arithmetic ----- + AddF { + a: NodeId, + b: NodeId, + }, + SubF { + a: NodeId, + b: NodeId, + }, + MulF { + a: NodeId, + b: NodeId, + }, + NegF { + a: NodeId, + }, + + // ----- Extension-field arithmetic ----- + AddEF { + a: NodeId, + b: NodeId, + }, + SubEF { + a: NodeId, + b: NodeId, + }, + MulEF { + a: NodeId, + b: NodeId, + }, + NegEF { + a: NodeId, + }, + + // ----- Mixed (EF, F) ----- + /// Lift base-field value `a` into the extension field. + EFFromF { + a: NodeId, + }, + /// `a: EF + b: F`. + EFAddF { + a: NodeId, + b: NodeId, + }, + /// `a: EF - b: F`. + EFSubF { + a: NodeId, + b: NodeId, + }, + /// `a: EF * b: F`. + EFMulF { + a: NodeId, + b: NodeId, + }, +} + +/// One top-level assertion. Produced by `assert_zero` / `assert_zero_ext`. +#[derive(Debug, Clone, Copy)] +pub struct ConstraintRef { + pub root: NodeId, + pub alpha_index: u32, + pub field: ConstraintField, +} + +/// Output of running a chip's `eval` against `DagBuilder`. +#[derive(Debug)] +pub struct ConstraintDag { + pub nodes: Vec, + pub constraints: Vec, + pub preprocessed_width: u32, + pub main_width: u32, +} + +impl ConstraintDag { + pub fn num_nodes(&self) -> usize { + self.nodes.len() + } + + pub fn num_constraints(&self) -> usize { + self.constraints.len() + } +} diff --git a/sp1-gpu/crates/air/src/ir/expr.rs b/sp1-gpu/crates/air/src/ir/expr.rs new file mode 100644 index 0000000000..5b604a6884 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/expr.rs @@ -0,0 +1,566 @@ +//! Expr types — `NodeId` wrappers for arbitrary DAG computations. +//! +//! `DagExprF` represents a base-field expression; `DagExprEF` represents an +//! extension-field expression. Both are thin newtypes around `NodeId`. All +//! operator overloads allocate new DAG nodes via the global state. + +use std::iter::{Product, Sum}; +use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}; + +use slop_algebra::{AbstractExtensionField, AbstractField}; + +use crate::ir::dag::{DagNode, NodeId}; +use crate::ir::state::with_state; +use crate::ir::var::{DagVarEF, DagVarF}; +use crate::{EF, F}; + +// ============================================================================ +// DagExprF +// ============================================================================ + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +#[repr(transparent)] +pub struct DagExprF(pub NodeId); + +impl DagExprF { + pub fn node_id(self) -> NodeId { + self.0 + } +} + +impl Default for DagExprF { + fn default() -> Self { + Self::zero() + } +} + +impl From for DagExprF { + fn from(f: F) -> Self { + let id = with_state(|s| s.intern_const_f(f)); + DagExprF(id) + } +} + +// ----- Add ----- +impl Add for DagExprF { + type Output = Self; + fn add(self, rhs: F) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_f(rhs); + s.alloc(DagNode::AddF { a, b }) + }); + DagExprF(id) + } +} + +impl Add for DagExprF { + type Output = Self; + fn add(self, rhs: DagVarF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::AddF { a, b })); + DagExprF(id) + } +} + +impl Add for DagExprF { + type Output = Self; + fn add(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::AddF { a, b })); + DagExprF(id) + } +} + +impl AddAssign for DagExprF { + fn add_assign(&mut self, rhs: Self) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::AddF { a, b })); + self.0 = id; + } +} + +// ----- Sub ----- +impl Sub for DagExprF { + type Output = Self; + fn sub(self, rhs: F) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_f(rhs); + s.alloc(DagNode::SubF { a, b }) + }); + DagExprF(id) + } +} + +impl Sub for DagExprF { + type Output = Self; + fn sub(self, rhs: DagVarF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::SubF { a, b })); + DagExprF(id) + } +} + +impl Sub for DagExprF { + type Output = Self; + fn sub(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::SubF { a, b })); + DagExprF(id) + } +} + +impl SubAssign for DagExprF { + fn sub_assign(&mut self, rhs: Self) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::SubF { a, b })); + self.0 = id; + } +} + +// ----- Mul ----- +impl Mul for DagExprF { + type Output = Self; + fn mul(self, rhs: F) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_f(rhs); + s.alloc(DagNode::MulF { a, b }) + }); + DagExprF(id) + } +} + +impl Mul for DagExprF { + type Output = Self; + fn mul(self, rhs: DagVarF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::MulF { a, b })); + DagExprF(id) + } +} + +impl Mul for DagExprF { + type Output = Self; + fn mul(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::MulF { a, b })); + DagExprF(id) + } +} + +impl MulAssign for DagExprF { + fn mul_assign(&mut self, rhs: Self) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::MulF { a, b })); + self.0 = id; + } +} + +// ----- Neg ----- +impl Neg for DagExprF { + type Output = Self; + fn neg(self) -> Self::Output { + let a = self.0; + let id = with_state(|s| s.alloc(DagNode::NegF { a })); + DagExprF(id) + } +} + +impl Sum for DagExprF { + fn sum>(iter: I) -> Self { + iter.fold(Self::zero(), |acc, x| acc + x) + } +} + +impl Product for DagExprF { + fn product>(iter: I) -> Self { + iter.fold(Self::one(), |acc, x| acc * x) + } +} + +impl AbstractField for DagExprF { + type F = F; + + fn zero() -> Self { + let id = with_state(|s| s.intern_const_f(F::zero())); + DagExprF(id) + } + fn one() -> Self { + let id = with_state(|s| s.intern_const_f(F::one())); + DagExprF(id) + } + fn two() -> Self { + let id = with_state(|s| s.intern_const_f(F::two())); + DagExprF(id) + } + fn neg_one() -> Self { + let id = with_state(|s| s.intern_const_f(F::neg_one())); + DagExprF(id) + } + fn from_f(f: Self::F) -> Self { + let id = with_state(|s| s.intern_const_f(f)); + DagExprF(id) + } + fn from_bool(b: bool) -> Self { + let id = with_state(|s| s.intern_const_f(F::from_bool(b))); + DagExprF(id) + } + fn from_canonical_u8(n: u8) -> Self { + let id = with_state(|s| s.intern_const_f(F::from_canonical_u8(n))); + DagExprF(id) + } + fn from_canonical_u16(n: u16) -> Self { + let id = with_state(|s| s.intern_const_f(F::from_canonical_u16(n))); + DagExprF(id) + } + fn from_canonical_u32(n: u32) -> Self { + let id = with_state(|s| s.intern_const_f(F::from_canonical_u32(n))); + DagExprF(id) + } + fn from_canonical_u64(n: u64) -> Self { + let id = with_state(|s| s.intern_const_f(F::from_canonical_u64(n))); + DagExprF(id) + } + fn from_canonical_usize(n: usize) -> Self { + let id = with_state(|s| s.intern_const_f(F::from_canonical_usize(n))); + DagExprF(id) + } + fn from_wrapped_u32(n: u32) -> Self { + let id = with_state(|s| s.intern_const_f(F::from_wrapped_u32(n))); + DagExprF(id) + } + fn from_wrapped_u64(n: u64) -> Self { + let id = with_state(|s| s.intern_const_f(F::from_wrapped_u64(n))); + DagExprF(id) + } + fn generator() -> Self { + let id = with_state(|s| s.intern_const_f(F::generator())); + DagExprF(id) + } +} + +// ============================================================================ +// DagExprEF +// ============================================================================ + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +#[repr(transparent)] +pub struct DagExprEF(pub NodeId); + +impl DagExprEF { + pub fn node_id(self) -> NodeId { + self.0 + } +} + +impl Default for DagExprEF { + fn default() -> Self { + Self::zero() + } +} + +impl From for DagExprEF { + fn from(f: EF) -> Self { + let id = with_state(|s| s.intern_const_ef(f)); + DagExprEF(id) + } +} + +// ----- Add ----- +impl Add for DagExprEF { + type Output = Self; + fn add(self, rhs: EF) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_ef(rhs); + s.alloc(DagNode::AddEF { a, b }) + }); + DagExprEF(id) + } +} + +impl Add for DagExprEF { + type Output = Self; + fn add(self, rhs: DagVarEF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::AddEF { a, b })); + DagExprEF(id) + } +} + +impl Add for DagExprEF { + type Output = Self; + fn add(self, rhs: DagExprEF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::AddEF { a, b })); + DagExprEF(id) + } +} + +impl AddAssign for DagExprEF { + fn add_assign(&mut self, rhs: Self) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::AddEF { a, b })); + self.0 = id; + } +} + +// ----- Sub ----- +impl Sub for DagExprEF { + type Output = Self; + fn sub(self, rhs: EF) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_ef(rhs); + s.alloc(DagNode::SubEF { a, b }) + }); + DagExprEF(id) + } +} + +impl Sub for DagExprEF { + type Output = Self; + fn sub(self, rhs: DagVarEF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::SubEF { a, b })); + DagExprEF(id) + } +} + +impl Sub for DagExprEF { + type Output = Self; + fn sub(self, rhs: DagExprEF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::SubEF { a, b })); + DagExprEF(id) + } +} + +impl SubAssign for DagExprEF { + fn sub_assign(&mut self, rhs: Self) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::SubEF { a, b })); + self.0 = id; + } +} + +// ----- Mul ----- +impl Mul for DagExprEF { + type Output = Self; + fn mul(self, rhs: EF) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_ef(rhs); + s.alloc(DagNode::MulEF { a, b }) + }); + DagExprEF(id) + } +} + +impl Mul for DagExprEF { + type Output = Self; + fn mul(self, rhs: DagExprEF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::MulEF { a, b })); + DagExprEF(id) + } +} + +impl MulAssign for DagExprEF { + fn mul_assign(&mut self, rhs: Self) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::MulEF { a, b })); + self.0 = id; + } +} + +// ----- Neg ----- +impl Neg for DagExprEF { + type Output = Self; + fn neg(self) -> Self::Output { + let a = self.0; + let id = with_state(|s| s.alloc(DagNode::NegEF { a })); + DagExprEF(id) + } +} + +impl Sum for DagExprEF { + fn sum>(iter: I) -> Self { + iter.fold(Self::zero(), |acc, x| acc + x) + } +} + +impl Product for DagExprEF { + fn product>(iter: I) -> Self { + iter.fold(Self::one(), |acc, x| acc * x) + } +} + +impl AbstractField for DagExprEF { + type F = EF; + + fn zero() -> Self { + let id = with_state(|s| s.intern_const_ef(EF::zero())); + DagExprEF(id) + } + fn one() -> Self { + let id = with_state(|s| s.intern_const_ef(EF::one())); + DagExprEF(id) + } + fn two() -> Self { + let id = with_state(|s| s.intern_const_ef(EF::two())); + DagExprEF(id) + } + fn neg_one() -> Self { + let id = with_state(|s| s.intern_const_ef(EF::neg_one())); + DagExprEF(id) + } + fn from_f(f: Self::F) -> Self { + let id = with_state(|s| s.intern_const_ef(f)); + DagExprEF(id) + } + fn from_bool(b: bool) -> Self { + let id = with_state(|s| s.intern_const_ef(EF::from_bool(b))); + DagExprEF(id) + } + fn from_canonical_u8(n: u8) -> Self { + let id = with_state(|s| s.intern_const_ef(EF::from_canonical_u8(n))); + DagExprEF(id) + } + fn from_canonical_u16(n: u16) -> Self { + let id = with_state(|s| s.intern_const_ef(EF::from_canonical_u16(n))); + DagExprEF(id) + } + fn from_canonical_u32(n: u32) -> Self { + let id = with_state(|s| s.intern_const_ef(EF::from_canonical_u32(n))); + DagExprEF(id) + } + fn from_canonical_u64(n: u64) -> Self { + let id = with_state(|s| s.intern_const_ef(EF::from_canonical_u64(n))); + DagExprEF(id) + } + fn from_canonical_usize(n: usize) -> Self { + let id = with_state(|s| s.intern_const_ef(EF::from_canonical_usize(n))); + DagExprEF(id) + } + fn from_wrapped_u32(n: u32) -> Self { + let id = with_state(|s| s.intern_const_ef(EF::from_wrapped_u32(n))); + DagExprEF(id) + } + fn from_wrapped_u64(n: u64) -> Self { + let id = with_state(|s| s.intern_const_ef(EF::from_wrapped_u64(n))); + DagExprEF(id) + } + fn generator() -> Self { + let id = with_state(|s| s.intern_const_ef(EF::generator())); + DagExprEF(id) + } +} + +// ----- Mixed (DagExprEF, DagExprF) ----- + +impl From for DagExprEF { + fn from(v: DagExprF) -> Self { + let a = v.0; + let id = with_state(|s| s.alloc(DagNode::EFFromF { a })); + DagExprEF(id) + } +} + +impl Add for DagExprEF { + type Output = Self; + fn add(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::EFAddF { a, b })); + DagExprEF(id) + } +} + +impl AddAssign for DagExprEF { + fn add_assign(&mut self, rhs: DagExprF) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::EFAddF { a, b })); + self.0 = id; + } +} + +impl Sub for DagExprEF { + type Output = Self; + fn sub(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::EFSubF { a, b })); + DagExprEF(id) + } +} + +impl SubAssign for DagExprEF { + fn sub_assign(&mut self, rhs: DagExprF) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::EFSubF { a, b })); + self.0 = id; + } +} + +impl Mul for DagExprEF { + type Output = DagExprEF; + fn mul(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::EFMulF { a, b })); + DagExprEF(id) + } +} + +impl MulAssign for DagExprEF { + fn mul_assign(&mut self, rhs: DagExprF) { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::EFMulF { a, b })); + self.0 = id; + } +} + +impl AbstractExtensionField for DagExprEF { + const D: usize = 4; + + fn from_base(value: DagExprF) -> Self { + let a = value.0; + let id = with_state(|s| s.alloc(DagNode::EFFromF { a })); + DagExprEF(id) + } + + fn from_base_slice(_: &[DagExprF]) -> Self { + todo!() + } + + fn from_base_fn DagExprF>(_: Func) -> Self { + todo!() + } + + fn as_base_slice(&self) -> &[DagExprF] { + todo!() + } +} diff --git a/sp1-gpu/crates/air/src/ir/lowering.rs b/sp1-gpu/crates/air/src/ir/lowering.rs new file mode 100644 index 0000000000..a306d451ed --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/lowering.rs @@ -0,0 +1,284 @@ +//! Lowering plans per chunk. +//! +//! A `Lowering` describes one way to execute a chunk on the GPU. The IR +//! itself is lowering-agnostic — these plans are derived from a chunk plus +//! its constraint DAGs, pre-computed once per (chunk, chip-set) and reused +//! across sumcheck rounds. +//! +//! Currently ships Sequential + ColumnTile. ListScheduled and TermExpanded +//! are sketched as variants but their plan builders are not yet implemented. + +use crate::ir::analysis::{ConstraintInfo, ConstraintShape}; +use crate::ir::chunker::Chunk; +use crate::ir::dag::{ConstraintDag, NodeId}; + +/// One executable plan for a chunk. The scheduler picks one of these per +/// `(chunk, round)` based on the lane budget at that round. +#[derive(Debug, Clone)] +pub enum Lowering { + /// One thread per `(row, eval)` lane. Walks the DAG sequentially in + /// topological order. Cache `(zero, one)` leaves in shared memory at + /// CTA preamble; each lane evaluates the same DAG in its own register + /// file. Work-optimal — pays nothing for parallelism beyond the lane + /// axis. + Sequential(SequentialPlan), + + /// Lanes vary along `(col, row, eval)`. Each lane reads one column's + /// `(zero, one)` direct from global (no shared cache). Only applies + /// when the chunk's DAG is `Σ_i (coeff_i · leaf_i)` over a contiguous + /// column range. + ColumnTile(ColumnTilePlan), + + /// Reserved for Phase 5. Lanes cooperate across DAG levels within a warp. + ListScheduled, + + /// Reserved for Phase 5. Lanes vary along `(term, row, eval)` after + /// expanding the chunk into monomials. + TermExpanded, +} + +#[derive(Debug, Clone)] +pub struct SequentialPlan { + /// Topologically-sorted node IDs across the chunk's constituent DAGs. + /// `topo_order[0]` is evaluated first; `topo_order.last()` is one of + /// the constraint roots. + pub topo_order: Vec, + /// Number of distinct DAG values live at any one point. The kernel + /// allocates this many per-thread registers (or shared slots, per the + /// downstream allocator). + pub max_live: u32, + /// Total arithmetic ops counted on the deduplicated topo order — reflects + /// actual emission cost after CSE. + pub topo_work: u32, +} + +#[derive(Debug, Clone)] +pub struct ColumnTilePlan { + /// Flat term list across the chunk's constituent constraints. + pub terms: Vec, +} + +#[derive(Debug, Clone, Copy)] +pub struct ColumnTilePlanTerm { + /// Node ID of the coefficient (must be a constant or public). + pub coeff_node: NodeId, + /// Node ID of the column leaf. + pub leaf_node: NodeId, + /// `α^k` index for the constraint that owns this term. + pub alpha_idx: u32, +} + +/// Enumerate the lowerings that *apply* to this chunk. Always includes +/// Sequential; includes ColumnTile only if the chunk's shape allows. +/// +/// Note: this returns plans, not a single choice. The scheduler picks one +/// at dispatch time given the round's lane budget. +pub fn enumerate_lowerings( + chunk: &Chunk, + constraints: &[ConstraintInfo], + dag: &ConstraintDag, +) -> Vec { + let mut out = Vec::new(); + out.push(Lowering::Sequential(build_sequential(chunk, constraints, dag))); + if matches!(chunk.shape, ConstraintShape::LinearWeightedSum) { + if let Some(plan) = build_column_tile(chunk, constraints, dag) { + out.push(Lowering::ColumnTile(plan)); + } + } + out +} + +/// Topological sort of the union of constraint subgraphs in the chunk. +fn build_sequential( + chunk: &Chunk, + constraints: &[ConstraintInfo], + dag: &ConstraintDag, +) -> SequentialPlan { + use std::collections::HashSet; + + let roots: Vec = + chunk.constraint_indices.iter().map(|&ci| constraints[ci].root).collect(); + + let mut visited: HashSet = HashSet::new(); + let mut topo_order: Vec = Vec::new(); + for &root in &roots { + post_order(dag, root, &mut visited, &mut topo_order); + } + + let topo_work = + topo_order.iter().filter(|&&n| is_arithmetic(&dag.nodes[n as usize])).count() as u32; + + // Liveness-bound: walk in topo order, track which already-emitted nodes + // still have un-emitted parents. Conservative O(N²) sweep — fine for + // chunk sizes today. + let max_live = compute_max_live(&topo_order, dag); + + SequentialPlan { topo_order, max_live, topo_work } +} + +fn post_order( + dag: &ConstraintDag, + node_id: NodeId, + visited: &mut std::collections::HashSet, + out: &mut Vec, +) { + if !visited.insert(node_id) { + return; + } + for child in children(&dag.nodes[node_id as usize]).into_iter().flatten() { + post_order(dag, child, visited, out); + } + out.push(node_id); +} + +fn children(node: &crate::ir::dag::DagNode) -> [Option; 2] { + use crate::ir::dag::DagNode::*; + match *node { + InputLeaf { .. } + | PublicValue { .. } + | GlobalCumulativeSum { .. } + | ConstF { .. } + | ConstEF { .. } + | IsFirstRow + | IsLastRow + | IsTransition => [None, None], + AddF { a, b } + | SubF { a, b } + | MulF { a, b } + | AddEF { a, b } + | SubEF { a, b } + | MulEF { a, b } + | EFAddF { a, b } + | EFSubF { a, b } + | EFMulF { a, b } => [Some(a), Some(b)], + NegF { a } | NegEF { a } | EFFromF { a } => [Some(a), None], + } +} + +fn is_arithmetic(node: &crate::ir::dag::DagNode) -> bool { + use crate::ir::dag::DagNode::*; + !matches!( + node, + InputLeaf { .. } + | PublicValue { .. } + | GlobalCumulativeSum { .. } + | ConstF { .. } + | ConstEF { .. } + | IsFirstRow + | IsLastRow + | IsTransition + ) +} + +fn compute_max_live(topo: &[NodeId], dag: &ConstraintDag) -> u32 { + // For each node, find the last position it's used. live(i) = count of + // nodes whose first-emit ≤ i and last-use ≥ i. + let pos_of: std::collections::HashMap = + topo.iter().enumerate().map(|(i, &n)| (n, i)).collect(); + + let mut last_use: std::collections::HashMap = + topo.iter().map(|&n| (n, 0)).collect(); + for (i, &n) in topo.iter().enumerate() { + for c in children(&dag.nodes[n as usize]).into_iter().flatten() { + if let Some(&p) = pos_of.get(&c) { + let _ = p; // suppress unused + let e = last_use.entry(c).or_insert(0); + if i > *e { + *e = i; + } + } + } + } + + let mut max_live: u32 = 0; + let mut live: u32 = 0; + let mut end_at: std::collections::HashMap> = Default::default(); + for (i, &n) in topo.iter().enumerate() { + live += 1; + end_at.entry(*last_use.get(&n).unwrap_or(&i)).or_default().push(n); + max_live = max_live.max(live); + if let Some(ending) = end_at.get(&i) { + live = live.saturating_sub(ending.len() as u32); + } + } + max_live +} + +/// Try to materialize a column-tile plan: requires every constraint in +/// the chunk to be `LinearWeightedSum`. Each constraint's terms get tagged +/// with its `alpha_index` so the kernel can weight them correctly. +fn build_column_tile( + chunk: &Chunk, + constraints: &[ConstraintInfo], + dag: &ConstraintDag, +) -> Option { + let mut terms: Vec = Vec::new(); + for &ci in &chunk.constraint_indices { + let c = &constraints[ci]; + if !matches!(c.shape, ConstraintShape::LinearWeightedSum) { + return None; + } + let mut raw: Vec<(NodeId, NodeId)> = Vec::new(); + flatten_linear(dag, c.root, &mut raw)?; + for (coeff, leaf) in raw { + terms.push(ColumnTilePlanTerm { + coeff_node: coeff, + leaf_node: leaf, + alpha_idx: c.alpha_index, + }); + } + } + Some(ColumnTilePlan { terms }) +} + +/// Walk an Add/Sub-of-Mul tree, pushing each `(coefficient, column_leaf)` pair. +fn flatten_linear( + dag: &ConstraintDag, + node_id: NodeId, + out: &mut Vec<(NodeId, NodeId)>, +) -> Option<()> { + use crate::ir::dag::DagNode::*; + match dag.nodes[node_id as usize] { + AddF { a, b } | SubF { a, b } => { + flatten_linear(dag, a, out)?; + flatten_linear(dag, b, out)?; + Some(()) + } + MulF { a, b } => { + let a_is_coeff = is_coefficient(dag, a); + let b_is_coeff = is_coefficient(dag, b); + match (a_is_coeff, b_is_coeff) { + (true, false) => { + out.push((a, b)); + Some(()) + } + (false, true) => { + out.push((b, a)); + Some(()) + } + _ => None, + } + } + InputLeaf { .. } => { + // Bare leaf at the spine — synthesize a coefficient-of-one node ID. + // We can't allocate here; punt by skipping this case for now. + // (Linear-sum chunks with bare-leaf terms are rare; revisit if seen.) + None + } + _ => None, + } +} + +fn is_coefficient(dag: &ConstraintDag, node_id: NodeId) -> bool { + use crate::ir::dag::DagNode::*; + matches!( + dag.nodes[node_id as usize], + ConstF { .. } + | ConstEF { .. } + | PublicValue { .. } + | GlobalCumulativeSum { .. } + | IsFirstRow + | IsLastRow + | IsTransition + ) +} diff --git a/sp1-gpu/crates/air/src/ir/mod.rs b/sp1-gpu/crates/air/src/ir/mod.rs new file mode 100644 index 0000000000..0854b85460 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/mod.rs @@ -0,0 +1,35 @@ +//! DAG-native AIR constraint IR. +//! +//! Each chip's `eval` produces a single `ConstraintDag` with explicit +//! cross-constraint sharing, which is then chunked, lowered, and compiled to +//! flat bytecode that the fused zerocheck kernels interpret. +//! +//! Entry point: [`builder::build_dag`]. + +pub mod analysis; +pub mod builder; +pub mod bytecode; +pub mod chunker; +pub mod column_tile_bytecode; +pub mod dag; +pub mod expr; +pub mod lowering; +pub mod scheduler; +mod state; +pub mod var; + +pub use analysis::{analyze_constraints, ColumnLeaf, ConstraintInfo, ConstraintShape}; +pub use builder::{build_dag, DagBuilder}; +pub use bytecode::{lower_sequential, BcOp, ChunkBytecode, DagInstr, LeafRef}; +pub use chunker::{chunk_dag, Chunk, ChunkBudget}; +pub use column_tile_bytecode::{ + lower_column_tile, synthesize_gkr_chunk, ColumnTermEntry, ColumnTileBytecode, COEFF_KIND_CONST, + COEFF_KIND_PUBLIC, COEFF_KIND_RUNTIME, +}; +pub use dag::{ConstraintDag, ConstraintField, ConstraintRef, DagNode, NodeId, TraceSource}; +pub use expr::{DagExprEF, DagExprF}; +pub use lowering::{ + enumerate_lowerings, ColumnTilePlan, ColumnTilePlanTerm, Lowering, SequentialPlan, +}; +pub use scheduler::{pick_lowering, GpuCaps, LaneBudget, LoweringPick}; +pub use var::{DagVarEF, DagVarF}; diff --git a/sp1-gpu/crates/air/src/ir/scheduler.rs b/sp1-gpu/crates/air/src/ir/scheduler.rs new file mode 100644 index 0000000000..fe96fe3b46 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/scheduler.rs @@ -0,0 +1,93 @@ +//! Per-round lowering selection. +//! +//! Given a chunk's pre-computed lowerings and the round's lane budget, picks +//! one lowering via a heuristic cost model. The model is deterministic and +//! intentionally stateless — replace with a measurement-fitted version later +//! if profiling shows mispredictions. + +use crate::ir::analysis::ConstraintShape; +use crate::ir::chunker::Chunk; +use crate::ir::lowering::Lowering; + +/// GPU caps that bound lowering choice. Today we only use shared-memory +/// budget; the others are placeholders for the Phase 2+ kernel code. +#[derive(Debug, Clone, Copy)] +pub struct GpuCaps { + pub shmem_bytes_per_cta: u32, + pub warp_size: u32, +} + +impl GpuCaps { + pub fn ampere_default() -> Self { + Self { shmem_bytes_per_cta: 48 * 1024, warp_size: 32 } + } +} + +/// Lanes available along the data axes (row × eval_point) for this +/// `(chunk, round)`. Computed by the scheduler from per-chip row counts. +#[derive(Debug, Clone, Copy)] +pub struct LaneBudget { + pub log_rows: u32, + pub eval_points: u32, +} + +impl LaneBudget { + pub fn total_lanes(&self) -> u32 { + (1u32 << self.log_rows).saturating_mul(self.eval_points) + } +} + +/// Decision returned by `pick_lowering`. Carries both the index into +/// `chunk.lowerings` and a brief reason string for telemetry. +#[derive(Debug, Clone)] +pub struct LoweringPick { + pub lowering_idx: usize, + pub reason: &'static str, +} + +/// Heuristic cost model. +/// +/// Rules, in order: +/// 1. If a ColumnTile plan exists, use it. Structurally optimal for +/// `Σ β_i · leaf_i` chunks; no shared cache, lanes vary along the +/// column axis so wide chunks fill warps even at row_tile=1. +/// 2. If lanes fill at least one warp on the row × eval axis, use +/// Sequential. Adding extra parallelism via other lowerings wastes +/// work without filling more lanes. +/// 3. Otherwise (lane-starved): fall through to Sequential anyway. +/// Phase 5 will plug in ListScheduled / TermExpanded here. +pub fn pick_lowering( + chunk: &Chunk, + lowerings: &[Lowering], + lane_budget: LaneBudget, + _caps: GpuCaps, +) -> LoweringPick { + // Rule 1: structural — column-tile if the chunk shape allows. + if matches!(chunk.shape, ConstraintShape::LinearWeightedSum) { + for (i, l) in lowerings.iter().enumerate() { + if matches!(l, Lowering::ColumnTile(_)) { + return LoweringPick { + lowering_idx: i, + reason: "column-tile (linear-weighted-sum shape)", + }; + } + } + } + + // Rule 2: lane budget fills a warp → sequential is work-optimal. + let total_lanes = lane_budget.total_lanes(); + let reason = if total_lanes >= _caps.warp_size { + "sequential (lanes fill warp)" + } else { + // Rule 3: lane-starved; Sequential is the fallback. A future + // ListScheduled / TermExpanded lowering would route in here. + "sequential (lane-starved fallback)" + }; + for (i, l) in lowerings.iter().enumerate() { + if matches!(l, Lowering::Sequential(_)) { + return LoweringPick { lowering_idx: i, reason }; + } + } + // Should not reach here — Sequential is always emitted. + LoweringPick { lowering_idx: 0, reason: "fallback (no sequential found)" } +} diff --git a/sp1-gpu/crates/air/src/ir/state.rs b/sp1-gpu/crates/air/src/ir/state.rs new file mode 100644 index 0000000000..784b5d2c79 --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/state.rs @@ -0,0 +1,160 @@ +//! Global DAG-builder state, mirroring the existing `CUDA_P3_EVAL_*` pattern. +//! +//! The state is process-global, serialized by a guard lock so chip eval can +//! be done one-at-a-time. The high-level entry (`build_dag`) acquires the lock, +//! resets state, runs `chip.eval`, snapshots the result, and releases. +//! +//! This is intentionally the same shape as the v1 globals — easier to reason +//! about, easier to migrate. Can be replaced by builder-owned state once the +//! v1 path is deleted. + +use std::collections::HashMap; +use std::sync::Mutex; + +use lazy_static::lazy_static; + +use crate::ir::dag::{ConstraintRef, DagNode, NodeId, TraceSource}; +use crate::{EF, F}; + +/// State accumulated during a single chip's `eval`. +#[derive(Default)] +pub struct DagState { + pub nodes: Vec, + pub constraints: Vec, + + // Interning tables for leaves and constants only. + pub leaf_intern: HashMap<(TraceSource, u32), NodeId>, + pub public_intern: HashMap, + pub gcs_intern: HashMap, + pub const_f_intern: HashMap, // keyed by F's u32 repr + pub const_ef_intern: HashMap<[u32; 4], NodeId>, + pub singleton_is_first_row: Option, + pub singleton_is_last_row: Option, + pub singleton_is_transition: Option, + + pub num_constraints: u32, +} + +impl DagState { + pub fn alloc(&mut self, node: DagNode) -> NodeId { + let id = self.nodes.len() as u32; + self.nodes.push(node); + id + } + + pub fn intern_leaf(&mut self, source: TraceSource, col: u32) -> NodeId { + if let Some(&id) = self.leaf_intern.get(&(source, col)) { + return id; + } + let id = self.alloc(DagNode::InputLeaf { source, col }); + self.leaf_intern.insert((source, col), id); + id + } + + pub fn intern_public(&mut self, idx: u32) -> NodeId { + if let Some(&id) = self.public_intern.get(&idx) { + return id; + } + let id = self.alloc(DagNode::PublicValue { idx }); + self.public_intern.insert(idx, id); + id + } + + pub fn intern_gcs(&mut self, idx: u32) -> NodeId { + if let Some(&id) = self.gcs_intern.get(&idx) { + return id; + } + let id = self.alloc(DagNode::GlobalCumulativeSum { idx }); + self.gcs_intern.insert(idx, id); + id + } + + pub fn intern_const_f(&mut self, value: F) -> NodeId { + let key = f_key(value); + if let Some(&id) = self.const_f_intern.get(&key) { + return id; + } + let id = self.alloc(DagNode::ConstF { value }); + self.const_f_intern.insert(key, id); + id + } + + pub fn intern_const_ef(&mut self, value: EF) -> NodeId { + let key = ef_key(value); + if let Some(&id) = self.const_ef_intern.get(&key) { + return id; + } + let id = self.alloc(DagNode::ConstEF { value }); + self.const_ef_intern.insert(key, id); + id + } + + pub fn intern_is_first_row(&mut self) -> NodeId { + if let Some(id) = self.singleton_is_first_row { + return id; + } + let id = self.alloc(DagNode::IsFirstRow); + self.singleton_is_first_row = Some(id); + id + } + + pub fn intern_is_last_row(&mut self) -> NodeId { + if let Some(id) = self.singleton_is_last_row { + return id; + } + let id = self.alloc(DagNode::IsLastRow); + self.singleton_is_last_row = Some(id); + id + } + + pub fn intern_is_transition(&mut self) -> NodeId { + if let Some(id) = self.singleton_is_transition { + return id; + } + let id = self.alloc(DagNode::IsTransition); + self.singleton_is_transition = Some(id); + id + } + + pub fn reset(&mut self) { + self.nodes.clear(); + self.constraints.clear(); + self.leaf_intern.clear(); + self.public_intern.clear(); + self.gcs_intern.clear(); + self.const_f_intern.clear(); + self.const_ef_intern.clear(); + self.singleton_is_first_row = None; + self.singleton_is_last_row = None; + self.singleton_is_transition = None; + self.num_constraints = 0; + } +} + +/// Stable key for an `F` value. Uses the canonical `u32` representation. +fn f_key(value: F) -> u32 { + use slop_algebra::PrimeField32; + value.as_canonical_u32() +} + +/// Stable key for an `EF` value. Each base coefficient mapped via `f_key`. +fn ef_key(value: EF) -> [u32; 4] { + use slop_algebra::AbstractExtensionField; + let slice: &[F] = value.as_base_slice(); + assert!(slice.len() == 4, "EF degree expected to be 4"); + [f_key(slice[0]), f_key(slice[1]), f_key(slice[2]), f_key(slice[3])] +} + +lazy_static! { + /// Outer guard. Acquired by `build_dag` for the duration of one chip's eval. + pub static ref DAG_BUILDER_LOCK: Mutex<()> = Mutex::new(()); + + /// The accumulating DAG state. Mutated by operator overloads and assert_zero. + pub static ref DAG_STATE: Mutex = Mutex::new(DagState::default()); +} + +/// Convenience: lock the state and run a closure with `&mut DagState`. +pub(crate) fn with_state(f: impl FnOnce(&mut DagState) -> R) -> R { + let mut guard = DAG_STATE.lock().unwrap(); + f(&mut guard) +} diff --git a/sp1-gpu/crates/air/src/ir/var.rs b/sp1-gpu/crates/air/src/ir/var.rs new file mode 100644 index 0000000000..67ebc2172d --- /dev/null +++ b/sp1-gpu/crates/air/src/ir/var.rs @@ -0,0 +1,213 @@ +//! Var types — opaque `NodeId` wrappers used as leaves in the DAG. +//! +//! Distinct from `DagExpr*` only to satisfy `AirBuilder`'s associated-type +//! constraints; both ultimately point at DAG nodes. Operations on `DagVar*` +//! promote them to `DagExpr*` via `Into`. + +use std::ops::{Add, Mul, Neg, Sub}; + +use crate::ir::dag::DagNode; +use crate::ir::dag::{NodeId, TraceSource}; +use crate::ir::expr::{DagExprEF, DagExprF}; +use crate::ir::state::with_state; +use crate::{EF, F}; + +/// Base-field variable. Wraps the `NodeId` of a leaf node. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct DagVarF(pub NodeId); + +impl DagVarF { + pub fn node_id(self) -> NodeId { + self.0 + } + + pub fn preprocessed_local(col: u32) -> Self { + let id = with_state(|s| s.intern_leaf(TraceSource::PreprocessedLocal, col)); + DagVarF(id) + } + + pub fn preprocessed_next(col: u32) -> Self { + let id = with_state(|s| s.intern_leaf(TraceSource::PreprocessedNext, col)); + DagVarF(id) + } + + pub fn main_local(col: u32) -> Self { + let id = with_state(|s| s.intern_leaf(TraceSource::MainLocal, col)); + DagVarF(id) + } + + pub fn main_next(col: u32) -> Self { + let id = with_state(|s| s.intern_leaf(TraceSource::MainNext, col)); + DagVarF(id) + } + + pub fn public_value(idx: u32) -> Self { + let id = with_state(|s| s.intern_public(idx)); + DagVarF(id) + } + + pub fn global_cumulative_sum(idx: u32) -> Self { + let id = with_state(|s| s.intern_gcs(idx)); + DagVarF(id) + } + + pub fn is_first_row() -> Self { + let id = with_state(|s| s.intern_is_first_row()); + DagVarF(id) + } + + pub fn is_last_row() -> Self { + let id = with_state(|s| s.intern_is_last_row()); + DagVarF(id) + } + + pub fn is_transition() -> Self { + let id = with_state(|s| s.intern_is_transition()); + DagVarF(id) + } +} + +impl From for DagExprF { + fn from(v: DagVarF) -> Self { + DagExprF(v.0) + } +} + +// ----- Add ----- +impl Add for DagVarF { + type Output = DagExprF; + fn add(self, rhs: F) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_f(rhs); + s.alloc(DagNode::AddF { a, b }) + }); + DagExprF(id) + } +} + +impl Add for DagVarF { + type Output = DagExprF; + fn add(self, rhs: DagVarF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::AddF { a, b })); + DagExprF(id) + } +} + +impl Add for DagVarF { + type Output = DagExprF; + fn add(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::AddF { a, b })); + DagExprF(id) + } +} + +// ----- Sub ----- +impl Sub for DagVarF { + type Output = DagExprF; + fn sub(self, rhs: F) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_f(rhs); + s.alloc(DagNode::SubF { a, b }) + }); + DagExprF(id) + } +} + +impl Sub for DagVarF { + type Output = DagExprF; + fn sub(self, rhs: DagVarF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::SubF { a, b })); + DagExprF(id) + } +} + +impl Sub for DagVarF { + type Output = DagExprF; + fn sub(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::SubF { a, b })); + DagExprF(id) + } +} + +// ----- Mul ----- +impl Mul for DagVarF { + type Output = DagExprF; + fn mul(self, rhs: F) -> Self::Output { + let a = self.0; + let id = with_state(|s| { + let b = s.intern_const_f(rhs); + s.alloc(DagNode::MulF { a, b }) + }); + DagExprF(id) + } +} + +impl Mul for DagVarF { + type Output = DagExprF; + fn mul(self, rhs: DagVarF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::MulF { a, b })); + DagExprF(id) + } +} + +impl Mul for DagVarF { + type Output = DagExprF; + fn mul(self, rhs: DagExprF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::MulF { a, b })); + DagExprF(id) + } +} + +impl Neg for DagVarF { + type Output = DagExprF; + fn neg(self) -> Self::Output { + let a = self.0; + let id = with_state(|s| s.alloc(DagNode::NegF { a })); + DagExprF(id) + } +} + +// ============================================================================ +// Extension-field variant. Mirrors `SymbolicVarEF`. For now this is a thin +// wrapper that holds an EF NodeId. SP1 chips today don't use EF vars at the +// `eval` level (the existing kernel rejects nontrivial EF var variants). +// We still define the type so trait bounds are satisfied. +// ============================================================================ + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct DagVarEF(pub NodeId); + +impl From for DagExprEF { + fn from(v: DagVarEF) -> Self { + DagExprEF(v.0) + } +} + +impl Mul for DagExprEF { + type Output = DagExprEF; + fn mul(self, rhs: DagVarEF) -> Self::Output { + let a = self.0; + let b = rhs.0; + let id = with_state(|s| s.alloc(DagNode::MulEF { a, b })); + DagExprEF(id) + } +} + +// Suppress unused-imports warning for EF in this file: +const _: fn() = || { + let _: EF; +}; diff --git a/sp1-gpu/crates/air/src/lib.rs b/sp1-gpu/crates/air/src/lib.rs index cad62939b3..42a4404867 100644 --- a/sp1-gpu/crates/air/src/lib.rs +++ b/sp1-gpu/crates/air/src/lib.rs @@ -1,199 +1,10 @@ #![allow(clippy::assign_op_pattern)] -pub mod air_block; -pub mod instruction; -pub mod optimizer; -pub mod symbolic_expr_ef; -pub mod symbolic_expr_f; -pub mod symbolic_var_ef; -pub mod symbolic_var_f; -use std::sync::Mutex; +pub mod ir; -use air_block::BlockAir; -use instruction::{Instruction16, Instruction32}; -use lazy_static::lazy_static; -use slop_air::{ - AirBuilder, AirBuilderWithPublicValues, ExtensionBuilder, PairBuilder, PermutationAirBuilder, -}; use slop_algebra::extension::BinomialExtensionField; -use slop_matrix::dense::RowMajorMatrixView; - -use sp1_core_machine::air::TrivialOperationBuilder; -use sp1_hypercube::air::EmptyMessageBuilder; -use sp1_hypercube::{AirOpenedValues, PROOF_MAX_NUM_PVS}; use sp1_primitives::SP1Field; -use symbolic_expr_ef::SymbolicExprEF; -use symbolic_expr_f::SymbolicExprF; -use symbolic_var_ef::SymbolicVarEF; -use symbolic_var_f::SymbolicVarF; pub type F = SP1Field; pub type EF = BinomialExtensionField; - -lazy_static! { - pub static ref CUDA_P3_EVAL_LOCK: Mutex<()> = Mutex::new(()); - pub static ref CUDA_P3_EVAL_CODE: Mutex> = Mutex::new(Vec::new()); - pub static ref CUDA_P3_EVAL_F_CONSTANTS: Mutex> = Mutex::new(Vec::new()); - pub static ref CUDA_P3_EVAL_EF_CONSTANTS: Mutex> = Mutex::new(Vec::new()); - pub static ref CUDA_P3_EVAL_EXPR_F_CTR: Mutex = Mutex::new(0); - pub static ref CUDA_P3_EVAL_EXPR_EF_CTR: Mutex = Mutex::new(0); -} - -pub struct SymbolicProverFolder<'a> { - pub preprocessed: RowMajorMatrixView<'a, SymbolicVarF>, - pub main: RowMajorMatrixView<'a, SymbolicVarF>, - pub public_values: &'a [SymbolicVarF], - pub num_constraints: u32, -} - -impl<'a> AirBuilder for SymbolicProverFolder<'a> { - type F = F; - type Var = SymbolicVarF; - type Expr = SymbolicExprF; - type M = RowMajorMatrixView<'a, SymbolicVarF>; - - fn main(&self) -> Self::M { - self.main - } - - fn is_first_row(&self) -> Self::Expr { - unimplemented!(); - } - - fn is_last_row(&self) -> Self::Expr { - unimplemented!(); - } - - fn is_transition_window(&self, _: usize) -> Self::Expr { - unimplemented!(); - } - - fn assert_zero>(&mut self, x: I) { - let x: Self::Expr = x.into(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assert_zero(x)); - self.num_constraints += 1; - drop(code); - } -} - -impl ExtensionBuilder for SymbolicProverFolder<'_> { - type EF = EF; - type ExprEF = SymbolicExprEF; - type VarEF = SymbolicVarEF; - - fn assert_zero_ext(&mut self, x: I) - where - I: Into, - { - let x: SymbolicExprEF = x.into(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assert_zero(x)); - self.num_constraints += 1; - drop(code); - } -} - -impl<'a> PermutationAirBuilder for SymbolicProverFolder<'a> { - type MP = RowMajorMatrixView<'a, SymbolicVarEF>; - type RandomVar = SymbolicVarEF; - - fn permutation(&self) -> Self::MP { - unimplemented!(); - } - fn permutation_randomness(&self) -> &[Self::RandomVar] { - unimplemented!(); - } -} - -impl PairBuilder for SymbolicProverFolder<'_> { - fn preprocessed(&self) -> Self::M { - self.preprocessed - } -} - -impl AirBuilderWithPublicValues for SymbolicProverFolder<'_> { - type PublicVar = SymbolicVarF; - - fn public_values(&self) -> &[Self::PublicVar] { - self.public_values - } -} - -impl EmptyMessageBuilder for SymbolicProverFolder<'_> {} - -impl TrivialOperationBuilder for SymbolicProverFolder<'_> {} - -/// Generates code in CUDA for evaluating the constraint polynomial on the device. -#[allow(clippy::type_complexity)] -pub fn codegen_cuda_eval( - air: &A, -) -> (Vec, Vec, Vec, Vec, Vec, Vec, Vec, u32, u32) -where - A: for<'a> BlockAir>, -{ - let preprocessed_width = air.preprocessed_width() as u32; - let width = air.width() as u32; - let preprocessed = AirOpenedValues { - local: (0..preprocessed_width).map(SymbolicVarF::preprocessed_local).collect(), - }; - let main = AirOpenedValues { local: (0..width).map(SymbolicVarF::main_local).collect() }; - let public_values = - (0..PROOF_MAX_NUM_PVS as u32).map(SymbolicVarF::public_value).collect::>(); - - let mut folder = SymbolicProverFolder { - preprocessed: preprocessed.view(), - main: main.view(), - public_values: &public_values, - num_constraints: 0, - }; - - let nb_block = air.num_blocks(); - let mut constraint_indices = Vec::new(); - let mut instructions = Vec::new(); - let mut block_indices = Vec::new(); - let mut f_constants = Vec::new(); - let mut f_constants_indices = Vec::new(); - let mut ef_constants = Vec::new(); - let mut ef_constants_indices = Vec::new(); - let mut f_ctr = 0; - let mut ef_ctr = 0; - for i in 0..nb_block { - constraint_indices.push(folder.num_constraints); - air.eval_block(&mut folder, i); - let code = CUDA_P3_EVAL_CODE.lock().unwrap().to_vec(); - let block_f_constants = CUDA_P3_EVAL_F_CONSTANTS.lock().unwrap().to_vec(); - let block_ef_constants = CUDA_P3_EVAL_EF_CONSTANTS.lock().unwrap().to_vec(); - CUDA_P3_EVAL_RESET(); - let (block_code, block_f_ctr, block_ef_ctr) = optimizer::optimize(code); - block_indices.push(instructions.len() as u32); - f_constants_indices.push(f_constants.len() as u32); - ef_constants_indices.push(ef_constants.len() as u32); - f_ctr = f_ctr.max(block_f_ctr); - ef_ctr = ef_ctr.max(block_ef_ctr); - instructions.extend(block_code); - f_constants.extend(block_f_constants); - ef_constants.extend(block_ef_constants); - } - - ( - constraint_indices, - instructions, - block_indices, - f_constants, - f_constants_indices, - ef_constants, - ef_constants_indices, - f_ctr as u32, - ef_ctr as u32, - ) -} - -#[allow(non_snake_case)] -pub fn CUDA_P3_EVAL_RESET() { - *CUDA_P3_EVAL_CODE.lock().unwrap() = Vec::new(); - *CUDA_P3_EVAL_EF_CONSTANTS.lock().unwrap() = Vec::new(); - *CUDA_P3_EVAL_EXPR_F_CTR.lock().unwrap() = 0; - *CUDA_P3_EVAL_EXPR_EF_CTR.lock().unwrap() = 0; -} diff --git a/sp1-gpu/crates/air/src/optimizer.rs b/sp1-gpu/crates/air/src/optimizer.rs deleted file mode 100644 index 1056fd0898..0000000000 --- a/sp1-gpu/crates/air/src/optimizer.rs +++ /dev/null @@ -1,199 +0,0 @@ -use std::collections::HashMap; - -use crate::instruction::{Instruction16, Instruction32, Opcode}; - -struct RegisterAllocator { - f_used: Vec, - ef_used: Vec, - f_vreg2phys_map: HashMap, - ef_vreg2phys_map: HashMap, - f_max: usize, - ef_max: usize, -} - -impl RegisterAllocator { - pub fn new() -> Self { - let mut f_used = vec![false; 2048]; - let mut ef_used = vec![false; 1024]; - - // Make %v0 always map to %p0. - f_used[0] = true; - ef_used[0] = true; - let f_vreg2phys_map = HashMap::new(); - let ef_vreg2phys_map = HashMap::new(); - - Self { f_used, ef_used, f_vreg2phys_map, ef_vreg2phys_map, f_max: 0, ef_max: 0 } - } - - pub fn f_vreg2phys(&mut self, vreg: u32) -> u32 { - if self.f_vreg2phys_map.contains_key(&vreg) { - return self.f_vreg2phys_map[&vreg]; - } - for i in 0..self.f_used.len() { - if !self.f_used[i] { - self.f_used[i] = true; - let phys = i as u32; - self.f_vreg2phys_map.insert(vreg, phys); - if i > self.f_max { - self.f_max = i; - } - return phys; - } - } - unreachable!() - } - - pub fn ef_vreg2phys(&mut self, vreg: u32) -> u32 { - if self.ef_vreg2phys_map.contains_key(&vreg) { - return self.ef_vreg2phys_map[&vreg]; - } - for i in 0..self.ef_used.len() { - if !self.ef_used[i] { - self.ef_used[i] = true; - let phys = i as u32; - self.ef_vreg2phys_map.insert(vreg, phys); - if i > self.ef_max { - self.ef_max = i; - } - return phys; - } - } - unreachable!() - } - - pub fn f_free(&mut self, vreg: u32) { - if self.f_vreg2phys_map.contains_key(&vreg) { - let phys = self.f_vreg2phys_map.remove(&vreg).unwrap(); - self.f_used[phys as usize] = false; - } - } - - pub fn ef_free(&mut self, vreg: u32) { - if self.ef_vreg2phys_map.contains_key(&vreg) { - let phys = self.ef_vreg2phys_map.remove(&vreg).unwrap(); - self.ef_used[phys as usize] = false; - } - } -} - -pub fn optimize(instructions: Vec) -> (Vec, usize, usize) { - let mut f_first_time_vreg_used: HashMap = HashMap::new(); - let mut f_last_time_vreg_used: HashMap = HashMap::new(); - let mut ef_first_time_vreg_used: HashMap = HashMap::new(); - let mut ef_last_time_vreg_used: HashMap = HashMap::new(); - - for (i, instr) in instructions.iter().enumerate() { - let i = i as u32; - let opcode = Opcode::from(instr.opcode); - - if opcode.is_f_assign() { - f_first_time_vreg_used.entry(instr.a).or_insert(i); - f_last_time_vreg_used.insert(instr.a, i); - } - - if opcode.is_f_arg1() { - f_first_time_vreg_used.entry(instr.b).or_insert(i); - f_last_time_vreg_used.insert(instr.b, i); - } - - if opcode.is_f_arg2() { - f_first_time_vreg_used.entry(instr.c).or_insert(i); - f_last_time_vreg_used.insert(instr.c, i); - } - - if opcode.is_e_assign() { - ef_first_time_vreg_used.entry(instr.a).or_insert(i); - ef_last_time_vreg_used.insert(instr.a, i); - } - - if opcode.is_e_arg1() { - ef_first_time_vreg_used.entry(instr.b).or_insert(i); - ef_last_time_vreg_used.insert(instr.b, i); - } - - if opcode.is_e_arg2() { - ef_first_time_vreg_used.entry(instr.c).or_insert(i); - ef_last_time_vreg_used.insert(instr.c, i); - } - } - - let mut optimized_instructions = Vec::new(); - let mut allocator = RegisterAllocator::new(); - for (i, instr) in instructions.iter().enumerate() { - let i = i as u32; - let opcode = Opcode::from(instr.opcode); - - let mut new_instr = *instr; - if opcode.is_f_assign() { - let phys_a = allocator.f_vreg2phys(instr.a); - new_instr.a = phys_a; - } - - if opcode.is_f_arg1() { - let phys_b = allocator.f_vreg2phys(instr.b); - new_instr.b = phys_b; - } - - if opcode.is_f_arg2() { - let phys_c = allocator.f_vreg2phys(instr.c); - new_instr.c = phys_c; - } - - if opcode.is_e_assign() { - let phys_a = allocator.ef_vreg2phys(instr.a); - new_instr.a = phys_a; - } - - if opcode.is_e_arg1() { - let phys_b = allocator.ef_vreg2phys(instr.b); - new_instr.b = phys_b; - } - - if opcode.is_e_arg2() { - let phys_c = allocator.ef_vreg2phys(instr.c); - new_instr.c = phys_c; - } - - optimized_instructions.push(new_instr); - - if opcode.is_f_assign() && f_last_time_vreg_used.get(&instr.a).unwrap() == &i { - allocator.f_free(instr.a); - } - - if opcode.is_f_arg1() && f_last_time_vreg_used.get(&instr.b).unwrap() == &i { - allocator.f_free(instr.b); - } - - if opcode.is_f_arg2() && f_last_time_vreg_used.get(&instr.c).unwrap() == &i { - allocator.f_free(instr.c); - } - - if opcode.is_e_assign() && ef_last_time_vreg_used.get(&instr.a).unwrap() == &i { - allocator.ef_free(instr.a); - } - - if opcode.is_e_arg1() && ef_last_time_vreg_used.get(&instr.b).unwrap() == &i { - allocator.ef_free(instr.b); - } - - if opcode.is_e_arg2() && ef_last_time_vreg_used.get(&instr.c).unwrap() == &i { - allocator.ef_free(instr.c); - } - } - - ( - optimized_instructions - .into_iter() - .map(|instr| Instruction16 { - opcode: instr.opcode, - b_variant: instr.b_variant, - c_variant: instr.c_variant, - a: instr.a as u16, - b: instr.b as u16, - c: instr.c as u16, - }) - .collect(), - allocator.f_max, - allocator.ef_max, - ) -} diff --git a/sp1-gpu/crates/air/src/symbolic_expr_ef.rs b/sp1-gpu/crates/air/src/symbolic_expr_ef.rs deleted file mode 100644 index 5aca313726..0000000000 --- a/sp1-gpu/crates/air/src/symbolic_expr_ef.rs +++ /dev/null @@ -1,484 +0,0 @@ -use std::{ - iter::{Product, Sum}, - ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}, -}; - -use slop_algebra::{AbstractExtensionField, AbstractField}; - -use crate::{ - instruction::Instruction32, symbolic_expr_f::SymbolicExprF, symbolic_var_ef::SymbolicVarEF, - CUDA_P3_EVAL_CODE, CUDA_P3_EVAL_EXPR_EF_CTR, EF, -}; - -#[derive(Debug, Copy, PartialEq, Eq, Hash)] -#[repr(C)] -pub struct SymbolicExprEF(pub u32); - -impl SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "Empty for SymbolicExprEF")] - pub fn empty() -> Self { - Self(u32::MAX) - } - - // #[instrument(skip_all, level = "trace", name = "Alloc for SymbolicExprEF")] - pub fn alloc() -> Self { - let mut tmp = CUDA_P3_EVAL_EXPR_EF_CTR.lock().unwrap(); - let id = *tmp; - *tmp += 1; - drop(tmp); - Self(id) - } - - pub fn variant(&self) -> u8 { - 0 - } - - pub fn data(&self) -> u32 { - self.0 - } -} - -impl Default for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "Default for SymbolicExprEF")] - fn default() -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::zero())); - drop(code); - output - } -} - -impl From for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from(f: EF) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, f)); - drop(code); - output - } -} - -impl Add for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicExprEF")] - fn add(self, rhs: EF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_add_ec(output, self, rhs)); - drop(code); - output - } -} - -impl Add for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicExprEF")] - fn add(self, rhs: SymbolicVarEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_add_ev(output, self, rhs)); - drop(code); - output - } -} - -impl Add for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicExprEF")] - fn add(self, rhs: SymbolicExprEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_add_ee(output, self, rhs)); - drop(code); - output - } -} - -impl AddAssign for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "AddAssign for SymbolicExprEF")] - fn add_assign(&mut self, rhs: Self) { - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_add_assign_e(*self, rhs)); - drop(code); - } -} - -impl Sub for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicExprEF")] - fn sub(self, rhs: EF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_sub_ec(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicExprEF")] - fn sub(self, rhs: SymbolicVarEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_sub_ev(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicExprEF")] - fn sub(self, rhs: SymbolicExprEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_sub_ee(output, self, rhs)); - drop(code); - output - } -} - -impl SubAssign for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "SubAssign for SymbolicExprEF")] - fn sub_assign(&mut self, rhs: Self) { - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_sub_assign_e(*self, rhs)); - drop(code); - } -} - -impl Mul for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicExprEF")] - fn mul(self, rhs: EF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_mul_ec(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicExprEF")] - fn mul(self, rhs: SymbolicVarEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_mul_ev(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicExprEF")] - fn mul(self, rhs: SymbolicExprEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_mul_ee(output, self, rhs)); - drop(code); - output - } -} - -impl MulAssign for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "MulAssign for SymbolicExprEF")] - fn mul_assign(&mut self, rhs: Self) { - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_mul_assign_e(*self, rhs)); - drop(code); - } -} - -impl Neg for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Neg for SymbolicExprEF")] - fn neg(self) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_neg_e(output, self)); - drop(code); - output - } -} - -impl Sum for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "Sum for SymbolicExprEF")] - fn sum>(iter: I) -> Self { - let mut output = SymbolicExprEF::zero(); - for item in iter { - output = output + item; - } - output - } -} - -impl Product for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "Product for SymbolicExprEF")] - fn product>(iter: I) -> Self { - let mut output = SymbolicExprEF::one(); - for item in iter { - output = output * item; - } - output - } -} - -impl Clone for SymbolicExprEF { - #[allow(clippy::non_canonical_clone_impl)] - // #[instrument(skip_all, level = "trace", name = "Clone for SymbolicExprEF")] - fn clone(&self) -> Self { - // let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - // let output = SymbolicExprEF::alloc(); - // code.push(Instruction32::e_assign_e(output, *self)); - // drop(code); - // output - *self - } -} - -impl AbstractField for SymbolicExprEF { - type F = EF; - - // #[instrument(skip_all, level = "trace", name = "Zero for SymbolicExprEF")] - fn zero() -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::zero())); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "One for SymbolicExprEF")] - fn one() -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::one())); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "Two for SymbolicExprEF")] - fn two() -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::two())); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "NegOne for SymbolicExprEF")] - fn neg_one() -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::neg_one())); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_f(f: Self::F) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, f)); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_bool(b: bool) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::from_bool(b))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_canonical_u8(n: u8) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::from_canonical_u8(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_canonical_u16(n: u16) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::from_canonical_u16(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_canonical_u32(n: u32) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::from_canonical_u32(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_canonical_u64(n: u64) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::from_canonical_u64(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_canonical_usize(n: usize) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::from_canonical_usize(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_wrapped_u32(n: u32) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::from_wrapped_u32(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from_wrapped_u64(n: u64) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::from_wrapped_u64(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "Generator for SymbolicExprEF")] - fn generator() -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_c(output, EF::generator())); - drop(code); - output - } -} - -impl From for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from(value: SymbolicExprF) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::ef_from_e(output, value)); - drop(code); - output - } -} - -impl Add for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicExprEF")] - fn add(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::ef_add_ee(output, self, rhs)); - drop(code); - output - } -} - -impl AddAssign for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "AddAssign for SymbolicExprEF")] - fn add_assign(&mut self, rhs: SymbolicExprF) { - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::ef_add_assign_e(*self, rhs)); - drop(code); - } -} - -impl Sub for SymbolicExprEF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicExprEF")] - fn sub(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::ef_sub_ee(output, self, rhs)); - drop(code); - output - } -} - -impl SubAssign for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "SubAssign for SymbolicExprEF")] - fn sub_assign(&mut self, rhs: SymbolicExprF) { - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::ef_sub_assign_e(*self, rhs)); - drop(code); - } -} - -impl Mul for SymbolicExprEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicExprEF")] - fn mul(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::ef_mul_ee(output, self, rhs)); - drop(code); - output - } -} - -impl MulAssign for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "MulAssign for SymbolicExprEF")] - fn mul_assign(&mut self, rhs: SymbolicExprF) { - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::ef_mul_assign_e(*self, rhs)); - drop(code); - } -} - -impl AbstractExtensionField for SymbolicExprEF { - const D: usize = 4; - - fn from_base(value: SymbolicExprF) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::ef_from_e(output, value)); - drop(code); - output - } - - fn from_base_slice(_: &[SymbolicExprF]) -> Self { - todo!() - } - - fn from_base_fn SymbolicExprF>(_: F) -> Self { - todo!() - } - - fn as_base_slice(&self) -> &[SymbolicExprF] { - todo!() - } -} diff --git a/sp1-gpu/crates/air/src/symbolic_expr_f.rs b/sp1-gpu/crates/air/src/symbolic_expr_f.rs deleted file mode 100644 index dfbfd3fddd..0000000000 --- a/sp1-gpu/crates/air/src/symbolic_expr_f.rs +++ /dev/null @@ -1,377 +0,0 @@ -use std::{ - iter::{Product, Sum}, - ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}, -}; - -use crate::{ - instruction::Instruction32, symbolic_var_f::SymbolicVarF, CUDA_P3_EVAL_CODE, - CUDA_P3_EVAL_EXPR_F_CTR, F, -}; - -use slop_algebra::AbstractField; - -#[derive(Debug, Copy, PartialEq, Eq, Hash)] -#[repr(C)] -pub struct SymbolicExprF(pub u32); - -impl SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "Empty for SymbolicExprF")] - pub fn empty() -> Self { - Self(u32::MAX) - } - - // #[instrument(skip_all, level = "trace", name = "Alloc for SymbolicExprF")] - pub fn alloc() -> Self { - let mut tmp = CUDA_P3_EVAL_EXPR_F_CTR.lock().unwrap(); - let id = *tmp; - *tmp += 1; - drop(tmp); - Self(id) - } - - pub fn variant(&self) -> u8 { - 0 - } - - pub fn data(&self) -> u32 { - self.0 - } -} - -impl Default for SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "Default for SymbolicExprF")] - fn default() -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::zero())); - drop(code); - output - } -} - -impl From for SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from(f: F) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, f)); - drop(code); - output - } -} - -impl Add for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicExprF")] - fn add(self, rhs: F) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_add_ec(output, self, rhs)); - drop(code); - output - } -} - -impl Add for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicExprF")] - fn add(self, rhs: SymbolicVarF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_add_ev(output, self, rhs)); - drop(code); - output - } -} - -impl Add for SymbolicExprF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicExprF")] - fn add(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_add_ee(output, self, rhs)); - drop(code); - output - } -} - -impl AddAssign for SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "AddAssign for SymbolicExprF")] - fn add_assign(&mut self, _: SymbolicExprF) { - unreachable!() - } -} - -impl Sub for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicExprF")] - fn sub(self, rhs: F) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_sub_ec(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicExprF")] - fn sub(self, rhs: SymbolicVarF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_sub_ev(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicExprF")] - fn sub(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_sub_ee(output, self, rhs)); - drop(code); - output - } -} - -impl SubAssign for SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "SubAssign for SymbolicExprF")] - fn sub_assign(&mut self, _: SymbolicExprF) { - unreachable!() - } -} - -impl Mul for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicExprF")] - fn mul(self, rhs: F) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_mul_ec(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicExprF")] - fn mul(self, rhs: SymbolicVarF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_mul_ev(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicExprF")] - fn mul(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_mul_ee(output, self, rhs)); - drop(code); - output - } -} - -impl MulAssign for SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "MulAssign for SymbolicExprF")] - fn mul_assign(&mut self, _: SymbolicExprF) { - unreachable!() - } -} - -impl Neg for SymbolicExprF { - type Output = Self; - - // #[instrument(skip_all, level = "trace", name = "Neg for SymbolicExprF")] - fn neg(self) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_neg_e(output, self)); - drop(code); - output - } -} - -impl Sum for SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "Sum for SymbolicExprF")] - fn sum>(iter: I) -> Self { - let mut output = SymbolicExprF::zero(); - for item in iter { - output = output + item; - } - output - } -} - -impl Product for SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "Product for SymbolicExprF")] - fn product>(iter: I) -> Self { - let mut output = SymbolicExprF::one(); - for item in iter { - output = output * item; - } - output - } -} - -impl Clone for SymbolicExprF { - #[allow(clippy::non_canonical_clone_impl)] - // #[instrument(skip_all, level = "trace", name = "Clone for SymbolicExprF")] - fn clone(&self) -> Self { - // let output = SymbolicExprF::alloc(); - // let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - // code.push(Instruction32::f_assign_e(output, *self)); - // drop(code); - // output - *self - } -} - -impl AbstractField for SymbolicExprF { - type F = F; - - // #[instrument(skip_all, level = "trace", name = "Zero for SymbolicExprF")] - fn zero() -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::zero())); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "One for SymbolicExprF")] - fn one() -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::one())); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "Two for SymbolicExprF")] - fn two() -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::two())); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "NegOne for SymbolicExprF")] - fn neg_one() -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::neg_one())); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_f(f: Self::F) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, f)); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_bool(b: bool) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::from_bool(b))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_canonical_u8(n: u8) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::from_canonical_u8(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_canonical_u16(n: u16) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::from_canonical_u16(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_canonical_u32(n: u32) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::from_canonical_u32(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_canonical_u64(n: u64) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::from_canonical_u64(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_canonical_usize(n: usize) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::from_canonical_usize(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_wrapped_u32(n: u32) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::from_wrapped_u32(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from_wrapped_u64(n: u64) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::from_wrapped_u64(n))); - drop(code); - output - } - - // #[instrument(skip_all, level = "trace", name = "Generator for SymbolicExprF")] - fn generator() -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_c(output, F::generator())); - drop(code); - output - } -} diff --git a/sp1-gpu/crates/air/src/symbolic_var_ef.rs b/sp1-gpu/crates/air/src/symbolic_var_ef.rs deleted file mode 100644 index 6864a5f1cb..0000000000 --- a/sp1-gpu/crates/air/src/symbolic_var_ef.rs +++ /dev/null @@ -1,188 +0,0 @@ -use std::fmt::Debug; -use std::ops::{Add, Mul, Sub}; - -use crate::{instruction::Instruction32, symbolic_expr_ef::SymbolicExprEF, CUDA_P3_EVAL_CODE, EF}; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SymbolicVarEF { - Empty, - PermutationLocal(u32), - PermutationNext(u32), - PermutationChallenge(u32), - CumulativeSum(u32), -} - -impl SymbolicVarEF { - // #[instrument(skip_all, level = "trace", name = "Empty for SymbolicVarEF")] - pub fn empty() -> Self { - Self::Empty - } - - // #[instrument(skip_all, level = "trace", name = "PermutationLocal for SymbolicVarEF")] - pub fn permutation_local(idx: u32) -> Self { - Self::PermutationLocal(idx) - } - - // #[instrument(skip_all, level = "trace", name = "PermutationNext for SymbolicVarEF")] - pub fn permutation_next(idx: u32) -> Self { - Self::PermutationNext(idx) - } - - // #[instrument(skip_all, level = "trace", name = "PermutationChallenge for SymbolicVarEF")] - pub fn permutation_challenge(idx: u32) -> Self { - Self::PermutationChallenge(idx) - } - - // #[instrument(skip_all, level = "trace", name = "CumulativeSum for SymbolicVarEF")] - pub fn cumulative_sum(idx: u32) -> Self { - Self::CumulativeSum(idx) - } - - pub fn variant(&self) -> u8 { - match self { - Self::Empty => 0x00, - Self::PermutationLocal(_) => 0x01, - Self::PermutationNext(_) => 0x02, - Self::PermutationChallenge(_) => 0x03, - Self::CumulativeSum(_) => 0x04, - } - } - - pub fn data(&self) -> u32 { - match self { - Self::Empty => 0, - Self::PermutationLocal(idx) => *idx, - Self::PermutationNext(idx) => *idx, - Self::PermutationChallenge(idx) => *idx, - Self::CumulativeSum(idx) => *idx, - } - } -} - -impl From for SymbolicExprEF { - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprEF")] - fn from(value: SymbolicVarEF) -> Self { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_assign_v(output, value)); - drop(code); - output - } -} - -impl Add for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicVarEF")] - fn add(self, rhs: EF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_add_vc(output, self, rhs)); - drop(code); - output - } -} - -impl Add for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicVarEF")] - fn add(self, rhs: SymbolicVarEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_add_vv(output, self, rhs)); - drop(code); - output - } -} - -impl Add for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicVarEF")] - fn add(self, rhs: SymbolicExprEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_add_ve(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicVarEF")] - fn sub(self, rhs: EF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_sub_vc(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicVarEF")] - fn sub(self, rhs: SymbolicVarEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_sub_vv(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicVarEF")] - fn sub(self, rhs: SymbolicExprEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_sub_ve(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicVarEF")] - fn mul(self, rhs: EF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_mul_vc(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicVarEF")] - fn mul(self, rhs: SymbolicVarEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_mul_vv(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicVarEF { - type Output = SymbolicExprEF; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicVarEF")] - fn mul(self, rhs: SymbolicExprEF) -> Self::Output { - let output = SymbolicExprEF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::e_mul_ve(output, self, rhs)); - drop(code); - output - } -} diff --git a/sp1-gpu/crates/air/src/symbolic_var_f.rs b/sp1-gpu/crates/air/src/symbolic_var_f.rs deleted file mode 100644 index bd3cf004de..0000000000 --- a/sp1-gpu/crates/air/src/symbolic_var_f.rs +++ /dev/null @@ -1,227 +0,0 @@ -use std::fmt::Debug; -use std::ops::{Add, Mul, Sub}; - -use crate::instruction::f_constant; -use crate::{instruction::Instruction32, symbolic_expr_f::SymbolicExprF, CUDA_P3_EVAL_CODE, F}; - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum SymbolicVarF { - Empty, - Constant(u32), - PreprocessedLocal(u32), - PreprocessedNext(u32), - MainLocal(u32), - MainNext(u32), - IsFirstRow, - IsLastRow, - IsTransition, - PublicValue(u32), - GlobalCumulativeSum(u32), -} - -impl SymbolicVarF { - pub fn empty() -> Self { - Self::Empty - } - - pub fn constant(f: F) -> Self { - let f = f_constant(f); - Self::Constant(f) - } - - pub fn preprocessed_local(idx: u32) -> Self { - Self::PreprocessedLocal(idx) - } - - pub fn preprocessed_next(idx: u32) -> Self { - Self::PreprocessedNext(idx) - } - - pub fn main_local(idx: u32) -> Self { - Self::MainLocal(idx) - } - - pub fn main_next(idx: u32) -> Self { - Self::MainNext(idx) - } - - pub fn is_first_row() -> Self { - Self::IsFirstRow - } - - pub fn is_last_row() -> Self { - Self::IsLastRow - } - - pub fn is_transition() -> Self { - Self::IsTransition - } - - pub fn public_value(idx: u32) -> Self { - Self::PublicValue(idx) - } - - pub fn global_cumulative_sum(idx: u32) -> Self { - Self::GlobalCumulativeSum(idx) - } - - pub fn variant(&self) -> u8 { - match self { - Self::Empty => 0x00, - Self::Constant(_) => 0x01, - Self::PreprocessedLocal(_) => 0x02, - Self::PreprocessedNext(_) => 0x03, - Self::MainLocal(_) => 0x04, - Self::MainNext(_) => 0x05, - Self::IsFirstRow => 0x06, - Self::IsLastRow => 0x07, - Self::IsTransition => 0x08, - Self::PublicValue(_) => 0x09, - Self::GlobalCumulativeSum(_) => 0x0A, - } - } - - pub fn data(&self) -> u32 { - match self { - Self::Empty => 0, - Self::Constant(f) => *f, - Self::PreprocessedLocal(idx) => *idx, - Self::PreprocessedNext(idx) => *idx, - Self::MainLocal(idx) => *idx, - Self::MainNext(idx) => *idx, - Self::IsFirstRow => 0, - Self::IsLastRow => 0, - Self::IsTransition => 0, - Self::PublicValue(idx) => *idx, - Self::GlobalCumulativeSum(idx) => *idx, - } - } -} - -impl From for SymbolicExprF { - // #[instrument(skip_all, level = "trace", name = "From for SymbolicExprF")] - fn from(val: SymbolicVarF) -> Self { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_assign_v(output, val)); - drop(code); - output - } -} - -impl Add for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicVarF")] - fn add(self, rhs: F) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_add_vc(output, self, rhs)); - drop(code); - output - } -} - -impl Add for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicVarF")] - fn add(self, rhs: SymbolicVarF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_add_vv(output, self, rhs)); - drop(code); - output - } -} - -impl Add for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Add for SymbolicVarF")] - fn add(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_add_ve(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicVarF")] - fn sub(self, rhs: F) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_sub_vc(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicVarF")] - fn sub(self, rhs: SymbolicVarF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_sub_vv(output, self, rhs)); - drop(code); - output - } -} - -impl Sub for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Sub for SymbolicVarF")] - fn sub(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_sub_ve(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicVarF")] - fn mul(self, rhs: F) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_mul_vc(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicVarF")] - fn mul(self, rhs: SymbolicVarF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_mul_vv(output, self, rhs)); - drop(code); - output - } -} - -impl Mul for SymbolicVarF { - type Output = SymbolicExprF; - - // #[instrument(skip_all, level = "trace", name = "Mul for SymbolicVarF")] - fn mul(self, rhs: SymbolicExprF) -> Self::Output { - let output = SymbolicExprF::alloc(); - let mut code = CUDA_P3_EVAL_CODE.lock().unwrap(); - code.push(Instruction32::f_mul_ve(output, self, rhs)); - drop(code); - output - } -} diff --git a/sp1-gpu/crates/cuda/src/mle/restrict.rs b/sp1-gpu/crates/cuda/src/mle/restrict.rs index fa7dfa42b9..140faef8a9 100644 --- a/sp1-gpu/crates/cuda/src/mle/restrict.rs +++ b/sp1-gpu/crates/cuda/src/mle/restrict.rs @@ -3,12 +3,12 @@ use std::sync::Arc; use slop_algebra::{ExtensionField, Field}; use slop_multilinear::MleEval; use sp1_gpu_sys::{ - runtime::KernelPtr, - v2_kernels::{ + kernels::{ fix_last_variable_ext_ext_kernel, fix_last_variable_felt_ext_kernel, mle_fix_last_variable_koala_bear_base_extension_constant_padding, mle_fix_last_variable_koala_bear_ext_ext_constant_padding, }, + runtime::KernelPtr, }; use sp1_primitives::{SP1ExtensionField, SP1Field}; diff --git a/sp1-gpu/crates/jagged_sumcheck/src/hadamard.rs b/sp1-gpu/crates/jagged_sumcheck/src/hadamard.rs index 2fa1c7970e..a6debe74be 100644 --- a/sp1-gpu/crates/jagged_sumcheck/src/hadamard.rs +++ b/sp1-gpu/crates/jagged_sumcheck/src/hadamard.rs @@ -5,14 +5,14 @@ use slop_algebra::{AbstractExtensionField, UnivariatePolynomial}; use slop_challenger::FieldChallenger; use slop_multilinear::Mle; use slop_multilinear::MleBaseBackend; +use sp1_gpu_cudart::sys::kernels::hadamard_fix_last_variable_and_sum_as_poly_base_ext_kernel; +use sp1_gpu_cudart::sys::kernels::hadamard_fix_last_variable_and_sum_as_poly_ext_ext_kernel; +use sp1_gpu_cudart::sys::kernels::hadamard_sum_as_poly_base_ext_kernel; +use sp1_gpu_cudart::sys::kernels::hadamard_sum_as_poly_ext_ext_kernel; +use sp1_gpu_cudart::sys::kernels::mle_fix_last_variable_koala_bear_ext_ext_zero_padding; +use sp1_gpu_cudart::sys::kernels::padded_hadamard_fix_and_sum; use sp1_gpu_cudart::sys::runtime::Dim3; use sp1_gpu_cudart::sys::runtime::KernelPtr; -use sp1_gpu_cudart::sys::v2_kernels::hadamard_fix_last_variable_and_sum_as_poly_base_ext_kernel; -use sp1_gpu_cudart::sys::v2_kernels::hadamard_fix_last_variable_and_sum_as_poly_ext_ext_kernel; -use sp1_gpu_cudart::sys::v2_kernels::hadamard_sum_as_poly_base_ext_kernel; -use sp1_gpu_cudart::sys::v2_kernels::hadamard_sum_as_poly_ext_ext_kernel; -use sp1_gpu_cudart::sys::v2_kernels::mle_fix_last_variable_koala_bear_ext_ext_zero_padding; -use sp1_gpu_cudart::sys::v2_kernels::padded_hadamard_fix_and_sum; use sp1_gpu_cudart::TaskScope; use sp1_gpu_cudart::{args, DeviceTensor}; use sp1_gpu_utils::{Ext, Felt}; diff --git a/sp1-gpu/crates/jagged_sumcheck/src/sumcheck.rs b/sp1-gpu/crates/jagged_sumcheck/src/sumcheck.rs index 783a8c7b29..15892e8916 100644 --- a/sp1-gpu/crates/jagged_sumcheck/src/sumcheck.rs +++ b/sp1-gpu/crates/jagged_sumcheck/src/sumcheck.rs @@ -1,6 +1,6 @@ use sp1_gpu_cudart::{ args, - sys::v2_kernels::{ + sys::kernels::{ jagged_fix_and_sum, jagged_sum_as_poly, mle_fix_last_variable_koala_bear_ext_ext_zero_padding, padded_hadamard_fix_and_sum, }, diff --git a/sp1-gpu/crates/jagged_tracegen/src/lib.rs b/sp1-gpu/crates/jagged_tracegen/src/lib.rs index ae443138c4..ed1eb682f3 100644 --- a/sp1-gpu/crates/jagged_tracegen/src/lib.rs +++ b/sp1-gpu/crates/jagged_tracegen/src/lib.rs @@ -21,7 +21,7 @@ use slop_alloc::{Backend, Buffer, HasBackend, Slice}; use slop_challenger::IopCtx; use slop_jagged::JaggedProverData; use slop_multilinear::Mle; -use sp1_gpu_cudart::sys::v2_kernels::{ +use sp1_gpu_cudart::sys::kernels::{ count_and_add_kernel, fill_buffer, generate_col_index, generate_start_indices, sum_to_trace_kernel, }; @@ -1035,7 +1035,7 @@ mod tests { use serial_test::serial; use slop_algebra::AbstractField; use slop_alloc::{Buffer, GLOBAL_CPU_BACKEND}; - use sp1_gpu_cudart::sys::v2_kernels::jagged_eval_kernel_chunked_felt; + use sp1_gpu_cudart::sys::kernels::jagged_eval_kernel_chunked_felt; use sp1_gpu_cudart::{ run_in_place, DeviceBuffer, DevicePoint, DeviceTensor, PinnedBuffer, TaskScope, }; diff --git a/sp1-gpu/crates/jagged_tracegen/src/test_utils.rs b/sp1-gpu/crates/jagged_tracegen/src/test_utils.rs index 741bf6f108..b8b3f6c351 100644 --- a/sp1-gpu/crates/jagged_tracegen/src/test_utils.rs +++ b/sp1-gpu/crates/jagged_tracegen/src/test_utils.rs @@ -69,13 +69,19 @@ pub mod random { /// {"name": "Memory", "preprocessed_width": 0, "main_width": 32, "height": 512} /// ] /// ``` + /// + /// Entries are sorted by `name` on load. The downstream cluster is a + /// `BTreeSet` (chip-name order) and the prover's per-chip trace + /// offset computation walks it in that order; the synthesized trace must + /// match. Sorting here makes the loader robust to any input array order. pub fn read_layout_from_json( path: impl AsRef, ) -> std::io::Result { let file = std::fs::File::open(path)?; let reader = std::io::BufReader::new(file); - let entries: Vec = serde_json::from_reader(reader) + let mut entries: Vec = serde_json::from_reader(reader) .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + entries.sort_by(|a, b| a.name.cmp(&b.name)); Ok(AbstractChipLayoutWithHeights::new( entries .into_iter() @@ -600,6 +606,16 @@ pub mod bench_utils { let layout = read_layout_from_json(path).expect("failed to read JSON layout"); let machine = RiscvAir::::machine(); let cluster = cluster_from_json_layout(&machine, &layout); + // The synthesized trace lays columns out in `layout` order, while + // the prover walks the cluster (`BTreeSet`) order. They must + // agree or per-chip trace offsets are wrong. `read_layout_from_json` + // sorts by name to match `BTreeSet` iteration; assert it here so a + // future change to `Chip`'s ordering fails loudly instead of + // silently corrupting the trace. + assert!( + cluster.iter().map(|c| c.name()).eq(layout.chip_names()), + "JSON layout column order must match cluster (BTreeSet) iteration order", + ); let device_mle = random_jagged_trace_mle_from_layout::(rng, &layout, LOG_STACKING_HEIGHT) .into_device(scope); diff --git a/sp1-gpu/crates/logup_gkr/src/execution.rs b/sp1-gpu/crates/logup_gkr/src/execution.rs index e1de75eca1..436894499c 100644 --- a/sp1-gpu/crates/logup_gkr/src/execution.rs +++ b/sp1-gpu/crates/logup_gkr/src/execution.rs @@ -1,7 +1,7 @@ use slop_alloc::{Buffer, HasBackend}; use sp1_gpu_cudart::{ args, - sys::v2_kernels::{ + sys::kernels::{ logup_gkr_circuit_transition, logup_gkr_extract_output, logup_gkr_first_layer_transition, }, DeviceBuffer, DeviceMle, diff --git a/sp1-gpu/crates/logup_gkr/src/sumcheck.rs b/sp1-gpu/crates/logup_gkr/src/sumcheck.rs index 26f7763a69..d313b940fe 100644 --- a/sp1-gpu/crates/logup_gkr/src/sumcheck.rs +++ b/sp1-gpu/crates/logup_gkr/src/sumcheck.rs @@ -13,7 +13,7 @@ use slop_tensor::Tensor; use sp1_gpu_cudart::DevicePoint; use sp1_gpu_cudart::{ args, - sys::v2_kernels::{ + sys::kernels::{ logup_gkr_first_sum_as_poly_circuit_layer as first_sum_as_poly_layer_circuit_layer_kernel, logup_gkr_fix_and_sum_circuit_layer as fix_and_sum_circuit_layer_kernel, logup_gkr_fix_and_sum_first_layer as fix_and_sum_first_layer_kernel, diff --git a/sp1-gpu/crates/logup_gkr/src/tracegen.rs b/sp1-gpu/crates/logup_gkr/src/tracegen.rs index 5b7e7bf246..0a02a03c30 100644 --- a/sp1-gpu/crates/logup_gkr/src/tracegen.rs +++ b/sp1-gpu/crates/logup_gkr/src/tracegen.rs @@ -9,8 +9,7 @@ use slop_alloc::{Buffer, HasBackend}; use slop_multilinear::Point; use slop_tensor::Tensor; use sp1_gpu_cudart::{ - args, sys::v2_kernels::logup_gkr_populate_last_circuit_layer, DeviceBuffer, DevicePoint, - TaskScope, + args, sys::kernels::logup_gkr_populate_last_circuit_layer, DeviceBuffer, DevicePoint, TaskScope, }; use sp1_hypercube::{air::MachineAir, Chip}; use tracing::instrument; diff --git a/sp1-gpu/crates/prover_components/Cargo.toml b/sp1-gpu/crates/prover_components/Cargo.toml index 54bb71d366..748410e7a4 100644 --- a/sp1-gpu/crates/prover_components/Cargo.toml +++ b/sp1-gpu/crates/prover_components/Cargo.toml @@ -21,6 +21,7 @@ sp1-gpu-merkle-tree = { workspace = true } sp1-gpu-jagged-tracegen = { workspace = true } # slop +slop-air = { workspace = true } slop-koala-bear = { workspace = true } slop-algebra = { workspace = true } slop-bn254 = { workspace = true } diff --git a/sp1-gpu/crates/prover_components/src/components.rs b/sp1-gpu/crates/prover_components/src/components.rs index 3528c3043e..a4b84fa4f8 100644 --- a/sp1-gpu/crates/prover_components/src/components.rs +++ b/sp1-gpu/crates/prover_components/src/components.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, sync::Arc}; -use sp1_gpu_air::{air_block::BlockAir, codegen_cuda_eval, SymbolicProverFolder}; +use sp1_gpu_air::ir::DagBuilder; use sp1_gpu_challenger::{DuplexChallenger, MultiField32Challenger}; use sp1_gpu_cudart::{PinnedBuffer, TaskScope}; @@ -68,18 +68,12 @@ where PC: CudaShardProverComponents, PC::P: CudaTcsProver, PC::Air: CudaTracegenAir - + for<'a> BlockAir> + + for<'a> slop_air::Air> + ZerocheckAir + std::fmt::Debug, { let machine = verifier.machine().clone(); - let mut cache = BTreeMap::new(); - for chip in machine.chips() { - let result = codegen_cuda_eval(chip.air.as_ref()); - cache.insert(chip.air.name().to_string(), result); - } - let log_stacking_height = verifier.log_stacking_height(); let max_log_row_count = verifier.max_log_row_count(); @@ -116,7 +110,6 @@ where max_trace_size, scope, all_interactions, - cache, recompute_first_layer, recompute_first_layer, ) diff --git a/sp1-gpu/crates/shard_prover/Cargo.toml b/sp1-gpu/crates/shard_prover/Cargo.toml index 8d696170d6..f834330236 100644 --- a/sp1-gpu/crates/shard_prover/Cargo.toml +++ b/sp1-gpu/crates/shard_prover/Cargo.toml @@ -25,6 +25,7 @@ sp1-gpu-jagged-tracegen = { workspace = true } sp1-gpu-logup-gkr = { workspace = true } # slop +slop-air = { workspace = true } slop-alloc = { workspace = true } slop-algebra = { workspace = true } slop-multilinear = { workspace = true } @@ -43,6 +44,7 @@ sp1-hypercube = { workspace = true } thiserror = "1.0" tracing = { workspace = true } tokio = { workspace = true } +serde_json = "1.0" [features] default = [] diff --git a/sp1-gpu/crates/shard_prover/benches/prove_trusted_evaluations.rs b/sp1-gpu/crates/shard_prover/benches/prove_trusted_evaluations.rs index 01cbc1f670..34736c4c73 100644 --- a/sp1-gpu/crates/shard_prover/benches/prove_trusted_evaluations.rs +++ b/sp1-gpu/crates/shard_prover/benches/prove_trusted_evaluations.rs @@ -18,7 +18,6 @@ use slop_commit::Rounds; use slop_futures::queue::WorkerQueue; use slop_multilinear::{MleEval, MultilinearPcsChallenger}; use sp1_core_machine::riscv::RiscvAir; -use sp1_gpu_air::codegen_cuda_eval; use sp1_gpu_basefold::FriCudaProver; use sp1_gpu_commit::commit_multilinears; use sp1_gpu_cudart::{DeviceTensor, PinnedBuffer, TaskScope}; @@ -88,12 +87,6 @@ fn run_prove_trusted_evaluations( all_interactions.insert(chip.name().to_string(), Arc::new(device_interactions)); } - let mut all_zerocheck_programs = BTreeMap::new(); - for chip in machine.chips().iter() { - let result = codegen_cuda_eval(chip.air.as_ref()); - all_zerocheck_programs.insert(chip.name().to_string(), result); - } - let trace_buffers = Arc::new(WorkerQueue::new(vec![PinnedBuffer::::with_capacity( CORE_MAX_TRACE_SIZE as usize, )])); @@ -106,7 +99,6 @@ fn run_prove_trusted_evaluations( CORE_MAX_TRACE_SIZE as usize, scope.clone(), all_interactions, - all_zerocheck_programs, false, false, ); diff --git a/sp1-gpu/crates/shard_prover/src/prover.rs b/sp1-gpu/crates/shard_prover/src/prover.rs index 5b7b428506..7af5fc699e 100644 --- a/sp1-gpu/crates/shard_prover/src/prover.rs +++ b/sp1-gpu/crates/shard_prover/src/prover.rs @@ -9,8 +9,7 @@ use slop_jagged::{ JaggedProverError, PrefixSumsMaxLogRowCount, }; use slop_multilinear::{MleEval, MultilinearPcsVerifier, Point}; -use sp1_gpu_air::air_block::BlockAir; -use sp1_gpu_air::SymbolicProverFolder; +use sp1_gpu_air::ir::{ChunkBudget, DagBuilder}; use sp1_gpu_basefold::{CudaStackedPcsProverData, DeviceGrindingChallenger, FriCudaProver}; use sp1_gpu_challenger::FromHostChallengerSync; use sp1_gpu_cudart::PinnedBuffer; @@ -22,8 +21,7 @@ use sp1_gpu_logup_gkr::{prove_logup_gkr, CudaLogUpGkrOptions, Interactions}; use sp1_gpu_merkle_tree::{CudaTcsProver, SingleLayerMerkleTreeProverError}; use sp1_gpu_tracegen::CudaTracegenAir; use sp1_gpu_utils::{Ext, Felt, JaggedTraceMle}; -use sp1_gpu_zerocheck::zerocheck; -use sp1_gpu_zerocheck::CudaEvalResult; +use sp1_gpu_zerocheck::prover::{upload_machine_bytecode, zerocheck, MachineBytecode}; use sp1_hypercube::prover::ZerocheckAir; use sp1_hypercube::{ air::{MachineAir, MachineProgram}, @@ -31,7 +29,7 @@ use sp1_hypercube::{ Machine, MachineVerifyingKey, ShardProof, }; use sp1_hypercube::{SP1PcsProof, ShardContextImpl}; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::iter::once; use std::vec; use std::{marker::PhantomData, sync::Arc}; @@ -43,7 +41,7 @@ pub trait CudaShardProverComponents: Send + Sync + 'static { type P: CudaTcsProver; type Air: CudaTracegenAir + ZerocheckAir - + for<'a> BlockAir>; + + for<'a> slop_air::Air>; type C: MultilinearPcsVerifier + Send + Sync; /// The device challenger type used for GPU-based challenger operations. type DeviceChallenger: sp1_gpu_jagged_assist::AsMutRawChallenger @@ -64,7 +62,7 @@ impl> Clone for CudaShardProver> CudaShardProver { +impl, PC: CudaShardProverComponents> CudaShardProver { #[allow(clippy::too_many_arguments)] pub fn new( trace_buffers: Arc>>, @@ -74,10 +72,15 @@ impl> CudaShardProver { max_trace_size: usize, backend: TaskScope, all_interactions: BTreeMap>>, - all_zerocheck_programs: BTreeMap, recompute_first_layer: bool, drop_ldes: bool, ) -> Self { + // Compile + upload the whole machine's zerocheck bytecode once. + // It's machine-stable, so every shard reuses this single upload. + let machine_bytecode = { + let chip_set: BTreeSet<_> = machine.chips().iter().cloned().collect(); + Arc::new(upload_machine_bytecode(&chip_set, ChunkBudget::recommended(), &backend)) + }; Self { inner: Arc::new(CudaShardProverInner { trace_buffers, @@ -87,7 +90,7 @@ impl> CudaShardProver { max_trace_size, backend, all_interactions, - all_zerocheck_programs, + machine_bytecode, recompute_first_layer, drop_ldes, _marker: PhantomData, @@ -139,7 +142,11 @@ pub(crate) struct CudaShardProverInner>>, - pub all_zerocheck_programs: BTreeMap, + /// The whole machine's DAG-native bytecode, compiled and uploaded to + /// the GPU once at prover construction. The bytecode is machine-stable + /// (cluster-independent — see `compile_chips`), so every shard reuses + /// this single upload instead of re-compiling + re-uploading per shard. + pub machine_bytecode: Arc, pub recompute_first_layer: bool, pub drop_ldes: bool, pub _marker: PhantomData, @@ -667,12 +674,15 @@ impl, PC: CudaShardProverComponents> // Get the challenge for batching the evaluations from the GKR proof. let gkr_opening_batch_challenge = challenger.sample_ext_element::(); - // Generate the zerocheck proof. + // Generate the zerocheck proof via the DAG-native fused-kernel path. + // The bytecode is compiled + uploaded once per machine at prover + // construction (see `machine_bytecode`); `zerocheck` just selects + // this shard's chips from it. let (shard_open_values, zerocheck_partial_sumcheck_proof) = tracing::debug_span!("zerocheck").in_scope(|| { zerocheck( shard_chips, - &self.all_zerocheck_programs, + &self.machine_bytecode, traces, batching_challenge, gkr_opening_batch_challenge, @@ -763,7 +773,6 @@ mod tests { use slop_tensor::Tensor; use sp1_core_machine::io::SP1Stdin; use sp1_core_machine::riscv::RiscvAir; - use sp1_gpu_air::codegen_cuda_eval; use sp1_gpu_cudart::run_in_place; use sp1_gpu_jagged_tracegen::test_utils::tracegen_setup::{ self, CORE_MAX_LOG_ROW_COUNT, LOG_STACKING_HEIGHT, @@ -827,12 +836,6 @@ mod tests { all_interactions.insert(chip.name().to_string(), Arc::new(device_interactions)); } - let mut cache = BTreeMap::new(); - for chip in machine.chips().iter() { - let result = codegen_cuda_eval(chip.air.as_ref()); - cache.insert(chip.name().to_string(), result); - } - let num_workers = 1; let mut trace_buffers = Vec::with_capacity(num_workers); for _ in 0..num_workers { @@ -840,11 +843,15 @@ mod tests { trace_buffers.push(buffer); } + let machine_bytecode = { + let chip_set: BTreeSet<_> = machine.chips().iter().cloned().collect(); + Arc::new(upload_machine_bytecode(&chip_set, ChunkBudget::recommended(), &scope)) + }; let shard_prover_inner: CudaShardProverInner = CudaShardProverInner { trace_buffers: Arc::new(WorkerQueue::new(trace_buffers)), all_interactions, - all_zerocheck_programs: cache, + machine_bytecode, max_log_row_count: CORE_MAX_LOG_ROW_COUNT, basefold_prover, max_trace_size: CORE_MAX_TRACE_SIZE as usize, diff --git a/sp1-gpu/crates/sys/include/sum_and_reduce/reduce.cuh b/sp1-gpu/crates/sys/include/sum_and_reduce/reduce.cuh index 66bfdb04a0..d3f627fb42 100644 --- a/sp1-gpu/crates/sys/include/sum_and_reduce/reduce.cuh +++ b/sp1-gpu/crates/sys/include/sum_and_reduce/reduce.cuh @@ -18,8 +18,32 @@ partialBlockReduce(const TyBlock& block, const TyTile& tile, F val, F* shared) { // Synchronize after warp-level reduction block.sync(); + // Tree-based reduction over the N = block.size() / tile.size() warp partials. + // + // The simple `for (stride = N/2; stride > 0; stride /= 2)` loop is only + // correct when N is a power of 2; for e.g. N = 3 it folds slot 1 into 0 and + // then exits with `shared[2]` never added in (silent undercount). + // + // Fix (smallest blast radius — does not change the caller-visible shmem + // contract of N slots): on the first pass fold the "tail" elements + // `[pow2, N)` into `[0, N - pow2)`, leaving `pow2` valid partials, then run + // the regular power-of-2 tree reduction. Since pow2 is the largest + // power-of-2 <= N we have N - pow2 < pow2, so the writes don't overlap. + const int n = block.size() / tile.size(); + int pow2 = 1; + while ((pow2 << 1) <= n) { + pow2 <<= 1; + } + if (pow2 < n) { + const int tail = n - pow2; + if (block.thread_rank() < tail) { + shared[block.thread_rank()] += shared[block.thread_rank() + pow2]; + } + block.sync(); + } + // Perform tree-based reduction on shared memory - for (int stride = (block.size() / tile.size()) / 2; stride > 0; stride /= 2) { + for (int stride = pow2 / 2; stride > 0; stride /= 2) { if (block.thread_rank() < stride) { shared[block.thread_rank()] += shared[block.thread_rank() + stride]; } @@ -32,4 +56,9 @@ partialBlockReduce(const TyBlock& block, const TyTile& tile, F val, F* shared) { // A reduction kernel for Felt extern "C" void* reduce_kernel_felt(); // A reduction kernel for Ext -extern "C" void* reduce_kernel_ext(); \ No newline at end of file +extern "C" void* reduce_kernel_ext(); + +// Test-only kernel exercising `partialBlockReduce` with one block. Each thread +// contributes `input[threadIdx.x]` and the per-block sum is written to +// `output[0]`. Used to verify correctness for non-power-of-2 warp counts. +extern "C" void* partial_block_reduce_test_kernel_felt(); \ No newline at end of file diff --git a/sp1-gpu/crates/sys/include/sum_and_reduce/reduction.cuh b/sp1-gpu/crates/sys/include/sum_and_reduce/reduction.cuh index de6bac5f1b..b59dcca9c1 100644 --- a/sp1-gpu/crates/sys/include/sum_and_reduce/reduction.cuh +++ b/sp1-gpu/crates/sys/include/sum_and_reduce/reduction.cuh @@ -51,8 +51,24 @@ partialBlockReduce(const TyBlock& block, const TyTile& tile, F val, F* shared, T } block.sync(); // Synchronize after warp-level reduction + // See `reduce.cuh::partialBlockReduce` for a discussion of the + // non-power-of-2 fold-tail trick used here. Shared-memory contract is + // unchanged: caller still allocates N = warps-per-block slots. + const int n = block.size() / tile.size(); + int pow2 = 1; + while ((pow2 << 1) <= n) { + pow2 <<= 1; + } + if (pow2 < n) { + const int tail = n - pow2; + if (block.thread_rank() < tail) { + op.evalAssign(shared[block.thread_rank()], shared[block.thread_rank() + pow2]); + } + block.sync(); + } + // Perform tree-based reduction on shared memory - for (int stride = (block.size() / tile.size()) / 2; stride > 0; stride /= 2) { + for (int stride = pow2 / 2; stride > 0; stride /= 2) { if (block.thread_rank() < stride) { op.evalAssign(shared[block.thread_rank()], shared[block.thread_rank() + stride]); } diff --git a/sp1-gpu/crates/sys/include/zerocheck/column_tile.cuh b/sp1-gpu/crates/sys/include/zerocheck/column_tile.cuh new file mode 100644 index 0000000000..6b97291349 --- /dev/null +++ b/sp1-gpu/crates/sys/include/zerocheck/column_tile.cuh @@ -0,0 +1,26 @@ +// ColumnTile lowering for v2 zerocheck. +// +// Handles `Σ_k α^k · (Σ_i coeff_{k,i} · leaf_{k,i})` chunks. One lane per +// `(term, row, eval)` tuple. No shared-mem cache — lane variation IS the +// program; each lane reads its own `(zero, one)` directly from global. + +#pragma once + +#include "config.cuh" +#include "zerocheck/sequential.cuh" // re-uses LeafRef +#include + +// Must match `ColumnTermEntry` in column_tile_bytecode.rs. +struct ColumnTermEntry { + uint32_t leaf_idx; + uint32_t coeff_kind; // 0 = const, 1 = public + uint32_t coeff_idx; + uint32_t alpha_idx; +}; + +constexpr uint32_t COEFF_KIND_CONST = 0; +constexpr uint32_t COEFF_KIND_PUBLIC = 1; +constexpr uint32_t COEFF_KIND_RUNTIME = 2; + +extern "C" void* zerocheck_column_tile_kb_kernel(); +extern "C" void* zerocheck_column_tile_ext_kernel(); diff --git a/sp1-gpu/crates/sys/include/zerocheck/sequential.cuh b/sp1-gpu/crates/sys/include/zerocheck/sequential.cuh new file mode 100644 index 0000000000..062785e9ad --- /dev/null +++ b/sp1-gpu/crates/sys/include/zerocheck/sequential.cuh @@ -0,0 +1,114 @@ +// Sequential lowering for v2 zerocheck. +// +// Interprets a chunk's bytecode per row, computing 3 eval-point partial sums +// per CTA. Matches the host-side `ChunkBytecode` layout from +// sp1-gpu-air/src/v2/bytecode.rs. +// +// Phase 2 scope: base-field DAG nodes only. Output is unweighted (no +// partial_lagrange / lambda multiplication). The host weighting + final +// reduction happens in a follow-up phase. + +#pragma once + +#include "config.cuh" +#include + +// Must match `DagInstr` in v2/bytecode.rs. +struct DagInstr { + uint8_t opcode; + uint8_t _pad; + uint16_t out; + uint16_t a; + uint16_t b; +}; + +// Must match `LeafRef` in v2/bytecode.rs. +struct LeafRef { + uint8_t source; // 2=PrepLocal, 3=PrepNext, 4=MainLocal, 5=MainNext + uint8_t _pad; + uint32_t col; +}; + +// Opcodes — must match `BcOp` in v2/bytecode.rs. +enum BcOp : uint8_t { + BC_LOAD_LEAF = 0, + BC_LOAD_CONST = 1, + BC_LOAD_PUBLIC = 2, + BC_ADD_F = 3, + BC_SUB_F = 4, + BC_MUL_F = 5, + BC_NEG_F = 6, + BC_ASSERT_F = 7, +}; + +// Entry points used by Rust via `KernelPtr`. Returns the device function +// pointer; the actual kernels are `zerocheck_sequential`. +// - kb_kernel: K = felt_t (base-field trace, round 0 of sumcheck) +// - ext_kernel: K = ext_t (extension-field trace, rounds 1+) +extern "C" void* zerocheck_sequential_kb_kernel(); +extern "C" void* zerocheck_sequential_ext_kernel(); + +// Tiered variants by per-chunk MAX_REGS. The chunk's `max_reg` selects the +// smallest tier whose MAX_REGS >= max_reg. Tight sizing matters for perf: +// `K regs[MAX_REGS][3]` is a per-thread stack array that spills to local +// memory, so an over-sized MAX_REGS pays a real load/store cost per row. +extern "C" void* zerocheck_sequential_kb_32_kernel(); +extern "C" void* zerocheck_sequential_kb_64_kernel(); +extern "C" void* zerocheck_sequential_kb_128_kernel(); +extern "C" void* zerocheck_sequential_kb_256_kernel(); +extern "C" void* zerocheck_sequential_ext_32_kernel(); +extern "C" void* zerocheck_sequential_ext_64_kernel(); +extern "C" void* zerocheck_sequential_ext_128_kernel(); +extern "C" void* zerocheck_sequential_ext_256_kernel(); + +// Per-chunk metadata for the fused dispatch kernel. One ChunkMeta entry per +// Sequential chunk that the round wants to evaluate. Must match +// `ChunkMetaC` in v2.rs (layout-compat). +// +// All per-chunk buffer pointers are device pointers into separate +// (per-chunk) buffers — we don't concatenate at this stage, because each +// chunk's bytecode/leaves are uploaded once and reused across rounds. +struct ChunkMeta { + const DagInstr* instrs; // 8 + const LeafRef* leaves; // 8 + const void* consts; // 8 — cast to felt_t* in kernel + const uint32_t* publics; // 8 + const uint16_t* assert_regs; // 8 + const uint32_t* assert_alphas; // 8 + uint64_t preprocessed_ptr; // 8 + uint64_t main_ptr; // 8 + uint32_t n_instrs; // 4 + uint32_t n_asserts; // 4 + uint32_t chip_idx; // 4 + uint32_t gkr_main_width; // 4 + uint32_t gkr_prep_width; // 4 + uint32_t height; // 4 + uint32_t row_count; // 4 + uint32_t chip_alpha_offset; // 4 — added to chip-relative alpha idx + uint32_t geq_threshold; // 4 — applied iff gkr_main_width != 0 + ext_t geq_eq_coefficient; // 16 + ext_t padded_row_adjustment; // 16 +}; + +// Fused dispatch kernel. One launch handles every Sequential chunk across +// every chip. Each thread binary-searches `row_starts` to find which chunk +// its `idx` belongs to, then runs that chunk's bytecode. +// +// Tiered variants by MAX_REGS — the launcher partitions chunks into tiers +// and launches one kernel per non-empty tier so each kernel's per-thread +// register array is sized to its tier's worst case (not the entire +// workload's worst case). +extern "C" void* zerocheck_fused_sequential_kb_kernel(); +extern "C" void* zerocheck_fused_sequential_ext_kernel(); +extern "C" void* zerocheck_fused_sequential_kb_32_kernel(); +extern "C" void* zerocheck_fused_sequential_kb_64_kernel(); +extern "C" void* zerocheck_fused_sequential_kb_128_kernel(); +extern "C" void* zerocheck_fused_sequential_kb_256_kernel(); +extern "C" void* zerocheck_fused_sequential_kb_512_kernel(); +extern "C" void* zerocheck_fused_sequential_kb_1024_kernel(); +extern "C" void* zerocheck_fused_sequential_ext_32_kernel(); +extern "C" void* zerocheck_fused_sequential_ext_64_kernel(); +extern "C" void* zerocheck_fused_sequential_ext_128_kernel(); +extern "C" void* zerocheck_fused_sequential_ext_256_kernel(); +extern "C" void* zerocheck_fused_sequential_ext_512_kernel(); +extern "C" void* zerocheck_fused_sequential_ext_1024_kernel(); diff --git a/sp1-gpu/crates/sys/include/zerocheck/zerocheck.cuh b/sp1-gpu/crates/sys/include/zerocheck/zerocheck.cuh deleted file mode 100644 index b71dd8cb89..0000000000 --- a/sp1-gpu/crates/sys/include/zerocheck/zerocheck.cuh +++ /dev/null @@ -1,6 +0,0 @@ -#pragma once - -extern "C" void* zerocheck_sum_as_poly_koala_bear_base_ext_kernel(); -extern "C" void* zerocheck_sum_as_poly_koala_bear_ext_ext_kernel(); -extern "C" void* zerocheck_fix_last_variable_and_sum_as_poly_base_ext_kernel(); -extern "C" void* zerocheck_fix_last_variable_and_sum_as_poly_ext_ext_kernel(); diff --git a/sp1-gpu/crates/sys/include/zerocheck/zerocheck_eval.cuh b/sp1-gpu/crates/sys/include/zerocheck/zerocheck_eval.cuh deleted file mode 100644 index f11988a2c0..0000000000 --- a/sp1-gpu/crates/sys/include/zerocheck/zerocheck_eval.cuh +++ /dev/null @@ -1,127 +0,0 @@ -#pragma once - -#include "config.cuh" -#include - -struct Instruction { - unsigned char opcode; - unsigned char b_variant; - unsigned char c_variant; - unsigned short a; - unsigned short b; - unsigned short c; -}; - -template -struct JaggedConstraintFolder { - public: - const K* data; - size_t preprocessed_ptr; - size_t main_ptr; - uint32_t height; - const felt_t* publicValues; - const ext_t* powersOfAlpha; - uint32_t constraintIndex; - ext_t accumulator; - uint32_t rowIdx; - unsigned char eval_point; - - public: - __device__ JaggedConstraintFolder() {} - - __inline__ __device__ K var_f(unsigned char variant, unsigned int idx) { - switch (variant) { - case 0: - return K::zero(); - case 1: - return K(idx); - case 2: - K zeroPrepVal = K::load(data, preprocessed_ptr + idx * height + (rowIdx << 1)); - K onePrepVal = K::load(data, preprocessed_ptr + idx * height + (rowIdx << 1 | 1)); - K result = zeroPrepVal; - K multi_diff; - switch (eval_point) { - case 0: - break; - case 2: - multi_diff = onePrepVal - zeroPrepVal; - multi_diff = multi_diff + multi_diff; - result += multi_diff; - break; - case 4: - multi_diff = onePrepVal - zeroPrepVal; - multi_diff = multi_diff + multi_diff; - multi_diff = multi_diff + multi_diff; - result += multi_diff; - break; - default: - assert(0); - break; - } - return result; - case 4: - K zeroMainVal = K::load(data, main_ptr + idx * height + (rowIdx << 1)); - K oneMainVal = K::load(data, main_ptr + idx * height + (rowIdx << 1 | 1)); - result = zeroMainVal; - switch (eval_point) { - case 0: - break; - case 2: - multi_diff = oneMainVal - zeroMainVal; - multi_diff = multi_diff + multi_diff; - result += multi_diff; - break; - case 4: - multi_diff = oneMainVal - zeroMainVal; - multi_diff = multi_diff + multi_diff; - multi_diff = multi_diff + multi_diff; - result += multi_diff; - break; - default: - assert(0); - break; - } - return result; - case 9: - return K(felt_t::load(publicValues, idx)); - default: - // Case 3: next row for for preprocessed trace for univariate. - // Case 5: next row for for main trace for univariate. - // Case 6: isFirstRow for univariate. - // Case 7: isLastRow for univariate. - // Case 8: isTransition for univariate. - // Case 10: globalCumulativeSum for univariate. - assert(0); - return K::zero(); - } - } - - __inline__ __device__ ext_t var_ef(unsigned char variant, unsigned int idx) { - switch (variant) { - case 0: - return ext_t::zero(); - default: - // Case 1: Permutation trace row for univariate. - // Case 2: Permutation trace next row for multivariate. - // Case 3: Permutation challenge for univariate. - // Case 4: Local cumulative sum for univariate. - assert(0); - return ext_t::zero(); - } - } -}; - - -extern "C" void* jagged_constraint_poly_eval_32_koala_bear_kernel(); -extern "C" void* jagged_constraint_poly_eval_64_koala_bear_kernel(); -extern "C" void* jagged_constraint_poly_eval_128_koala_bear_kernel(); -extern "C" void* jagged_constraint_poly_eval_256_koala_bear_kernel(); -extern "C" void* jagged_constraint_poly_eval_512_koala_bear_kernel(); -extern "C" void* jagged_constraint_poly_eval_1024_koala_bear_kernel(); - -extern "C" void* jagged_constraint_poly_eval_32_koala_bear_extension_kernel(); -extern "C" void* jagged_constraint_poly_eval_64_koala_bear_extension_kernel(); -extern "C" void* jagged_constraint_poly_eval_128_koala_bear_extension_kernel(); -extern "C" void* jagged_constraint_poly_eval_256_koala_bear_extension_kernel(); -extern "C" void* jagged_constraint_poly_eval_512_koala_bear_extension_kernel(); -extern "C" void* jagged_constraint_poly_eval_1024_koala_bear_extension_kernel(); diff --git a/sp1-gpu/crates/sys/lib/sum_and_reduce/reduce.cu b/sp1-gpu/crates/sys/lib/sum_and_reduce/reduce.cu index 47a5a754ef..e4943fcd51 100644 --- a/sp1-gpu/crates/sys/lib/sum_and_reduce/reduce.cu +++ b/sp1-gpu/crates/sys/lib/sum_and_reduce/reduce.cu @@ -48,4 +48,29 @@ __global__ void reduceKernel(F* input, F* output, size_t width, size_t height) { extern "C" void* reduce_kernel_felt() { return reinterpret_cast(reduceKernel); } -extern "C" void* reduce_kernel_ext() { return reinterpret_cast(reduceKernel); } \ No newline at end of file +extern "C" void* reduce_kernel_ext() { return reinterpret_cast(reduceKernel); } + +// Test-only kernel: launched as a single block with `len` threads (where +// `len <= blockDim.x`). Each thread loads `input[threadIdx.x]` (or zero if out +// of range) and `partialBlockReduce` reduces them. Result goes to `output[0]`. +// +// Used by `examples/partial_block_reduce_test.rs` to exercise non-power-of-2 +// warp counts (e.g. 96 threads = 3 warps, 160 = 5, 224 = 7, 288 = 9). +__global__ void partialBlockReduceTestKernel(const felt_t* input, felt_t* output, uint32_t len) { + auto block = cg::this_thread_block(); + auto tile = cg::tiled_partition<32>(block); + + extern __shared__ unsigned char memory[]; + felt_t* shared = reinterpret_cast(memory); + + felt_t val = (threadIdx.x < len) ? input[threadIdx.x] : felt_t::zero(); + felt_t block_sum = partialBlockReduce(block, tile, val, shared); + + if (threadIdx.x == 0) { + output[0] = block_sum; + } +} + +extern "C" void* partial_block_reduce_test_kernel_felt() { + return reinterpret_cast(partialBlockReduceTestKernel); +} \ No newline at end of file diff --git a/sp1-gpu/crates/sys/lib/zerocheck/CMakeLists.txt b/sp1-gpu/crates/sys/lib/zerocheck/CMakeLists.txt index 01fc2c9cbe..2e3ea67f07 100644 --- a/sp1-gpu/crates/sys/lib/zerocheck/CMakeLists.txt +++ b/sp1-gpu/crates/sys/lib/zerocheck/CMakeLists.txt @@ -1,6 +1,6 @@ add_library(zerocheck_objs OBJECT jagged_mle.cu - zerocheck.cu - zerocheck_eval.cu + sequential.cu + column_tile.cu ) target_link_libraries(zerocheck_objs PRIVATE sp1_gpu_common) diff --git a/sp1-gpu/crates/sys/lib/zerocheck/column_tile.cu b/sp1-gpu/crates/sys/lib/zerocheck/column_tile.cu new file mode 100644 index 0000000000..cd970e16c1 --- /dev/null +++ b/sp1-gpu/crates/sys/lib/zerocheck/column_tile.cu @@ -0,0 +1,133 @@ +// v2 ColumnTile lowering. +// +// Each thread = one `(term, row)` pair. Per thread: +// 1. Look up term's coefficient (constant, public, or runtime ext_t). +// 2. Read `(zero, one)` for the term's leaf at this row (K-typed). +// 3. Compute interp at all 3 eval points (eval-point caching). +// 4. Multiply by `α^{term.alpha_idx} · coeff · eq[row] · λ_chip`. +// 5. Block-reduce per eval point into one partial per CTA per eval. +// +// Templated on `K` ∈ {felt_t, ext_t} for the trace element type. Constants +// and publics stay base-field; runtime coeffs and the per-row weighting +// stay extension-field; the accumulator is always ext_t. + +#include "zerocheck/column_tile.cuh" +#include "zerocheck/sequential.cuh" +#include "config.cuh" +#include "sum_and_reduce/reduce.cuh" + +#include +#include + +namespace cg = cooperative_groups; + +namespace { + +__device__ __forceinline__ felt_t eval_point(int i) { + return felt_t::from_canonical_u32(2u * i); +} + +template +__global__ void zerocheck_column_tile( + const ColumnTermEntry* __restrict__ terms, + uint32_t n_terms, + const LeafRef* __restrict__ leaves, + const felt_t* __restrict__ consts, + const uint32_t* __restrict__ publics, + const ext_t* __restrict__ runtime_coeffs, + const K* __restrict__ trace_data, + size_t preprocessed_ptr, + size_t main_ptr, + uint32_t height, + const felt_t* __restrict__ public_values, + const ext_t* __restrict__ powers_of_alpha, + const ext_t* __restrict__ partial_lagrange, + const ext_t* __restrict__ powers_of_lambda, + uint32_t chip_idx, + uint32_t rest_point_dim, + uint32_t row_start, + uint32_t row_count, + ext_t* __restrict__ partials +) { + const uint64_t total = (uint64_t)n_terms * (uint64_t)row_count; + const uint64_t stride = (uint64_t)blockDim.x * (uint64_t)gridDim.x; + const ext_t lambda = ext_t::load(powers_of_lambda, chip_idx); + const uint32_t row_limit = 1u << rest_point_dim; + + ext_t thread_acc[3] = { ext_t::zero(), ext_t::zero(), ext_t::zero() }; + + // Grid-stride loop. Each lane covers one (term, row) tuple per iter. + for (uint64_t lane = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + lane < total; + lane += stride) + { + const uint32_t term_idx = (uint32_t)(lane / row_count); + const uint32_t local_row = (uint32_t)(lane - (uint64_t)term_idx * row_count); + const uint32_t row = row_start + local_row; + + ColumnTermEntry t = terms[term_idx]; + + LeafRef leaf = leaves[t.leaf_idx]; + size_t base = (leaf.source == 4 || leaf.source == 5) ? main_ptr : preprocessed_ptr; + K z = K::load(trace_data, base + leaf.col * height + (row << 1)); + K o = K::load(trace_data, base + leaf.col * height + (row << 1 | 1)); + K diff = o - z; + + ext_t alpha = ext_t::load(powers_of_alpha, t.alpha_idx); + + K v0 = z; + K v1 = z + eval_point(1) * diff; + K v2 = z + eval_point(2) * diff; + + ext_t a0, a1, a2; + if (t.coeff_kind == COEFF_KIND_RUNTIME) { + ext_t coeff = ext_t::load(runtime_coeffs, t.coeff_idx); + a0 = alpha * (coeff * v0); + a1 = alpha * (coeff * v1); + a2 = alpha * (coeff * v2); + } else { + felt_t coeff; + if (t.coeff_kind == COEFF_KIND_CONST) { + coeff = consts[t.coeff_idx]; + } else { + uint32_t pv_idx = publics[t.coeff_idx]; + coeff = felt_t::load(public_values, pv_idx); + } + a0 = alpha * (coeff * v0); + a1 = alpha * (coeff * v1); + a2 = alpha * (coeff * v2); + } + + if (row < row_limit) { + ext_t eq = ext_t::load(partial_lagrange, row); + ext_t w = eq * lambda; + thread_acc[0] += a0 * w; + thread_acc[1] += a1 * w; + thread_acc[2] += a2 * w; + } + } + + extern __shared__ unsigned char smem[]; + ext_t* shared = reinterpret_cast(smem); + + auto block = cg::this_thread_block(); + auto tile_warp = cg::tiled_partition<32>(block); + + for (int e = 0; e < 3; e++) { + ext_t block_sum = partialBlockReduce(block, tile_warp, thread_acc[e], shared); + if (threadIdx.x == 0) { + ext_t::store(partials, blockIdx.x * 3 + e, block_sum); + } + __syncthreads(); + } +} + +} // namespace + +extern "C" void* zerocheck_column_tile_kb_kernel() { + return (void*)zerocheck_column_tile; +} + +extern "C" void* zerocheck_column_tile_ext_kernel() { + return (void*)zerocheck_column_tile; +} diff --git a/sp1-gpu/crates/sys/lib/zerocheck/sequential.cu b/sp1-gpu/crates/sys/lib/zerocheck/sequential.cu new file mode 100644 index 0000000000..1739b34515 --- /dev/null +++ b/sp1-gpu/crates/sys/lib/zerocheck/sequential.cu @@ -0,0 +1,478 @@ +// v2 Sequential lowering — DAG bytecode interpreter with eval-point caching. +// +// Each CTA processes a tile of rows. Each thread = one row. Per-thread state +// is a register file of `MAX_REGS` × 3 K values (one slot per eval point), +// so leaf (zero, one) loads happen ONCE per thread per leaf and feed all +// three eval points without re-loading. +// +// Templated on `K` ∈ {felt_t, ext_t}: +// - Round 0 of sumcheck uses K = felt_t (base-field trace). +// - Rounds 1+ use K = ext_t (the trace has been folded with extension-field +// challenges). +// Constants and public values are always base-field; runtime coeffs and the +// per-row weighting/accumulator are always ext_t. + +#include "zerocheck/sequential.cuh" +#include "config.cuh" +#include "sum_and_reduce/reduce.cuh" + +#include +#include + +namespace cg = cooperative_groups; + +namespace { + +__device__ __forceinline__ felt_t eval_point(int i) { + return felt_t::from_canonical_u32(2u * i); +} + +template +__global__ void zerocheck_sequential( + const DagInstr* __restrict__ instrs, + uint32_t n_instrs, + const LeafRef* __restrict__ leaves, + const felt_t* __restrict__ consts, + const uint32_t* __restrict__ publics, + const uint16_t* __restrict__ assert_regs, + const uint32_t* __restrict__ assert_alphas, + uint32_t n_asserts, + const K* __restrict__ trace_data, + size_t preprocessed_ptr, + size_t main_ptr, + uint32_t height, + const felt_t* __restrict__ public_values, + const ext_t* __restrict__ powers_of_alpha, + const ext_t* __restrict__ partial_lagrange, + const ext_t* __restrict__ powers_of_lambda, + uint32_t chip_idx, + uint32_t rest_point_dim, + uint32_t row_start, + uint32_t row_count, + // GKR fusion: if both widths are non-zero, the kernel appends a per-row + // column sweep `Σ_i gkr_powers[i] · col_i(row)` after the bytecode + + // asserts pass, sharing column loads with the constraint reads. + // gkr_powers[0..main_w-1] weight main cols, gkr_powers[main_w..main_w+prep_w-1] weight prep cols. + const ext_t* __restrict__ gkr_powers, + uint32_t gkr_main_width, + uint32_t gkr_prep_width, + ext_t* __restrict__ partials +) { + K regs[MAX_REGS][3]; + ext_t thread_acc[3] = { ext_t::zero(), ext_t::zero(), ext_t::zero() }; + + const uint32_t stride = blockDim.x * gridDim.x; + const ext_t lambda = ext_t::load(powers_of_lambda, chip_idx); + const uint32_t row_limit = 1u << rest_point_dim; + + // Grid-stride loop: amortize launch + reduce overhead across many rows. + for (uint32_t lane = blockIdx.x * blockDim.x + threadIdx.x; + lane < row_count; + lane += stride) + { + const uint32_t row_idx = row_start + lane; + ext_t acc[3] = { ext_t::zero(), ext_t::zero(), ext_t::zero() }; + + for (uint32_t i = 0; i < n_instrs; i++) { + DagInstr instr = instrs[i]; + switch (instr.opcode) { + case BC_LOAD_LEAF: { + LeafRef leaf = leaves[instr.a]; + size_t base = (leaf.source == 4 || leaf.source == 5) + ? main_ptr + : preprocessed_ptr; + K z = K::load(trace_data, base + leaf.col * height + (row_idx << 1)); + K o = K::load(trace_data, base + leaf.col * height + (row_idx << 1 | 1)); + K diff = o - z; + regs[instr.out][0] = z; + regs[instr.out][1] = z + eval_point(1) * diff; + regs[instr.out][2] = z + eval_point(2) * diff; + break; + } + case BC_LOAD_CONST: { + K c = K(consts[instr.a]); + regs[instr.out][0] = c; + regs[instr.out][1] = c; + regs[instr.out][2] = c; + break; + } + case BC_LOAD_PUBLIC: { + uint32_t pv_idx = publics[instr.a]; + K v = K(felt_t::load(public_values, pv_idx)); + regs[instr.out][0] = v; + regs[instr.out][1] = v; + regs[instr.out][2] = v; + break; + } + case BC_ADD_F: { + regs[instr.out][0] = regs[instr.a][0] + regs[instr.b][0]; + regs[instr.out][1] = regs[instr.a][1] + regs[instr.b][1]; + regs[instr.out][2] = regs[instr.a][2] + regs[instr.b][2]; + break; + } + case BC_SUB_F: { + regs[instr.out][0] = regs[instr.a][0] - regs[instr.b][0]; + regs[instr.out][1] = regs[instr.a][1] - regs[instr.b][1]; + regs[instr.out][2] = regs[instr.a][2] - regs[instr.b][2]; + break; + } + case BC_MUL_F: { + regs[instr.out][0] = regs[instr.a][0] * regs[instr.b][0]; + regs[instr.out][1] = regs[instr.a][1] * regs[instr.b][1]; + regs[instr.out][2] = regs[instr.a][2] * regs[instr.b][2]; + break; + } + case BC_NEG_F: { + regs[instr.out][0] = K::zero() - regs[instr.a][0]; + regs[instr.out][1] = K::zero() - regs[instr.a][1]; + regs[instr.out][2] = K::zero() - regs[instr.a][2]; + break; + } + default: + break; + } + } + + for (uint32_t i = 0; i < n_asserts; i++) { + uint16_t reg = assert_regs[i]; + uint32_t alpha_idx = assert_alphas[i]; + ext_t alpha = ext_t::load(powers_of_alpha, alpha_idx); + acc[0] += alpha * regs[reg][0]; + acc[1] += alpha * regs[reg][1]; + acc[2] += alpha * regs[reg][2]; + } + + // GKR sweep — appended in-thread to share row's column reads with + // the constraint pass. With main_w + prep_w columns and one mult-add + // per col per eval point, this is what the synthesized GKR ColumnTile + // chunk would otherwise do in a separate launch. + if (gkr_main_width != 0 || gkr_prep_width != 0) { + for (uint32_t i = 0; i < gkr_main_width; i++) { + K z = K::load(trace_data, main_ptr + i * height + (row_idx << 1)); + K o = K::load(trace_data, main_ptr + i * height + (row_idx << 1 | 1)); + K diff = o - z; + ext_t bp = ext_t::load(gkr_powers, i); + acc[0] += bp * z; + acc[1] += bp * (z + eval_point(1) * diff); + acc[2] += bp * (z + eval_point(2) * diff); + } + for (uint32_t i = 0; i < gkr_prep_width; i++) { + K z = K::load(trace_data, preprocessed_ptr + i * height + (row_idx << 1)); + K o = K::load(trace_data, preprocessed_ptr + i * height + (row_idx << 1 | 1)); + K diff = o - z; + ext_t bp = ext_t::load(gkr_powers, gkr_main_width + i); + acc[0] += bp * z; + acc[1] += bp * (z + eval_point(1) * diff); + acc[2] += bp * (z + eval_point(2) * diff); + } + } + + if (row_idx < row_limit) { + ext_t eq = ext_t::load(partial_lagrange, row_idx); + ext_t w = eq * lambda; + thread_acc[0] += acc[0] * w; + thread_acc[1] += acc[1] * w; + thread_acc[2] += acc[2] * w; + } + } + + extern __shared__ unsigned char smem[]; + ext_t* shared = reinterpret_cast(smem); + + auto block = cg::this_thread_block(); + auto tile_warp = cg::tiled_partition<32>(block); + + for (int e = 0; e < 3; e++) { + ext_t block_sum = partialBlockReduce(block, tile_warp, thread_acc[e], shared); + if (threadIdx.x == 0) { + ext_t::store(partials, blockIdx.x * 3 + e, block_sum); + } + __syncthreads(); + } +} + +// ============================================================================ +// Fused dispatch kernel: one launch handles every Sequential chunk in the +// round. Each thread binary-searches `row_starts` to find its chunk, then +// runs that chunk's bytecode. +// +// Output: one ext_t[3] per block (block-reduced across all rows the block +// touched, regardless of which chunks). The chip-specific weighting +// (eq * λ_chip) is applied INSIDE the per-row body, so the cross-chunk sum +// is unambiguous — chip totals are NOT separately tracked. +// ============================================================================ + +__device__ __forceinline__ uint32_t +upper_bound_u32_local(const uint32_t* arr, uint32_t n, uint32_t target) { + uint32_t lo = 0, hi = n; + while (lo < hi) { + uint32_t mid = (lo + hi) >> 1; + if (arr[mid] <= target) lo = mid + 1; + else hi = mid; + } + return lo; +} + +template +__global__ void zerocheck_fused_sequential( + const ChunkMeta* __restrict__ chunk_meta, + const uint32_t* __restrict__ row_starts, // n_chunks + 1 entries + uint32_t n_chunks, + uint32_t total_rows, + const K* __restrict__ trace_data, + const felt_t* __restrict__ public_values, + const ext_t* __restrict__ powers_of_alpha, + const ext_t* __restrict__ partial_lagrange, + const ext_t* __restrict__ powers_of_lambda, + const ext_t* __restrict__ gkr_powers, + uint32_t rest_point_dim, + ext_t* __restrict__ partials +) { + // The eval-point axis lives on grid_z (matches v1's design): each thread + // handles ONE eval point and the per-thread register file is + // single-dimensional. Stashing all 3 eval points in regs[MAX_REGS][3] + // tripled register pressure and forced heavy local-memory spilling + // (profiling showed 5.9M local loads + 3.3M stores per launch). With + // grid_z=3 the same total work runs across 3x more blocks, each block + // has 1/3 the per-thread state, and the GPU schedules them concurrently. + K regs[MAX_REGS]; + ext_t thread_acc = ext_t::zero(); + + const uint32_t stride = blockDim.x * gridDim.x; + const uint32_t row_limit = 1u << rest_point_dim; + const int e = blockIdx.z; + const felt_t ep_v = (e == 0) ? felt_t::zero() + : felt_t::from_canonical_u32(2u * (uint32_t)e); + + for (uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; + idx < total_rows; + idx += stride) + { + // Find chunk index for this idx. + uint32_t chunk_idx = upper_bound_u32_local(row_starts, n_chunks + 1, idx) - 1; + uint32_t row_idx = idx - row_starts[chunk_idx]; + ChunkMeta cm = chunk_meta[chunk_idx]; + const felt_t* consts = reinterpret_cast(cm.consts); + + ext_t acc = ext_t::zero(); + + for (uint32_t i = 0; i < cm.n_instrs; i++) { + DagInstr instr = cm.instrs[i]; + switch (instr.opcode) { + case BC_LOAD_LEAF: { + LeafRef leaf = cm.leaves[instr.a]; + size_t base = (leaf.source == 4 || leaf.source == 5) + ? cm.main_ptr + : cm.preprocessed_ptr; + K z = K::load(trace_data, base + leaf.col * cm.height + (row_idx << 1)); + if (e == 0) { + regs[instr.out] = z; + } else { + K o = K::load(trace_data, + base + leaf.col * cm.height + (row_idx << 1 | 1)); + regs[instr.out] = z + ep_v * (o - z); + } + break; + } + case BC_LOAD_CONST: { + regs[instr.out] = K(consts[instr.a]); + break; + } + case BC_LOAD_PUBLIC: { + uint32_t pv_idx = cm.publics[instr.a]; + regs[instr.out] = K(felt_t::load(public_values, pv_idx)); + break; + } + case BC_ADD_F: { + regs[instr.out] = regs[instr.a] + regs[instr.b]; + break; + } + case BC_SUB_F: { + regs[instr.out] = regs[instr.a] - regs[instr.b]; + break; + } + case BC_MUL_F: { + regs[instr.out] = regs[instr.a] * regs[instr.b]; + break; + } + case BC_NEG_F: { + regs[instr.out] = K::zero() - regs[instr.a]; + break; + } + default: + break; + } + } + + for (uint32_t i = 0; i < cm.n_asserts; i++) { + uint16_t reg = cm.assert_regs[i]; + // Bytecode stores chip-relative alpha indices; shift into the + // cluster's powers_of_alpha table here. + uint32_t alpha_idx = cm.chip_alpha_offset + cm.assert_alphas[i]; + ext_t alpha = ext_t::load(powers_of_alpha, alpha_idx); + acc += alpha * regs[reg]; + } + + if (cm.gkr_main_width != 0 || cm.gkr_prep_width != 0) { + for (uint32_t i = 0; i < cm.gkr_main_width; i++) { + K z = K::load(trace_data, cm.main_ptr + i * cm.height + (row_idx << 1)); + K v; + if (e == 0) { + v = z; + } else { + K o = K::load(trace_data, cm.main_ptr + i * cm.height + (row_idx << 1 | 1)); + v = z + ep_v * (o - z); + } + ext_t bp = ext_t::load(gkr_powers, i); + acc += bp * v; + } + for (uint32_t i = 0; i < cm.gkr_prep_width; i++) { + K z = K::load(trace_data, + cm.preprocessed_ptr + i * cm.height + (row_idx << 1)); + K v; + if (e == 0) { + v = z; + } else { + K o = K::load(trace_data, + cm.preprocessed_ptr + i * cm.height + (row_idx << 1 | 1)); + v = z + ep_v * (o - z); + } + ext_t bp = ext_t::load(gkr_powers, cm.gkr_main_width + i); + acc += bp * v; + } + + // Geq correction — subtract `geq(row, eval_pt) * padded_row_adjustment` + // from acc. Carrier chunk (gkr_main_width > 0) is the natural place + // since each chip's geq fires exactly once. Moving this from a host + // loop into the kernel was a big win: the host loop iterated + // row_count rows × 9 ext ops per row per chip per round. + uint32_t z_idx = row_idx << 1; + uint32_t o_idx = (row_idx << 1) | 1; + ext_t geq_z, geq_o; + if (z_idx < cm.geq_threshold) { + geq_z = ext_t::zero(); + } else if (z_idx == cm.geq_threshold) { + geq_z = ext_t::one() + cm.geq_eq_coefficient; + } else { + geq_z = ext_t::one(); + } + if (o_idx < cm.geq_threshold) { + geq_o = ext_t::zero(); + } else if (o_idx == cm.geq_threshold) { + geq_o = ext_t::one() + cm.geq_eq_coefficient; + } else { + geq_o = ext_t::one(); + } + ext_t geq_v = (e == 0) ? geq_z : (geq_z + ep_v * (geq_o - geq_z)); + acc -= geq_v * cm.padded_row_adjustment; + } + + if (row_idx < row_limit) { + ext_t eq = ext_t::load(partial_lagrange, row_idx); + ext_t lambda = ext_t::load(powers_of_lambda, cm.chip_idx); + thread_acc += acc * (eq * lambda); + } + } + + extern __shared__ unsigned char smem[]; + ext_t* shared = reinterpret_cast(smem); + + auto block = cg::this_thread_block(); + auto tile_warp = cg::tiled_partition<32>(block); + + ext_t block_sum = partialBlockReduce(block, tile_warp, thread_acc, shared); + if (threadIdx.x == 0) { + // Output layout: blockIdx.z slot lives at (block.x * 3 + e). Matches + // the existing host aggregation that sums 3 ext per block. + ext_t::store(partials, blockIdx.x * 3 + (uint32_t)e, block_sum); + } +} + +} // namespace + +// Tiered register-file sizes. Each chunk picks the smallest tier whose +// MAX_REGS >= chunk.max_reg. Larger MAX_REGS means more local-memory +// spilling (compile-time-sized stack array) so picking tightly matters a lot. + +extern "C" void* zerocheck_sequential_kb_32_kernel() { + return (void*)zerocheck_sequential; +} +extern "C" void* zerocheck_sequential_kb_64_kernel() { + return (void*)zerocheck_sequential; +} +extern "C" void* zerocheck_sequential_kb_128_kernel() { + return (void*)zerocheck_sequential; +} +extern "C" void* zerocheck_sequential_kb_256_kernel() { + return (void*)zerocheck_sequential; +} + +extern "C" void* zerocheck_sequential_ext_32_kernel() { + return (void*)zerocheck_sequential; +} +extern "C" void* zerocheck_sequential_ext_64_kernel() { + return (void*)zerocheck_sequential; +} +extern "C" void* zerocheck_sequential_ext_128_kernel() { + return (void*)zerocheck_sequential; +} +extern "C" void* zerocheck_sequential_ext_256_kernel() { + return (void*)zerocheck_sequential; +} + +// Back-compat default (used by code that doesn't tier yet). +extern "C" void* zerocheck_sequential_kb_kernel() { + return (void*)zerocheck_sequential; +} +extern "C" void* zerocheck_sequential_ext_kernel() { + return (void*)zerocheck_sequential; +} + +// Fused dispatch entry points, tiered by MAX_REGS so the per-thread local +// memory footprint matches the largest chunk in the tier. The launcher +// partitions seq_meta into tiers and launches one fused kernel per non-empty +// tier (typically 1-2 launches per round on real workloads). +extern "C" void* zerocheck_fused_sequential_kb_32_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_kb_64_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_kb_128_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_kb_256_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_kb_512_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_kb_1024_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_ext_32_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_ext_64_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_ext_128_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_ext_256_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_ext_512_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_ext_1024_kernel() { + return (void*)zerocheck_fused_sequential; +} + +// Back-compat entry points (deprecated: use the tiered variants). +extern "C" void* zerocheck_fused_sequential_kb_kernel() { + return (void*)zerocheck_fused_sequential; +} +extern "C" void* zerocheck_fused_sequential_ext_kernel() { + return (void*)zerocheck_fused_sequential; +} diff --git a/sp1-gpu/crates/sys/lib/zerocheck/zerocheck.cu b/sp1-gpu/crates/sys/lib/zerocheck/zerocheck.cu deleted file mode 100644 index b4a9b647f0..0000000000 --- a/sp1-gpu/crates/sys/lib/zerocheck/zerocheck.cu +++ /dev/null @@ -1,194 +0,0 @@ -#include "config.cuh" -#include "sum_and_reduce/reduce.cuh" -#include "zerocheck/zerocheck.cuh" - -#include -#include - -namespace cg = cooperative_groups; - -// see crates/prover-clea/src/zerocheck.rs -template -__device__ inline EF zerocheckEval(EF a, F b) { - EF six = EF::two() + EF::two() + EF::two(); // six = 6 using F's two() method - return a * a * b + a * b * b + six * b * b * b; -} - -template -__global__ void zerocheckFixLastVariableAndSumAsPoly( - const F* base_input, - const EF* ext_input, - EF* __restrict base_output, - EF* __restrict ext_output, - EF alpha, - EF* univariate_result, - size_t numPolys, - size_t inputHeight) { - size_t outputHeight = (inputHeight + 1) >> 1; - bool padding = inputHeight & 1; - EF evalZero = EF::zero(); - EF evalTwo = EF::zero(); - EF evalFour = EF::zero(); - - auto block = cg::this_thread_block(); - auto tile = cg::tiled_partition<2>(block); - - for (size_t j = blockDim.y * blockIdx.y + threadIdx.y; j < numPolys; - j += blockDim.y * gridDim.y) { - for (size_t i = blockDim.x * blockIdx.x + threadIdx.x; i < outputHeight; - i += blockDim.x * gridDim.x) { - size_t startingIdx = j * inputHeight + (i << 1); - - F baseZeroValue = F::load(base_input, startingIdx); - F baseOneValue; - if (padding && i >= outputHeight - 1) { - baseOneValue = F::zero(); - } else { - baseOneValue = F::load(base_input, startingIdx + 1); - } - // Compute value = zeroValue * (1 - alpha) + oneValue * alpha - EF baseValue = alpha * baseOneValue + (F::one() - alpha) * baseZeroValue; - EF::store(base_output, j * outputHeight + i, baseValue); - - EF extZeroValue = EF::load(ext_input, startingIdx); - EF extOneValue; - if (padding && i >= outputHeight - 1) { - extOneValue = EF::zero(); - } else { - extOneValue = EF::load(ext_input, startingIdx + 1); - } - - // Compute value = zeroValue * (1 - alpha) + oneValue * alpha - EF extValue = alpha * extOneValue + (EF::one() - alpha) * extZeroValue; - EF::store(ext_output, j * outputHeight + i, extValue); - - EF prevBaseValue = tile.shfl(baseValue, 0); - EF prevExtValue = tile.shfl(extValue, 0); - - bool amEven = (tile.thread_rank() & 1) == 0; - if (amEven) { - // The even threads sum in evalZero. - evalZero += zerocheckEval(extValue, baseValue); - } else { - EF extSlope = extValue - prevExtValue; - EF extSlopeTimesTwo = extSlope + extSlope; - EF extSlopeTimesFour = extSlopeTimesTwo + extSlopeTimesTwo; - EF extEvalAtTwo = prevExtValue + extSlopeTimesTwo; - EF extEvalAtFour = prevExtValue + extSlopeTimesFour; - - EF baseSlope = baseValue - prevBaseValue; - EF baseSlopeTimesTwo = baseSlope + baseSlope; - EF baseSlopeTimesFour = baseSlopeTimesTwo + baseSlopeTimesTwo; - EF baseEvalAtTwo = prevBaseValue + baseSlopeTimesTwo; - EF baseEvalAtFour = prevBaseValue + baseSlopeTimesFour; - // The odd threads read from the even threads and sum in evalHalf. - evalTwo += zerocheckEval(extEvalAtTwo, baseEvalAtTwo); - evalFour += zerocheckEval(extEvalAtFour, baseEvalAtFour); - } - } - } - // __syncthreads(); - - // Allocate shared memory - extern __shared__ unsigned char memory[]; - EF* shared = reinterpret_cast(memory); - - auto reduce_tile = cg::tiled_partition<32>(block); - EF evalZeroblockSum = partialBlockReduce(block, reduce_tile, evalZero, shared); - EF evalTwoblockSum = partialBlockReduce(block, reduce_tile, evalTwo, shared); - EF evalFourblockSum = partialBlockReduce(block, reduce_tile, evalFour, shared); - - if (threadIdx.x == 0) { - EF::store(univariate_result, gridDim.x * blockIdx.y + blockIdx.x, evalZeroblockSum); - EF::store( - univariate_result, - gridDim.x * gridDim.y + gridDim.x * blockIdx.y + blockIdx.x, - evalTwoblockSum); - EF::store( - univariate_result, - 2 * gridDim.x * gridDim.y + gridDim.x * blockIdx.y + blockIdx.x, - evalFourblockSum); - } -} - -template -__global__ void zerocheckSumAsPoly( - EF* __restrict__ result, - const F* __restrict__ base_mle, - const EF* __restrict__ ext_mle, - size_t numVariablesMinusOne, - size_t numPolys) { - size_t height = 1 << (numVariablesMinusOne); - size_t inputHeight = height << 1; - EF evalZero = EF::zero(); - EF evalTwo = EF::zero(); - EF evalFour = EF::zero(); - for (size_t i = blockDim.x * blockIdx.x + threadIdx.x; i < height; - i += blockDim.x * gridDim.x) { - for (size_t j = blockDim.y * blockIdx.y + threadIdx.y; j < numPolys; - j += blockDim.y * gridDim.y) { - size_t evenIdx = j * inputHeight + (i << 1); - size_t oddIdx = evenIdx + 1; - EF prevExtValue = EF::load(ext_mle, evenIdx); - EF extValue = EF::load(ext_mle, oddIdx); - - EF extSlope = extValue - prevExtValue; - EF extSlopeTimesTwo = extSlope + extSlope; - EF extSlopeTimesFour = extSlopeTimesTwo + extSlopeTimesTwo; - EF extEvalAtTwo = prevExtValue + extSlopeTimesTwo; - EF extEvalAtFour = prevExtValue + extSlopeTimesFour; - - F prevBaseValue = F::load(base_mle, evenIdx); - F baseValue = F::load(base_mle, oddIdx); - - F baseSlope = baseValue - prevBaseValue; - F baseSlopeTimesTwo = baseSlope + baseSlope; - F baseSlopeTimesFour = baseSlopeTimesTwo + baseSlopeTimesTwo; - F baseEvalAtTwo = prevBaseValue + baseSlopeTimesTwo; - F baseEvalAtFour = prevBaseValue + baseSlopeTimesFour; - - evalZero += zerocheckEval(prevExtValue, prevBaseValue); - evalTwo += zerocheckEval(extEvalAtTwo, baseEvalAtTwo); - evalFour += zerocheckEval(extEvalAtFour, baseEvalAtFour); - } - } - - // Allocate shared memory - extern __shared__ unsigned char memory[]; - EF* shared = reinterpret_cast(memory); - - auto block = cg::this_thread_block(); - auto tile = cg::tiled_partition<32>(block); - - EF evalZeroBlockSum = partialBlockReduce(block, tile, evalZero, shared); - EF evalTwoBlockSum = partialBlockReduce(block, tile, evalTwo, shared); - EF evalFourBlockSum = partialBlockReduce(block, tile, evalFour, shared); - - if (threadIdx.x == 0) { - EF::store(result, gridDim.x * blockIdx.y + blockIdx.x, evalZeroBlockSum); - EF::store( - result, - gridDim.x * gridDim.y + gridDim.x * blockIdx.y + blockIdx.x, - evalTwoBlockSum); - EF::store( - result, - 2 * gridDim.x * gridDim.y + gridDim.x * blockIdx.y + blockIdx.x, - evalFourBlockSum); - } -} - -extern "C" void* zerocheck_sum_as_poly_base_ext_kernel() { - return (void*)zerocheckSumAsPoly; -} - -extern "C" void* zerocheck_sum_as_poly_ext_ext_kernel() { - return (void*)zerocheckSumAsPoly; -} - -extern "C" void* zerocheck_fix_last_variable_and_sum_as_poly_base_ext_kernel() { - return (void*)zerocheckFixLastVariableAndSumAsPoly; -} - -extern "C" void* zerocheck_fix_last_variable_and_sum_as_poly_ext_ext_kernel() { - return (void*)zerocheckFixLastVariableAndSumAsPoly; -} diff --git a/sp1-gpu/crates/sys/lib/zerocheck/zerocheck_eval.cu b/sp1-gpu/crates/sys/lib/zerocheck/zerocheck_eval.cu deleted file mode 100644 index 6f5e227b47..0000000000 --- a/sp1-gpu/crates/sys/lib/zerocheck/zerocheck_eval.cu +++ /dev/null @@ -1,367 +0,0 @@ -#include - -#include "zerocheck/jagged_mle.cuh" -#include "zerocheck/zerocheck_eval.cuh" -#include "config.cuh" -#include "sum_and_reduce/reduce.cuh" -#include "tracegen/jagged_tracegen/jagged.cuh" - -#include -#include - -namespace cg = cooperative_groups; - -#define DEBUG_FLAG 0 // Set this to 0 or 1 - -#if DEBUG_FLAG == 1 -#define DEBUG(...) printf(__VA_ARGS__) -#else -#define DEBUG(...) // Do nothing -#endif - -__device__ inline unsigned char get_input_point(size_t idx) { - return (unsigned char)(2 * idx); -} - -__device__ inline ext_t geq_eval(size_t idx, uint32_t threshold, ext_t eq_coefficient) { - if (idx < threshold) { - return ext_t::zero(); - } else if (idx == threshold) { - return ext_t::one() + eq_coefficient; - } else { - return ext_t::one(); - } -} - -template -__global__ void jaggedConstraintPolyEval( - const uint32_t* __restrict__ constraintIndices, - const Instruction* evalProgram, - const uint32_t* __restrict__ evalProgramIndices, - const kb31_t* evalConstantsF, - const uint32_t* __restrict__ evalConstantsFIndices, - const ext_t* evalConstantsEF, - const uint32_t* __restrict__ evalConstantsEFIndices, - const JaggedMle> inputJaggedMle, - const JaggedMle inputJaggedInfo, - const ext_t* __restrict__ partialLagrange, - const uint32_t* __restrict__ geq_thresholds, - const ext_t* __restrict__ eq_coefficients, - uint32_t totalLen, - const ext_t* __restrict__ paddedRowAdjustment, - const felt_t* __restrict__ publicValues, - const ext_t* __restrict__ powersOfAlpha, - const ext_t* __restrict__ batchingPowers, - const ext_t* __restrict__ powersOfLambda, - const uint32_t* __restrict__ preprocessed_column, - const uint32_t* __restrict__ main_column, - uint32_t total_num_preprocessed_column, - ext_t* __restrict__ constraintValues, - uint32_t rest_point_dim -) { - K expr_f[MEMORY_SIZE]; - JaggedConstraintFolder folder = JaggedConstraintFolder(); - - ext_t thread_sum = ext_t::zero(); - - // This kernel assumes that a single block deals with a single `xValueIdx`. - size_t xValueIdx = blockDim.z * blockIdx.z + threadIdx.z; - unsigned char eval_point = get_input_point(xValueIdx); - - for (size_t idx = blockDim.x * blockIdx.x + threadIdx.x; idx < totalLen; idx += blockDim.x * gridDim.x) { - uint64_t packed_info = inputJaggedInfo.denseData.data[idx << 1]; - assert(packed_info == inputJaggedInfo.denseData.data[idx << 1 | 1]); - - uint32_t chip_idx = (packed_info >> 1) & 0x7FFF; - uint32_t preprocessed_idx = (packed_info >> 16) & 0xFFFF; - uint32_t main_idx = (packed_info >> 32) & 0xFFFFFFFF; - uint32_t num_preprocessed_columns = preprocessed_column[chip_idx]; - uint32_t num_main_columns = main_column[chip_idx]; - - uint32_t airBlockIdx = inputJaggedInfo.colIndex[idx]; - bool is_first_air_block = ((packed_info & 1) == 1); - - uint32_t rowIdx = idx - inputJaggedInfo.startIndices[airBlockIdx]; - uint32_t height = 2 * (inputJaggedInfo.startIndices[airBlockIdx + 1] - inputJaggedInfo.startIndices[airBlockIdx]); - - size_t constraint_offset = constraintIndices[airBlockIdx]; - size_t program_start_idx = evalProgramIndices[airBlockIdx]; - size_t program_end_idx = evalProgramIndices[airBlockIdx + 1]; - size_t f_constant_offset = evalConstantsFIndices[airBlockIdx]; - size_t ef_constant_offset = evalConstantsEFIndices[airBlockIdx]; - - for (size_t i = 0; i < MEMORY_SIZE; i++) { - expr_f[i] = K::zero(); - } - - folder.data = inputJaggedMle.denseData.data; - folder.preprocessed_ptr = inputJaggedMle.startIndices[preprocessed_idx] << 1; - folder.main_ptr = inputJaggedMle.startIndices[total_num_preprocessed_column + 1 + main_idx] << 1; - folder.height = height; - folder.publicValues = publicValues; - folder.powersOfAlpha = powersOfAlpha; - folder.constraintIndex = constraint_offset; - folder.accumulator = ext_t::zero(); - folder.rowIdx = rowIdx; - folder.eval_point = eval_point; - - for (size_t i = program_start_idx; i < program_end_idx; i++) { - Instruction instr = evalProgram[i]; - switch (instr.opcode) { - case 0: - DEBUG("EMPTY\n"); - break; - - case 1: - DEBUG("FAssignC: %d <- %d\n", instr.a, instr.b); - expr_f[instr.a] = evalConstantsF[f_constant_offset + instr.b]; - break; - case 2: - DEBUG("FAssignV: %d <- (%d, %d)\n", instr.a, instr.b_variant, instr.b); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b); - break; - case 3: - DEBUG("FAssignE: %d <- %d\n", instr.a, instr.b); - expr_f[instr.a] = expr_f[instr.b]; - break; - case 4: - DEBUG("FAddVC: %d <- %d + %d\n", instr.a, instr.b_variant, instr.b); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) + - evalConstantsF[f_constant_offset + instr.c]; - break; - case 5: - DEBUG( - "FAddVV: %d <- (%d, %d) + (%d, %d)\n", - instr.a, - instr.b_variant, - instr.b, - instr.c_variant, - instr.c); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) + - folder.var_f(instr.c_variant, instr.c); - break; - case 6: - DEBUG( - "FAddVE: %d <- (%d, %d) + %d\n", - instr.a, - instr.b_variant, - instr.b, - instr.c); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) + expr_f[instr.c]; - break; - - case 7: - DEBUG("FAddEC: %d <- %d + %d\n", instr.a, instr.b_variant, instr.b); - expr_f[instr.a] = expr_f[instr.b] + evalConstantsF[f_constant_offset + instr.c]; - break; - case 8: - DEBUG( - "FAddEV: %d <- %d + (%d, %d)\n", - instr.a, - instr.b, - instr.c_variant, - instr.c); - expr_f[instr.a] = expr_f[instr.b] + folder.var_f(instr.c_variant, instr.c); - break; - case 9: - DEBUG("FAddEE: %d <- %d + %d\n", instr.a, instr.b, instr.c); - expr_f[instr.a] = expr_f[instr.b] + expr_f[instr.c]; - break; - case 10: - DEBUG("FAddAssignE: %d <- %d\n", instr.a, instr.b); - expr_f[instr.a] += expr_f[instr.b]; - break; - - case 11: - DEBUG("FSubVC: %d <- %d - %d\n", instr.a, instr.b_variant, instr.b); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) - - evalConstantsF[f_constant_offset + instr.c]; - break; - case 12: - DEBUG( - "FSubVV: %d <- (%d, %d) - (%d, %d)\n", - instr.a, - instr.b_variant, - instr.b, - instr.c_variant, - instr.c); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) - - folder.var_f(instr.c_variant, instr.c); - break; - case 13: - DEBUG( - "FSubVE: %d <- (%d, %d) - %d\n", - instr.a, - instr.b_variant, - instr.b, - instr.c); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) - expr_f[instr.c]; - break; - - case 14: - DEBUG("FSubEC: %d <- %d - %d\n", instr.a, instr.b, instr.c); - expr_f[instr.a] = expr_f[instr.b] - evalConstantsF[f_constant_offset + instr.c]; - break; - case 15: - DEBUG( - "FSubEV: %d <- %d - (%d, %d)\n", - instr.a, - instr.b, - instr.c_variant, - instr.c); - expr_f[instr.a] = expr_f[instr.b] - folder.var_f(instr.c_variant, instr.c); - break; - case 16: - DEBUG("FSubEE: %d <- %d - %d\n", instr.a, instr.b, instr.c); - expr_f[instr.a] = expr_f[instr.b] - expr_f[instr.c]; - break; - case 17: - DEBUG("FSubAssignE: %d <- %d\n", instr.a, instr.b); - expr_f[instr.a] -= expr_f[instr.b]; - break; - - case 18: - DEBUG("FMulVC: %d <- %d * %d\n", instr.a, instr.b_variant, instr.b); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) * - evalConstantsF[f_constant_offset + instr.c]; - break; - case 19: - DEBUG( - "FMulVV: %d <- (%d, %d) * (%d, %d)\n", - instr.a, - instr.b_variant, - instr.b, - instr.c_variant, - instr.c); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) * - folder.var_f(instr.c_variant, instr.c); - break; - case 20: - DEBUG( - "FMulVE: %d <- (%d, %d) * %d\n", - instr.a, - instr.b_variant, - instr.b, - instr.c); - expr_f[instr.a] = folder.var_f(instr.b_variant, instr.b) * expr_f[instr.c]; - break; - - case 21: - DEBUG("FMulEC: %d <- %d * %d\n", instr.a, instr.b_variant, instr.b); - expr_f[instr.a] = expr_f[instr.b] * evalConstantsF[f_constant_offset + instr.c]; - break; - case 22: - DEBUG( - "FMulEV: %d <- %d * (%d, %d)\n", - instr.a, - instr.b, - instr.c_variant, - instr.c); - expr_f[instr.a] = expr_f[instr.b] * folder.var_f(instr.c_variant, instr.c); - break; - case 23: - DEBUG("FMulEE: %d <- %d * %d\n", instr.a, instr.b, instr.c); - DEBUG("FMulEE Input: %d, %d\n", expr_f[instr.b], expr_f[instr.c]); - expr_f[instr.a] = expr_f[instr.b] * expr_f[instr.c]; - DEBUG("FMulEE Output: %d\n", expr_f[instr.a]); - break; - case 24: - DEBUG("FMulAssignE: %d <- %d\n", instr.a, instr.b); - expr_f[instr.a] *= expr_f[instr.b]; - break; - - case 25: - DEBUG("FNegE: %d <- -%d\n", instr.a, instr.b); - expr_f[instr.a] = -expr_f[instr.b]; - break; - - case 59: - DEBUG("FAssertZero: %d\n", instr.a); - folder.accumulator += - (folder.powersOfAlpha[folder.constraintIndex] * - expr_f[instr.a]); - folder.constraintIndex++; - break; - } - } - - ext_t gkr_correction = ext_t::zero(); - ext_t geq_correction = ext_t::zero(); - - if (is_first_air_block) { - for (size_t i = 0; i < num_main_columns; i++) { - gkr_correction += batchingPowers[i] * folder.var_f(4, i); - } - for (size_t i = 0; i < num_preprocessed_columns ; i++) { - gkr_correction += batchingPowers[num_main_columns + i] * folder.var_f(2, i); - } - ext_t zeroVal = geq_eval(rowIdx << 1, geq_thresholds[chip_idx], eq_coefficients[chip_idx]); - ext_t oneVal = geq_eval(rowIdx << 1 | 1, geq_thresholds[chip_idx], eq_coefficients[chip_idx]); - geq_correction = (zeroVal + eval_point * (oneVal - zeroVal)) * paddedRowAdjustment[chip_idx]; - } - - if (rowIdx < (1 << rest_point_dim)) { - ext_t eq = ext_t::load(partialLagrange, rowIdx); - thread_sum += (folder.accumulator + gkr_correction - geq_correction) * eq * powersOfLambda[chip_idx]; - } - } - - extern __shared__ unsigned char memory[]; - ext_t* shared = reinterpret_cast(memory); - - auto block = cg::this_thread_block(); - auto tile = cg::tiled_partition<32>(block); - ext_t thread_block_sum = partialBlockReduce(block, tile, thread_sum, shared); - - if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) { - ext_t::store(constraintValues, xValueIdx * gridDim.x + blockIdx.x, thread_block_sum); - } -} - -extern "C" void* jagged_constraint_poly_eval_32_koala_bear_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_64_koala_bear_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_128_koala_bear_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_256_koala_bear_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_512_koala_bear_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_1024_koala_bear_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_32_koala_bear_extension_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_64_koala_bear_extension_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_128_koala_bear_extension_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_256_koala_bear_extension_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_512_koala_bear_extension_kernel() { - return (void*)jaggedConstraintPolyEval; -} - -extern "C" void* jagged_constraint_poly_eval_1024_koala_bear_extension_kernel() { - return (void*)jaggedConstraintPolyEval; -} diff --git a/sp1-gpu/crates/sys/src/v2_kernels.rs b/sp1-gpu/crates/sys/src/kernels.rs similarity index 76% rename from sp1-gpu/crates/sys/src/v2_kernels.rs rename to sp1-gpu/crates/sys/src/kernels.rs index f2609c0682..43a91d40a5 100644 --- a/sp1-gpu/crates/sys/src/v2_kernels.rs +++ b/sp1-gpu/crates/sys/src/kernels.rs @@ -31,6 +31,42 @@ extern "C" { pub fn fix_last_variable_jagged_felt() -> KernelPtr; pub fn fix_last_variable_jagged_ext() -> KernelPtr; + // zerocheck (DAG-native): Sequential lowering kernels. + // _kb: K = felt_t (base-field trace, sumcheck round 0) + // _ext: K = ext_t (extension-field trace, sumcheck rounds 1+) + pub fn zerocheck_sequential_kb_kernel() -> KernelPtr; + pub fn zerocheck_sequential_ext_kernel() -> KernelPtr; + // Tiered by per-chunk MAX_REGS (kernel selection by max_reg). + pub fn zerocheck_sequential_kb_32_kernel() -> KernelPtr; + pub fn zerocheck_sequential_kb_64_kernel() -> KernelPtr; + pub fn zerocheck_sequential_kb_128_kernel() -> KernelPtr; + pub fn zerocheck_sequential_kb_256_kernel() -> KernelPtr; + pub fn zerocheck_sequential_ext_32_kernel() -> KernelPtr; + pub fn zerocheck_sequential_ext_64_kernel() -> KernelPtr; + pub fn zerocheck_sequential_ext_128_kernel() -> KernelPtr; + pub fn zerocheck_sequential_ext_256_kernel() -> KernelPtr; + + // Fused dispatch: one launch handles every Sequential chunk in a round. + // Per-thread idx → chunk_idx via binary search on `row_starts`. + pub fn zerocheck_fused_sequential_kb_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_ext_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_kb_32_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_kb_64_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_kb_128_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_kb_256_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_kb_512_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_kb_1024_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_ext_32_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_ext_64_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_ext_128_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_ext_256_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_ext_512_kernel() -> KernelPtr; + pub fn zerocheck_fused_sequential_ext_1024_kernel() -> KernelPtr; + + // zerocheck (DAG-native): ColumnTile lowering kernels. + pub fn zerocheck_column_tile_kb_kernel() -> KernelPtr; + pub fn zerocheck_column_tile_ext_kernel() -> KernelPtr; + // Jagged Zerocheck Kernels pub fn jagged_constraint_poly_eval_32_koala_bear_kernel() -> KernelPtr; pub fn jagged_constraint_poly_eval_64_koala_bear_kernel() -> KernelPtr; diff --git a/sp1-gpu/crates/sys/src/lib.rs b/sp1-gpu/crates/sys/src/lib.rs index 481335b62f..e654401d51 100644 --- a/sp1-gpu/crates/sys/src/lib.rs +++ b/sp1-gpu/crates/sys/src/lib.rs @@ -3,6 +3,7 @@ pub mod basefold; pub mod challenger; pub mod dft; pub mod jagged; +pub mod kernels; pub mod logup_gkr; pub mod merkle_tree; pub mod mle; @@ -12,5 +13,3 @@ pub mod scan; pub mod sumcheck; pub mod tracegen; pub mod transpose; -pub mod v2_kernels; -pub mod zerocheck; diff --git a/sp1-gpu/crates/sys/src/zerocheck.rs b/sp1-gpu/crates/sys/src/zerocheck.rs deleted file mode 100644 index 3cb366f861..0000000000 --- a/sp1-gpu/crates/sys/src/zerocheck.rs +++ /dev/null @@ -1,18 +0,0 @@ -use crate::runtime::KernelPtr; - -extern "C" { - pub fn interpolate_row_koala_bear_kernel() -> KernelPtr; - pub fn interpolate_row_koala_bear_extension_kernel() -> KernelPtr; - pub fn constraint_poly_eval_32_koala_bear_kernel() -> KernelPtr; - pub fn constraint_poly_eval_64_koala_bear_kernel() -> KernelPtr; - pub fn constraint_poly_eval_128_koala_bear_kernel() -> KernelPtr; - pub fn constraint_poly_eval_256_koala_bear_kernel() -> KernelPtr; - pub fn constraint_poly_eval_512_koala_bear_kernel() -> KernelPtr; - pub fn constraint_poly_eval_1024_koala_bear_kernel() -> KernelPtr; - pub fn constraint_poly_eval_32_koala_bear_extension_kernel() -> KernelPtr; - pub fn constraint_poly_eval_64_koala_bear_extension_kernel() -> KernelPtr; - pub fn constraint_poly_eval_128_koala_bear_extension_kernel() -> KernelPtr; - pub fn constraint_poly_eval_256_koala_bear_extension_kernel() -> KernelPtr; - pub fn constraint_poly_eval_512_koala_bear_extension_kernel() -> KernelPtr; - pub fn constraint_poly_eval_1024_koala_bear_extension_kernel() -> KernelPtr; -} diff --git a/sp1-gpu/crates/zerocheck/benches/zerocheck.rs b/sp1-gpu/crates/zerocheck/benches/zerocheck.rs index 8b68105ae7..acd2bb8d11 100644 --- a/sp1-gpu/crates/zerocheck/benches/zerocheck.rs +++ b/sp1-gpu/crates/zerocheck/benches/zerocheck.rs @@ -1,8 +1,6 @@ -//! Bench `zerocheck`. The zerocheck *prover* runs on any trace data — only verification cares -//! about constraint satisfaction — so source selection (random / JSON / real) goes through -//! [`with_trace_source`] with [`FullKind`]. For random and JSON sources the helper synthesizes -//! `cluster` (default: the machine's `core` cluster; override with `random:N,cluster=all-chips`) -//! and a zero-filled `public_values` of the right length. See [`benches/README.md`] for details. +//! Bench `zerocheck` — the DAG-native lowering path. Trace source selection +//! (random / JSON / real) goes through [`with_trace_source`] with +//! [`FullKind`]; see [`benches/README.md`] for details. use std::collections::BTreeMap; use std::sync::Arc; @@ -14,16 +12,16 @@ use slop_algebra::AbstractField; use slop_challenger::{CanObserve, CanSample, FieldChallenger, IopCtx}; use slop_multilinear::{MleEval, Point}; use slop_tensor::Tensor; -use sp1_gpu_air::codegen_cuda_eval; +use sp1_gpu_air::ir::ChunkBudget; use sp1_gpu_cudart::TaskScope; use sp1_gpu_jagged_tracegen::test_utils::bench_utils::{ with_trace_source, FullKind, RealTraceData, }; -use sp1_gpu_jagged_tracegen::test_utils::tracegen_setup::CORE_MAX_LOG_ROW_COUNT; use sp1_gpu_utils::{Ext, Felt, TestGC}; use sp1_gpu_zerocheck::primitives::evaluate_jagged_columns; -use sp1_gpu_zerocheck::zerocheck; +use sp1_gpu_zerocheck::prover::{upload_machine_bytecode, zerocheck}; use sp1_hypercube::air::MachineAir; +use sp1_hypercube::log2_ceil_usize; use sp1_hypercube::{ChipEvaluation, LogUpEvaluations}; fn run_zerocheck( @@ -36,11 +34,9 @@ fn run_zerocheck( let RealTraceData { machine: _, cluster, public_values, device_mle } = data; let chips = &cluster; - let mut cache = BTreeMap::new(); - for chip in chips.iter() { - let result = codegen_cuda_eval(chip.air.as_ref()); - cache.insert(chip.name().to_string(), result); - } + // Compile + upload the machine's v2 bytecode once. + let machine_bytecode = + Arc::new(upload_machine_bytecode(chips, ChunkBudget::recommended(), scope)); let trace_mle = Arc::new(device_mle); @@ -55,9 +51,15 @@ fn run_zerocheck( let mut challenger_prover = challenger.clone(); let batching_challenge = challenger_prover.sample_ext_element(); let gkr_opening_batch_randomness = challenger_prover.sample_ext_element(); - let max_log_row_count = CORE_MAX_LOG_ROW_COUNT; - - let zeta = Point::::rand(rng, CORE_MAX_LOG_ROW_COUNT); + // Derive max_log_row_count from the actual chip heights. The bench's old + // `CORE_MAX_LOG_ROW_COUNT = 22` only covers `core_2^25`-sized random + // traces — bigger areas can put more than 2^22 rows in a single chip and + // trip `VirtualGeq::new`'s assert. + let max_chip_height = + trace_mle.dense_data.main_table_index.values().map(|o| o.poly_size).max().unwrap_or(1); + let max_log_row_count = log2_ceil_usize(max_chip_height).max(1); + + let zeta = Point::::rand(rng, max_log_row_count as u32); let individual_column_evals = evaluate_jagged_columns(trace_mle.as_ref(), zeta.clone()); let mut preprocessed_ptr: usize = 0; @@ -101,14 +103,14 @@ fn run_zerocheck( |pv| { let result = zerocheck( chips, - &cache, + &machine_bytecode, trace_mle.as_ref(), batching_challenge, gkr_opening_batch_randomness, &logup_evaluations, pv, &mut challenger, - max_log_row_count, + max_log_row_count as u32, ); scope.synchronize_blocking().unwrap(); black_box(result) diff --git a/sp1-gpu/crates/zerocheck/examples/column_tile_equiv.rs b/sp1-gpu/crates/zerocheck/examples/column_tile_equiv.rs new file mode 100644 index 0000000000..70d0b6c55f --- /dev/null +++ b/sp1-gpu/crates/zerocheck/examples/column_tile_equiv.rs @@ -0,0 +1,329 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! Equivalence test for the v2 ColumnTile kernel. +//! +//! Constructs a synthetic `LinearWeightedSum` DAG (every constraint is a +//! linear combination of column leaves with constant or public coefficients), +//! lowers it via `lower_column_tile`, runs both a CPU oracle and the GPU +//! kernel, and verifies the 3 per-eval-point partial sums match bitwise. + +use rand::{rngs::StdRng, Rng, SeedableRng}; +use slop_algebra::{AbstractExtensionField, AbstractField}; +use sp1_gpu_air::ir::{ + analyze_constraints, chunk_dag, enumerate_lowerings, lower_column_tile, ChunkBudget, + ColumnTileBytecode, ConstraintDag, ConstraintField, ConstraintRef, ConstraintShape, DagNode, + Lowering, TraceSource, COEFF_KIND_CONST, +}; +use sp1_gpu_air::{EF, F}; +use sp1_gpu_cudart::sys::kernels::zerocheck_column_tile_kb_kernel; +use sp1_gpu_cudart::{args, run_sync_in_place, DeviceBuffer}; + +// LinearWeightedSum DAG with constants AND publics as coefficients: +// c0: 3*l0 + 5*l1 + 7*l2 (constant coeffs) +// c1: 11*l3 + 13*l4 (constant coeffs) +// c2: pv0 * l0 + pv1 * l1 (public coeffs) +fn build_linear_dag() -> ConstraintDag { + let mut nodes = vec![]; + // 0..5: column leaves + for c in 0..5 { + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: c }); + } + // 5..10: constants (3, 5, 7, 11, 13) + for k in [3u32, 5, 7, 11, 13] { + nodes.push(DagNode::ConstF { value: F::from_canonical_u32(k) }); + } + // 10..12: public values 0, 1 + nodes.push(DagNode::PublicValue { idx: 0 }); + nodes.push(DagNode::PublicValue { idx: 1 }); + + // c0: 3*l0 + 5*l1 + 7*l2 + let m0 = nodes.len() as u32; + nodes.push(DagNode::MulF { a: 5, b: 0 }); // 3 * l0 + let m1 = nodes.len() as u32; + nodes.push(DagNode::MulF { a: 6, b: 1 }); // 5 * l1 + let s0 = nodes.len() as u32; + nodes.push(DagNode::AddF { a: m0, b: m1 }); + let m2 = nodes.len() as u32; + nodes.push(DagNode::MulF { a: 7, b: 2 }); // 7 * l2 + let c0_root = nodes.len() as u32; + nodes.push(DagNode::AddF { a: s0, b: m2 }); + + // c1: 11*l3 + 13*l4 + let m3 = nodes.len() as u32; + nodes.push(DagNode::MulF { a: 8, b: 3 }); // 11 * l3 + let m4 = nodes.len() as u32; + nodes.push(DagNode::MulF { a: 9, b: 4 }); // 13 * l4 + let c1_root = nodes.len() as u32; + nodes.push(DagNode::AddF { a: m3, b: m4 }); + + // c2: pv0 * l0 + pv1 * l1 + let m5 = nodes.len() as u32; + nodes.push(DagNode::MulF { a: 10, b: 0 }); + let m6 = nodes.len() as u32; + nodes.push(DagNode::MulF { a: 11, b: 1 }); + let c2_root = nodes.len() as u32; + nodes.push(DagNode::AddF { a: m5, b: m6 }); + + let constraints = vec![ + ConstraintRef { root: c0_root, alpha_index: 0, field: ConstraintField::Base }, + ConstraintRef { root: c1_root, alpha_index: 1, field: ConstraintField::Base }, + ConstraintRef { root: c2_root, alpha_index: 2, field: ConstraintField::Base }, + ]; + + ConstraintDag { nodes, constraints, preprocessed_width: 0, main_width: 5 } +} + +const N_COLS: usize = 5; +const N_PARTIAL_PER_BLOCK: usize = 3; + +fn main() { + let dag = build_linear_dag(); + let infos = analyze_constraints(&dag); + + // Sanity: every constraint should be LinearWeightedSum. + for (i, info) in infos.iter().enumerate() { + assert!( + matches!(info.shape, ConstraintShape::LinearWeightedSum), + "constraint {} not detected as LinearWeightedSum", + i + ); + } + + let chunks = chunk_dag(&infos, &ChunkBudget::recommended()); + assert_eq!(chunks.len(), 1, "should chunk into one linear chunk"); + let chunk = &chunks[0]; + assert!(matches!(chunk.shape, ConstraintShape::LinearWeightedSum)); + + let lowerings = enumerate_lowerings(chunk, &infos, &dag); + let plan = lowerings + .iter() + .find_map(|l| match l { + Lowering::ColumnTile(p) => Some(p), + _ => None, + }) + .expect("ColumnTile plan should exist for a LinearWeightedSum chunk"); + let bc = + lower_column_tile(chunk, &infos, &dag, plan).expect("lower_column_tile should succeed"); + + println!( + "column-tile bytecode: {} terms, {} leaves, {} consts, {} publics", + bc.terms.len(), + bc.leaves.len(), + bc.consts.len(), + bc.publics.len() + ); + for (i, t) in bc.terms.iter().enumerate() { + let kind = if t.coeff_kind == COEFF_KIND_CONST { "const" } else { "public" }; + println!( + " term {}: leaf_idx={} {}#{} alpha_idx={}", + i, t.leaf_idx, kind, t.coeff_idx, t.alpha_idx + ); + } + + // Run multiple scenarios to exercise grid bounds. n_rows == 2^rest_point_dim. + let scenarios: &[(&str, usize, u32, u64)] = + &[("single block, full row tile", 8, 3, 0xAA), ("multi block, aligned tail", 32, 5, 0xCC)]; + + let mut all_ok = true; + for &(label, n_rows, rest_point_dim, seed) in scenarios { + let mut rng = StdRng::seed_from_u64(seed); + let height = 2 * n_rows; + let trace: Vec = (0..N_COLS * height).map(|_| rand_f(&mut rng)).collect(); + let powers_of_alpha: Vec = + (0..dag.constraints.len()).map(|_| rand_ef(&mut rng)).collect(); + let public_values: Vec = (0..16).map(|_| rand_f(&mut rng)).collect(); + let partial_lagrange: Vec = + (0..(1u32 << rest_point_dim)).map(|_| rand_ef(&mut rng)).collect(); + let powers_of_lambda: Vec = vec![rand_ef(&mut rng); 4]; + let chip_idx: u32 = 2; + + let host = host_oracle( + &bc, + &trace, + height, + &powers_of_alpha, + &public_values, + &partial_lagrange, + &powers_of_lambda, + chip_idx, + rest_point_dim, + n_rows, + ); + let gpu = run_gpu( + &bc, + &trace, + height, + &powers_of_alpha, + &public_values, + &partial_lagrange, + &powers_of_lambda, + chip_idx, + rest_point_dim, + n_rows, + ); + + let ok = (0..N_PARTIAL_PER_BLOCK).all(|e| host[e] == gpu[e]); + println!("[{}] n_rows={} : {}", label, n_rows, if ok { "match" } else { "MISMATCH" }); + if !ok { + for e in 0..N_PARTIAL_PER_BLOCK { + println!(" e{}: host={:?} gpu={:?}", e, host[e], gpu[e]); + } + all_ok = false; + } + } + if all_ok { + println!("\nv2 ColumnTile kernel matches host oracle bitwise across all scenarios"); + } else { + std::process::exit(1); + } +} + +fn rand_f(rng: &mut StdRng) -> F { + F::from_wrapped_u32(rng.gen()) +} + +fn rand_ef(rng: &mut StdRng) -> EF { + EF::from_base_slice(&[rand_f(rng), rand_f(rng), rand_f(rng), rand_f(rng)]) +} + +fn host_oracle( + bc: &ColumnTileBytecode, + trace: &[F], + height: usize, + powers_of_alpha: &[EF], + public_values: &[F], + partial_lagrange: &[EF], + powers_of_lambda: &[EF], + chip_idx: u32, + rest_point_dim: u32, + n_rows: usize, +) -> [EF; 3] { + let mut sum = [EF::zero(); 3]; + let eval_pts: [F; 3] = + [F::from_canonical_u32(0), F::from_canonical_u32(2), F::from_canonical_u32(4)]; + let domain = 1u32 << rest_point_dim; + let lambda = powers_of_lambda[chip_idx as usize]; + for row in 0..n_rows { + // Each term in a column-tile chunk contributes one weighted value per + // (row, eval). The kernel multiplies each term's contribution by + // eq[row] · λ; the host mirrors that exactly per-term. + let weight = + if (row as u32) < domain { partial_lagrange[row] * lambda } else { EF::zero() }; + for t in &bc.terms { + let leaf = bc.leaves[t.leaf_idx as usize]; + let off = (leaf.col as usize) * height + (row << 1); + let z = trace[off]; + let o = trace[off + 1]; + let coeff = if t.coeff_kind == COEFF_KIND_CONST { + bc.consts[t.coeff_idx as usize] + } else { + public_values[bc.publics[t.coeff_idx as usize] as usize] + }; + let alpha = powers_of_alpha[t.alpha_idx as usize]; + for e in 0..3 { + let v = z + eval_pts[e] * (o - z); + let contribution = alpha * EF::from_base(coeff * v); + sum[e] += contribution * weight; + } + } + } + sum +} + +fn run_gpu( + bc: &ColumnTileBytecode, + trace: &[F], + height: usize, + powers_of_alpha: &[EF], + public_values: &[F], + partial_lagrange: &[EF], + powers_of_lambda: &[EF], + chip_idx: u32, + rest_point_dim: u32, + n_rows: usize, +) -> [EF; 3] { + // 1D block; each thread = one (term, row) pair, computing all 3 eval points. + const BLOCK_SIZE: u32 = 64; + let total = (bc.terms.len() as u32) * (n_rows as u32); + let grid_x = (total + BLOCK_SIZE - 1) / BLOCK_SIZE; + let n_warps = BLOCK_SIZE / 32; + let shmem_bytes = (n_warps as usize) * std::mem::size_of::(); + + let result = run_sync_in_place(|scope| { + let d_terms = DeviceBuffer::from_host_slice(&bc.terms, &scope).unwrap(); + let d_leaves = DeviceBuffer::from_host_slice(&bc.leaves, &scope).unwrap(); + let d_consts = DeviceBuffer::from_host_slice(&bc.consts, &scope).unwrap(); + let d_publics = DeviceBuffer::from_host_slice(&bc.publics, &scope).unwrap(); + // No COEFF_KIND_RUNTIME terms in this test; pass a 1-element dummy so + // the kernel pointer is non-null. + let d_runtime: DeviceBuffer = + DeviceBuffer::from_host_slice(&[EF::zero()], &scope).unwrap(); + let d_trace = DeviceBuffer::from_host_slice(trace, &scope).unwrap(); + let d_public_values = DeviceBuffer::from_host_slice(public_values, &scope).unwrap(); + let d_powers = DeviceBuffer::from_host_slice(powers_of_alpha, &scope).unwrap(); + let d_partial_lagrange = DeviceBuffer::from_host_slice(partial_lagrange, &scope).unwrap(); + let d_powers_of_lambda = DeviceBuffer::from_host_slice(powers_of_lambda, &scope).unwrap(); + let mut d_partials: DeviceBuffer = + DeviceBuffer::with_capacity_in(grid_x as usize * 3, scope.clone()); + unsafe { + d_partials.assume_init(); + } + + let n_terms: u32 = bc.terms.len() as u32; + let preprocessed_ptr: u64 = 0; + let main_ptr: u64 = 0; + let height_u: u32 = height as u32; + let row_start: u32 = 0; + let row_count: u32 = n_rows as u32; + + unsafe { + let a = args!( + d_terms.as_ptr(), + n_terms, + d_leaves.as_ptr(), + d_consts.as_ptr(), + d_publics.as_ptr(), + d_runtime.as_ptr(), + d_trace.as_ptr(), + preprocessed_ptr, + main_ptr, + height_u, + d_public_values.as_ptr(), + d_powers.as_ptr(), + d_partial_lagrange.as_ptr(), + d_powers_of_lambda.as_ptr(), + chip_idx, + rest_point_dim, + row_start, + row_count, + d_partials.as_mut_ptr() + ); + scope + .launch_kernel( + zerocheck_column_tile_kb_kernel(), + (grid_x, 1, 1), + (BLOCK_SIZE, 1, 1), + &a, + shmem_bytes, + ) + .unwrap(); + } + + d_partials.to_host().unwrap().to_vec() + }) + .unwrap(); + + let mut out = [EF::zero(); 3]; + for (i, &p) in result.iter().enumerate() { + out[i % 3] += p; + } + out +} diff --git a/sp1-gpu/crates/zerocheck/examples/gkr_equiv.rs b/sp1-gpu/crates/zerocheck/examples/gkr_equiv.rs new file mode 100644 index 0000000000..f19ef4b20b --- /dev/null +++ b/sp1-gpu/crates/zerocheck/examples/gkr_equiv.rs @@ -0,0 +1,257 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! Equivalence test for GKR-correction chunks via ColumnTile. +//! +//! Uses `synthesize_gkr_chunk` to construct a column-tile bytecode whose +//! coefficients are `COEFF_KIND_RUNTIME` ext_t values supplied at launch +//! (the per-proof batching powers). Verifies the GPU kernel matches the +//! host oracle bitwise. + +use rand::{rngs::StdRng, Rng, SeedableRng}; +use slop_algebra::{AbstractExtensionField, AbstractField}; +use sp1_gpu_air::ir::{synthesize_gkr_chunk, ColumnTileBytecode, COEFF_KIND_RUNTIME}; +use sp1_gpu_air::{EF, F}; +use sp1_gpu_cudart::sys::kernels::zerocheck_column_tile_kb_kernel; +use sp1_gpu_cudart::{args, run_sync_in_place, DeviceBuffer}; + +const N_PARTIAL_PER_BLOCK: usize = 3; + +fn main() { + // Synthesize a GKR-correction chunk: 5 main + 2 preprocessed columns. + // Total 7 terms, all COEFF_KIND_RUNTIME, all alpha_idx=0. + let main_width: u32 = 5; + let prep_width: u32 = 2; + let bc = synthesize_gkr_chunk(main_width, prep_width); + println!( + "synthesized GKR chunk: main={} prep={} → {} terms, {} leaves", + main_width, + prep_width, + bc.terms.len(), + bc.leaves.len() + ); + // Sanity: every term should be COEFF_KIND_RUNTIME. + for t in &bc.terms { + assert_eq!(t.coeff_kind, COEFF_KIND_RUNTIME); + } + // Main first, then prep — matches v1's ordering. + for (i, t) in bc.terms.iter().enumerate() { + let leaf = bc.leaves[t.leaf_idx as usize]; + let kind = if leaf.source == 4 { "main" } else { "prep" }; + println!( + " term {}: {}#{} runtime_coeff[{}] alpha_idx={}", + i, kind, leaf.col, t.coeff_idx, t.alpha_idx + ); + } + + let scenarios: &[(&str, usize, u32, u64)] = + &[("single block, small row tile", 8, 3, 0xA1), ("multi block, aligned tail", 64, 6, 0xC1)]; + + let mut all_ok = true; + for &(label, n_rows, rest_point_dim, seed) in scenarios { + let mut rng = StdRng::seed_from_u64(seed); + let n_cols_total = (main_width + prep_width) as usize; + let height = 2 * n_rows; + + let trace: Vec = (0..n_cols_total * height).map(|_| rand_f(&mut rng)).collect(); + let runtime_coeffs: Vec = (0..n_cols_total).map(|_| rand_ef(&mut rng)).collect(); + let powers_of_alpha: Vec = vec![rand_ef(&mut rng); 1]; + let partial_lagrange: Vec = + (0..(1u32 << rest_point_dim)).map(|_| rand_ef(&mut rng)).collect(); + let powers_of_lambda: Vec = vec![rand_ef(&mut rng); 4]; + let chip_idx: u32 = 3; + + let host = host_oracle( + &bc, + &trace, + height, + &runtime_coeffs, + &powers_of_alpha, + &partial_lagrange, + &powers_of_lambda, + chip_idx, + rest_point_dim, + main_width, + n_rows, + ); + let gpu = run_gpu( + &bc, + &trace, + height, + &runtime_coeffs, + &powers_of_alpha, + &partial_lagrange, + &powers_of_lambda, + chip_idx, + rest_point_dim, + main_width, + n_rows, + ); + + let ok = (0..N_PARTIAL_PER_BLOCK).all(|e| host[e] == gpu[e]); + println!("[{}] n_rows={} : {}", label, n_rows, if ok { "match" } else { "MISMATCH" }); + if !ok { + for e in 0..N_PARTIAL_PER_BLOCK { + println!(" e{}: host={:?} gpu={:?}", e, host[e], gpu[e]); + } + all_ok = false; + } + } + if all_ok { + println!("\nv2 GKR via ColumnTile (runtime coefficients) matches host oracle bitwise"); + } else { + std::process::exit(1); + } +} + +fn rand_f(rng: &mut StdRng) -> F { + F::from_wrapped_u32(rng.gen()) +} + +fn rand_ef(rng: &mut StdRng) -> EF { + EF::from_base_slice(&[rand_f(rng), rand_f(rng), rand_f(rng), rand_f(rng)]) +} + +/// CPU port of the kernel logic for a runtime-coefficient chunk. Trace layout: +/// columns 0..main_width are "main" (source byte 4), preprocessed_ptr offsets +/// main_width * height ahead. Both share the same `trace` buffer. +fn host_oracle( + bc: &ColumnTileBytecode, + trace: &[F], + height: usize, + runtime_coeffs: &[EF], + powers_of_alpha: &[EF], + partial_lagrange: &[EF], + powers_of_lambda: &[EF], + chip_idx: u32, + rest_point_dim: u32, + main_width: u32, + n_rows: usize, +) -> [EF; 3] { + let mut sum = [EF::zero(); 3]; + let eval_pts: [F; 3] = + [F::from_canonical_u32(0), F::from_canonical_u32(2), F::from_canonical_u32(4)]; + let main_ptr: usize = 0; + let preprocessed_ptr: usize = (main_width as usize) * height; + let domain = 1u32 << rest_point_dim; + let lambda = powers_of_lambda[chip_idx as usize]; + for row in 0..n_rows { + let weight = + if (row as u32) < domain { partial_lagrange[row] * lambda } else { EF::zero() }; + for t in &bc.terms { + let leaf = bc.leaves[t.leaf_idx as usize]; + let base = + if leaf.source == 4 || leaf.source == 5 { main_ptr } else { preprocessed_ptr }; + let off = base + (leaf.col as usize) * height + (row << 1); + let z = trace[off]; + let o = trace[off + 1]; + let diff = o - z; + + let coeff = runtime_coeffs[t.coeff_idx as usize]; + let alpha = powers_of_alpha[t.alpha_idx as usize]; + + for e in 0..3 { + let v = z + eval_pts[e] * diff; + let contribution = alpha * (coeff * EF::from_base(v)); + sum[e] += contribution * weight; + } + } + } + sum +} + +fn run_gpu( + bc: &ColumnTileBytecode, + trace: &[F], + height: usize, + runtime_coeffs: &[EF], + powers_of_alpha: &[EF], + partial_lagrange: &[EF], + powers_of_lambda: &[EF], + chip_idx: u32, + rest_point_dim: u32, + main_width: u32, + n_rows: usize, +) -> [EF; 3] { + const BLOCK_SIZE: u32 = 64; + let total = (bc.terms.len() as u32) * (n_rows as u32); + let grid_x = (total + BLOCK_SIZE - 1) / BLOCK_SIZE; + let n_warps = BLOCK_SIZE / 32; + let shmem_bytes = (n_warps as usize) * std::mem::size_of::(); + + let result = run_sync_in_place(|scope| { + let d_terms = DeviceBuffer::from_host_slice(&bc.terms, &scope).unwrap(); + let d_leaves = DeviceBuffer::from_host_slice(&bc.leaves, &scope).unwrap(); + // GKR chunk has no const-pool or public-pool entries; pass 1-elem dummies. + let d_consts: DeviceBuffer = + DeviceBuffer::from_host_slice(&[F::zero()], &scope).unwrap(); + let d_publics: DeviceBuffer = DeviceBuffer::from_host_slice(&[0u32], &scope).unwrap(); + let d_runtime = DeviceBuffer::from_host_slice(runtime_coeffs, &scope).unwrap(); + let d_trace = DeviceBuffer::from_host_slice(trace, &scope).unwrap(); + let d_public_values = DeviceBuffer::from_host_slice(&vec![F::zero(); 4], &scope).unwrap(); + let d_powers = DeviceBuffer::from_host_slice(powers_of_alpha, &scope).unwrap(); + let d_partial_lagrange = DeviceBuffer::from_host_slice(partial_lagrange, &scope).unwrap(); + let d_powers_of_lambda = DeviceBuffer::from_host_slice(powers_of_lambda, &scope).unwrap(); + let mut d_partials: DeviceBuffer = + DeviceBuffer::with_capacity_in(grid_x as usize * 3, scope.clone()); + unsafe { + d_partials.assume_init(); + } + + let n_terms: u32 = bc.terms.len() as u32; + let preprocessed_ptr: u64 = (main_width as u64) * (height as u64); + let main_ptr: u64 = 0; + let height_u: u32 = height as u32; + let row_start: u32 = 0; + let row_count: u32 = n_rows as u32; + + unsafe { + let a = args!( + d_terms.as_ptr(), + n_terms, + d_leaves.as_ptr(), + d_consts.as_ptr(), + d_publics.as_ptr(), + d_runtime.as_ptr(), + d_trace.as_ptr(), + preprocessed_ptr, + main_ptr, + height_u, + d_public_values.as_ptr(), + d_powers.as_ptr(), + d_partial_lagrange.as_ptr(), + d_powers_of_lambda.as_ptr(), + chip_idx, + rest_point_dim, + row_start, + row_count, + d_partials.as_mut_ptr() + ); + scope + .launch_kernel( + zerocheck_column_tile_kb_kernel(), + (grid_x, 1, 1), + (BLOCK_SIZE, 1, 1), + &a, + shmem_bytes, + ) + .unwrap(); + } + + d_partials.to_host().unwrap().to_vec() + }) + .unwrap(); + + let mut out = [EF::zero(); 3]; + for (i, &p) in result.iter().enumerate() { + out[i % 3] += p; + } + out +} diff --git a/sp1-gpu/crates/zerocheck/examples/kernel_equiv.rs b/sp1-gpu/crates/zerocheck/examples/kernel_equiv.rs new file mode 100644 index 0000000000..3fd3c4b53c --- /dev/null +++ b/sp1-gpu/crates/zerocheck/examples/kernel_equiv.rs @@ -0,0 +1,364 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! Equivalence test for the v2 Sequential kernel. +//! +//! Constructs a tiny synthetic `ConstraintDag` (no AIR-side machinery), lowers +//! it to bytecode, then evaluates it both ways: +//! (a) on the host via a faithful CPU port of the kernel logic +//! (b) on the GPU via `zerocheck_sequential_kb_kernel` +//! +//! Both inputs are identical synthetic random data. The 3 per-eval-point +//! partial sums must match bitwise. + +use rand::{rngs::StdRng, Rng, SeedableRng}; +use slop_algebra::{AbstractExtensionField, AbstractField}; +use sp1_gpu_air::ir::{ + analyze_constraints, chunk_dag, enumerate_lowerings, lower_sequential, ChunkBudget, + ConstraintDag, ConstraintField, ConstraintRef, DagNode, Lowering, TraceSource, +}; +use sp1_gpu_air::{EF, F}; +use sp1_gpu_cudart::sys::kernels::zerocheck_sequential_kb_kernel; +use sp1_gpu_cudart::{args, run_sync_in_place, DeviceBuffer}; + +// Synthetic DAG covering every opcode the Sequential kernel handles: +// +// c0(row): l0 + l1 +// c1(row): (l0 * l2) - 5 +// c2(row): (l0 + l1) + l2 // shares the (l0 + l1) sub-DAG with c0 +// c3(row): -((l0 + l1) * l3) // exercises NegF +// c4(row): l4 * 7 - (l0 + l1) // exercises Mul-by-const, deep reuse +// +// Cross-constraint sharing on node 4 (l0+l1) is the key sanity property. +fn build_synthetic_dag() -> ConstraintDag { + let mut nodes = vec![]; + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 0 }); // 0 + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 1 }); // 1 + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 2 }); // 2 + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 3 }); // 3 + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 4 }); // 4 + nodes.push(DagNode::ConstF { value: F::from_canonical_u32(5) }); // 5 + nodes.push(DagNode::ConstF { value: F::from_canonical_u32(7) }); // 6 + nodes.push(DagNode::AddF { a: 0, b: 1 }); // 7 — c0 root, also referenced by c2 and c3 and c4 + nodes.push(DagNode::MulF { a: 0, b: 2 }); // 8 + nodes.push(DagNode::SubF { a: 8, b: 5 }); // 9 — c1 root + nodes.push(DagNode::AddF { a: 7, b: 2 }); // 10 — c2 root + nodes.push(DagNode::MulF { a: 7, b: 3 }); // 11 + nodes.push(DagNode::NegF { a: 11 }); // 12 — c3 root + nodes.push(DagNode::MulF { a: 4, b: 6 }); // 13 + nodes.push(DagNode::SubF { a: 13, b: 7 }); // 14 — c4 root + + let constraints = vec![ + ConstraintRef { root: 7, alpha_index: 0, field: ConstraintField::Base }, + ConstraintRef { root: 9, alpha_index: 1, field: ConstraintField::Base }, + ConstraintRef { root: 10, alpha_index: 2, field: ConstraintField::Base }, + ConstraintRef { root: 12, alpha_index: 3, field: ConstraintField::Base }, + ConstraintRef { root: 14, alpha_index: 4, field: ConstraintField::Base }, + ]; + + ConstraintDag { nodes, constraints, preprocessed_width: 0, main_width: 5 } +} + +const N_COLS: usize = 5; +const N_PARTIAL_PER_BLOCK: usize = 3; + +fn main() { + let dag = build_synthetic_dag(); + let infos = analyze_constraints(&dag); + let chunks = chunk_dag(&infos, &ChunkBudget::recommended()); + // Shape-aware chunker may split this DAG into multiple chunks (one per + // shape category). Test all of them. + println!("DAG produced {} chunks (shape-aware split)", chunks.len()); + + // Lower each chunk's Sequential plan to bytecode. + let bcs: Vec<_> = chunks + .iter() + .map(|chunk| { + let lowerings = enumerate_lowerings(chunk, &infos, &dag); + let plan = lowerings + .iter() + .find_map(|l| match l { + Lowering::Sequential(p) => Some(p), + _ => None, + }) + .unwrap(); + let bc = lower_sequential(chunk, &infos, &dag, plan); + println!( + " chunk: {} instrs, {} leaves, {} consts, {} asserts, max_reg={}", + bc.instrs.len(), + bc.leaves.len(), + bc.consts.len(), + bc.asserts.len(), + bc.max_reg + ); + bc + }) + .collect(); + + // n_rows == 2^rest_point_dim so every row contributes. + let scenarios: &[(&str, usize, u32, u64)] = + &[("single-warp, full block", 32, 5, 0xAA), ("multi-block, aligned tail", 128, 7, 0xCC)]; + let mut all_ok = true; + for &(label, n_rows, rest_point_dim, seed) in scenarios { + let mut rng = StdRng::seed_from_u64(seed); + let height = 2 * n_rows; + let trace: Vec = (0..N_COLS * height).map(|_| rand_f(&mut rng)).collect(); + let powers_of_alpha: Vec = + (0..dag.constraints.len()).map(|_| rand_ef(&mut rng)).collect(); + let public_values: Vec = vec![F::zero(); 16]; + // partial_lagrange has 2^rest_point_dim slots. + let partial_lagrange: Vec = + (0..(1u32 << rest_point_dim)).map(|_| rand_ef(&mut rng)).collect(); + let powers_of_lambda: Vec = vec![rand_ef(&mut rng); 4]; + let chip_idx: u32 = 1; + + let mut ok = true; + for (i, bc) in bcs.iter().enumerate() { + let host = host_oracle( + bc, + &trace, + height, + &powers_of_alpha, + &public_values, + &partial_lagrange, + &powers_of_lambda, + chip_idx, + rest_point_dim, + n_rows, + ); + let gpu = run_gpu( + bc, + &trace, + height, + &powers_of_alpha, + &public_values, + &partial_lagrange, + &powers_of_lambda, + chip_idx, + rest_point_dim, + n_rows, + ); + for e in 0..N_PARTIAL_PER_BLOCK { + if host[e] != gpu[e] { + ok = false; + eprintln!(" chunk {} e{}: host={:?} gpu={:?}", i, e, host[e], gpu[e]); + } + } + } + println!( + "[{}] n_rows={} dim={} : {}", + label, + n_rows, + rest_point_dim, + if ok { "match" } else { "MISMATCH" } + ); + if !ok { + all_ok = false; + } + } + if all_ok { + println!("\nv2 Sequential kernel matches host oracle bitwise across all scenarios"); + } else { + std::process::exit(1); + } +} + +fn rand_f(rng: &mut StdRng) -> F { + F::from_wrapped_u32(rng.gen()) +} + +fn rand_ef(rng: &mut StdRng) -> EF { + EF::from_base_slice(&[rand_f(rng), rand_f(rng), rand_f(rng), rand_f(rng)]) +} + +/// Faithful CPU port of the kernel logic. Per row, per eval point: +/// - lerp every leaf reference +/// - walk bytecode +/// - sum `α^k · constraint_k` +/// Then sum across rows. +fn host_oracle( + bc: &sp1_gpu_air::ir::ChunkBytecode, + trace: &[F], + height: usize, + powers_of_alpha: &[EF], + public_values: &[F], + partial_lagrange: &[EF], + powers_of_lambda: &[EF], + chip_idx: u32, + rest_point_dim: u32, + n_rows: usize, +) -> [EF; 3] { + let mut sum = [EF::zero(); 3]; + let eval_pts: [F; 3] = + [F::from_canonical_u32(0), F::from_canonical_u32(2), F::from_canonical_u32(4)]; + let domain = 1u32 << rest_point_dim; + let lambda = powers_of_lambda[chip_idx as usize]; + for row in 0..n_rows { + let mut regs: Vec<[F; 3]> = vec![[F::zero(); 3]; bc.max_reg as usize]; + + for instr in &bc.instrs { + let op = instr.opcode; + let out = instr.out as usize; + let a = instr.a as usize; + let b = instr.b as usize; + match op { + 0 => { + let leaf = bc.leaves[a]; + let off = (leaf.col as usize) * height + (row << 1); + let z = trace[off]; + let o = trace[off + 1]; + let diff = o - z; + regs[out][0] = z + eval_pts[0] * diff; + regs[out][1] = z + eval_pts[1] * diff; + regs[out][2] = z + eval_pts[2] * diff; + } + 1 => { + let c = bc.consts[a]; + regs[out] = [c, c, c]; + } + 2 => { + let pv_idx = bc.publics[a] as usize; + let v = public_values[pv_idx]; + regs[out] = [v, v, v]; + } + 3 => { + for e in 0..3 { + regs[out][e] = regs[a][e] + regs[b][e]; + } + } + 4 => { + for e in 0..3 { + regs[out][e] = regs[a][e] - regs[b][e]; + } + } + 5 => { + for e in 0..3 { + regs[out][e] = regs[a][e] * regs[b][e]; + } + } + 6 => { + for e in 0..3 { + regs[out][e] = -regs[a][e]; + } + } + _ => panic!("unhandled opcode {}", op), + } + } + + // Per-row constraint-weighted sum. + let mut row_acc = [EF::zero(); 3]; + for &(reg, alpha_idx) in &bc.asserts { + let alpha = powers_of_alpha[alpha_idx as usize]; + for e in 0..3 { + row_acc[e] += alpha * EF::from_base(regs[reg as usize][e]); + } + } + // Apply eq[row] · λ_chip if row is in the partial-Lagrange domain. + if (row as u32) < domain { + let eq = partial_lagrange[row]; + let w = eq * lambda; + for e in 0..3 { + sum[e] += row_acc[e] * w; + } + } + } + sum +} + +fn run_gpu( + bc: &sp1_gpu_air::ir::ChunkBytecode, + trace: &[F], + height: usize, + powers_of_alpha: &[EF], + public_values: &[F], + partial_lagrange: &[EF], + powers_of_lambda: &[EF], + chip_idx: u32, + rest_point_dim: u32, + n_rows: usize, +) -> [EF; 3] { + let assert_regs: Vec = bc.asserts.iter().map(|&(r, _)| r).collect(); + let assert_alphas: Vec = bc.asserts.iter().map(|&(_, a)| a).collect(); + + let block_size: u32 = 32; + let n_blocks: u32 = ((n_rows as u32) + block_size - 1) / block_size; + let shmem_bytes: usize = (block_size as usize / 32) * std::mem::size_of::(); + + let result = run_sync_in_place(|scope| { + let d_instrs = DeviceBuffer::from_host_slice(&bc.instrs, &scope).unwrap(); + let d_leaves = DeviceBuffer::from_host_slice(&bc.leaves, &scope).unwrap(); + let d_consts = DeviceBuffer::from_host_slice(&bc.consts, &scope).unwrap(); + let d_publics = DeviceBuffer::from_host_slice(&bc.publics, &scope).unwrap(); + let d_assert_regs = DeviceBuffer::from_host_slice(&assert_regs, &scope).unwrap(); + let d_assert_alphas = DeviceBuffer::from_host_slice(&assert_alphas, &scope).unwrap(); + let d_trace = DeviceBuffer::from_host_slice(trace, &scope).unwrap(); + let d_public_values = DeviceBuffer::from_host_slice(public_values, &scope).unwrap(); + let d_powers = DeviceBuffer::from_host_slice(powers_of_alpha, &scope).unwrap(); + let d_partial_lagrange = DeviceBuffer::from_host_slice(partial_lagrange, &scope).unwrap(); + let d_powers_of_lambda = DeviceBuffer::from_host_slice(powers_of_lambda, &scope).unwrap(); + let mut d_partials: DeviceBuffer = + DeviceBuffer::with_capacity_in(n_blocks as usize * 3, scope.clone()); + unsafe { + d_partials.assume_init(); + } + + let n_instrs: u32 = bc.instrs.len() as u32; + let n_asserts: u32 = bc.asserts.len() as u32; + let preprocessed_ptr: u64 = 0; + let main_ptr: u64 = 0; + let height_u: u32 = height as u32; + let row_start: u32 = 0; + let row_count: u32 = n_rows as u32; + + unsafe { + let a = args!( + d_instrs.as_ptr(), + n_instrs, + d_leaves.as_ptr(), + d_consts.as_ptr(), + d_publics.as_ptr(), + d_assert_regs.as_ptr(), + d_assert_alphas.as_ptr(), + n_asserts, + d_trace.as_ptr(), + preprocessed_ptr, + main_ptr, + height_u, + d_public_values.as_ptr(), + d_powers.as_ptr(), + d_partial_lagrange.as_ptr(), + d_powers_of_lambda.as_ptr(), + chip_idx, + rest_point_dim, + row_start, + row_count, + d_partials.as_mut_ptr() + ); + scope + .launch_kernel( + zerocheck_sequential_kb_kernel(), + (n_blocks, 1, 1), + (block_size, 1, 1), + &a, + shmem_bytes, + ) + .unwrap(); + } + + d_partials.to_host().unwrap().to_vec() + }) + .unwrap(); + + let mut out = [EF::zero(); 3]; + for (i, &p) in result.iter().enumerate() { + out[i % 3] += p; + } + out +} diff --git a/sp1-gpu/crates/zerocheck/examples/pipeline_equiv.rs b/sp1-gpu/crates/zerocheck/examples/pipeline_equiv.rs new file mode 100644 index 0000000000..02f274824c --- /dev/null +++ b/sp1-gpu/crates/zerocheck/examples/pipeline_equiv.rs @@ -0,0 +1,524 @@ +#![allow( + clippy::print_stdout, + clippy::print_stderr, + clippy::too_many_arguments, + clippy::needless_range_loop, + clippy::vec_init_then_push, + clippy::useless_vec, + clippy::manual_div_ceil, + clippy::doc_lazy_continuation +)] +//! End-to-end pipeline equivalence test. +//! +//! Drives one round's worth of v2 zerocheck for a synthetic chip: +//! 1. Build a `ConstraintDag` with constraints of both shapes (General + +//! LinearWeightedSum), plus an injected GKR-correction chunk. +//! 2. Run the full pipeline: analyze → chunk → lowerings → bytecode. +//! 3. Launch the right kernel per chunk (Sequential or ColumnTile). +//! 4. Sum all partials across all chunks per eval point. +//! 5. Compute the same quantity host-side as ground truth. +//! Verify bitwise match. + +use rand::{rngs::StdRng, Rng, SeedableRng}; +use slop_algebra::{AbstractExtensionField, AbstractField}; +use sp1_gpu_air::ir::{ + analyze_constraints, chunk_dag, enumerate_lowerings, lower_column_tile, lower_sequential, + synthesize_gkr_chunk, ChunkBudget, ChunkBytecode, ColumnTileBytecode, ConstraintDag, + ConstraintField, ConstraintRef, DagNode, Lowering, TraceSource, COEFF_KIND_CONST, + COEFF_KIND_PUBLIC, COEFF_KIND_RUNTIME, +}; +use sp1_gpu_air::{EF, F}; +use sp1_gpu_cudart::sys::kernels::{ + zerocheck_column_tile_kb_kernel, zerocheck_sequential_kb_kernel, +}; +use sp1_gpu_cudart::{args, run_sync_in_place, DeviceBuffer, TaskScope}; + +// Same DAG as v2_kernel_equiv: mix of LinearWeightedSum and General constraints. +fn build_chip_dag() -> ConstraintDag { + let mut nodes = vec![]; + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 0 }); // 0 + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 1 }); // 1 + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 2 }); // 2 + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 3 }); // 3 + nodes.push(DagNode::InputLeaf { source: TraceSource::MainLocal, col: 4 }); // 4 + nodes.push(DagNode::ConstF { value: F::from_canonical_u32(5) }); // 5 + nodes.push(DagNode::ConstF { value: F::from_canonical_u32(7) }); // 6 + nodes.push(DagNode::AddF { a: 0, b: 1 }); // 7 + nodes.push(DagNode::MulF { a: 0, b: 2 }); // 8 + nodes.push(DagNode::SubF { a: 8, b: 5 }); // 9 + nodes.push(DagNode::AddF { a: 7, b: 2 }); // 10 + nodes.push(DagNode::MulF { a: 7, b: 3 }); // 11 + nodes.push(DagNode::NegF { a: 11 }); // 12 + nodes.push(DagNode::MulF { a: 4, b: 6 }); // 13 + nodes.push(DagNode::SubF { a: 13, b: 7 }); // 14 + + let constraints = vec![ + ConstraintRef { root: 7, alpha_index: 0, field: ConstraintField::Base }, + ConstraintRef { root: 9, alpha_index: 1, field: ConstraintField::Base }, + ConstraintRef { root: 10, alpha_index: 2, field: ConstraintField::Base }, + ConstraintRef { root: 12, alpha_index: 3, field: ConstraintField::Base }, + ConstraintRef { root: 14, alpha_index: 4, field: ConstraintField::Base }, + ]; + ConstraintDag { nodes, constraints, preprocessed_width: 0, main_width: 5 } +} + +const MAIN_WIDTH: u32 = 5; +const PREP_WIDTH: u32 = 0; +// NOTE: `synthesize_gkr_chunk` now emits `alpha_idx = 0` and flags the chunk +// `is_gkr_carrier`; the real launcher (`launch_chunk_into`) points +// `powers_of_alpha` at the `1` slot for such chunks. This standalone example +// still models the older "reserved slot" convention — its GKR equivalence +// check is stale and needs reworking to mirror the launcher. +const GKR_ALPHA_IDX: u32 = 5; // beyond the constraint alphas +const N_ALPHAS: usize = 6; + +enum CompiledChunk { + Sequential(ChunkBytecode), + ColumnTile(ColumnTileBytecode), +} + +fn main() { + let dag = build_chip_dag(); + let infos = analyze_constraints(&dag); + let chunks = chunk_dag(&infos, &ChunkBudget::recommended()); + + // Compile every chunk to its appropriate bytecode. + let mut compiled: Vec = Vec::new(); + for chunk in &chunks { + let lowerings = enumerate_lowerings(chunk, &infos, &dag); + // Prefer ColumnTile when it applies; otherwise Sequential. + let column_tile = lowerings.iter().find_map(|l| match l { + Lowering::ColumnTile(p) => Some(p), + _ => None, + }); + if let Some(plan) = column_tile { + if let Some(bc) = lower_column_tile(chunk, &infos, &dag, plan) { + compiled.push(CompiledChunk::ColumnTile(bc)); + continue; + } + } + let plan = lowerings + .iter() + .find_map(|l| match l { + Lowering::Sequential(p) => Some(p), + _ => None, + }) + .unwrap(); + let bc = lower_sequential(chunk, &infos, &dag, plan); + compiled.push(CompiledChunk::Sequential(bc)); + } + // Append the synthesized GKR chunk. + let gkr_bc = synthesize_gkr_chunk(MAIN_WIDTH, PREP_WIDTH); + compiled.push(CompiledChunk::ColumnTile(gkr_bc)); + + println!("compiled {} chunks:", compiled.len()); + for (i, c) in compiled.iter().enumerate() { + match c { + CompiledChunk::Sequential(bc) => println!( + " [{}] Sequential: {} instrs, {} asserts, max_reg={}", + i, + bc.instrs.len(), + bc.asserts.len(), + bc.max_reg + ), + CompiledChunk::ColumnTile(bc) => println!( + " [{}] ColumnTile: {} terms ({} consts/publics, {} runtime)", + i, + bc.terms.len(), + bc.terms.iter().filter(|t| t.coeff_kind != COEFF_KIND_RUNTIME).count(), + bc.terms.iter().filter(|t| t.coeff_kind == COEFF_KIND_RUNTIME).count(), + ), + } + } + + let scenarios: &[(&str, usize, u32, u64)] = + &[("medium chip, n_rows=32", 32, 5, 0xE0), ("larger chip, n_rows=128", 128, 7, 0xE1)]; + + let mut all_ok = true; + for &(label, n_rows, rest_point_dim, seed) in scenarios { + let mut rng = StdRng::seed_from_u64(seed); + let height = 2 * n_rows; + let trace: Vec = (0..(MAIN_WIDTH as usize) * height).map(|_| rand_f(&mut rng)).collect(); + let powers_of_alpha: Vec = (0..N_ALPHAS).map(|_| rand_ef(&mut rng)).collect(); + // Reserve a slot at GKR_ALPHA_IDX = 5, set it to one() so GKR is + // additive (matches v1 semantics). + let mut powers_of_alpha = powers_of_alpha; + powers_of_alpha[GKR_ALPHA_IDX as usize] = EF::one(); + let partial_lagrange: Vec = + (0..(1u32 << rest_point_dim)).map(|_| rand_ef(&mut rng)).collect(); + let powers_of_lambda: Vec = vec![rand_ef(&mut rng); 4]; + let runtime_coeffs: Vec = + (0..(MAIN_WIDTH + PREP_WIDTH) as usize).map(|_| rand_ef(&mut rng)).collect(); + let public_values: Vec = vec![F::zero(); 16]; + let chip_idx: u32 = 2; + + let host = full_host_oracle( + &dag, + &trace, + height, + &powers_of_alpha, + &public_values, + &partial_lagrange, + &powers_of_lambda, + chip_idx, + rest_point_dim, + n_rows, + &runtime_coeffs, + MAIN_WIDTH, + ); + let gpu = full_pipeline_gpu( + &compiled, + &trace, + height, + &powers_of_alpha, + &public_values, + &partial_lagrange, + &powers_of_lambda, + chip_idx, + rest_point_dim, + n_rows, + &runtime_coeffs, + MAIN_WIDTH, + ); + + let ok = (0..3).all(|e| host[e] == gpu[e]); + println!("[{}] dim={} : {}", label, rest_point_dim, if ok { "match" } else { "MISMATCH" }); + if !ok { + for e in 0..3 { + println!(" e{}: host={:?} gpu={:?}", e, host[e], gpu[e]); + } + all_ok = false; + } + } + + if all_ok { + println!("\nv2 full pipeline matches host oracle bitwise across all scenarios"); + } else { + std::process::exit(1); + } +} + +fn rand_f(rng: &mut StdRng) -> F { + F::from_wrapped_u32(rng.gen()) +} + +fn rand_ef(rng: &mut StdRng) -> EF { + EF::from_base_slice(&[rand_f(rng), rand_f(rng), rand_f(rng), rand_f(rng)]) +} + +/// Full host oracle: evaluates the entire DAG plus the synthesized GKR +/// correction, weighted by eq · λ, summed across rows. The kernel pipeline +/// should produce the same value. +fn full_host_oracle( + dag: &ConstraintDag, + trace: &[F], + height: usize, + powers_of_alpha: &[EF], + _public_values: &[F], + partial_lagrange: &[EF], + powers_of_lambda: &[EF], + chip_idx: u32, + rest_point_dim: u32, + n_rows: usize, + runtime_coeffs: &[EF], + main_width: u32, +) -> [EF; 3] { + let mut sum = [EF::zero(); 3]; + let eval_pts: [F; 3] = + [F::from_canonical_u32(0), F::from_canonical_u32(2), F::from_canonical_u32(4)]; + let domain = 1u32 << rest_point_dim; + let lambda = powers_of_lambda[chip_idx as usize]; + let gkr_alpha = powers_of_alpha[GKR_ALPHA_IDX as usize]; + + for row in 0..n_rows { + // Evaluate the DAG at this row for each eval point. + let mut vals: Vec<[F; 3]> = vec![[F::zero(); 3]; dag.nodes.len()]; + for (i, node) in dag.nodes.iter().enumerate() { + match *node { + DagNode::InputLeaf { source, col } => { + assert!(matches!(source, TraceSource::MainLocal)); + let off = (col as usize) * height + (row << 1); + let z = trace[off]; + let o = trace[off + 1]; + let diff = o - z; + vals[i][0] = z + eval_pts[0] * diff; + vals[i][1] = z + eval_pts[1] * diff; + vals[i][2] = z + eval_pts[2] * diff; + } + DagNode::ConstF { value } => vals[i] = [value, value, value], + DagNode::AddF { a, b } => { + for e in 0..3 { + vals[i][e] = vals[a as usize][e] + vals[b as usize][e]; + } + } + DagNode::SubF { a, b } => { + for e in 0..3 { + vals[i][e] = vals[a as usize][e] - vals[b as usize][e]; + } + } + DagNode::MulF { a, b } => { + for e in 0..3 { + vals[i][e] = vals[a as usize][e] * vals[b as usize][e]; + } + } + DagNode::NegF { a } => { + for e in 0..3 { + vals[i][e] = -vals[a as usize][e]; + } + } + _ => panic!("oracle: unhandled node {:?}", node), + } + } + + // Per-row weighted constraint sum. + let weight = + if (row as u32) < domain { partial_lagrange[row] * lambda } else { EF::zero() }; + + // Constraint contributions. + let mut row_acc = [EF::zero(); 3]; + for c in &dag.constraints { + let alpha = powers_of_alpha[c.alpha_index as usize]; + for e in 0..3 { + row_acc[e] += alpha * EF::from_base(vals[c.root as usize][e]); + } + } + + // GKR contribution: Σ_i runtime_coeffs[i] · col_i(row, eval) · α_gkr. + // Same column ordering as `synthesize_gkr_chunk`: main first, then prep. + for col in 0..main_width { + let off = (col as usize) * height + (row << 1); + let z = trace[off]; + let o = trace[off + 1]; + let diff = o - z; + let coeff = runtime_coeffs[col as usize]; + for e in 0..3 { + let v = z + eval_pts[e] * diff; + row_acc[e] += gkr_alpha * (coeff * EF::from_base(v)); + } + } + // (No preprocessed columns in this test.) + + for e in 0..3 { + sum[e] += row_acc[e] * weight; + } + } + sum +} + +fn full_pipeline_gpu( + compiled: &[CompiledChunk], + trace: &[F], + height: usize, + powers_of_alpha: &[EF], + public_values: &[F], + partial_lagrange: &[EF], + powers_of_lambda: &[EF], + chip_idx: u32, + rest_point_dim: u32, + n_rows: usize, + runtime_coeffs: &[EF], + _main_width: u32, +) -> [EF; 3] { + run_sync_in_place(|scope| { + // Upload all the per-round buffers once. + let d_trace = DeviceBuffer::from_host_slice(trace, &scope).unwrap(); + let d_public_values = DeviceBuffer::from_host_slice(public_values, &scope).unwrap(); + let d_powers = DeviceBuffer::from_host_slice(powers_of_alpha, &scope).unwrap(); + let d_partial_lagrange = DeviceBuffer::from_host_slice(partial_lagrange, &scope).unwrap(); + let d_powers_of_lambda = DeviceBuffer::from_host_slice(powers_of_lambda, &scope).unwrap(); + let d_runtime_coeffs = DeviceBuffer::from_host_slice(runtime_coeffs, &scope).unwrap(); + + let mut total = [EF::zero(); 3]; + for chunk in compiled { + let partials = match chunk { + CompiledChunk::Sequential(bc) => launch_sequential( + &scope, + bc, + &d_trace, + &d_public_values, + &d_powers, + &d_partial_lagrange, + &d_powers_of_lambda, + chip_idx, + rest_point_dim, + height, + n_rows, + ), + CompiledChunk::ColumnTile(bc) => launch_column_tile( + &scope, + bc, + &d_trace, + &d_public_values, + &d_powers, + &d_partial_lagrange, + &d_powers_of_lambda, + &d_runtime_coeffs, + chip_idx, + rest_point_dim, + height, + n_rows, + ), + }; + for (i, &p) in partials.iter().enumerate() { + total[i % 3] += p; + } + } + total + }) + .unwrap() +} + +fn launch_sequential( + scope: &TaskScope, + bc: &ChunkBytecode, + d_trace: &DeviceBuffer, + d_public_values: &DeviceBuffer, + d_powers: &DeviceBuffer, + d_partial_lagrange: &DeviceBuffer, + d_powers_of_lambda: &DeviceBuffer, + chip_idx: u32, + rest_point_dim: u32, + height: usize, + n_rows: usize, +) -> Vec { + let d_instrs = DeviceBuffer::from_host_slice(&bc.instrs, scope).unwrap(); + let d_leaves = DeviceBuffer::from_host_slice(&bc.leaves, scope).unwrap(); + let d_consts = DeviceBuffer::from_host_slice(&bc.consts, scope).unwrap(); + let d_publics = DeviceBuffer::from_host_slice(&bc.publics, scope).unwrap(); + let assert_regs: Vec = bc.asserts.iter().map(|&(r, _)| r).collect(); + let assert_alphas: Vec = bc.asserts.iter().map(|&(_, a)| a).collect(); + let d_assert_regs = DeviceBuffer::from_host_slice(&assert_regs, scope).unwrap(); + let d_assert_alphas = DeviceBuffer::from_host_slice(&assert_alphas, scope).unwrap(); + + let block_size: u32 = 32; + let n_blocks: u32 = (n_rows as u32).div_ceil(block_size); + let shmem_bytes = (block_size as usize / 32) * std::mem::size_of::(); + let mut d_partials: DeviceBuffer = + DeviceBuffer::with_capacity_in(n_blocks as usize * 3, scope.clone()); + unsafe { + d_partials.assume_init(); + } + + let n_instrs: u32 = bc.instrs.len() as u32; + let n_asserts: u32 = bc.asserts.len() as u32; + let preprocessed_ptr: u64 = 0; + let main_ptr: u64 = 0; + let height_u: u32 = height as u32; + let row_start: u32 = 0; + let row_count: u32 = n_rows as u32; + + unsafe { + let a = args!( + d_instrs.as_ptr(), + n_instrs, + d_leaves.as_ptr(), + d_consts.as_ptr(), + d_publics.as_ptr(), + d_assert_regs.as_ptr(), + d_assert_alphas.as_ptr(), + n_asserts, + d_trace.as_ptr(), + preprocessed_ptr, + main_ptr, + height_u, + d_public_values.as_ptr(), + d_powers.as_ptr(), + d_partial_lagrange.as_ptr(), + d_powers_of_lambda.as_ptr(), + chip_idx, + rest_point_dim, + row_start, + row_count, + d_partials.as_mut_ptr() + ); + scope + .launch_kernel( + zerocheck_sequential_kb_kernel(), + (n_blocks, 1, 1), + (block_size, 1, 1), + &a, + shmem_bytes, + ) + .unwrap(); + } + d_partials.to_host().unwrap().to_vec() +} + +fn launch_column_tile( + scope: &TaskScope, + bc: &ColumnTileBytecode, + d_trace: &DeviceBuffer, + d_public_values: &DeviceBuffer, + d_powers: &DeviceBuffer, + d_partial_lagrange: &DeviceBuffer, + d_powers_of_lambda: &DeviceBuffer, + d_runtime_coeffs: &DeviceBuffer, + chip_idx: u32, + rest_point_dim: u32, + height: usize, + n_rows: usize, +) -> Vec { + let d_terms = DeviceBuffer::from_host_slice(&bc.terms, scope).unwrap(); + let d_leaves = DeviceBuffer::from_host_slice(&bc.leaves, scope).unwrap(); + // Empty const/public pools can't be zero-length (gives null ptr); pad with 1. + let consts = if bc.consts.is_empty() { vec![F::zero()] } else { bc.consts.clone() }; + let publics = if bc.publics.is_empty() { vec![0u32] } else { bc.publics.clone() }; + let d_consts = DeviceBuffer::from_host_slice(&consts, scope).unwrap(); + let d_publics = DeviceBuffer::from_host_slice(&publics, scope).unwrap(); + + const BLOCK_SIZE: u32 = 64; + let total = (bc.terms.len() as u32) * (n_rows as u32); + let grid_x = total.div_ceil(BLOCK_SIZE); + let n_warps = BLOCK_SIZE / 32; + let shmem_bytes = (n_warps as usize) * std::mem::size_of::(); + let mut d_partials: DeviceBuffer = + DeviceBuffer::with_capacity_in(grid_x as usize * 3, scope.clone()); + unsafe { + d_partials.assume_init(); + } + + let n_terms: u32 = bc.terms.len() as u32; + let preprocessed_ptr: u64 = 0; + let main_ptr: u64 = 0; + let height_u: u32 = height as u32; + let row_start: u32 = 0; + let row_count: u32 = n_rows as u32; + + // Sanity check: kinds present. + let _has_const = bc.terms.iter().any(|t| t.coeff_kind == COEFF_KIND_CONST); + let _has_pub = bc.terms.iter().any(|t| t.coeff_kind == COEFF_KIND_PUBLIC); + + unsafe { + let a = args!( + d_terms.as_ptr(), + n_terms, + d_leaves.as_ptr(), + d_consts.as_ptr(), + d_publics.as_ptr(), + d_runtime_coeffs.as_ptr(), + d_trace.as_ptr(), + preprocessed_ptr, + main_ptr, + height_u, + d_public_values.as_ptr(), + d_powers.as_ptr(), + d_partial_lagrange.as_ptr(), + d_powers_of_lambda.as_ptr(), + chip_idx, + rest_point_dim, + row_start, + row_count, + d_partials.as_mut_ptr() + ); + scope + .launch_kernel( + zerocheck_column_tile_kb_kernel(), + (grid_x, 1, 1), + (BLOCK_SIZE, 1, 1), + &a, + shmem_bytes, + ) + .unwrap(); + } + d_partials.to_host().unwrap().to_vec() +} diff --git a/sp1-gpu/crates/zerocheck/src/lib.rs b/sp1-gpu/crates/zerocheck/src/lib.rs index bf929185e1..3b501b958f 100644 --- a/sp1-gpu/crates/zerocheck/src/lib.rs +++ b/sp1-gpu/crates/zerocheck/src/lib.rs @@ -1,526 +1,11 @@ -use data::JaggedDenseInfo; use itertools::Itertools; -use primitives::{ - evaluate_jagged_fix_last_variable, evaluate_jagged_info_fix_last_variable, - initialize_jagged_dense_info, JaggedFixLastVariableKernel, -}; -use slop_air::BaseAir; -use slop_algebra::{ - interpolate_univariate_polynomial, AbstractExtensionField, AbstractField, ExtensionField, - Field, UnivariatePolynomial, -}; -use slop_alloc::{Backend, Buffer, CpuBackend, HasBackend}; -use slop_challenger::{FieldChallenger, VariableLengthChallenger}; -use slop_matrix::dense::RowMajorMatrixView; -use slop_multilinear::{Point, VirtualGeq}; -use slop_sumcheck::PartialSumcheckProof; -use slop_tensor::Tensor; -use sp1_gpu_air::instruction::Instruction16; -use sp1_gpu_air::{air_block::BlockAir, SymbolicProverFolder}; -use sp1_gpu_cudart::sys::runtime::KernelPtr; -use sp1_gpu_cudart::sys::v2_kernels::{ - jagged_constraint_poly_eval_1024_koala_bear_extension_kernel, - jagged_constraint_poly_eval_1024_koala_bear_kernel, - jagged_constraint_poly_eval_128_koala_bear_extension_kernel, - jagged_constraint_poly_eval_128_koala_bear_kernel, - jagged_constraint_poly_eval_256_koala_bear_extension_kernel, - jagged_constraint_poly_eval_256_koala_bear_kernel, - jagged_constraint_poly_eval_32_koala_bear_extension_kernel, - jagged_constraint_poly_eval_32_koala_bear_kernel, - jagged_constraint_poly_eval_512_koala_bear_extension_kernel, - jagged_constraint_poly_eval_512_koala_bear_kernel, - jagged_constraint_poly_eval_64_koala_bear_extension_kernel, - jagged_constraint_poly_eval_64_koala_bear_kernel, -}; -use sp1_gpu_cudart::{args, DeviceBuffer, DevicePoint, DeviceTensor, TaskScope}; -use sp1_gpu_utils::{Ext, Felt, JaggedTraceMle}; -use sp1_hypercube::air::MachineAir; -use sp1_hypercube::prover::ZerocheckAir; -use sp1_hypercube::{ - AirOpenedValues, Chip, ChipEvaluation, ChipOpenedValues, ConstraintSumcheckFolder, - LogUpEvaluations, ShardOpenedValues, -}; -use std::borrow::Cow; -use std::collections::{BTreeMap, BTreeSet}; +use slop_algebra::{AbstractExtensionField, UnivariatePolynomial}; +use slop_challenger::FieldChallenger; +use sp1_gpu_utils::{Ext, Felt}; pub mod data; pub mod primitives; - -pub struct EvalProgramInfo { - pub constraint_indices: Buffer, - pub operations: Buffer, - pub operations_indices: Buffer, - pub f_ctr: u32, - pub f_constants: Buffer, - pub f_constants_indices: Buffer, - pub ef_constants: Buffer, - pub ef_constants_indices: Buffer, -} - -pub type CudaEvalResult = - (Vec, Vec, Vec, Vec, Vec, Vec, Vec, u32, u32); - -pub struct ZeroCheckJaggedPoly<'a, K: Field> { - /// The program for **all chips**. - pub program: EvalProgramInfo, - /// The data in a `JaggedTraceMle` form. - pub data: Cow<'a, JaggedTraceMle>, - /// The information in a `JaggedDenseInfo` form. - pub info: JaggedDenseInfo, - /// The `VirtualGeq` for each table. - pub virtual_geq: Vec>, - /// The `eq_adjustment`, identical to the `ZeroCheckPoly`. - pub eq_adjustment: Ext, - /// The `padded_row_adjustment` for each table. - pub padded_row_adjustment: Buffer, - /// The random evaluation point, from the GKR. - pub zeta: Point, - /// The claimed evaluation. - pub claim: Ext, - /// The number of preprocessed columns. - pub total_num_preprocessed_column: u32, - /// The public values. - pub public_values: Buffer, - /// The powers of alpha. - pub powers_of_alpha: Buffer, - /// The gkr powers. - pub gkr_powers: Buffer, - /// The powers of lambda. - pub powers_of_lambda: Buffer, - /// The number of preprocessed column per chip. - pub preprocessed_column: Buffer, - /// The number of main column per chips. - pub main_column: Buffer, - /// The chips. - pub chips_info: Vec<(u32, u32, u32)>, - /// The total length of the info data. - pub total_len: usize, -} - -pub fn initialize_program_cpu( - chips: &BTreeSet>, - zerocheck_programs: &BTreeMap, -) -> EvalProgramInfo -where - A: for<'a> BlockAir>, -{ - let mut constraint_indices: Vec = vec![]; - let mut operations: Vec = vec![]; - let mut operations_indices: Vec = vec![]; - let mut f_ctr: u32 = 0; - let mut ef_ctr: u32 = 0; - let mut f_constants: Vec = vec![]; - let mut f_constants_indices: Vec = vec![]; - let mut ef_constants: Vec = vec![]; - let mut ef_constants_indices: Vec = vec![]; - - let max_num_constraints = - itertools::max(chips.iter().map(|chip| chip.num_constraints)).unwrap(); - - for chip in chips { - let ( - chip_constraint_indices, - chip_operations, - chip_operations_indices, - chip_f_constants, - chip_f_constants_indices, - chip_ef_constants, - chip_ef_constants_indices, - chip_f_ctr, - chip_ef_ctr, - ) = zerocheck_programs.get(chip.air.name()).unwrap_or_else(|| { - panic!("Chip name {} not found in CUDA eval cache", chip.air.name()); - }); - - for constraint_index in chip_constraint_indices { - constraint_indices - .push((max_num_constraints - chip.num_constraints) as u32 + constraint_index); - } - - let current_instruction_len = operations.len() as u32; - operations.extend(chip_operations); - for idx in chip_operations_indices.iter() { - operations_indices.push(current_instruction_len + idx); - } - - let current_f_constants_len = f_constants.len() as u32; - f_constants.extend(chip_f_constants); - for idx in chip_f_constants_indices.iter() { - f_constants_indices.push(current_f_constants_len + idx); - } - - let current_ef_constants_len = ef_constants.len() as u32; - ef_constants.extend(chip_ef_constants); - for idx in chip_ef_constants_indices.iter() { - ef_constants_indices.push(current_ef_constants_len + idx); - } - - f_ctr = f_ctr.max(*chip_f_ctr); - ef_ctr = ef_ctr.max(*chip_ef_ctr); - } - operations_indices.push(operations.len() as u32); - - EvalProgramInfo { - constraint_indices: Buffer::from(constraint_indices), - operations: Buffer::from(operations), - operations_indices: Buffer::from(operations_indices), - f_ctr, - f_constants: Buffer::from(f_constants), - f_constants_indices: Buffer::from(f_constants_indices), - ef_constants: Buffer::from(ef_constants), - ef_constants_indices: Buffer::from(ef_constants_indices), - } -} - -pub fn initialize_program( - chips: &BTreeSet>, - zerocheck_programs: &BTreeMap, - scope: &TaskScope, -) -> EvalProgramInfo -where - A: for<'a> BlockAir>, -{ - let cpu_info = initialize_program_cpu(chips, zerocheck_programs); - - let constraint_indices_device = - DeviceBuffer::from_host(&cpu_info.constraint_indices, scope).unwrap().into_inner(); - let operations_device = - DeviceBuffer::from_host(&cpu_info.operations, scope).unwrap().into_inner(); - let operations_indices_device = - DeviceBuffer::from_host(&cpu_info.operations_indices, scope).unwrap().into_inner(); - let f_constants_device = - DeviceBuffer::from_host(&cpu_info.f_constants, scope).unwrap().into_inner(); - let f_constants_indices_device = - DeviceBuffer::from_host(&cpu_info.f_constants_indices, scope).unwrap().into_inner(); - let ef_constants_device = - DeviceBuffer::from_host(&cpu_info.ef_constants, scope).unwrap().into_inner(); - let ef_constants_indices_device = - DeviceBuffer::from_host(&cpu_info.ef_constants_indices, scope).unwrap().into_inner(); - - EvalProgramInfo { - constraint_indices: constraint_indices_device, - operations: operations_device, - operations_indices: operations_indices_device, - f_constants: f_constants_device, - f_constants_indices: f_constants_indices_device, - f_ctr: cpu_info.f_ctr, - ef_constants: ef_constants_device, - ef_constants_indices: ef_constants_indices_device, - } -} - -/// The format is `is_first || chip_idx || prep_start || main_start` in little endian. -/// `is_first` is 1 bit, `chip_idx` is 15 bits, `prep_start` is 16 bits, and `main_start` is 32 bits. -pub fn pack_info(is_first: bool, chip_idx: u32, prep_start: u32, main_start: u32) -> u64 { - (is_first as u64) - + ((chip_idx as u64) << 1) - + ((prep_start as u64) << 16) - + ((main_start as u64) << 32) -} - -/// The `packed_info` is `is_first || chip_idx || prep_start || main_start` in little endian. -/// `is_first` is 1 bit, `chip_idx` is 15 bits, `prep_start` is 16 bits, and `main_start` is 32 bits. -pub fn unpack_info(packed_info: u64) -> (u32, u32, u32, u32) { - ( - (packed_info & 1) as u32, - ((packed_info >> 1) & 0x7FFF) as u32, - ((packed_info >> 16) & 0xFFFF) as u32, - ((packed_info >> 32) & 0xFFFFFFFF) as u32, - ) -} - -#[allow(clippy::type_complexity)] -pub fn initialize_dense_info( - chips: &BTreeSet>, - initial_heights: &[u32], - backend: &TaskScope, -) -> (JaggedDenseInfo, u32, usize, Vec<(u32, u32, u32)>) -where - A: for<'a> BlockAir>, -{ - let mut tot_len: usize = 0; - let mut total_num_block: usize = 0; - let mut chips_info = vec![]; - for (chip, height) in chips.iter().zip_eq(initial_heights.iter()) { - let num_air_blocks = chip.air.clone().num_blocks(); - assert_eq!(height % 4, 0, "invariant height % 4 == 0 invalidated"); - tot_len += (*height as usize) * num_air_blocks; - total_num_block += num_air_blocks; - chips_info.push(( - num_air_blocks as u32, - chip.preprocessed_width() as u32, - chip.width() as u32, - )); - } - - let mut info_data = vec![0; total_num_block]; - let mut info_heights = vec![0; total_num_block]; - let mut air_block_idx: usize = 0; - let mut total_preprocessed: u32 = 0; - let mut total_main: u32 = 0; - for (chip_idx, (chip, height)) in chips.iter().zip_eq(initial_heights.iter()).enumerate() { - let num_air_blocks = chip.air.clone().num_blocks(); - for air_block in 0..num_air_blocks { - let data = pack_info(air_block == 0, chip_idx as u32, total_preprocessed, total_main); - info_data[air_block_idx] = data; - info_heights[air_block_idx] = height / 2; - air_block_idx += 1; - } - total_preprocessed += chip.air.clone().preprocessed_width() as u32; - total_main += chip.air.clone().width() as u32; - } - - let info = initialize_jagged_dense_info(info_heights, info_data, backend); - (info, total_preprocessed, tot_len, chips_info) -} - -#[allow(clippy::too_many_arguments)] -pub fn initialize_zerocheck_poly<'b, A>( - data: &'b JaggedTraceMle, - chips: &BTreeSet>, - zerocheck_programs: &BTreeMap, - initial_heights: Vec, - public_values: Vec, - powers_of_alpha: Vec, - gkr_powers: Vec, - powers_of_lambda: Vec, - padded_row_adjustment: Vec, - zeta: Point, - claim: Ext, -) -> ZeroCheckJaggedPoly<'b, Felt> -where - A: for<'a> BlockAir>, -{ - let scope = data.dense().backend(); - let program = initialize_program(chips, zerocheck_programs, scope); - let mut virtual_geq: Vec> = vec![]; - let mut info_input_heights = vec![]; - - for (chip, height) in chips.iter().zip_eq(initial_heights.iter()) { - virtual_geq.push(VirtualGeq::new( - *height, - Ext::one(), - Ext::zero(), - zeta.dimension() as u32, - )); - let num_air_blocks = chip.air.clone().num_blocks(); - for _ in 0..num_air_blocks { - info_input_heights.push(height / 2); - } - } - - let preprocessed_columns = - chips.iter().map(|chip| chip.preprocessed_width() as u32).collect::>(); - let main_columns = chips.iter().map(|chip| chip.width() as u32).collect::>(); - - let (info, total_preprocessed, tot_len, chips_info) = - initialize_dense_info(chips, &initial_heights, scope); - - let padded_row_adjustment = - DeviceBuffer::from_host(&Buffer::from(padded_row_adjustment), scope).unwrap().into_inner(); - let public_values_device = - DeviceBuffer::from_host(&Buffer::from(public_values), scope).unwrap().into_inner(); - let powers_of_alpha_device = - DeviceBuffer::from_host(&Buffer::from(powers_of_alpha), scope).unwrap().into_inner(); - let gkr_powers_device = - DeviceBuffer::from_host(&Buffer::from(gkr_powers), scope).unwrap().into_inner(); - let powers_of_lambda_device = - DeviceBuffer::from_host(&Buffer::from(powers_of_lambda), scope).unwrap().into_inner(); - let preprocessed_columns_device = - DeviceBuffer::from_host(&preprocessed_columns, scope).unwrap().into_inner(); - let main_columns_device = DeviceBuffer::from_host(&main_columns, scope).unwrap().into_inner(); - - ZeroCheckJaggedPoly { - program, - data: Cow::Borrowed(data), - info, - virtual_geq, - eq_adjustment: Ext::one(), - padded_row_adjustment, - zeta, - claim, - total_num_preprocessed_column: total_preprocessed, - public_values: public_values_device, - powers_of_alpha: powers_of_alpha_device, - gkr_powers: gkr_powers_device, - powers_of_lambda: powers_of_lambda_device, - preprocessed_column: preprocessed_columns_device, - main_column: main_columns_device, - chips_info, - total_len: tot_len, - } -} - -pub trait JaggedConstraintPolyEvalKernel { - fn jagged_constraint_poly_eval_kernel(memory_size: usize) -> KernelPtr; -} - -impl JaggedConstraintPolyEvalKernel for TaskScope { - fn jagged_constraint_poly_eval_kernel(memory_size: usize) -> KernelPtr { - match memory_size { - 0..=32 => unsafe { jagged_constraint_poly_eval_32_koala_bear_kernel() }, - 33..=64 => unsafe { jagged_constraint_poly_eval_64_koala_bear_kernel() }, - 65..=128 => unsafe { jagged_constraint_poly_eval_128_koala_bear_kernel() }, - 129..=256 => unsafe { jagged_constraint_poly_eval_256_koala_bear_kernel() }, - 257..=512 => unsafe { jagged_constraint_poly_eval_512_koala_bear_kernel() }, - 513..=1024 => unsafe { jagged_constraint_poly_eval_1024_koala_bear_kernel() }, - _ => unreachable!(), - } - } -} - -impl JaggedConstraintPolyEvalKernel for TaskScope { - fn jagged_constraint_poly_eval_kernel(memory_size: usize) -> KernelPtr { - match memory_size { - 0..=32 => unsafe { jagged_constraint_poly_eval_32_koala_bear_extension_kernel() }, - 33..=64 => unsafe { jagged_constraint_poly_eval_64_koala_bear_extension_kernel() }, - 65..=128 => unsafe { jagged_constraint_poly_eval_128_koala_bear_extension_kernel() }, - 129..=256 => unsafe { jagged_constraint_poly_eval_256_koala_bear_extension_kernel() }, - 257..=512 => unsafe { jagged_constraint_poly_eval_512_koala_bear_extension_kernel() }, - 513..=1024 => unsafe { jagged_constraint_poly_eval_1024_koala_bear_extension_kernel() }, - _ => unreachable!(), - } - } -} - -pub fn evaluate_zerocheck<'b, K: Field>( - input: &'b ZeroCheckJaggedPoly<'b, K>, -) -> UnivariatePolynomial -where - TaskScope: JaggedConstraintPolyEvalKernel, -{ - let backend = input.data.backend(); - const BLOCK_SIZE: usize = 256; - const NUM_EVAL_POINT: usize = 3; - - let n_chunks = input.total_len.div_ceil(1 << 12); - let grid_size_x = n_chunks.max(256); - let grid_size = (grid_size_x, 1, NUM_EVAL_POINT); - - let num_tiles = BLOCK_SIZE.div_ceil(32); - let shared_mem = num_tiles * std::mem::size_of::(); - - let mut output: Tensor = - Tensor::with_sizes_in([NUM_EVAL_POINT, grid_size_x], backend.clone()); - - let (rest, last) = input.zeta.split_at(input.zeta.dimension() - 1); - let last = *last[0]; - let thresholds = input.virtual_geq.iter().map(|geq| geq.threshold).collect::>(); - let eq_coefficients = - input.virtual_geq.iter().map(|geq| geq.eq_coefficient).collect::>(); - - let rest_point = DevicePoint::from_host(&rest, backend).unwrap(); - let thresholds = DeviceBuffer::from_host(&thresholds, backend).unwrap().into_inner(); - let eq_coefficients = DeviceBuffer::from_host(&eq_coefficients, backend).unwrap().into_inner(); - - let partial_lagrange = rest_point.partial_lagrange(); - let rest_point_dim = rest.dimension() as u32; - - unsafe { - output.assume_init(); - let args = args!( - input.program.constraint_indices.as_ptr(), - input.program.operations, - input.program.operations_indices.as_ptr(), - input.program.f_constants.as_ptr(), - input.program.f_constants_indices.as_ptr(), - input.program.ef_constants.as_ptr(), - input.program.ef_constants_indices.as_ptr(), - input.data.as_raw(), - input.info.as_raw(), - partial_lagrange.as_ptr(), - thresholds.as_ptr(), - eq_coefficients.as_ptr(), - (input.total_len / 2) as u32, - input.padded_row_adjustment.as_ptr(), - input.public_values.as_ptr(), - input.powers_of_alpha.as_ptr(), - input.gkr_powers.as_ptr(), - input.powers_of_lambda.as_ptr(), - input.preprocessed_column.as_ptr(), - input.main_column.as_ptr(), - input.total_num_preprocessed_column, - output.as_mut_ptr(), - rest_point_dim - ); - backend - .launch_kernel( - >::jagged_constraint_poly_eval_kernel( - input.program.f_ctr as usize, - ), - grid_size, - (BLOCK_SIZE, 1, 1), - &args, - shared_mem, - ) - .unwrap(); - } - - let output_eval = DeviceTensor::from_raw(output).sum_dim(1); - let result = output_eval.to_host().unwrap().into_buffer().into_vec(); - - let mut xs = - vec![Ext::from_canonical_u32(0), Ext::from_canonical_u32(2), Ext::from_canonical_u32(4)]; - - let mut ys = result - .iter() - .zip_eq(xs.iter()) - .map(|(&result, &x)| { - let last_var_eq = (Ext::one() - x) * (Ext::one() - last) + x * last; - result * last_var_eq * input.eq_adjustment - }) - .collect::>(); - - xs.push(Ext::from_canonical_u32(1)); - ys.push(input.claim - ys[0]); - - xs.push((last - Ext::one()) / (last + last - Ext::one())); - ys.push(Ext::zero()); - - interpolate_univariate_polynomial(&xs, &ys) -} - -pub fn zerocheck_fix_last_variable<'b, K: Field>( - input: ZeroCheckJaggedPoly<'b, K>, - point: Ext, - claim: Ext, -) -> ZeroCheckJaggedPoly<'b, Ext> -where - TaskScope: JaggedFixLastVariableKernel, - Ext: ExtensionField, -{ - let (rest, last) = input.zeta.split_at(input.zeta.dimension() - 1); - let last = *last[0]; - - let new_data = evaluate_jagged_fix_last_variable(&input.data, point); - let new_info = evaluate_jagged_info_fix_last_variable(input.info); - - let virtual_geq = - input.virtual_geq.iter().map(|geq| geq.fix_last_variable(point)).collect::>(); - let eq = (Ext::one() - last) * (Ext::one() - point) + last * point; - let eq_adjustment = input.eq_adjustment * eq; - let total_len = 2 * new_info.column_heights.iter().sum::() as usize; - - ZeroCheckJaggedPoly { - program: input.program, - data: Cow::Owned(new_data), - info: new_info, - virtual_geq, - eq_adjustment, - padded_row_adjustment: input.padded_row_adjustment, - zeta: rest, - claim, - total_num_preprocessed_column: input.total_num_preprocessed_column, - public_values: input.public_values, - powers_of_alpha: input.powers_of_alpha, - gkr_powers: input.gkr_powers, - powers_of_lambda: input.powers_of_lambda, - preprocessed_column: input.preprocessed_column, - main_column: input.main_column, - chips_info: input.chips_info, - total_len, - } -} +pub mod prover; pub fn challenger_update( input_poly: &UnivariatePolynomial, @@ -537,187 +22,13 @@ where (point, claim) } -#[allow(clippy::too_many_arguments)] -pub fn zerocheck( - chips: &BTreeSet>, - zerocheck_programs: &BTreeMap, - trace_mle: &JaggedTraceMle, - batching_challenge: Ext, - gkr_opening_batch_randomness: Ext, - logup_evaluations: &LogUpEvaluations, - public_values: Vec, - challenger: &mut C, - max_log_row_count: u32, -) -> (ShardOpenedValues, PartialSumcheckProof) -where - A: ZerocheckAir + for<'a> BlockAir>, - C: FieldChallenger, -{ - let data_input_heights = &trace_mle.column_heights; - let initial_heights = trace_mle - .dense_data - .main_table_index - .values() - .map(|trace_offset| trace_offset.poly_size as u32) - .collect::>(); - - let max_num_constraints = - itertools::max(chips.iter().map(|chip| chip.num_constraints)).unwrap(); - let max_columns = - itertools::max(chips.iter().map(|chip| chip.preprocessed_width() + chip.width())).unwrap(); - let total_preprocessed_columns = trace_mle.dense().preprocessed_cols; - let mut powers_of_challenge = - batching_challenge.powers().take(max_num_constraints).collect::>(); - powers_of_challenge.reverse(); - let num_chips = chips.len(); - - let mut padded_row_adjustment = vec![Ext::zero(); num_chips]; - - for (i, chip) in chips.iter().enumerate() { - let prep_len = chip.preprocessed_width(); - let main_len = chip.width(); - let prep_zero = vec![Felt::zero(); prep_len]; - let main_zero = vec![Felt::zero(); main_len]; - let mut folder = ConstraintSumcheckFolder { - preprocessed: RowMajorMatrixView::new_row(&prep_zero), - main: RowMajorMatrixView::new_row(&main_zero), - accumulator: Ext::zero(), - public_values: &public_values, - constraint_index: 0, - powers_of_alpha: &powers_of_challenge - [powers_of_challenge.len() - chip.num_constraints..], - }; - chip.air.eval(&mut folder); - padded_row_adjustment[i] = folder.accumulator; - } - - let gkr_powers = - gkr_opening_batch_randomness.powers().skip(1).take(max_columns).collect::>(); - - let lambda: Ext = challenger.sample_ext_element(); - let powers_of_lambda = - lambda.powers().take(num_chips).collect_vec().into_iter().rev().collect(); - - let mut claim = Ext::zero(); - - let LogUpEvaluations { point: gkr_point, chip_openings } = logup_evaluations; - - for chip in chips.iter() { - let ChipEvaluation { - main_trace_evaluations: main_opening, - preprocessed_trace_evaluations: prep_opening, - } = chip_openings.get(chip.name()).unwrap(); - - claim *= lambda; - - let addend = main_opening - .evaluations() - .as_slice() - .iter() - .chain( - prep_opening - .as_ref() - .map_or_else(Vec::new, |mle| mle.evaluations().as_slice().to_vec()) - .iter(), - ) - .zip(gkr_powers.iter()) - .map(|(opening, power)| *opening * *power) - .sum::(); - - claim += addend; - } - - let main_poly = initialize_zerocheck_poly( - trace_mle, - chips, - zerocheck_programs, - initial_heights.clone(), - public_values, - powers_of_challenge, - gkr_powers, - powers_of_lambda, - padded_row_adjustment, - gkr_point.clone(), - claim, - ); - - let mut univariate_polys = vec![]; - let mut jagged_point: Point = Point::from(vec![]); - let mut result = evaluate_zerocheck(&main_poly); - let (mut point, mut next_claim) = challenger_update(&result, challenger); - univariate_polys.push(result); - jagged_point.add_dimension(point); - let mut next_poly = zerocheck_fix_last_variable(main_poly, point, next_claim); - for _ in 0..max_log_row_count - 1 { - result = evaluate_zerocheck(&next_poly); - (point, next_claim) = challenger_update(&result, challenger); - univariate_polys.push(result); - jagged_point.add_dimension(point); - next_poly = zerocheck_fix_last_variable(next_poly, point, next_claim); - } - - let final_jagged_data = - unsafe { next_poly.data.as_ref().dense_data.dense.copy_into_host_vec() }; - - let mut idx = 0; - let mut individual_column_evals = vec![Ext::zero(); data_input_heights.len()]; - for i in 0..data_input_heights.len() { - if data_input_heights[i] != 0 { - individual_column_evals[i] = final_jagged_data[idx]; - idx += 4; - } - } - - let mut preprocessed_ptr = 0; - let mut main_ptr = total_preprocessed_columns; - let mut opened_values: BTreeMap> = BTreeMap::new(); - challenger.observe(Felt::from_canonical_usize(chips.len())); - for (i, chip) in chips.iter().enumerate() { - let preprocessed_width = chip.preprocessed_width(); - let preprocessed = AirOpenedValues { - local: individual_column_evals[preprocessed_ptr..preprocessed_ptr + preprocessed_width] - .to_vec(), - }; - challenger.observe_variable_length_extension_slice(&preprocessed.local); - preprocessed_ptr += preprocessed_width; - - let width = chip.width(); - - let main = - AirOpenedValues { local: individual_column_evals[main_ptr..main_ptr + width].to_vec() }; - challenger.observe_variable_length_extension_slice(&main.local); - main_ptr += width; - - opened_values.insert( - chip.air.name().to_string(), - ChipOpenedValues { - preprocessed, - main, - degree: Point::from_usize( - initial_heights[i] as usize, - (max_log_row_count + 1) as usize, - ), - }, - ); - } - - let partial_sumcheck_proof = PartialSumcheckProof { - univariate_polys, - claimed_sum: claim, - point_and_eval: (jagged_point, next_claim), - }; - - let shard_open_values = ShardOpenedValues { chips: opened_values }; - - (shard_open_values, partial_sumcheck_proof) -} - #[cfg(test)] pub mod tests { + #![allow(clippy::print_stdout)] use itertools::Itertools; use rand::Rng; use serial_test::serial; - use slop_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder}; + use slop_air::{Air, BaseAir, PairBuilder}; use slop_algebra::{AbstractField, PrimeField32}; use slop_alloc::{Buffer, CpuBackend}; use slop_challenger::{ @@ -732,11 +43,10 @@ pub mod tests { use sp1_hypercube::air::{MachineAir, SP1AirBuilder}; use sp1_hypercube::prover::ZerocheckAir; use sp1_hypercube::{ - prover::ProverSemaphore, Chip, ChipEvaluation, ChipOpenedValues, ConstraintSumcheckFolder, - LogUpEvaluations, ShardOpenedValues, VerifierConstraintFolder, + prover::ProverSemaphore, Chip, ChipEvaluation, ChipOpenedValues, LogUpEvaluations, + ShardOpenedValues, VerifierConstraintFolder, }; - use sp1_gpu_air::codegen_cuda_eval; use sp1_primitives::SP1Field; use std::collections::{BTreeMap, BTreeSet}; @@ -753,18 +63,12 @@ pub mod tests { use sp1_gpu_utils::{Ext, Felt, JaggedTraceMle, TestGC, TraceDenseData, TraceOffset}; use super::primitives::evaluate_jagged_columns; - use super::{ - challenger_update, evaluate_zerocheck, initialize_zerocheck_poly, zerocheck, - zerocheck_fix_last_variable, - }; + use super::prover::{compile_chips, upload_compiled_bytecode, zerocheck, CompiledChip}; + use sp1_gpu_air::ir::ChunkBudget; use core::{borrow::Borrow, mem::size_of}; use sp1_core_executor::{ExecutionRecord, Program}; use sp1_derive::AlignedBorrow; - use sp1_gpu_air::{ - air_block::BlockAir, symbolic_expr_f::SymbolicExprF, symbolic_var_f::SymbolicVarF, - SymbolicProverFolder, - }; #[derive(Debug)] pub enum ZerocheckTestChip { @@ -841,24 +145,6 @@ pub mod tests { } } - impl<'a> BlockAir> for ZerocheckTestChip { - fn eval_block(&self, builder: &mut SymbolicProverFolder<'a>, index: usize) { - match self { - Self::Chip1(chip) => chip.eval_block(builder, index), - Self::Chip2(chip) => chip.eval_block(builder, index), - Self::Chip3(chip) => chip.eval_block(builder, index), - } - } - - fn num_blocks(&self) -> usize { - match self { - Self::Chip1(chip) => chip.num_blocks(), - Self::Chip2(chip) => chip.num_blocks(), - Self::Chip3(chip) => chip.num_blocks(), - } - } - } - #[derive(Default, Clone)] pub struct ZerocheckTestChip1; @@ -886,37 +172,6 @@ pub mod tests { } } - impl<'a> BlockAir> for ZerocheckTestChip1 { - fn num_blocks(&self) -> usize { - 2 - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder, index: usize) { - let main = builder.main(); - let local = main.row_slice(0); - let local: &ZerocheckTestCols1 = (*local).borrow(); - - match index { - 0 => { - builder.assert_zero( - local.op_a + SymbolicExprF::from_canonical_u32(3) * local.op_b - - (local.op_b + local.op_c + SymbolicExprF::one()) - * (local.op_b + local.op_c + SymbolicExprF::two()) - * (local.op_b - local.op_c + SymbolicExprF::from_canonical_u32(8)), - ); - } - 1 => { - builder.assert_zero( - local.op_d - * (local.op_d - SymbolicExprF::one()) - * (local.op_d - SymbolicExprF::two()), - ); - } - _ => unreachable!(), - } - } - } - impl MachineAir for ZerocheckTestChip1 { type Record = ExecutionRecord; type Program = Program; @@ -1001,49 +256,6 @@ pub mod tests { } } - impl<'a> BlockAir> for ZerocheckTestChip2 { - fn num_blocks(&self) -> usize { - 2 - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder, index: usize) { - let main = builder.main(); - let local = main.row_slice(0); - let local: &ZerocheckTestCols2 = (*local).borrow(); - - match index { - 0 => { - builder.assert_zero( - local.op_a + SymbolicExprF::from_canonical_u32(5) * local.op_b - - (local.op_b + local.op_c + SymbolicExprF::two()) - * (local.op_b - + SymbolicExprF::from_canonical_u32(3) * local.op_c - + SymbolicExprF::one()) - * (local.op_b - local.op_c + SymbolicExprF::from_canonical_u32(10)), - ); - } - 1 => { - builder.assert_zero( - (local.op_d + SymbolicExprF::one()) - * (local.op_d - SymbolicExprF::one()) - * (local.op_d - SymbolicExprF::two()), - ); - - builder.assert_zero( - local.op_b - - local.op_c * local.op_d * SymbolicExprF::from_canonical_u32(5) - - SymbolicExprF::from_canonical_u32(8) - * local.op_d - * local.op_d - * local.op_d - - SymbolicExprF::one(), - ); - } - _ => unreachable!(), - } - } - } - impl MachineAir for ZerocheckTestChip2 { type Record = ExecutionRecord; type Program = Program; @@ -1148,56 +360,6 @@ pub mod tests { } } - impl<'a> BlockAir> for ZerocheckTestChip3 { - fn num_blocks(&self) -> usize { - 3 - } - - fn eval_block(&self, builder: &mut SymbolicProverFolder, index: usize) { - let main = builder.main(); - let local = main.row_slice(0); - let local: &ZerocheckTestCols3 = (*local).borrow(); - - let prep = builder.preprocessed(); - let prep = prep.row_slice(0); - let prep: &ZerocheckTestPrepCols3 = (*prep).borrow(); - - let pv = builder.public_values(); - let pv_0 = pv[0]; - let pv_1 = pv[1]; - - match index { - 0 => { - builder.assert_zero( - prep.prep_a - - (local.op_a * local.op_a * local.op_b - + SymbolicExprF::one() - + SymbolicExprF::from_canonical_u32(3) * pv_0 * local.op_c), - ); - } - 1 => { - builder.assert_zero( - prep.prep_b - - (SymbolicExprF::from_canonical_u32(8) * prep.prep_a * local.op_c - + pv_0 * local.op_a * local.op_b - + SymbolicExprF::from_canonical_u32(17)), - ); - } - 2 => { - builder.assert_zero( - local.op_a - - (local.op_b * local.op_c * SymbolicExprF::from_canonical_u32(8) - + local.op_b * local.op_b * local.op_b - + local.op_c * local.op_c * local.op_c - + SymbolicExprF::from_canonical_u32(178) - + pv_1), - ); - } - _ => unreachable!(), - } - } - } - impl MachineAir for ZerocheckTestChip3 { type Record = ExecutionRecord; type Program = Program; @@ -1351,7 +513,7 @@ pub mod tests { challenger: &mut C, max_log_row_count: usize, ) where - A: MachineAir + ZerocheckAir + for<'a> BlockAir>, + A: MachineAir + ZerocheckAir, C: FieldChallenger, { // Get the random challenge to merge the constraints. @@ -1597,7 +759,7 @@ pub mod tests { public_values: &[Felt], ) -> JaggedTraceMle where - A: for<'a> BlockAir>, + A: MachineAir, { let mut rng = rand::thread_rng(); let total_main = @@ -1728,190 +890,40 @@ pub mod tests { } } + /// Scaling sanity check for the future "machine with thousands of chips" + /// goal: replicate the compiled test chips up to 5000 and confirm the + /// once-per-machine flat upload (`upload_compiled_bytecode`) handles it. #[test] #[serial] - fn test_zerocheck() { + fn test_machine_bytecode_scaling() { let mut chips: BTreeSet> = BTreeSet::new(); chips.insert(Chip::new(ZerocheckTestChip::Chip1(ZerocheckTestChip1))); chips.insert(Chip::new(ZerocheckTestChip::Chip2(ZerocheckTestChip2))); chips.insert(Chip::new(ZerocheckTestChip::Chip3(ZerocheckTestChip3))); - - let mut cache = BTreeMap::new(); - for chip in chips.iter() { - let result = codegen_cuda_eval(chip.air.as_ref()); - cache.insert(chip.name().to_string(), result); + let base = compile_chips(&chips, ChunkBudget::recommended()); + + // Replicate the compiled chips to a machine-sized chip count, each + // with a distinct name (the chunker output is identical — we only + // exercise the flatten + upload path at scale). + const TARGET: usize = 5000; + let mut big: Vec = Vec::with_capacity(TARGET); + for i in 0..TARGET { + let mut c = base[i % base.len()].clone(); + c.chip_idx = i as u32; + c.name = format!("{}_{}", c.name, i); + big.push(c); } - let chips_vec = chips.iter().cloned().collect::>(); - let num_chips = chips_vec.len(); - let row_variables = 22; run_sync_in_place(move |t| { - let input_size = get_input_sizes(); - let mut rng = rand::thread_rng(); - let public_values = vec![random_felt(&mut rng), random_felt(&mut rng)]; - - let trace_mle = get_input(&input_size, &chips_vec, &public_values); - let trace_mle = Arc::new(trace_mle.into_device(&t)); - let initial_heights = input_size.clone(); - - let mut challenger = TestGC::default_challenger(); - let _lambda: Ext = challenger.sample(); - - let alpha = challenger.sample(); - let beta = challenger.sample(); - let lambda: Ext = challenger.sample(); - - let max_num_constraints = - itertools::max(chips.iter().map(|chip| chip.num_constraints)).unwrap(); - let mut powers_of_alpha = vec![Ext::one(); max_num_constraints]; - for i in 1..max_num_constraints { - powers_of_alpha[i] = powers_of_alpha[i - 1] * alpha; - } - powers_of_alpha.reverse(); - - let mut gkr_powers = vec![Ext::zero(); 1024]; - gkr_powers[0] = beta; - for i in 1..1024 { - gkr_powers[i] = gkr_powers[i - 1] * beta; - } - - let mut powers_of_lambda = vec![Ext::one(); num_chips]; - for i in 1..num_chips { - powers_of_lambda[i] = powers_of_lambda[i - 1] * lambda; - } - powers_of_lambda.reverse(); - - let mut padded_row_adjustment = vec![Ext::zero(); num_chips]; - - for i in 0..num_chips { - let prep_len = chips_vec[i].preprocessed_width(); - let main_len = chips_vec[i].width(); - let prep_zero = vec![Felt::zero(); prep_len]; - let main_zero = vec![Felt::zero(); main_len]; - let mut folder = ConstraintSumcheckFolder { - preprocessed: RowMajorMatrixView::new_row(&prep_zero), - main: RowMajorMatrixView::new_row(&main_zero), - accumulator: Ext::zero(), - public_values: &public_values, - constraint_index: 0, - powers_of_alpha: &powers_of_alpha - [powers_of_alpha.len() - chips_vec[i].num_constraints..], - }; - chips_vec[i].air.eval(&mut folder); - padded_row_adjustment[i] = folder.accumulator; - } - - let zeta = Point::::rand(&mut rng, row_variables); - let individual_column_evals = evaluate_jagged_columns(&trace_mle, zeta.clone()); - let mut claim = Ext::zero(); - let mut preprocessed_ptr: usize = 0; - let mut main_ptr = chips_vec.iter().map(|x| x.preprocessed_width()).sum::() + 1; - - for chip in chips_vec.iter() { - let preprocessed_width = chip.preprocessed_width(); - let main_width = chip.width(); - claim *= lambda; - for idx in 0..preprocessed_width { - claim += gkr_powers[main_width + idx] - * individual_column_evals[preprocessed_ptr + idx]; - } - for idx in 0..main_width { - claim += gkr_powers[idx] * individual_column_evals[main_ptr + idx]; - } - preprocessed_ptr += preprocessed_width; - main_ptr += main_width; - } - - let main_poly = initialize_zerocheck_poly( - trace_mle.as_ref(), - &chips, - &cache, - initial_heights.clone(), - public_values.clone(), - powers_of_alpha.clone(), - gkr_powers.clone(), - powers_of_lambda.clone(), - padded_row_adjustment.clone(), - zeta.clone(), - claim, + let start = std::time::Instant::now(); + let mb = upload_compiled_bytecode(big, &t); + println!( + "upload_compiled_bytecode: {} chips uploaded in {:?}", + mb.chips.len(), + start.elapsed(), ); - - let mut jagged_point = vec![]; - let mut result = evaluate_zerocheck(&main_poly); - let (mut point, mut claim) = challenger_update(&result, &mut challenger); - jagged_point.insert(0, point); - let mut next_poly = zerocheck_fix_last_variable(main_poly, point, claim); - for _ in 0..21 { - result = evaluate_zerocheck(&next_poly); - (point, claim) = challenger_update(&result, &mut challenger); - jagged_point.insert(0, point); - next_poly = zerocheck_fix_last_variable(next_poly, point, claim); - } - let result = unsafe { next_poly.data.as_ref().dense_data.dense.copy_into_host_vec() }; - let mut idx = 0; - let data_input_heights = &trace_mle.column_heights; - let mut individual_column_evals = vec![Ext::zero(); data_input_heights.len()]; - for i in 0..data_input_heights.len() { - if data_input_heights[i] != 0 { - individual_column_evals[i] = result[idx]; - idx += 4; - } - } - - let mut jagged_point = Point::from(jagged_point); - - let mut expected_final_claim = Ext::zero(); - let mut preprocessed_ptr: usize = 0; - let mut main_ptr = chips_vec.iter().map(|x| x.preprocessed_width()).sum::() + 1; - - let eq_mul = Mle::full_lagrange_eval(&zeta, &jagged_point); - jagged_point.add_dimension(Ext::zero()); - - for (i, chip) in chips_vec.iter().enumerate() { - let preprocessed_width = chip.preprocessed_width(); - let main_width = chip.width(); - expected_final_claim *= lambda; - for idx in 0..preprocessed_width { - expected_final_claim += gkr_powers[main_width + idx] - * individual_column_evals[preprocessed_ptr + idx]; - } - for idx in 0..main_width { - expected_final_claim += - gkr_powers[idx] * individual_column_evals[main_ptr + idx]; - } - - let mut folder = VerifierConstraintFolder:: { - preprocessed: RowMajorMatrixView::new_row( - &individual_column_evals - [preprocessed_ptr..preprocessed_ptr + preprocessed_width], - ), - main: RowMajorMatrixView::new_row( - &individual_column_evals[main_ptr..main_ptr + main_width], - ), - alpha, - accumulator: Ext::zero(), - public_values: &public_values, - _marker: PhantomData, - }; - chip.air.eval(&mut folder); - expected_final_claim += folder.accumulator; - - expected_final_claim -= padded_row_adjustment[i] - * full_geq( - &Point::::from_usize( - initial_heights[i] as usize, - row_variables as usize + 1, - ), - &jagged_point, - ); - preprocessed_ptr += preprocessed_width; - main_ptr += main_width; - } - - expected_final_claim *= eq_mul; - assert_eq!(claim, expected_final_claim); - - t.synchronize_blocking().unwrap(); + assert_eq!(mb.chips.len(), TARGET); + assert_eq!(mb.chip_index.len(), TARGET); }) .unwrap(); } @@ -1924,14 +936,20 @@ pub mod tests { chips.insert(Chip::new(ZerocheckTestChip::Chip2(ZerocheckTestChip2))); chips.insert(Chip::new(ZerocheckTestChip::Chip3(ZerocheckTestChip3))); - let mut cache = BTreeMap::new(); - for chip in chips.iter() { - let result = codegen_cuda_eval(chip.air.as_ref()); - cache.insert(chip.name().to_string(), result); - } let chips_vec = chips.iter().cloned().collect::>(); let row_variables = 22; + run_sync_in_place(move |t| { + // Build machine bytecode as a strict SUPERSET of the shard: a + // fake extra chip is prepended so the real chips sit at machine + // indices 1,2,3 while the shard sees them at 0,1,2. Exercises + // the machine-⊋-shard selection path that e2e proving hits. + let mut machine_compiled = compile_chips(&chips, ChunkBudget::recommended()); + let mut fake = machine_compiled[0].clone(); + fake.name = "ZZZ_fake_extra_chip".to_string(); + let mut superset = vec![fake]; + superset.append(&mut machine_compiled); + let machine_bytecode = Arc::new(upload_compiled_bytecode(superset, &t)); let input_size = get_input_sizes(); let mut rng = rand::thread_rng(); let public_values = vec![random_felt(&mut rng), random_felt(&mut rng)]; @@ -1963,7 +981,6 @@ pub mod tests { for chip in chips_vec.iter() { let preprocessed_width = chip.preprocessed_width(); let main_width = chip.width(); - let chip_eval = ChipEvaluation { preprocessed_trace_evaluations: match preprocessed_width { 0 => None, @@ -1977,7 +994,6 @@ pub mod tests { individual_column_evals[main_ptr..main_ptr + main_width].to_vec(), )), }; - chip_openings.insert( >::name( &chip.air, @@ -1993,7 +1009,7 @@ pub mod tests { let (opened_values, zerocheck_proof) = zerocheck( &chips, - &cache, + &machine_bytecode, trace_mle.as_ref(), batching_challenge, gkr_opening_batch_randomness, @@ -2017,6 +1033,11 @@ pub mod tests { .unwrap(); } + /// Zerocheck on real RISC-V traces. Builds the machine bytecode from the + /// WHOLE machine (machine ⊋ shard cluster) — mirroring what the e2e + /// prover does — proves the shard cluster against it, and verifies. Also + /// proves against a 5000-chip padded machine to confirm per-shard work is + /// pay-per-use (independent of machine size). #[tokio::test] #[serial] async fn test_zerocheck_real_traces() { @@ -2045,39 +1066,39 @@ pub mod tests { ) .await; let chips = machine.smallest_cluster(&chips).unwrap(); - let mut cache = BTreeMap::new(); - for chip in chips.iter() { - let result = codegen_cuda_eval(chip.air.as_ref()); - cache.insert(chip.name().to_string(), result); - } - let trace_mle = Arc::new(trace_mle); - let mut challenger = TestGC::default_challenger(); - challenger.observe(Felt::from_canonical_u32(0x2013)); - challenger.observe(Felt::from_canonical_u32(0x2015)); - challenger.observe(Felt::from_canonical_u32(0x2016)); - challenger.observe(Felt::from_canonical_u32(0x2023)); - challenger.observe(Felt::from_canonical_u32(0x2024)); - - let _lambda: Ext = challenger.sample(); - - let mut challenger_prover = challenger.clone(); - let batching_challenge = challenger_prover.sample_ext_element(); - let gkr_opening_batch_randomness = challenger_prover.sample_ext_element(); - let max_log_row_count = CORE_MAX_LOG_ROW_COUNT; + // Build two machine bytecodes from the whole machine: one sized + // to the real chip set, and one padded to 5000 chips with + // distinct-named clones (simulating an auto-generated many-chip + // machine). The shard cluster is selected by name from each. + let machine_chip_set: BTreeSet<_> = machine.chips().iter().cloned().collect(); + let base_compiled = compile_chips(&machine_chip_set, ChunkBudget::recommended()); + let real_n = base_compiled.len(); + let mut padded = base_compiled.clone(); + let mut k = 0; + while padded.len() < 5000 { + let mut c = padded[k % real_n].clone(); + c.name = format!("{}__pad{}", c.name, padded.len()); + padded.push(c); + k += 1; + } + println!( + "5000-chip pay-per-use test: cluster={} chips, machine sizes {} and {}", + chips.len(), + real_n, + padded.len(), + ); + // The verifier's chip openings, derived from the committed trace. let zeta = Point::::rand(&mut rng, CORE_MAX_LOG_ROW_COUNT); let individual_column_evals = evaluate_jagged_columns(&trace_mle, zeta.clone()); - let mut preprocessed_ptr: usize = 0; let mut main_ptr = chips.iter().map(|x| x.preprocessed_width()).sum::() + 1; - let mut chip_openings: BTreeMap> = BTreeMap::new(); for chip in chips.iter() { let preprocessed_width = chip.preprocessed_width(); let main_width = chip.width(); - let chip_eval = ChipEvaluation { preprocessed_trace_evaluations: match preprocessed_width { 0 => None, @@ -2091,35 +1112,65 @@ pub mod tests { individual_column_evals[main_ptr..main_ptr + main_width].to_vec(), )), }; - chip_openings.insert(chip.air.name().to_string(), chip_eval); preprocessed_ptr += preprocessed_width; main_ptr += main_width; } - let logup_evaluations = LogUpEvaluations { point: zeta, chip_openings }; - let (opened_values, zerocheck_proof) = zerocheck( - chips, - &cache, - trace_mle.as_ref(), - batching_challenge, - gkr_opening_batch_randomness, - &logup_evaluations, - public_values.clone(), - &mut challenger_prover, - max_log_row_count, + let mut challenger = TestGC::default_challenger(); + challenger.observe(Felt::from_canonical_u32(0x2013)); + challenger.observe(Felt::from_canonical_u32(0x2015)); + challenger.observe(Felt::from_canonical_u32(0x2016)); + + // Prove the shard cluster against a given machine bytecode and + // verify; returns the wall time of the `zerocheck` call. + let run = |machine_compiled: Vec| { + let mb = Arc::new(upload_compiled_bytecode(machine_compiled, &t)); + let mut challenger_prover = challenger.clone(); + let batching_challenge = challenger_prover.sample_ext_element(); + let gkr_opening_batch_randomness = challenger_prover.sample_ext_element(); + + let start = std::time::Instant::now(); + let (opened_values, zerocheck_proof) = zerocheck( + chips, + &mb, + trace_mle.as_ref(), + batching_challenge, + gkr_opening_batch_randomness, + &logup_evaluations, + public_values.clone(), + &mut challenger_prover, + CORE_MAX_LOG_ROW_COUNT, + ); + let elapsed = start.elapsed(); + + let mut challenger_verifier = challenger.clone(); + crate::tests::verify_zerocheck( + chips, + &opened_values, + &logup_evaluations, + zerocheck_proof, + &public_values, + &mut challenger_verifier, + CORE_MAX_LOG_ROW_COUNT as usize, + ); + elapsed + }; + + let t_small = run(base_compiled); + let t_5000 = run(padded); + println!( + " zerocheck: {}-chip machine = {:?}, 5000-chip machine = {:?}", + real_n, t_small, t_5000, ); - let mut challenger_verifier = challenger.clone(); - crate::tests::verify_zerocheck( - chips, - &opened_values, - &logup_evaluations, - zerocheck_proof, - &public_values, - &mut challenger_verifier, - max_log_row_count as usize, + // Pay-per-use: the 5000-chip machine must not slow the per-shard + // proof down meaningfully vs the small machine. + let slower = t_5000.as_secs_f64() / t_small.as_secs_f64(); + assert!( + slower < 1.5, + "per-shard work scaled with machine size ({slower:.2}x) — not pay-per-use", ); }) .await; diff --git a/sp1-gpu/crates/zerocheck/src/primitives.rs b/sp1-gpu/crates/zerocheck/src/primitives.rs index a3acf70e0b..2abfd16063 100644 --- a/sp1-gpu/crates/zerocheck/src/primitives.rs +++ b/sp1-gpu/crates/zerocheck/src/primitives.rs @@ -6,11 +6,11 @@ use slop_commit::Rounds; use slop_multilinear::{MleEval, Point}; use slop_tensor::{Tensor, TensorView}; -use sp1_gpu_cudart::sys::runtime::KernelPtr; -use sp1_gpu_cudart::sys::v2_kernels::{ +use sp1_gpu_cudart::sys::kernels::{ fix_last_variable_jagged_ext, fix_last_variable_jagged_felt, fix_last_variable_jagged_info, initialize_jagged_info, }; +use sp1_gpu_cudart::sys::runtime::KernelPtr; use sp1_gpu_cudart::{ args, dot_along_dim_view, DeviceBuffer, DevicePoint, DeviceTensor, TaskScope, }; @@ -411,7 +411,7 @@ mod tests { use slop_multilinear::Mle; use slop_multilinear::Point; use sp1_gpu_cudart::run_sync_in_place; - use sp1_gpu_cudart::sys::v2_kernels::jagged_eval_kernel_chunked_felt; + use sp1_gpu_cudart::sys::kernels::jagged_eval_kernel_chunked_felt; use sp1_gpu_cudart::{DeviceBuffer, DevicePoint}; use sp1_hypercube::log2_ceil_usize; use sp1_primitives::SP1Field; diff --git a/sp1-gpu/crates/zerocheck/src/prover.rs b/sp1-gpu/crates/zerocheck/src/prover.rs new file mode 100644 index 0000000000..b2bbe902fa --- /dev/null +++ b/sp1-gpu/crates/zerocheck/src/prover.rs @@ -0,0 +1,1547 @@ +//! The GPU zerocheck prover. +//! +//! Constraint evaluation runs through the DAG-native IR in `sp1-gpu-air::ir`: +//! each chip's AIR is compiled once to flat bytecode (`compile_chips` + +//! `upload_machine_bytecode`), and the per-round sumcheck (`zerocheck` → +//! `evaluate_zerocheck` / `zerocheck_fix_last_variable`) interprets that +//! bytecode in fused CUDA kernels (`sp1-gpu-sys` `zerocheck/` kernels). + +use std::borrow::Cow; +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use itertools::Itertools; +use slop_air::{Air, BaseAir}; +use slop_algebra::{ + interpolate_univariate_polynomial, AbstractField, ExtensionField, Field, UnivariatePolynomial, +}; +use slop_alloc::{Buffer, HasBackend}; +use slop_challenger::{FieldChallenger, VariableLengthChallenger}; +use slop_matrix::dense::RowMajorMatrixView; +use slop_multilinear::{Point, VirtualGeq}; +use slop_sumcheck::PartialSumcheckProof; +use slop_tensor::Tensor; + +use sp1_gpu_air::ir::{ + analyze_constraints, build_dag, chunk_dag, enumerate_lowerings, lower_column_tile, + lower_sequential, synthesize_gkr_chunk, ChunkBudget, ChunkBytecode, ColumnTileBytecode, + DagBuilder, Lowering, +}; +use sp1_gpu_cudart::sys::kernels::{ + zerocheck_column_tile_ext_kernel, zerocheck_column_tile_kb_kernel, + zerocheck_fused_sequential_ext_1024_kernel, zerocheck_fused_sequential_ext_128_kernel, + zerocheck_fused_sequential_ext_256_kernel, zerocheck_fused_sequential_ext_32_kernel, + zerocheck_fused_sequential_ext_512_kernel, zerocheck_fused_sequential_ext_64_kernel, + zerocheck_fused_sequential_ext_kernel, zerocheck_fused_sequential_kb_1024_kernel, + zerocheck_fused_sequential_kb_128_kernel, zerocheck_fused_sequential_kb_256_kernel, + zerocheck_fused_sequential_kb_32_kernel, zerocheck_fused_sequential_kb_512_kernel, + zerocheck_fused_sequential_kb_64_kernel, zerocheck_fused_sequential_kb_kernel, + zerocheck_sequential_ext_128_kernel, zerocheck_sequential_ext_256_kernel, + zerocheck_sequential_ext_32_kernel, zerocheck_sequential_ext_64_kernel, + zerocheck_sequential_ext_kernel, zerocheck_sequential_kb_128_kernel, + zerocheck_sequential_kb_256_kernel, zerocheck_sequential_kb_32_kernel, + zerocheck_sequential_kb_64_kernel, zerocheck_sequential_kb_kernel, +}; +use sp1_gpu_cudart::sys::runtime::KernelPtr; +use sp1_gpu_cudart::{args, DeviceBuffer, DevicePoint, TaskScope}; +use sp1_gpu_utils::{Ext, Felt, JaggedTraceMle}; +use sp1_hypercube::air::MachineAir; +use sp1_hypercube::prover::ZerocheckAir; +use sp1_hypercube::{ + AirOpenedValues, Chip, ChipEvaluation, ChipOpenedValues, ConstraintSumcheckFolder, + LogUpEvaluations, ShardOpenedValues, +}; + +use crate::challenger_update; +use crate::primitives::{evaluate_jagged_fix_last_variable, JaggedFixLastVariableKernel}; + +// ============================================================================ +// Compiled per-chip data — built once per session, reused across all rounds +// and shards. +// ============================================================================ + +/// One chunk's bytecode + a discriminator for which kernel runs it. +#[derive(Debug, Clone)] +pub enum CompiledChunk { + Sequential(ChunkBytecode), + ColumnTile(ColumnTileBytecode), +} + +/// Per-chip compiled program: a list of chunks (Sequential + ColumnTile) +/// plus a final synthesized GKR-correction chunk. +#[derive(Debug, Clone)] +pub struct CompiledChip { + pub chip_idx: u32, + pub name: String, + pub main_width: u32, + pub prep_width: u32, + pub chunks: Vec, +} + +/// Index of the synthesized GKR chunk inside `chunks`. Used by the launcher +/// to wire `runtime_coeffs` only to that one chunk's launch. +pub fn gkr_chunk_idx(chip: &CompiledChip) -> usize { + // GKR is always the last chunk we append in `compile_chip`. + chip.chunks.len() - 1 +} + +/// Compile a chip set to per-chip v2 chunks. +/// +/// The emitted bytecode is **machine-stable**: it depends only on each chip's +/// AIR, not on the cluster it lands in. In particular, assertion alpha +/// indices are stored *chip-relative* (`0 .. chip.num_constraints`), NOT +/// shifted into the cluster's `powers_of_alpha` table. The cluster-dependent +/// shift (`max_num_constraints - chip.num_constraints`) is instead applied at +/// kernel-launch time via `ChunkMetaC::chip_alpha_offset` (fused Sequential +/// kernel) or by offsetting the `powers_of_alpha` pointer (ColumnTile). This +/// lets the compiled+uploaded bytecode be cached once per machine and reused +/// across every shard and cluster. +/// +/// The synthesized GKR chunk (only for chips with no Sequential carrier) +/// stores the chip-relative index `num_constraints - 1`, which the same +/// runtime shift maps onto the `powers_of_alpha` slot holding `EF::one()`. +pub fn compile_chips(chips: &BTreeSet>, budget: ChunkBudget) -> Vec +where + A: MachineAir + for<'a> Air>, +{ + let t_compile = std::time::Instant::now(); + + let mut out = Vec::with_capacity(chips.len()); + for (i, chip) in chips.iter().enumerate() { + let air: &A = chip.air.as_ref(); + let dag = build_dag(air); + let infos = analyze_constraints(&dag); + let chunks_meta = chunk_dag(&infos, &budget); + + let mut compiled_chunks: Vec = Vec::new(); + for chunk in &chunks_meta { + let lowerings = enumerate_lowerings(chunk, &infos, &dag); + // Prefer ColumnTile if it applies; fall back to Sequential. + let mut placed = false; + if let Some(plan) = lowerings.iter().find_map(|l| match l { + Lowering::ColumnTile(p) => Some(p), + _ => None, + }) { + if let Some(bc) = lower_column_tile(chunk, &infos, &dag, plan) { + // `bc.terms[*].alpha_idx` stays chip-relative; the + // cluster shift is applied at launch. + compiled_chunks.push(CompiledChunk::ColumnTile(bc)); + placed = true; + } + } + if !placed { + let plan = lowerings + .iter() + .find_map(|l| match l { + Lowering::Sequential(p) => Some(p), + _ => None, + }) + .expect("every chunk must have a Sequential lowering"); + let bc = lower_sequential(chunk, &infos, &dag, plan); + // `bc.asserts[*].1` (alpha index) stays chip-relative; the + // cluster shift is applied at launch. + if std::env::var("SP1_GPU_DEBUG_MAXREG").is_ok() { + eprintln!( + "compile chip={} max_reg={} n_instrs={} n_asserts={}", + air.name(), + bc.max_reg, + bc.instrs.len(), + bc.asserts.len() + ); + } + compiled_chunks.push(CompiledChunk::Sequential(bc)); + } + } + + // GKR carrier selection: the synthesized GKR sweep is + // `Σ_i bp_i · col_i` over all main + prep cols. We bake it into the + // FIRST Sequential chunk of the chip (it then does the column sweep + // after its bytecode + asserts, sharing column reads with the + // constraint pass in L1 — same memory pattern as v1). + // + // Chips with no Sequential chunk (would be ColumnTile-only — rare in + // practice for SP1 today) fall back to the legacy synthesized + // ColumnTile GKR chunk so they still get GKR. + let main_width = air.width() as u32; + let prep_width = air.preprocessed_width() as u32; + let carrier = compiled_chunks.iter_mut().find_map(|c| match c { + CompiledChunk::Sequential(bc) => Some(bc), + _ => None, + }); + if let Some(carrier_bc) = carrier { + carrier_bc.gkr_main_width = main_width; + carrier_bc.gkr_prep_width = prep_width; + } else { + // No Sequential carrier — synthesize a ColumnTile GKR chunk. It + // is flagged `is_gkr_carrier`; the launcher weights it by the + // `powers_of_alpha` slot holding `1` directly (see + // `synthesize_gkr_chunk` / `launch_chunk_into`). + let gkr_bc = synthesize_gkr_chunk(main_width, prep_width); + compiled_chunks.push(CompiledChunk::ColumnTile(gkr_bc)); + } + + out.push(CompiledChip { + chip_idx: i as u32, + name: air.name().to_string(), + main_width, + prep_width, + chunks: compiled_chunks, + }); + } + if std::env::var("SP1_GPU_ZEROCHECK_TIMING").is_ok() { + tracing::info!("compile_chips: {} chips in {:?}", out.len(), t_compile.elapsed()); + } + out +} + +// ============================================================================ +// Device-uploaded per-chunk buffers — built once per shard. +// ============================================================================ + +/// Device-side per-chunk metadata for the fused dispatch kernel. Layout +/// must match `struct ChunkMeta` in `include/zerocheck/sequential.cuh`. +/// +/// The "carrier" chunk (one per chip, gkr_main_width > 0) also carries the +/// chip's per-round geq parameters (`padded_row_adjustment`, `geq_threshold`, +/// `geq_eq_coefficient`). The kernel subtracts the geq correction +/// in-thread, eliminating the per-round host loop that used to dominate +/// runtime for chips with large row counts. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct ChunkMetaC { + pub instrs: *const sp1_gpu_air::ir::DagInstr, + pub leaves: *const sp1_gpu_air::ir::LeafRef, + pub consts: *const Felt, + pub publics: *const u32, + pub assert_regs: *const u16, + pub assert_alphas: *const u32, + pub preprocessed_ptr: u64, + pub main_ptr: u64, + pub n_instrs: u32, + pub n_asserts: u32, + pub chip_idx: u32, + pub gkr_main_width: u32, + pub gkr_prep_width: u32, + pub height: u32, + pub row_count: u32, + /// Cluster-dependent shift added to every chip-relative alpha index in + /// this chunk's bytecode before indexing `powers_of_alpha`. + pub chip_alpha_offset: u32, + pub geq_threshold: u32, + pub geq_eq_coefficient: Ext, + pub padded_row_adjustment: Ext, +} + +// SAFETY: ChunkMetaC contains raw device pointers but the kernel will only +// dereference them on the GPU after we copy the struct over. Send/Sync is +// fine for our usage — we never share these across threads on the host. +unsafe impl Send for ChunkMetaC {} +unsafe impl Sync for ChunkMetaC {} + +/// A per-chunk view into the flat machine-wide bytecode buffers +/// (`MachineBytecode`). Holds raw device pointers, not owned allocations — +/// the backing memory lives in the `MachineBytecode` that contains this +/// struct, so a view is valid exactly as long as that machine bytecode is. +#[derive(Clone, Copy)] +pub struct ChunkDeviceBufs { + pub kind: ChunkKind, + // Common + pub leaves: *const sp1_gpu_air::ir::LeafRef, + pub consts: *const Felt, + pub publics: *const u32, + // Sequential-only (dummy-but-valid pointer + zero counts for ColumnTile) + pub instrs: *const sp1_gpu_air::ir::DagInstr, + pub assert_regs: *const u16, + pub assert_alphas: *const u32, + pub max_reg: u16, + pub n_instrs: u32, + pub n_asserts: u32, + /// Sequential chunks that carry the chip's GKR sweep. When > 0, the + /// kernel appends `Σ_i gkr_powers[i] · col_i` over (main_w main cols, + /// then prep_w prep cols) after the bytecode + asserts pass, sharing + /// column reads with the constraint bytecode. + pub gkr_main_width: u32, + pub gkr_prep_width: u32, + // ColumnTile-only (dummy-but-valid pointer + zero count for Sequential) + pub terms: *const sp1_gpu_air::ir::ColumnTermEntry, + pub n_terms: u32, + // Marker for GKR (uses the per-shard runtime_coeffs buffer) + pub is_gkr: bool, +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum ChunkKind { + Sequential, + ColumnTile, +} + +/// Per-chip device views into the flat machine bytecode. +#[derive(Clone)] +pub struct CompiledChipDevice { + pub chip_idx: u32, + pub name: String, + pub main_width: u32, + pub prep_width: u32, + pub chunks: Vec, +} + +/// The whole machine's compiled v2 bytecode, uploaded once (at prover +/// construction) into a small fixed number of flat device buffers. +/// +/// Every chunk of every chip is concatenated, per array type, into one of +/// the seven `flat_*` buffers; each chunk records pointers into them via +/// `ChunkDeviceBufs`. This replaces the old per-shard, per-chunk upload +/// (~7 tiny allocations × every chunk × every shard) with 7 allocations +/// uploaded exactly once per machine. +pub struct MachineBytecode { + // Flat buffers — own the device memory the `ChunkDeviceBufs` views point + // into. Never read directly; kept alive for the views' sake. + _flat_instrs: Buffer, + _flat_leaves: Buffer, + _flat_consts: Buffer, + _flat_publics: Buffer, + _flat_assert_regs: Buffer, + _flat_assert_alphas: Buffer, + _flat_terms: Buffer, + /// One entry per machine chip, in machine chip order. + pub chips: Vec, + /// Chip name → index into `chips`, for per-shard subset selection. + pub chip_index: BTreeMap, +} + +// SAFETY: `ChunkDeviceBufs`/`CompiledChipDevice` hold raw device pointers +// into the `_flat_*` buffers owned by the same `MachineBytecode`. They are +// only ever dereferenced on the GPU after being copied across, and the +// backing buffers share this struct's lifetime. We never mutate through the +// pointers on the host. +unsafe impl Send for MachineBytecode {} +unsafe impl Sync for MachineBytecode {} + +/// Compile + upload the entire machine's bytecode once. Call at prover +/// construction; the result is reused for every shard and every cluster. +pub fn upload_machine_bytecode( + chips: &BTreeSet>, + budget: ChunkBudget, + scope: &TaskScope, +) -> MachineBytecode +where + A: MachineAir + for<'a> Air>, +{ + upload_compiled_bytecode(compile_chips(chips, budget), scope) +} + +/// Flatten + upload an already-compiled set of chips. Split out of +/// `upload_machine_bytecode` so scaling tests can feed a synthetically large +/// chip list without paying repeated AIR compilation. +pub fn upload_compiled_bytecode(compiled: Vec, scope: &TaskScope) -> MachineBytecode { + // ---- Pass 1: concatenate every chunk's arrays into flat host vecs. ---- + let mut flat_instrs: Vec = Vec::new(); + let mut flat_leaves: Vec = Vec::new(); + let mut flat_consts: Vec = Vec::new(); + let mut flat_publics: Vec = Vec::new(); + let mut flat_assert_regs: Vec = Vec::new(); + let mut flat_assert_alphas: Vec = Vec::new(); + let mut flat_terms: Vec = Vec::new(); + + // Per-chunk (offset, len) into each flat vec, recorded in pass 1 and + // resolved to pointers in pass 2 (after the flat buffers are uploaded). + struct ChunkOffsets { + kind: ChunkKind, + leaves: (usize, usize), + consts: (usize, usize), + publics: (usize, usize), + instrs: (usize, usize), + assert_regs: (usize, usize), + assert_alphas: (usize, usize), + terms: (usize, usize), + max_reg: u16, + gkr_main_width: u32, + gkr_prep_width: u32, + is_gkr: bool, + } + let mut chip_offsets: Vec> = Vec::with_capacity(compiled.len()); + + // Append `src` to `dst`, returning the (offset, len) of the appended run. + fn extend_flat(dst: &mut Vec, src: &[T]) -> (usize, usize) { + let off = dst.len(); + dst.extend_from_slice(src); + (off, src.len()) + } + + for chip in &compiled { + let mut chunks = Vec::with_capacity(chip.chunks.len()); + for c in chip.chunks.iter() { + chunks.push(match c { + CompiledChunk::Sequential(bc) => { + let regs: Vec = bc.asserts.iter().map(|&(r, _)| r).collect(); + let alphas: Vec = bc.asserts.iter().map(|&(_, a)| a).collect(); + ChunkOffsets { + kind: ChunkKind::Sequential, + leaves: extend_flat(&mut flat_leaves, &bc.leaves), + consts: extend_flat(&mut flat_consts, &bc.consts), + publics: extend_flat(&mut flat_publics, &bc.publics), + instrs: extend_flat(&mut flat_instrs, &bc.instrs), + assert_regs: extend_flat(&mut flat_assert_regs, ®s), + assert_alphas: extend_flat(&mut flat_assert_alphas, &alphas), + terms: (flat_terms.len(), 0), + max_reg: bc.max_reg, + gkr_main_width: bc.gkr_main_width, + gkr_prep_width: bc.gkr_prep_width, + is_gkr: false, + } + } + CompiledChunk::ColumnTile(bc) => ChunkOffsets { + kind: ChunkKind::ColumnTile, + leaves: extend_flat(&mut flat_leaves, &bc.leaves), + consts: extend_flat(&mut flat_consts, &bc.consts), + publics: extend_flat(&mut flat_publics, &bc.publics), + instrs: (flat_instrs.len(), 0), + assert_regs: (flat_assert_regs.len(), 0), + assert_alphas: (flat_assert_alphas.len(), 0), + terms: extend_flat(&mut flat_terms, &bc.terms), + max_reg: 0, + gkr_main_width: 0, + gkr_prep_width: 0, + is_gkr: bc.is_gkr_carrier, + }, + }); + } + chip_offsets.push(chunks); + } + + // ---- Upload the seven flat buffers (≥1 element so the base pointer is + // never null even for an array type no chunk uses). ---- + fn upload_flat(v: &mut Vec, scope: &TaskScope) -> Buffer { + if v.is_empty() { + // POD `#[repr(C)]` bytecode structs — an all-zero element is a + // valid (never-dereferenced) placeholder so the base pointer of + // an unused array type is non-null. + v.push(unsafe { std::mem::zeroed() }); + } + DeviceBuffer::from_host_slice(v, scope).unwrap().into_inner() + } + if std::env::var("SP1_GPU_ZEROCHECK_TIMING").is_ok() { + let mb = |n: usize, sz: usize| (n * sz) as f64 / (1024.0 * 1024.0); + tracing::info!( + "upload_compiled_bytecode: {} chips, flat bytes — instrs={:.1}M leaves={:.1}M \ + consts={:.1}M publics={:.1}M assert_regs={:.1}M assert_alphas={:.1}M terms={:.1}M", + compiled.len(), + mb(flat_instrs.len(), std::mem::size_of::()), + mb(flat_leaves.len(), std::mem::size_of::()), + mb(flat_consts.len(), std::mem::size_of::()), + mb(flat_publics.len(), 4), + mb(flat_assert_regs.len(), 2), + mb(flat_assert_alphas.len(), 4), + mb(flat_terms.len(), std::mem::size_of::()), + ); + } + let flat_instrs_buf = upload_flat(&mut flat_instrs, scope); + let flat_leaves_buf = upload_flat(&mut flat_leaves, scope); + let flat_consts_buf = upload_flat(&mut flat_consts, scope); + let flat_publics_buf = upload_flat(&mut flat_publics, scope); + let flat_assert_regs_buf = upload_flat(&mut flat_assert_regs, scope); + let flat_assert_alphas_buf = upload_flat(&mut flat_assert_alphas, scope); + let flat_terms_buf = upload_flat(&mut flat_terms, scope); + + // ---- Pass 2: resolve offsets to device pointers. ---- + let instrs_base = flat_instrs_buf.as_ptr(); + let leaves_base = flat_leaves_buf.as_ptr(); + let consts_base = flat_consts_buf.as_ptr(); + let publics_base = flat_publics_buf.as_ptr(); + let assert_regs_base = flat_assert_regs_buf.as_ptr(); + let assert_alphas_base = flat_assert_alphas_buf.as_ptr(); + let terms_base = flat_terms_buf.as_ptr(); + + let mut device_chips = Vec::with_capacity(compiled.len()); + let mut chip_index = BTreeMap::new(); + for (chip, offsets) in compiled.iter().zip(chip_offsets.iter()) { + let chunks = offsets + .iter() + .map(|o| ChunkDeviceBufs { + kind: o.kind, + leaves: unsafe { leaves_base.add(o.leaves.0) }, + consts: unsafe { consts_base.add(o.consts.0) }, + publics: unsafe { publics_base.add(o.publics.0) }, + instrs: unsafe { instrs_base.add(o.instrs.0) }, + assert_regs: unsafe { assert_regs_base.add(o.assert_regs.0) }, + assert_alphas: unsafe { assert_alphas_base.add(o.assert_alphas.0) }, + terms: unsafe { terms_base.add(o.terms.0) }, + max_reg: o.max_reg, + n_instrs: o.instrs.1 as u32, + n_asserts: o.assert_regs.1 as u32, + gkr_main_width: o.gkr_main_width, + gkr_prep_width: o.gkr_prep_width, + n_terms: o.terms.1 as u32, + is_gkr: o.is_gkr, + }) + .collect(); + chip_index.insert(chip.name.clone(), device_chips.len()); + device_chips.push(CompiledChipDevice { + chip_idx: chip.chip_idx, + name: chip.name.clone(), + main_width: chip.main_width, + prep_width: chip.prep_width, + chunks, + }); + } + + MachineBytecode { + _flat_instrs: flat_instrs_buf, + _flat_leaves: flat_leaves_buf, + _flat_consts: flat_consts_buf, + _flat_publics: flat_publics_buf, + _flat_assert_regs: flat_assert_regs_buf, + _flat_assert_alphas: flat_assert_alphas_buf, + _flat_terms: flat_terms_buf, + chips: device_chips, + chip_index, + } +} + +// ============================================================================ +// Per-round poly state (parallel to v1's ZeroCheckJaggedPoly). +// ============================================================================ + +pub struct ZeroCheckJaggedPoly<'b, K: Field> { + pub data: Cow<'b, JaggedTraceMle>, + /// This shard's chips, as pointer-views into `machine_bytecode`. + pub compiled: Vec, + /// The machine-wide flat bytecode the `compiled` views point into. Held + /// here so the device buffers outlive every round of this shard. + pub machine_bytecode: Arc, + pub virtual_geq: Vec>, + pub eq_adjustment: Ext, + pub zeta: Point, + pub claim: Ext, + pub initial_heights: Vec, + pub padded_row_adjustment_host: Vec, + pub public_values: Buffer, + pub powers_of_alpha: Buffer, + pub gkr_powers: Buffer, + pub powers_of_lambda: Buffer, + pub powers_of_lambda_host: Vec, + pub chip_main_ptrs: Vec, + pub chip_preprocessed_ptrs: Vec, + pub chip_heights: Vec, + /// Per-chip shift into the cluster's reversed `powers_of_alpha` table: + /// `max_num_constraints - chip.num_constraints`. The compiled bytecode + /// stores chip-relative alpha indices; this shift is applied at launch + /// (see `compile_chips`). Indexed by `chip_idx`. + pub chip_alpha_offset: Vec, +} + +/// Per-trace-element-type kernel selection. Round 0 uses `K = Felt`; rounds +/// 1+ use `K = Ext` after the trace is folded into the extension field. +pub trait EvalKernels { + fn sequential_kernel() -> KernelPtr; + /// Smallest tier whose MAX_REGS >= `max_reg`. `K regs[MAX_REGS][3]` is + /// a per-thread stack array that spills to local memory, so tight + /// sizing avoids paying for unused slots on every row. + fn sequential_kernel_for(max_reg: u16) -> KernelPtr; + fn column_tile_kernel() -> KernelPtr; + /// Fused dispatch kernel — one launch handles every Sequential chunk + /// across every chip in a round. + fn fused_sequential_kernel() -> KernelPtr; + /// Tiered fused dispatch kernel. The launcher partitions chunks into + /// tiers by their `max_reg` and launches one kernel per non-empty + /// tier so each kernel's local register array is sized to its tier's + /// worst case. + fn fused_sequential_kernel_for(max_reg_in_tier: u16) -> KernelPtr; +} + +impl EvalKernels for TaskScope { + fn sequential_kernel() -> KernelPtr { + unsafe { zerocheck_sequential_kb_kernel() } + } + fn sequential_kernel_for(max_reg: u16) -> KernelPtr { + unsafe { + if max_reg <= 32 { + zerocheck_sequential_kb_32_kernel() + } else if max_reg <= 64 { + zerocheck_sequential_kb_64_kernel() + } else if max_reg <= 128 { + zerocheck_sequential_kb_128_kernel() + } else { + zerocheck_sequential_kb_256_kernel() + } + } + } + fn column_tile_kernel() -> KernelPtr { + unsafe { zerocheck_column_tile_kb_kernel() } + } + fn fused_sequential_kernel() -> KernelPtr { + unsafe { zerocheck_fused_sequential_kb_kernel() } + } + fn fused_sequential_kernel_for(max_reg_in_tier: u16) -> KernelPtr { + unsafe { + if max_reg_in_tier <= 32 { + zerocheck_fused_sequential_kb_32_kernel() + } else if max_reg_in_tier <= 64 { + zerocheck_fused_sequential_kb_64_kernel() + } else if max_reg_in_tier <= 128 { + zerocheck_fused_sequential_kb_128_kernel() + } else if max_reg_in_tier <= 256 { + zerocheck_fused_sequential_kb_256_kernel() + } else if max_reg_in_tier <= 512 { + zerocheck_fused_sequential_kb_512_kernel() + } else { + zerocheck_fused_sequential_kb_1024_kernel() + } + } + } +} + +impl EvalKernels for TaskScope { + fn sequential_kernel() -> KernelPtr { + unsafe { zerocheck_sequential_ext_kernel() } + } + fn sequential_kernel_for(max_reg: u16) -> KernelPtr { + unsafe { + if max_reg <= 32 { + zerocheck_sequential_ext_32_kernel() + } else if max_reg <= 64 { + zerocheck_sequential_ext_64_kernel() + } else if max_reg <= 128 { + zerocheck_sequential_ext_128_kernel() + } else { + zerocheck_sequential_ext_256_kernel() + } + } + } + fn column_tile_kernel() -> KernelPtr { + unsafe { zerocheck_column_tile_ext_kernel() } + } + fn fused_sequential_kernel() -> KernelPtr { + unsafe { zerocheck_fused_sequential_ext_kernel() } + } + fn fused_sequential_kernel_for(max_reg_in_tier: u16) -> KernelPtr { + unsafe { + if max_reg_in_tier <= 32 { + zerocheck_fused_sequential_ext_32_kernel() + } else if max_reg_in_tier <= 64 { + zerocheck_fused_sequential_ext_64_kernel() + } else if max_reg_in_tier <= 128 { + zerocheck_fused_sequential_ext_128_kernel() + } else if max_reg_in_tier <= 256 { + zerocheck_fused_sequential_ext_256_kernel() + } else if max_reg_in_tier <= 512 { + zerocheck_fused_sequential_ext_512_kernel() + } else { + zerocheck_fused_sequential_ext_1024_kernel() + } + } + } +} + +// ============================================================================ +// Build the initial poly state for round 0. +// ============================================================================ + +#[allow(clippy::too_many_arguments)] +pub fn initialize_zerocheck_poly<'b, A>( + data: &'b JaggedTraceMle, + chips: &BTreeSet>, + compiled_chips_dev: Vec, + machine_bytecode: Arc, + initial_heights: Vec, + public_values: Vec, + powers_of_alpha: Vec, + gkr_powers: Vec, + powers_of_lambda: Vec, + padded_row_adjustment: Vec, + zeta: Point, + claim: Ext, +) -> ZeroCheckJaggedPoly<'b, Felt> +where + A: MachineAir, +{ + let scope = data.dense().backend(); + + let mut virtual_geq: Vec> = vec![]; + for (chip, height) in chips.iter().zip_eq(initial_heights.iter()) { + virtual_geq.push(VirtualGeq::new( + *height, + Ext::one(), + Ext::zero(), + zeta.dimension() as u32, + )); + let _ = chip; + } + + let (chip_main_ptrs, chip_preprocessed_ptrs, chip_heights) = compute_chip_offsets(data, chips); + + // Per-chip launch-time shift into the reversed `powers_of_alpha` table. + let max_num_constraints = + chips.iter().map(|c| c.num_constraints).max().unwrap_or(1).max(1) as u32; + let chip_alpha_offset: Vec = + chips.iter().map(|c| max_num_constraints - c.num_constraints as u32).collect(); + + let public_values_device = + DeviceBuffer::from_host(&Buffer::from(public_values), scope).unwrap().into_inner(); + let powers_of_alpha_device = + DeviceBuffer::from_host(&Buffer::from(powers_of_alpha), scope).unwrap().into_inner(); + let gkr_powers_device = + DeviceBuffer::from_host(&Buffer::from(gkr_powers), scope).unwrap().into_inner(); + let powers_of_lambda_host = powers_of_lambda.clone(); + let powers_of_lambda_device = + DeviceBuffer::from_host(&Buffer::from(powers_of_lambda), scope).unwrap().into_inner(); + + ZeroCheckJaggedPoly { + data: Cow::Borrowed(data), + compiled: compiled_chips_dev, + machine_bytecode, + virtual_geq, + eq_adjustment: Ext::one(), + zeta, + claim, + initial_heights, + padded_row_adjustment_host: padded_row_adjustment, + public_values: public_values_device, + powers_of_alpha: powers_of_alpha_device, + gkr_powers: gkr_powers_device, + powers_of_lambda: powers_of_lambda_device, + powers_of_lambda_host, + chip_main_ptrs, + chip_preprocessed_ptrs, + chip_heights, + chip_alpha_offset, + } +} + +/// Compute the per-chip main/preprocessed trace pointers within the jagged +/// dense buffer, using the jagged structure's `start_indices` and +/// `column_heights` (in pair units) as the source of truth: +/// main_ptr = startIndices[total_prep_cols + has_prep_padding + main_idx] << 1 +/// +/// The dense_offset values from `update_offset` (used by JaggedTraceMle's +/// table indices after a fold) silently drift from the actual jagged layout +/// when fold padding introduces extra elements at column ends — using +/// start_indices avoids that whole class of bugs. +/// +/// Returns (main_ptrs, preprocessed_ptrs, heights), all in element units of +/// the underlying buffer type (Felt before round 0, Ext after). +fn compute_chip_offsets( + data: &JaggedTraceMle, + chips: &BTreeSet>, +) -> (Vec, Vec, Vec) +where + A: MachineAir, +{ + let column_heights = &data.0.column_heights; + // start_indices in pair units derived from column_heights (cumulative). + let mut starts: Vec = Vec::with_capacity(column_heights.len() + 1); + starts.push(0); + let mut acc: u64 = 0; + for &h in column_heights.iter() { + acc += h as u64; + starts.push(acc); + } + + let total_prep_widths: usize = chips.iter().map(|c| c.preprocessed_width()).sum(); + // Match v1's kernel convention: always assume a single prep-padding + // column sits between prep cols and main cols. v1's evaluate_zerocheck + // hardcodes `total_num_preprocessed_column + 1 + main_idx` for the same + // reason. In practice padded prep is always rounded up so this column + // exists. + let main_section_start_col: usize = total_prep_widths + 1; + + let mut prep_ptrs = Vec::with_capacity(chips.len()); + let mut main_ptrs = Vec::with_capacity(chips.len()); + let mut heights = Vec::with_capacity(chips.len()); + + let mut cum_prep: usize = 0; + let mut cum_main: usize = 0; + for chip in chips.iter() { + let prep_w = chip.preprocessed_width(); + let main_w = chip.width(); + + let prep_col_idx = cum_prep; + let main_col_idx = main_section_start_col + cum_main; + + let prep_ptr = if prep_w > 0 { starts[prep_col_idx] * 2 } else { 0 }; + let main_ptr = if main_w > 0 { starts[main_col_idx] * 2 } else { 0 }; + + // Column stride in element units. Prefer main if present (chip has + // main cols), else prep. Chips with neither shouldn't reach here. + let height = if main_w > 0 { + column_heights[main_col_idx] * 2 + } else if prep_w > 0 { + column_heights[prep_col_idx] * 2 + } else { + 0 + }; + + prep_ptrs.push(prep_ptr); + main_ptrs.push(main_ptr); + heights.push(height); + + cum_prep += prep_w; + cum_main += main_w; + } + + (main_ptrs, prep_ptrs, heights) +} + +// ============================================================================ +// Per-round kernel launcher. +// ============================================================================ + +pub fn evaluate_zerocheck<'b, K: Field>( + poly: &ZeroCheckJaggedPoly<'b, K>, +) -> UnivariatePolynomial +where + TaskScope: EvalKernels, +{ + let backend = poly.data.backend(); + // Three evaluation points per univariate round (degree-2 polynomial + // recovered by Lagrange interpolation downstream). + const NUM_EVAL_POINT: usize = 3; + // Cap on grid x-dim. Beyond this, blocks oversubscribe SMs and the per- + // block reduce overhead dominates the actual per-row work. Measured: 4096 + // saturates an SM-rich GPU (Ada/Hopper); higher values are flat or worse. + const MAX_GRID: u32 = 4096; + // Two-tier launch threshold. Sequential chunks with `max_reg > TIER_SPLIT` + // run in a separate larger-template kernel so the bulk of chunks (which + // sit at `max_reg ≤ 128`) don't pay for the bigger `regs[]` array. We + // only actually tier-split when the high-`max_reg` chunks are a small + // minority (see `do_tier_split` below); otherwise the launch + // fragmentation cost outweighs the per-thread footprint win. + const TIER_SPLIT: u16 = 256; + + let (rest, last) = poly.zeta.split_at(poly.zeta.dimension() - 1); + let last = *last[0]; + let rest_point = DevicePoint::from_host(&rest, backend).unwrap(); + let partial_lagrange = rest_point.partial_lagrange(); + let rest_point_dim = rest.dimension() as u32; + + let trace_ptr = poly.data.as_ref().dense_data.dense.as_ptr(); + + // ---- Pass 1: partition Sequential chunks into 2 tiers by max_reg ---- + // + // We tier-split ONLY when tier 1 (the large-register-pressure tier) has + // a SMALL number of chunks relative to tier 0 — i.e. when there are a + // few outliers forcing the whole kernel into a large MAX_REGS template. + // This catches real RSP shards (1-2 outliers at max_reg=340 force + // MAX_REGS=512 for the entire kernel; tiering them out lets ~95% of + // chunks run at MAX_REGS=128 with much less spill) without hurting + // all-chips, where many chunks have moderately high max_reg and the + // unified kernel is actually more efficient than 2 launches. + // + // Decide the tier assignment in two passes: first pre-scan to count + // chunks per tier; if tier 1's count is non-trivial, collapse everything + // into tier 0 (the unified single-launch path). + // + // ColumnTile chunks (only when a chip had no Sequential chunk to carry + // GKR — empty for current SP1 workloads) launch individually. + let mut tier1_candidate_count = 0usize; + let mut total_seq_count = 0usize; + for chip in poly.compiled.iter() { + // Skip empty chips — they aren't launched, so they must not skew the + // tier-split ratio below. + if poly.chip_heights[chip.chip_idx as usize] / 2 == 0 { + continue; + } + for chunk in chip.chunks.iter() { + if matches!(chunk.kind, ChunkKind::Sequential) { + total_seq_count += 1; + if chunk.max_reg > TIER_SPLIT { + tier1_candidate_count += 1; + } + } + } + } + // Heuristic: tier-split only when tier 1 is a small minority. The + // exact ratio comes from measurement — real RSP has ≤1/30 chunks in + // tier 1 and benefits hugely; all-chips has ~1/3 and regresses. + let do_tier_split = total_seq_count > 0 + && tier1_candidate_count > 0 + && tier1_candidate_count * 10 <= total_seq_count; + let mut seq_meta_tiers: [Vec; 2] = [Vec::new(), Vec::new()]; + let mut row_starts_tiers: [Vec; 2] = [vec![0], vec![0]]; + let mut max_reg_tiers: [u16; 2] = [0, 0]; + let mut ct_launches: Vec<(u32, &ChunkDeviceBufs, u64, u64, u32, u32)> = Vec::new(); + + for chip in poly.compiled.iter() { + let chip_idx = chip.chip_idx; + let height = poly.chip_heights[chip_idx as usize]; + let row_count = height / 2; + // An empty chip (0 rows) contributes nothing to the zerocheck sum: + // its constraint pass and GKR sweep both run over 0 rows. Skip it + // entirely so the per-round `ChunkMetaC` array and `row_starts` table + // stay O(non-empty chips) rather than O(total chips) — this matters + // when the machine has thousands of chips most of which are empty in + // any given shard. + if row_count == 0 { + continue; + } + let main_ptr = poly.chip_main_ptrs[chip_idx as usize]; + let preprocessed_ptr = poly.chip_preprocessed_ptrs[chip_idx as usize]; + + for chunk in chip.chunks.iter() { + match chunk.kind { + ChunkKind::Sequential => { + let geq = &poly.virtual_geq[chip_idx as usize]; + let pad_adj = poly.padded_row_adjustment_host[chip_idx as usize]; + let tier: usize = + if do_tier_split && chunk.max_reg > TIER_SPLIT { 1 } else { 0 }; + max_reg_tiers[tier] = max_reg_tiers[tier].max(chunk.max_reg); + seq_meta_tiers[tier].push(ChunkMetaC { + instrs: chunk.instrs, + leaves: chunk.leaves, + consts: chunk.consts, + publics: chunk.publics, + assert_regs: chunk.assert_regs, + assert_alphas: chunk.assert_alphas, + preprocessed_ptr, + main_ptr, + n_instrs: chunk.n_instrs, + n_asserts: chunk.n_asserts, + chip_idx, + gkr_main_width: chunk.gkr_main_width, + gkr_prep_width: chunk.gkr_prep_width, + height, + row_count, + chip_alpha_offset: poly.chip_alpha_offset[chip_idx as usize], + geq_threshold: geq.threshold, + geq_eq_coefficient: geq.eq_coefficient, + padded_row_adjustment: pad_adj, + }); + let last = *row_starts_tiers[tier].last().unwrap(); + row_starts_tiers[tier].push(last + row_count); + } + ChunkKind::ColumnTile => { + ct_launches.push(( + chip_idx, + chunk, + preprocessed_ptr, + main_ptr, + height, + row_count, + )); + } + } + } + } + + // Block size per tier, keyed off the tier's worst-case `max_reg`: above + // 128 the per-thread `regs[]` array crosses into a bigger MAX_REGS + // template and we need a smaller block so enough warps still fit per SM + // for latency hiding. + const BLOCK_SIZE_LOW_REG: u32 = 256; + const BLOCK_SIZE_HIGH_REG: u32 = 64; + let block_size_for = |tier: usize| -> u32 { + if max_reg_tiers[tier] > 128 { + BLOCK_SIZE_HIGH_REG + } else { + BLOCK_SIZE_LOW_REG + } + }; + + // Target a few grid-stride iterations per thread so per-block reduce + // overhead is amortised by real work. Larger values under-saturate SMs; + // smaller values over-saturate and spill local memory under the per-block + // reduce. Measured plateau in [2, 8]; 4 sits comfortably in the middle. + const GRID_STRIDE_ROWS_PER_THREAD: u64 = 4; + let mut tier_n_blocks: [u32; 2] = [0, 0]; + let mut tier_slot: [usize; 2] = [0, 0]; + let mut total_slots: usize = 0; + for t in 0..2 { + let bs = block_size_for(t); + let total_rows = *row_starts_tiers[t].last().unwrap(); + if !seq_meta_tiers[t].is_empty() && total_rows > 0 { + let target_threads = + (total_rows as u64).div_ceil(GRID_STRIDE_ROWS_PER_THREAD).max(bs as u64); + let blocks = target_threads.div_ceil(bs as u64); + tier_n_blocks[t] = blocks.min(MAX_GRID as u64).max(1) as u32; + } + tier_slot[t] = total_slots; + total_slots += (tier_n_blocks[t] as usize) * NUM_EVAL_POINT; + } + + let mut ct_slots: Vec<(usize, u32)> = Vec::with_capacity(ct_launches.len()); + let ct_block_size: u32 = 128; // unchanged for ColumnTile fallback + for &(_, chunk, _, _, _, row_count) in &ct_launches { + let total = chunk.n_terms as u64 * row_count as u64; + let n_blocks = if total == 0 { + 0 + } else { + total.div_ceil(ct_block_size as u64).min(MAX_GRID as u64).max(1) as u32 + }; + ct_slots.push((total_slots, n_blocks)); + total_slots += (n_blocks as usize) * NUM_EVAL_POINT; + } + + let mut shared_output: Tensor = + Tensor::with_sizes_in([total_slots.max(1)], backend.clone()); + unsafe { + shared_output.assume_init(); + } + let shared_output_ptr = shared_output.as_mut_ptr(); + + // Launch one fused kernel per non-empty tier. + let mut _seq_keepalive: Vec<(Buffer, Buffer)> = + Vec::with_capacity(2); + for t in 0..2 { + if tier_n_blocks[t] == 0 { + continue; + } + let bs = block_size_for(t); + let total_rows = *row_starts_tiers[t].last().unwrap(); + let chunk_meta_buf = + DeviceBuffer::from_host_slice(&seq_meta_tiers[t], backend).unwrap().into_inner(); + let row_starts_buf = + DeviceBuffer::from_host_slice(&row_starts_tiers[t], backend).unwrap().into_inner(); + let out_ptr = unsafe { shared_output_ptr.add(tier_slot[t]) }; + let shmem_bytes = (bs as usize / 32).max(1) * std::mem::size_of::(); + unsafe { + let args = args!( + chunk_meta_buf.as_ptr(), + row_starts_buf.as_ptr(), + (seq_meta_tiers[t].len() as u32), + total_rows, + trace_ptr, + poly.public_values.as_ptr(), + poly.powers_of_alpha.as_ptr(), + partial_lagrange.as_ptr(), + poly.powers_of_lambda.as_ptr(), + poly.gkr_powers.as_ptr(), + rest_point_dim, + out_ptr + ); + backend + .launch_kernel( + >::fused_sequential_kernel_for(max_reg_tiers[t]), + (tier_n_blocks[t], 1, 3), + (bs, 1, 1), + &args, + shmem_bytes, + ) + .unwrap(); + } + _seq_keepalive.push((chunk_meta_buf, row_starts_buf)); + } + + // Launch any remaining ColumnTile chunks individually (typically zero + // for current SP1 workloads). + for (i, &(chip_idx, chunk, preprocessed_ptr, main_ptr, height, row_count)) in + ct_launches.iter().enumerate() + { + let (slot, n_blocks) = ct_slots[i]; + if n_blocks == 0 { + continue; + } + let out_slot = unsafe { shared_output_ptr.add(slot) }; + launch_chunk_into::( + backend, + chunk, + trace_ptr, + preprocessed_ptr, + main_ptr, + height, + &poly.public_values, + &poly.powers_of_alpha, + poly.chip_alpha_offset[chip_idx as usize], + &poly.gkr_powers, + partial_lagrange.as_ptr(), + &poly.powers_of_lambda, + chip_idx, + rest_point_dim, + 0, + row_count, + n_blocks, + ct_block_size, + out_slot, + ); + } + + // ---- Single host sync + copy ---- + let host_partials: Vec = unsafe { shared_output.into_buffer().copy_into_host_vec() }; + + // ---- Pass 3: aggregate ---- + let mut totals: [Ext; NUM_EVAL_POINT] = [Ext::zero(); NUM_EVAL_POINT]; + + // Fused Sequential output: cross-tier, cross-chunk sum (λ applied per-row). + for t in 0..2 { + for b in 0..(tier_n_blocks[t] as usize) { + for e in 0..NUM_EVAL_POINT { + totals[e] += host_partials[tier_slot[t] + b * NUM_EVAL_POINT + e]; + } + } + } + for (i, &(_chip_idx, _chunk, _, _, _, _)) in ct_launches.iter().enumerate() { + let (slot, n_blocks) = ct_slots[i]; + for b in 0..(n_blocks as usize) { + for e in 0..NUM_EVAL_POINT { + totals[e] += host_partials[slot + b * NUM_EVAL_POINT + e]; + } + } + } + + // Geq correction is applied INSIDE the fused kernel (per-row, on the + // carrier chunk for each chip). Nothing to do here — it was the dominant + // host cost before being pushed into the kernel. + + // Apply last_var_eq and eq_adjustment (mirror v1's evaluate_zerocheck). + let mut xs = + vec![Ext::from_canonical_u32(0), Ext::from_canonical_u32(2), Ext::from_canonical_u32(4)]; + let mut ys: Vec = xs + .iter() + .zip(totals.iter()) + .map(|(&x, &t)| { + let last_var_eq = (Ext::one() - x) * (Ext::one() - last) + x * last; + t * last_var_eq * poly.eq_adjustment + }) + .collect(); + + xs.push(Ext::from_canonical_u32(1)); + ys.push(poly.claim - ys[0]); + + xs.push((last - Ext::one()) / (last + last - Ext::one())); + ys.push(Ext::zero()); + + interpolate_univariate_polynomial(&xs, &ys) +} + +/// Launch a chunk against a caller-provided device output pointer. +/// All launches in a round write into one shared buffer; the caller does a +/// single sync + copy_into_host at the end. +#[allow(clippy::too_many_arguments)] +fn launch_chunk_into( + scope: &TaskScope, + chunk: &ChunkDeviceBufs, + trace_ptr: *const K, + preprocessed_ptr: u64, + main_ptr: u64, + height: u32, + public_values: &Buffer, + powers_of_alpha: &Buffer, + chip_alpha_offset: u32, + gkr_powers: &Buffer, + partial_lagrange_ptr: *const Ext, + powers_of_lambda: &Buffer, + chip_idx: u32, + rest_point_dim: u32, + row_start: u32, + row_count: u32, + n_blocks: u32, + block_size: u32, + output_ptr: *mut Ext, +) where + TaskScope: EvalKernels, +{ + let shmem_bytes = (block_size as usize / 32) * std::mem::size_of::(); + match chunk.kind { + ChunkKind::Sequential => unsafe { + let args = args!( + chunk.instrs, + chunk.n_instrs, + chunk.leaves, + chunk.consts, + chunk.publics, + chunk.assert_regs, + chunk.assert_alphas, + chunk.n_asserts, + trace_ptr, + preprocessed_ptr, + main_ptr, + height, + public_values.as_ptr(), + powers_of_alpha.as_ptr(), + partial_lagrange_ptr, + powers_of_lambda.as_ptr(), + chip_idx, + rest_point_dim, + row_start, + row_count, + gkr_powers.as_ptr(), + chunk.gkr_main_width, + chunk.gkr_prep_width, + output_ptr + ); + scope + .launch_kernel( + >::sequential_kernel_for(chunk.max_reg), + (n_blocks, 1, 1), + (block_size, 1, 1), + &args, + shmem_bytes, + ) + .unwrap(); + }, + ChunkKind::ColumnTile => unsafe { + // Regular ColumnTile chunks store chip-relative `alpha_idx`; shift + // the `powers_of_alpha` base by the per-chip offset. The + // synthesized GKR carrier (`is_gkr`) instead wants weight `1`: + // its terms store `alpha_idx = 0`, so we point the base at the + // last slot (which holds `1` in the reversed table). The per-chip + // offset can't express this for a 0-constraint chip. + let shift = if chunk.is_gkr { + powers_of_alpha.len().saturating_sub(1) + } else { + chip_alpha_offset as usize + }; + let powers_of_alpha_shifted = powers_of_alpha.as_ptr().add(shift); + let args = args!( + chunk.terms, + chunk.n_terms, + chunk.leaves, + chunk.consts, + chunk.publics, + gkr_powers.as_ptr(), + trace_ptr, + preprocessed_ptr, + main_ptr, + height, + public_values.as_ptr(), + powers_of_alpha_shifted, + partial_lagrange_ptr, + powers_of_lambda.as_ptr(), + chip_idx, + rest_point_dim, + row_start, + row_count, + output_ptr + ); + scope + .launch_kernel( + >::column_tile_kernel(), + (n_blocks, 1, 1), + (block_size, 1, 1), + &args, + shmem_bytes, + ) + .unwrap(); + }, + } +} + +// ============================================================================ +// Fix-last-variable: fold trace data (reused from v1 path), update eq_adjustment. +// ============================================================================ + +pub fn zerocheck_fix_last_variable<'b, K: Field>( + input: ZeroCheckJaggedPoly<'b, K>, + point: Ext, + claim: Ext, +) -> ZeroCheckJaggedPoly<'b, Ext> +where + TaskScope: JaggedFixLastVariableKernel, + Ext: ExtensionField, +{ + let (rest, last) = input.zeta.split_at(input.zeta.dimension() - 1); + let last = *last[0]; + + let new_data = evaluate_jagged_fix_last_variable(&input.data, point); + let virtual_geq = + input.virtual_geq.iter().map(|geq| geq.fix_last_variable(point)).collect::>(); + let eq = (Ext::one() - last) * (Ext::one() - point) + last * point; + let eq_adjustment = input.eq_adjustment * eq; + + // After fold, the jagged layout may pad some columns (fold kernel adds 2 + // ext slots when `column_height_pairs % 4 != 0`). The `update_offset` + // path used by `evaluate_jagged_fix_last_variable` to rewrite the + // `dense_offset` ranges silently DRIFTS from the real column stride in + // those cases — `update_offset` uses `poly_size.div_ceil(4) * 2`, the + // actual stride uses `column_height_pairs.div_ceil(4) * 4`. Compute + // ptrs/heights from the new jagged structure's `column_heights` instead. + let column_heights = &new_data.0.column_heights; + let mut starts: Vec = Vec::with_capacity(column_heights.len() + 1); + starts.push(0); + let mut acc: u64 = 0; + for &h in column_heights.iter() { + acc += h as u64; + starts.push(acc); + } + + let total_prep_widths: usize = input.compiled.iter().map(|c| c.prep_width as usize).sum(); + // See `compute_chip_offsets`: match v1's `+1` convention. + let main_section_start_col: usize = total_prep_widths + 1; + + let mut chip_main_ptrs = Vec::with_capacity(input.compiled.len()); + let mut chip_preprocessed_ptrs = Vec::with_capacity(input.compiled.len()); + let mut chip_heights: Vec = Vec::with_capacity(input.compiled.len()); + + let mut cum_prep: usize = 0; + let mut cum_main: usize = 0; + for c in input.compiled.iter() { + let prep_w = c.prep_width as usize; + let main_w = c.main_width as usize; + let prep_col_idx = cum_prep; + let main_col_idx = main_section_start_col + cum_main; + let prep_ptr = if prep_w > 0 { starts[prep_col_idx] * 2 } else { 0 }; + let main_ptr = if main_w > 0 { starts[main_col_idx] * 2 } else { 0 }; + let height = if main_w > 0 { + (column_heights[main_col_idx]) * 2 + } else if prep_w > 0 { + (column_heights[prep_col_idx]) * 2 + } else { + 0 + }; + chip_preprocessed_ptrs.push(prep_ptr); + chip_main_ptrs.push(main_ptr); + chip_heights.push(height); + cum_prep += prep_w; + cum_main += main_w; + } + let initial_heights = chip_heights.clone(); + + ZeroCheckJaggedPoly { + data: Cow::Owned(new_data), + compiled: input.compiled, + machine_bytecode: input.machine_bytecode, + virtual_geq, + eq_adjustment, + zeta: rest, + claim, + initial_heights, + padded_row_adjustment_host: input.padded_row_adjustment_host, + public_values: input.public_values, + powers_of_alpha: input.powers_of_alpha, + gkr_powers: input.gkr_powers, + powers_of_lambda: input.powers_of_lambda, + powers_of_lambda_host: input.powers_of_lambda_host, + chip_main_ptrs, + chip_preprocessed_ptrs, + chip_heights, + chip_alpha_offset: input.chip_alpha_offset, + } +} + +// ============================================================================ +// Outer driver — parallel to v1's `zerocheck`. +// ============================================================================ + +#[allow(clippy::too_many_arguments)] +pub fn zerocheck( + chips: &BTreeSet>, + machine_bytecode: &Arc, + trace_mle: &JaggedTraceMle, + batching_challenge: Ext, + gkr_opening_batch_randomness: Ext, + logup_evaluations: &LogUpEvaluations, + public_values: Vec, + challenger: &mut C, + max_log_row_count: u32, +) -> (ShardOpenedValues, PartialSumcheckProof) +where + A: ZerocheckAir, + C: FieldChallenger, +{ + let data_input_heights = &trace_mle.column_heights; + let initial_heights = trace_mle + .dense_data + .main_table_index + .values() + .map(|trace_offset| trace_offset.poly_size as u32) + .collect::>(); + + let max_num_constraints = + itertools::max(chips.iter().map(|chip| chip.num_constraints)).unwrap(); + let max_columns = + itertools::max(chips.iter().map(|chip| chip.preprocessed_width() + chip.width())).unwrap(); + let total_preprocessed_columns = trace_mle.dense().preprocessed_cols; + let mut powers_of_challenge = + batching_challenge.powers().take(max_num_constraints).collect::>(); + powers_of_challenge.reverse(); + let num_chips = chips.len(); + let debug_timing = std::env::var("SP1_GPU_ZEROCHECK_TIMING").is_ok(); + + // padded_row_adjustment (CPU per-chip folder accumulator at all-zero trace). + let t_setup = std::time::Instant::now(); + let mut padded_row_adjustment = vec![Ext::zero(); num_chips]; + for (i, chip) in chips.iter().enumerate() { + let prep_len = chip.preprocessed_width(); + let main_len = chip.width(); + let prep_zero = vec![Felt::zero(); prep_len]; + let main_zero = vec![Felt::zero(); main_len]; + let mut folder = ConstraintSumcheckFolder { + preprocessed: RowMajorMatrixView::new_row(&prep_zero), + main: RowMajorMatrixView::new_row(&main_zero), + accumulator: Ext::zero(), + public_values: &public_values, + constraint_index: 0, + powers_of_alpha: &powers_of_challenge + [powers_of_challenge.len() - chip.num_constraints..], + }; + chip.air.eval(&mut folder); + padded_row_adjustment[i] = folder.accumulator; + } + + let gkr_powers = + gkr_opening_batch_randomness.powers().skip(1).take(max_columns).collect::>(); + + let lambda: Ext = challenger.sample_ext_element(); + let powers_of_lambda = + lambda.powers().take(num_chips).collect_vec().into_iter().rev().collect::>(); + + let mut claim = Ext::zero(); + let LogUpEvaluations { point: gkr_point, chip_openings } = logup_evaluations; + for chip in chips.iter() { + let ChipEvaluation { + main_trace_evaluations: main_opening, + preprocessed_trace_evaluations: prep_opening, + } = chip_openings.get(chip.name()).unwrap(); + claim *= lambda; + let addend = main_opening + .evaluations() + .as_slice() + .iter() + .chain( + prep_opening + .as_ref() + .map_or_else(Vec::new, |mle| mle.evaluations().as_slice().to_vec()) + .iter(), + ) + .zip(gkr_powers.iter()) + .map(|(opening, power)| *opening * *power) + .sum::(); + claim += addend; + } + + let t_pra_and_claim = t_setup.elapsed(); + + // Select this shard's chips from the machine-wide bytecode (uploaded once + // at prover construction). Cheap pointer-view clones — no device upload. + let t_select = std::time::Instant::now(); + let compiled_dev: Vec = chips + .iter() + .enumerate() + .map(|(shard_idx, chip)| { + let m = *machine_bytecode + .chip_index + .get(chip.name()) + .expect("shard chip not present in machine bytecode"); + let mut view = machine_bytecode.chips[m].clone(); + // Re-stamp to the shard-relative index used by the per-shard + // arrays (`chip_heights`, `chip_alpha_offset`, …). + view.chip_idx = shard_idx as u32; + view + }) + .collect(); + let t_select = t_select.elapsed(); + if debug_timing { + tracing::info!( + "zerocheck setup: num_chips={} pra+claim={:?} select={:?}", + num_chips, + t_pra_and_claim, + t_select, + ); + } + + let main_poly = initialize_zerocheck_poly( + trace_mle, + chips, + compiled_dev, + machine_bytecode.clone(), + initial_heights.clone(), + public_values, + powers_of_challenge, + gkr_powers, + powers_of_lambda, + padded_row_adjustment, + gkr_point.clone(), + claim, + ); + + let mut univariate_polys = vec![]; + let mut jagged_point: Point = Point::from(vec![]); + let t_eval_total = std::time::Instant::now(); + let mut total_fold = std::time::Duration::ZERO; + let mut total_eval = std::time::Duration::ZERO; + let mut total_chal = std::time::Duration::ZERO; + let t = std::time::Instant::now(); + let mut result = evaluate_zerocheck(&main_poly); + if debug_timing { + total_eval += t.elapsed(); + } + let t = std::time::Instant::now(); + let (mut point, mut next_claim) = challenger_update(&result, challenger); + if debug_timing { + total_chal += t.elapsed(); + } + univariate_polys.push(result); + jagged_point.add_dimension(point); + let t = std::time::Instant::now(); + let mut next_poly = zerocheck_fix_last_variable(main_poly, point, next_claim); + if debug_timing { + total_fold += t.elapsed(); + } + for _ in 0..max_log_row_count - 1 { + let t = std::time::Instant::now(); + result = evaluate_zerocheck(&next_poly); + if debug_timing { + total_eval += t.elapsed(); + } + let t = std::time::Instant::now(); + (point, next_claim) = challenger_update(&result, challenger); + if debug_timing { + total_chal += t.elapsed(); + } + univariate_polys.push(result); + jagged_point.add_dimension(point); + let t = std::time::Instant::now(); + next_poly = zerocheck_fix_last_variable(next_poly, point, next_claim); + if debug_timing { + total_fold += t.elapsed(); + } + } + if debug_timing { + tracing::info!( + "zerocheck: total={:?} eval={:?} fold={:?} chal={:?}", + t_eval_total.elapsed(), + total_eval, + total_fold, + total_chal + ); + } + + let final_jagged_data = + unsafe { next_poly.data.as_ref().dense_data.dense.copy_into_host_vec() }; + + let mut idx = 0; + let mut individual_column_evals = vec![Ext::zero(); data_input_heights.len()]; + for i in 0..data_input_heights.len() { + if data_input_heights[i] != 0 { + individual_column_evals[i] = final_jagged_data[idx]; + idx += 4; + } + } + + let mut preprocessed_ptr = 0; + let mut main_ptr = total_preprocessed_columns; + let mut opened_values: BTreeMap> = BTreeMap::new(); + challenger.observe(Felt::from_canonical_usize(chips.len())); + for (i, chip) in chips.iter().enumerate() { + let preprocessed_width = chip.preprocessed_width(); + let preprocessed = AirOpenedValues { + local: individual_column_evals[preprocessed_ptr..preprocessed_ptr + preprocessed_width] + .to_vec(), + }; + challenger.observe_variable_length_extension_slice(&preprocessed.local); + preprocessed_ptr += preprocessed_width; + let width = chip.width(); + let main = + AirOpenedValues { local: individual_column_evals[main_ptr..main_ptr + width].to_vec() }; + challenger.observe_variable_length_extension_slice(&main.local); + main_ptr += width; + opened_values.insert( + chip.air.name().to_string(), + ChipOpenedValues { + preprocessed, + main, + degree: Point::from_usize( + initial_heights[i] as usize, + (max_log_row_count + 1) as usize, + ), + }, + ); + } + + let partial_sumcheck_proof = PartialSumcheckProof { + univariate_polys, + claimed_sum: claim, + point_and_eval: (jagged_point, next_claim), + }; + let shard_open_values = ShardOpenedValues { chips: opened_values }; + (shard_open_values, partial_sumcheck_proof) +}