From 631f7139525b059b851848980d6336d6540e29a2 Mon Sep 17 00:00:00 2001 From: Gabriel Barreto Date: Fri, 31 Jul 2026 15:56:41 -0300 Subject: [PATCH 1/3] aiur: circuit-level function grouping (no source pragma) Several functions can now be proven by ONE circuit: the members are walked like branches of a single function - auxiliary columns and lookup slots are shared across members (the same save/restore sharing match arms already use), selector columns are laid out consecutively per member, and every member folds its selector-gated return message (carrying its own function index) into the shared lookup slot 0 against a single shared multiplicity column. One extra constraint enforces cross-member exclusivity: the sum of the members' top-block selectors must be boolean. Callers are untouched - calls still target function indices on the function channel - so grouping is invisible to execution, the query record, and the interpreter. Grouping is a CIRCUIT-level choice, not a property of the function library, so there is no source annotation: Source.Toplevel.compile builds the default singleton partition (Bytecode.Toplevel.circuits, one circuit per constrained function - behavior-identical to before), and CompiledToplevel.groupFunctions optionally regroups it by function NAME (validated: known, constrained, non-entry, no duplicates). The merged layout is max inputs, summed selectors, max auxiliaries, max lookups - so grouping fits rarely-called functions of similar shape: each (rare) row pays the group's selector count while the system sheds one circuit (vk entry, commitment matrix, verifier work) per absorbed member. Rust consumes the partition directly (bytecode Circuit via FFI; constraints/trace/synthesis iterate circuits, witness rows concatenate the members' queried rows in member order). The stage-2 lookup group size and the branchless raw-argument rule now key on the CIRCUIT layout: multi-member circuits are branching by construction, so their arguments are selector-superposed exactly like match arms. Tests: the aiur suite proves the same toplevel twice - ungrouped and with a 3-member test group (different arities, matches, cross-member call, recursion) - plus structural checks on the partition (members, merge-rule layout, every constrained function in exactly one circuit). All suites pass unchanged (ixvm FFT pins identical - the default partition is behavior-neutral); codegen is unaffected (execution ignores the partition). --- Ix/Aiur/Compiler.lean | 81 ++++++++++++++++++++++++++- Ix/Aiur/Compiler/Lower.lean | 2 +- Ix/Aiur/Stages/Bytecode.lean | 24 ++++++++ Ix/Aiur/Statistics.lean | 41 ++++++-------- Tests/Aiur/Aiur.lean | 79 ++++++++++++++++++++++++++ Tests/Aiur/Common.lean | 5 +- Tests/Main.lean | 12 +++- crates/aiur/src/bytecode.rs | 17 ++++++ crates/aiur/src/constraints.rs | 98 +++++++++++++++++++++------------ crates/aiur/src/synthesis.rs | 45 +++++++++------ crates/aiur/src/trace.rs | 84 +++++++++++++++++++++------- crates/ffi/src/aiur/toplevel.rs | 18 +++++- crates/ffi/src/lean.rs | 3 +- 13 files changed, 405 insertions(+), 104 deletions(-) diff --git a/Ix/Aiur/Compiler.lean b/Ix/Aiur/Compiler.lean index d85266632..752b2a878 100644 --- a/Ix/Aiur/Compiler.lean +++ b/Ix/Aiur/Compiler.lean @@ -42,6 +42,68 @@ def CompiledToplevel.getFuncIdx (ct : CompiledToplevel) (name : Lean.Name) : Option Bytecode.FunIdx := ct.nameMap[Global.mk name]? +/-- Regroup the circuit partition: each `(name, members)` in `groups` becomes +ONE circuit proving all the listed functions (branching on the member; see +`Bytecode.Circuit`), positioned where its first member's singleton circuit +was; every other constrained function keeps its singleton circuit. Grouping +is a circuit-level choice — the function "library", its bytecode, execution, +and the query record are untouched, and callers still target function +indices on the function channel. + +Grouping is favorable for RARELY-CALLED functions of similar shape: the +merged circuit sums the members' selector columns but takes the max of their +auxiliary columns, so each (rare) row pays the group's selector count, while +the system sheds one circuit (vk entry, commitment matrix, verifier work) +per absorbed member. + +Errors if a name is unknown, unconstrained (it has no circuit to group), an +entry function, listed twice, or if a group is empty. -/ +def CompiledToplevel.groupFunctions (ct : CompiledToplevel) + (groups : Array (String × Array Lean.Name)) : + Except String CompiledToplevel := do + let t := ct.bytecode + -- Resolve and validate the groups into member-index arrays. + let mut grouped : Std.HashMap Bytecode.FunIdx Nat := {} + let mut resolved : Array (String × Array Bytecode.FunIdx) := #[] + for (gname, names) in groups do + if names.isEmpty then + throw s!"group {gname} is empty" + let mut members := #[] + for name in names do + let some i := ct.getFuncIdx name + | throw s!"group {gname}: unknown function {name}" + let f := t.functions[i]! + unless f.constrained do + throw s!"group {gname}: {name} is unconstrained (it has no circuit)" + if f.entry then + throw s!"group {gname}: {name} is an entry function" + if grouped.contains i then + throw s!"group {gname}: {name} is already grouped" + grouped := grouped.insert i resolved.size + members := members.push i + resolved := resolved.push (gname, members) + -- Rebuild the partition in first-occurrence order over the existing + -- (singleton-ordered) circuits. + let mut circuits : Array Bytecode.Circuit := #[] + let mut placed : Array Bool := .replicate resolved.size false + for c in t.circuits do + let members := c.members + if h : members.size = 1 then + let i := members[0] + match grouped[i]? with + | none => circuits := circuits.push c + | some g => + unless placed[g]! do + placed := placed.set! g true + let (gname, ms) := resolved[g]! + let layout := ms.foldl (init := t.functions[ms[0]!]!.layout) + fun acc m => if m == ms[0]! then acc + else acc.merge t.functions[m]!.layout + circuits := circuits.push { name := gname, members := ms, layout } + else + throw "groupFunctions: partition already grouped; group from a freshly compiled toplevel" + pure { ct with bytecode := { t with circuits } } + /-- Termination helper for the `Block`/`Ctrl` traversal below. -/ private theorem Bytecode.Block.sizeOf_ctrl_lt'' (b : Bytecode.Block) : sizeOf b.ctrl < sizeOf b := by @@ -90,6 +152,18 @@ decreasing_by | (apply Prod.Lex.left; exact Bytecode.Block.sizeOf_ctrl_lt'' _) end +/-- The default circuit partition: one singleton circuit per constrained +function, in function-index order, named by `nameOf`. -/ +def Bytecode.Toplevel.singletonCircuits (t : Bytecode.Toplevel) + (nameOf : Bytecode.FunIdx → String) : Array Bytecode.Circuit := Id.run do + let mut circuits : Array Bytecode.Circuit := #[] + for h : i in [:t.functions.size] do + let f := t.functions[i] + if f.constrained then + circuits := circuits.push + { name := nameOf i, members := #[i], layout := f.layout } + pure circuits + /-- Compute which functions need a circuit. A function needs a circuit iff it is reachable from an entry point through a chain of constrained call edges. -/ def Bytecode.Toplevel.needsCircuit (t : Bytecode.Toplevel) : Array Bool := Id.run do @@ -118,11 +192,16 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev let (bytecodeRaw, preNameMap) ← concDecls.toBytecode let (bytecodeDedup, remap) := bytecodeRaw.deduplicate let needs := bytecodeDedup.needsCircuit - let bytecode := { bytecodeDedup with + let bytecode : Bytecode.Toplevel := { bytecodeDedup with functions := bytecodeDedup.functions.mapIdx fun i f => { f with constrained := needs[i]! } } let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx)) fun acc name idx => acc.insert name (remap idx) + -- Singleton circuits are labeled with (one of) the function's source names. + let reverseMap := nameMap.fold (init := (∅ : Std.HashMap Bytecode.FunIdx String)) + fun acc global idx => if acc.contains idx then acc else acc.insert idx (toString global) + let bytecode := { bytecode with + circuits := bytecode.singletonCircuits fun i => reverseMap[i]?.getD s!"" } pure (CompiledToplevel.mk t bytecode nameMap) /-- Progress helper: given success of the three `Except`-returning stages, diff --git a/Ix/Aiur/Compiler/Lower.lean b/Ix/Aiur/Compiler/Lower.lean index 4cc3dc33b..136df82ec 100644 --- a/Ix/Aiur/Compiler/Lower.lean +++ b/Ix/Aiur/Compiler/Lower.lean @@ -631,7 +631,7 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) : let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes pure (functions.push function, memSizes, nameMap) | _ => pure acc - pure (⟨functions, memSizes.toArray⟩, nameMap) + pure (⟨functions, memSizes.toArray, #[]⟩, nameMap) end Aiur diff --git a/Ix/Aiur/Stages/Bytecode.lean b/Ix/Aiur/Stages/Bytecode.lean index acd8e2579..21c52fd66 100644 --- a/Ix/Aiur/Stages/Bytecode.lean +++ b/Ix/Aiur/Stages/Bytecode.lean @@ -122,9 +122,33 @@ structure Function where constrained : Bool deriving Inhabited, Repr +/-- A circuit of the proving system, backing one or more functions. By +default every constrained function gets a singleton circuit named after it; +`CompiledToplevel.groupFunctions` can regroup several functions into one +circuit whose branching selects the member function. `layout` is the merged +layout: max `inputSize`, sum of `selectors`, max `auxiliaries` (which +includes the single shared multiplicity column), max `lookups` (slot 0 is +the shared return lookup). -/ +structure Circuit where + name : String + members : Array FunIdx + layout : FunctionLayout + deriving Inhabited, Repr + +/-- Merged layout of a group of functions (see `Circuit`). -/ +def FunctionLayout.merge (a b : FunctionLayout) : FunctionLayout where + inputSize := a.inputSize.max b.inputSize + selectors := a.selectors + b.selectors + auxiliaries := a.auxiliaries.max b.auxiliaries + lookups := a.lookups.max b.lookups + structure Toplevel where functions : Array Function memorySizes : Array Nat + /-- Circuit partition of the constrained functions, in first-occurrence + order. Built by `Source.Toplevel.compile` (singletons by default; see + `CompiledToplevel.groupFunctions`); empty on a freshly lowered toplevel. -/ + circuits : Array Circuit := #[] deriving Repr end Bytecode diff --git a/Ix/Aiur/Statistics.lean b/Ix/Aiur/Statistics.lean index aa91de166..095a1d2d8 100644 --- a/Ix/Aiur/Statistics.lean +++ b/Ix/Aiur/Statistics.lean @@ -81,44 +81,37 @@ def computeStats (compiled : CompiledToplevel) (queryCounts : Array QueryCount) (logBlowup : Nat := defaultCommitmentParameters.logBlowup) : ExecutionStats := let t := compiled.bytecode - -- Invert nameMap to get FunIdx → String - let reverseMap := compiled.nameMap.fold (init := (∅ : Std.HashMap Bytecode.FunIdx String)) - fun acc global idx => if !acc.contains idx then acc.insert idx (toString global) else acc let nAllFuns := t.functions.size - let nConstrained := t.functions.foldl (fun n f => if f.constrained then n + 1 else n) 0 - -- Shapes arrive in canonical system order: constrained functions - -- (ascending index), memories, `Bytes1`, `Bytes2`. A mismatch means the - -- shapes were built from a different toplevel; misindexing would silently - -- attribute costs to the wrong circuits. - if shapes.size != nConstrained + t.memorySizes.size + 2 then + -- Shapes arrive in canonical system order: function circuits (grouped; + -- singletons for ungrouped functions, in ascending member index), + -- memories, `Bytes1`, `Bytes2`. A mismatch means the shapes were built + -- from a different toplevel; misindexing would silently attribute costs + -- to the wrong circuits. + if shapes.size != t.circuits.size + t.memorySizes.size + 2 then panic! s!"computeStats: {shapes.size} circuit shapes for \ - {nConstrained} constrained functions + {t.memorySizes.size} memories + 2 gadgets" + {t.circuits.size} function circuits + {t.memorySizes.size} memories + 2 gadgets" else let mkStats (name : String) (shape : CircuitShape) (h hits : Nat) : CircuitStats := { name, width := shape.committedWidth, height := h, cacheHits := hits, fftCost := fftCost shape h logBlowup, uncachedFftCost := fftCost shape (h + hits) logBlowup } - let functionCircuits := Id.run do - let mut acc := #[] - let mut shapeIdx := 0 - for i in [:nAllFuns] do - if t.functions[i]!.constrained then - let shape := shapes[shapeIdx]! - shapeIdx := shapeIdx + 1 - let qc := queryCounts[i]! - let name := reverseMap[i]?.getD s!"" - acc := acc.push - (mkStats name shape qc.uniqueRows (qc.totalHits - qc.uniqueRows)) - acc + -- One row per function circuit: heights and cache hits are summed over + -- the circuit's member functions (singletons sum over one). + let functionCircuits := t.circuits.mapIdx fun cIdx c => + let shape := shapes[cIdx]! + let (h, hits) := c.members.foldl (init := (0, 0)) fun (h, hits) i => + let qc := queryCounts[i]! + (h + qc.uniqueRows, hits + (qc.totalHits - qc.uniqueRows)) + mkStats c.name shape h hits let memoryCircuits := t.memorySizes.mapIdx fun i size => - let shape := shapes[nConstrained + i]! + let shape := shapes[t.circuits.size + i]! let qc := queryCounts[nAllFuns + i]! mkStats s!"memory[{size}]" shape qc.uniqueRows (qc.totalHits - qc.uniqueRows) -- The byte gadgets commit full-table traces in every proof: their height -- is the (fixed) preprocessed height, independent of the query set, so -- they carry no cache-hit counterfactual. let gadgetCircuits := #["Bytes1", "Bytes2"].mapIdx fun i name => - let shape := shapes[nConstrained + t.memorySizes.size + i]! + let shape := shapes[t.circuits.size + t.memorySizes.size + i]! mkStats name shape shape.preprocessedHeight 0 let circuits := (functionCircuits ++ memoryCircuits ++ gadgetCircuits).qsort (·.fftCost > ·.fftCost) diff --git a/Tests/Aiur/Aiur.lean b/Tests/Aiur/Aiur.lean index d78fa2f6e..1e7b9fcc3 100644 --- a/Tests/Aiur/Aiur.lean +++ b/Tests/Aiur/Aiur.lean @@ -719,6 +719,38 @@ def toplevel := ⟦ let s5 = c[4] + c[0]; -- 255 s1 + s2 + 10 * s3 + s4 + s5 -- 1309 } + + --------------------------------------------------------------------------- + -- Grouped circuits (`CompiledToplevel.groupFunctions`): the test runner + -- groups these three into one circuit whose branching selects the member. + -- Grouping is a circuit-level choice, so there is NO source annotation: + -- the same functions also run ungrouped in the plain suite. Members + -- differ in arity, output and branch count, call each other (through the + -- shared circuit) and recurse. + --------------------------------------------------------------------------- + fn grouped_double(x: G) -> G { + x + x + } + + -- Different arity, a match (two selectors), calls a fellow group member. + fn grouped_pick(t: G, a: G, b: G) -> G { + match t { + 0 => grouped_double(a), + _ => b, + } + } + + -- Recursive group member: self-calls route through the shared circuit. + fn grouped_sum_range(n: G) -> G { + match n { + 0 => 0, + _ => n + grouped_sum_range(n - 1), + } + } + + pub fn calls_grouped(t: G, a: G, b: G) -> G { + grouped_pick(t, a, b) + grouped_sum_range(a) + } ⟧ /-- The PROVING suite: every case runs the full prove+verify pipeline @@ -843,6 +875,53 @@ def aiurTestCases : List AiurTestCase := [ -- Unconstrained g_to_bytes / g_inverse hints: all cases in one proof .prove `hint_test #[] #[1309], + + -- Grouped-circuit member functions, run UNGROUPED here (the grouped + -- variant runs in the grouped env; see `testGroups`). + -- t=0 → grouped_double(5) + Σ1..5 = 10 + 15 = 25; t≠0 → 9 + Σ1..3 = 15. + .prove `calls_grouped #[0, 5, 9] #[25] + (label := "calls_grouped(0,5,9)"), + .prove `calls_grouped #[1, 3, 9] #[15] + (label := "calls_grouped(1,3,9)"), ] +/-- The grouping the `aiur` runner applies for the grouped environment. -/ +def testGroups : Array (String × Array Lean.Name) := + #[("test_group", #[`grouped_double, `grouped_pick, `grouped_sum_range])] + +def groupedTestCases : List AiurTestCase := [ + .prove `calls_grouped #[0, 5, 9] #[25] + (label := "calls_grouped(0,5,9) [grouped]"), + .prove `calls_grouped #[1, 3, 9] #[15] + (label := "calls_grouped(1,3,9) [grouped]"), +] + +/-- Structural checks on the grouped partition: the grouped circuit exists, +holds exactly its members, its layout follows the merge rule (max inputs, +summed selectors, max auxiliaries, max lookups), and every constrained +function lands in exactly one circuit. -/ +def groupingStructureChecks (compiled : Aiur.CompiledToplevel) : TestSeq := + let t := compiled.bytecode + let memberOf := fun (name : Lean.Name) => compiled.getFuncIdx name |>.get! + let expectedMembers := + #[`grouped_double, `grouped_pick, `grouped_sum_range].map memberOf + match t.circuits.find? (·.name == "test_group") with + | none => test "test_group circuit exists" false + | some c => + let layouts := c.members.map (t.functions[·]!.layout) + let expected := layouts.foldl (init := (⟨0, 0, 0, 0⟩ : Aiur.Bytecode.FunctionLayout)) + Aiur.Bytecode.FunctionLayout.merge + let allCircuitMembers := t.circuits.flatMap (·.members) + let constrained := (Array.range t.functions.size).filter + (t.functions[·]!.constrained) + test "test_group circuit exists" true ++ + test "test_group members" (c.members == expectedMembers) ++ + test "test_group layout follows the merge rule" + (c.layout.inputSize == expected.inputSize && + c.layout.selectors == expected.selectors && + c.layout.auxiliaries == expected.auxiliaries && + c.layout.lookups == expected.lookups) ++ + test "every constrained function is in exactly one circuit" + (allCircuitMembers.qsort (· < ·) == constrained) + end diff --git a/Tests/Aiur/Common.lean b/Tests/Aiur/Common.lean index f90e64f66..701462489 100644 --- a/Tests/Aiur/Common.lean +++ b/Tests/Aiur/Common.lean @@ -72,10 +72,13 @@ structure AiurTestEnv where aiurSystem : Aiur.AiurSystem shapes : Array Aiur.CircuitShape -def AiurTestEnv.build (toplevelFn : Except Aiur.Global Aiur.Source.Toplevel) : +def AiurTestEnv.build (toplevelFn : Except Aiur.Global Aiur.Source.Toplevel) + (groups : Array (String × Array Lean.Name) := #[]) : Except String AiurTestEnv := do let toplevel ← toplevelFn.mapError toString let compiled ← toplevel.compile + let compiled ← if groups.isEmpty then pure compiled + else compiled.groupFunctions groups let decls ← toplevel.mkDecls.mapError toString let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters return ⟨compiled, decls, aiurSystem, aiurSystem.circuitShapes⟩ diff --git a/Tests/Main.lean b/Tests/Main.lean index 592033552..5a4cc9264 100644 --- a/Tests/Main.lean +++ b/Tests/Main.lean @@ -130,7 +130,17 @@ def primaryRunners : List (String × IO UInt32) := [ IO.println "aiur-prove" match AiurTestEnv.build (pure toplevel) with | .error e => IO.eprintln s!"Aiur setup failed: {e}"; return 1 - | .ok env => LSpec.lspecEachIO aiurTestCases fun tc => pure (env.runTestCase tc)), + | .ok env => do + let r1 ← LSpec.lspecEachIO aiurTestCases fun tc => pure (env.runTestCase tc) + -- The same toplevel with `testGroups` applied: the members share one + -- circuit, and the whole suite of grouped cases proves through it. + match AiurTestEnv.build (pure toplevel) testGroups with + | .error e => IO.eprintln s!"Aiur grouped setup failed: {e}"; return 1 + | .ok genv => do + let r2 ← LSpec.lspecEachIO groupedTestCases fun tc => pure (genv.runTestCase tc) + let r3 ← LSpec.lspecIO + (.ofList [("aiur-grouping", [groupingStructureChecks genv.compiled])]) [] + return if r1 == 0 && r2 == 0 && r3 == 0 then 0 else 1), ("aiur-hashes", do IO.println "aiur-hashes" let .ok blake3Env := AiurTestEnv.build (do diff --git a/crates/aiur/src/bytecode.rs b/crates/aiur/src/bytecode.rs index 57e426e82..3349906ff 100644 --- a/crates/aiur/src/bytecode.rs +++ b/crates/aiur/src/bytecode.rs @@ -5,6 +5,23 @@ use super::G; pub struct Toplevel { pub functions: Vec, pub memory_sizes: Vec, + /// Circuit partition of the constrained functions, in first-occurrence + /// order. Computed by the Lean compiler (singletons by default; + /// `CompiledToplevel.groupFunctions` regroups); every constrained + /// function appears in exactly one circuit. + pub circuits: Vec, +} + +/// A circuit of the proving system, backing one or more functions. Ungrouped +/// functions get a singleton circuit; grouped functions share one circuit +/// whose branching selects the member function. +/// +/// `layout` is the merged layout: max `input_size`, sum of `selectors`, max +/// `auxiliaries` (which includes the single shared multiplicity column), max +/// `lookups` (slot 0 is the shared return lookup). +pub struct Circuit { + pub members: Vec, + pub layout: FunctionLayout, } pub struct Function { diff --git a/crates/aiur/src/constraints.rs b/crates/aiur/src/constraints.rs index 26488ce50..8691198b2 100644 --- a/crates/aiur/src/constraints.rs +++ b/crates/aiur/src/constraints.rs @@ -6,7 +6,7 @@ use std::{array, ops::Range, sync::LazyLock}; use crate::{ FxIndexMap, G, - bytecode::{Block, Ctrl, Function, FunctionLayout, Op, Toplevel, ValIdx}, + bytecode::{Block, Ctrl, Op, Toplevel, ValIdx}, function_channel, gadgets::{ AiurGadget, @@ -65,11 +65,18 @@ pub struct Constraints { } struct ConstraintState { + /// Index of the circuit member currently being walked. function_index: G, - /// Exactly one selector: the function has a single leaf block (no - /// matches), so every lookup slot is written by exactly one branch. + /// Exactly one selector: the circuit backs a single function with a + /// single leaf block (no matches), so every lookup slot is written by + /// exactly one branch. branchless: bool, - layout: FunctionLayout, + /// Input size of the current member (inputs live in columns + /// `0..input_size` for every member; the circuit reserves the max). + input_size: usize, + /// Column of the current member's first selector: the circuit's input + /// block plus the selector counts of the members walked before it. + sel_base: usize, column: usize, lookup: usize, lookups: Vec>, @@ -88,7 +95,7 @@ struct SharedState { impl ConstraintState { fn selector_index(&self, sel: usize) -> usize { - sel + self.layout.input_size + sel + self.sel_base } /// Selector-gate a lookup argument. Lookup slots shared across branches @@ -131,34 +138,75 @@ impl ConstraintState { } impl Toplevel { + /// Build the constraints of one circuit. The circuit's members are walked + /// like branches of a single function: each walk restarts the auxiliary + /// column / lookup-slot counters (so members share those, like match arms + /// do), while selector columns are laid out consecutively per member. All + /// members fold their return message into the shared lookup slot 0, gated + /// by their own selectors and carrying their own function index, against + /// the single shared multiplicity column. pub fn build_constraints( &self, - function_index: usize, + circuit_index: usize, ) -> (Constraints, Vec>) { - let function = &self.functions[function_index]; + let circuit = &self.circuits[circuit_index]; + let layout = circuit.layout; let constraints = Constraints { zeros: vec![], - selectors: 0..0, - width: function.layout.width(), + selectors: layout.input_size..layout.input_size + layout.selectors, + width: layout.width(), }; let mut state = ConstraintState { - function_index: G::from_usize(function_index), - branchless: function.layout.selectors == 1, - layout: function.layout, + function_index: G::ZERO, + branchless: layout.selectors == 1, + input_size: 0, + sel_base: 0, column: 0, lookup: 0, map: vec![], - lookups: vec![empty_lookup(); function.layout.lookups], + lookups: vec![empty_lookup(); layout.lookups], constraints, yield_info: vec![], }; - function.build_constraints(&mut state); + // The shared multiplicity column: first auxiliary, right after the + // selectors. The return lookup occupies the first lookup slot. + let multiplicity = var(layout.input_size + layout.selectors); + state.lookups[0].multiplicity = -multiplicity; + let aux_start = layout.input_size + layout.selectors + 1; + let mut sel_base = layout.input_size; + let mut circuit_sel = Expr::from(G::ZERO); + for &member in &circuit.members { + let function = &self.functions[member]; + state.function_index = G::from_usize(member); + state.input_size = function.layout.input_size; + state.sel_base = sel_base; + state.column = aux_start; + state.lookup = 1; + state.map.clear(); + (0..function.layout.input_size).for_each(|i| state.map.push((var(i), 1))); + let body_sel = function.body.get_block_selector(&state); + circuit_sel = circuit_sel + body_sel.clone(); + function.body.collect_constraints(body_sel, &mut state); + debug_assert!(state.yield_info.is_empty()); + sel_base += function.layout.selectors; + } // The old `Air::eval` asserted each selector column boolean; the new // system compiles a constraint vector, so materialize those explicitly. for sel in state.constraints.selectors.clone() { let s = var(sel); state.constraints.zeros.push(s.clone() * (s - konst(G::ONE))); } + // Cross-member exclusivity: the circuit-level selector (the sum of the + // members' top-block selectors) must be boolean, so at most one member + // is active per row and the shared return lookup emits a single + // member's message. A singleton circuit already gets this from its top + // block's own boolean constraint. + if circuit.members.len() > 1 { + state + .constraints + .zeros + .push(circuit_sel.clone() * (Expr::from(G::ONE) - circuit_sel)); + } (state.constraints, state.lookups) } } @@ -167,26 +215,6 @@ fn empty_lookup() -> Lookup { Lookup { multiplicity: konst(G::ZERO), args: vec![] } } -impl Function { - fn build_constraints(&self, state: &mut ConstraintState) { - // the first columns are occupied by the input, which is also mapped - state.column += self.layout.input_size; - (0..self.layout.input_size).for_each(|i| state.map.push((var(i), 1))); - // then comes the selectors, which are not mapped - let init_sel = state.column; - let final_sel = state.column + self.layout.selectors; - state.constraints.selectors = init_sel..final_sel; - state.column = final_sel; - // the multiplicity occupies another column - let multiplicity = var(state.column); - state.column += 1; - // the return lookup occupies the first lookup slot - state.lookups[0].multiplicity = -multiplicity.clone(); - state.lookup += 1; - self.body.collect_constraints(self.body.get_block_selector(state), state); - } -} - impl Block { fn collect_constraints(&self, sel: Expr, state: &mut ConstraintState) { // Boolean constraint for this block's selector @@ -277,7 +305,7 @@ impl Ctrl { ]; // input args.extend( - (0..state.layout.input_size) + (0..state.input_size) .map(|arg| state.gate(&sel, state.map[arg].0.clone())), ); // output diff --git a/crates/aiur/src/synthesis.rs b/crates/aiur/src/synthesis.rs index 47d7ea37c..6099213ec 100644 --- a/crates/aiur/src/synthesis.rs +++ b/crates/aiur/src/synthesis.rs @@ -89,20 +89,17 @@ impl AiurSystem { }); }; - // Constrained functions (ascending index). - for i in 0..toplevel.functions.len() { - if !toplevel.functions[i].constrained { - continue; - } + // Function circuits, in partition order (singletons unless grouped). + for i in 0..toplevel.circuits.len() { let (constraints, lookups) = toplevel.build_constraints(i); - // A branchless function's lookup arguments are sent raw (degree 1; + // A branchless circuit's lookup arguments are sent raw (degree 1; // see `ConstraintState::gate`), so two lookups fit in one chained // accumulator step at degree 3 — within the degree the selector-gated - // constraints already pay for. Branching functions keep k = 1: their + // constraints already pay for. Branching circuits keep k = 1: their // superposed arguments are degree 2, and grouping would push the // logUp constraints past the quotient budget. let group_size = - if toplevel.functions[i].layout.selectors == 1 && lookups.len() >= 2 { + if toplevel.circuits[i].layout.selectors == 1 && lookups.len() >= 2 { 2 } else { 1 @@ -156,11 +153,8 @@ impl AiurSystem { /// order the circuits were chained in [`AiurSystem::build`], so index `i` /// of the returned `Vec` corresponds to `self.system.circuits[i]`. fn circuit_types(&self) -> Vec { - let functions = (0..self.toplevel.functions.len()).filter_map(|idx| { - self.toplevel.functions[idx] - .constrained - .then_some(CircuitType::Function { idx }) - }); + let functions = (0..self.toplevel.circuits.len()) + .map(|idx| CircuitType::Function { idx }); let memories = self .toplevel .memory_sizes @@ -384,6 +378,25 @@ mod tests { /// fresh auxiliary column pinned by `sel * (col - a*b)`. /// - `lookups = 1`: the function-provide (return) lookup in slot 0, which /// pulls the claim `[function_channel, fun_idx, a, b, a*b]`. + /// + /// Test-side singleton partition (production circuits come pre-built from + /// the Lean compiler). + fn with_singleton_circuits( + functions: Vec, + memory_sizes: Vec, + ) -> Toplevel { + let circuits = functions + .iter() + .enumerate() + .filter(|(_, f)| f.constrained) + .map(|(i, f)| crate::bytecode::Circuit { + members: vec![i], + layout: f.layout, + }) + .collect(); + Toplevel { functions, memory_sizes, circuits } + } + fn mul_toplevel() -> Toplevel { let body = Block { ops: vec![Op::Mul(0, 1)], ctrl: Ctrl::Return(0, vec![2]) }; @@ -398,7 +411,7 @@ mod tests { entry: true, constrained: true, }; - Toplevel { functions: vec![function], memory_sizes: vec![] } + with_singleton_circuits(vec![function], vec![]) } fn xor_splits_toplevel() -> Toplevel { @@ -417,7 +430,7 @@ mod tests { entry: true, constrained: true, }; - Toplevel { functions: vec![function], memory_sizes: vec![] } + with_singleton_circuits(vec![function], vec![]) } #[test] @@ -523,7 +536,7 @@ mod tests { constrained: true, }; - Toplevel { functions: vec![f, g], memory_sizes: vec![1] } + with_singleton_circuits(vec![f, g], vec![1]) } #[test] diff --git a/crates/aiur/src/trace.rs b/crates/aiur/src/trace.rs index e87dfe151..66e88cf92 100644 --- a/crates/aiur/src/trace.rs +++ b/crates/aiur/src/trace.rs @@ -12,7 +12,7 @@ use rayon::{ use crate::{ FxIndexMap, G, - bytecode::{Block, Ctrl, Function, Op, Toplevel}, + bytecode::{Block, Ctrl, Function, FunctionLayout, Op, Toplevel}, execute::{ IOBuffer, IOKeyInfo, QueryRecord, find_unconstrained_big_uint_div_mod, g_inverse_value, @@ -20,6 +20,7 @@ use crate::{ function_channel, gadgets::{bytes1::Bytes1, bytes2::Bytes2}, memory::Memory, + querymap::QueryRef, u8_add_channel, u8_and_channel, u8_bit_decomposition_channel, u8_less_than_channel, u8_mul_channel, u8_or_channel, u8_range_check_channel, u8_shift_left_channel, u8_shift_right_channel, u8_sub_channel, @@ -55,15 +56,23 @@ fn u32_sum(values: &[u64]) -> ([G; 4], G) { } impl<'a, 'b> ColumnMutSlice<'a, 'b> { + /// Slice a circuit row into the regions of one member function: the + /// member's inputs are a prefix of the circuit's input block, its + /// selectors a sub-range of the circuit's selector block at `sel_offset`, + /// and the auxiliary block is shared by all members. fn from_slice( function: &Function, + circuit_layout: &FunctionLayout, + sel_offset: usize, slice: &'a mut [G], lookups: &'a mut LookupRowMut<'b, G>, ) -> Self { - let (inputs, slice) = slice.split_at_mut(function.layout.input_size); - let (selectors, slice) = slice.split_at_mut(function.layout.selectors); - let (auxiliaries, slice) = slice.split_at_mut(function.layout.auxiliaries); - assert!(slice.is_empty()); + let (inputs, slice) = slice.split_at_mut(circuit_layout.input_size); + let (selectors, auxiliaries) = slice.split_at_mut(circuit_layout.selectors); + assert_eq!(auxiliaries.len(), circuit_layout.auxiliaries); + let inputs = &mut inputs[..function.layout.input_size]; + let selectors = + &mut selectors[sel_offset..sel_offset + function.layout.selectors]; Self { inputs, selectors, auxiliaries, lookups } } @@ -92,22 +101,49 @@ struct TraceContext<'a> { query_record: &'a QueryRecord, } +/// One row of a circuit trace: the member function it belongs to, the +/// member's selector offset within the circuit, its function index, and the +/// recorded query. +struct RowMeta<'a> { + function: &'a Function, + sel_offset: usize, + function_index: G, + inputs: &'a [G], + result: QueryRef<'a>, +} + impl Toplevel { pub fn witness_data( &self, - function_index: usize, + circuit_index: usize, query_record: &QueryRecord, io_buffer: &IOBuffer, slot_arg_widths: &[usize], ) -> (RowMajorMatrix, LookupValues) { - let func = &self.functions[function_index]; - let width = func.width(); - let unfiltered_queries = &query_record.function_queries[function_index]; - let queries = unfiltered_queries - .iter() - .filter(|(_, res)| !res.multiplicity.is_zero()) - .collect::>(); - let height_no_padding = queries.len(); + let circuit = &self.circuits[circuit_index]; + let layout = &circuit.layout; + let width = layout.width(); + // Concatenate the members' queried rows, in member order. + let mut rows_meta = Vec::new(); + let mut sel_offset = 0; + for &member in &circuit.members { + let function = &self.functions[member]; + let function_index = G::from_usize(member); + rows_meta.extend( + query_record.function_queries[member] + .iter() + .filter(|(_, res)| !res.multiplicity.is_zero()) + .map(|(inputs, result)| RowMeta { + function, + sel_offset, + function_index, + inputs, + result, + }), + ); + sel_offset += function.layout.selectors; + } + let height_no_padding = rows_meta.len(); // An unqueried circuit yields an EMPTY trace (not a padded height-1 one): // the prover deactivates it, so it is neither committed nor opened. let height = if height_no_padding == 0 { @@ -126,21 +162,27 @@ impl Toplevel { .zip(row_writers[..height_no_padding].par_iter_mut()) .enumerate() .for_each(|(i, (row, lookups))| { - let (inputs, result) = queries[i]; + let meta = &rows_meta[i]; let index = &mut ColumnIndex { auxiliary: 0, // we skip the first lookup, which is reserved for return lookup: 1, }; - let slice = &mut ColumnMutSlice::from_slice(func, row, lookups); + let slice = &mut ColumnMutSlice::from_slice( + meta.function, + layout, + meta.sel_offset, + row, + lookups, + ); let context = TraceContext { - function_index: G::from_usize(function_index), - inputs, - multiplicity: result.multiplicity, - output: result.output, + function_index: meta.function_index, + inputs: meta.inputs, + multiplicity: meta.result.multiplicity, + output: meta.result.output, query_record, }; - func.populate_row(index, slice, context, io_buffer); + meta.function.populate_row(index, slice, context, io_buffer); }); drop(row_writers); let trace = RowMajorMatrix::new(rows, width); diff --git a/crates/ffi/src/aiur/toplevel.rs b/crates/ffi/src/aiur/toplevel.rs index 4b88cdacd..fbead94f7 100644 --- a/crates/ffi/src/aiur/toplevel.rs +++ b/crates/ffi/src/aiur/toplevel.rs @@ -2,12 +2,15 @@ use multi_stark::p3_field::PrimeCharacteristicRing; use lean_ffi::object::{LeanBorrowed, LeanCtor, LeanRef}; +use crate::lean::LeanAiurCircuit; use crate::lean::LeanAiurFunction; use crate::lean::LeanAiurToplevel; use aiur::{ FxIndexMap, G, - bytecode::{Block, Ctrl, Function, FunctionLayout, Op, Toplevel, ValIdx}, + bytecode::{ + Block, Circuit, Ctrl, Function, FunctionLayout, Op, Toplevel, ValIdx, + }, }; use crate::aiur::{lean_unbox_g, lean_unbox_nat_as_usize}; @@ -294,14 +297,23 @@ fn decode_function(ctor: LeanCtor>) -> Function { Function { body, layout, entry, constrained } } +fn decode_circuit(ctor: LeanCtor>) -> Circuit { + let ctor = LeanAiurCircuit::from_ctor(ctor); + // Object field 0 is the circuit's display name (`String`), unused here. + let members = ctor.get_obj(1).as_array().map(|x| lean_unbox_nat_as_usize(&x)); + let layout = decode_function_layout(ctor.get_obj(2).as_ctor()); + Circuit { members, layout } +} + pub(crate) fn decode_toplevel( obj: &LeanAiurToplevel, ) -> Toplevel { let ctor = obj.as_ctor(); - let [functions_obj, memory_sizes_obj] = ctor.objs::<2>(); + let [functions_obj, memory_sizes_obj, circuits_obj] = ctor.objs::<3>(); let functions = functions_obj.as_array().map(|o| decode_function(o.as_ctor())); let memory_sizes = memory_sizes_obj.as_array().map(|x| lean_unbox_nat_as_usize(&x)); - Toplevel { functions, memory_sizes } + let circuits = circuits_obj.as_array().map(|o| decode_circuit(o.as_ctor())); + Toplevel { functions, memory_sizes, circuits } } diff --git a/crates/ffi/src/lean.rs b/crates/ffi/src/lean.rs index e9040c66f..0cd5b9ed3 100644 --- a/crates/ffi/src/lean.rs +++ b/crates/ffi/src/lean.rs @@ -270,8 +270,9 @@ lean_ffi::lean_inductive! { // --- Aiur types --- - LeanAiurToplevel [ { num_obj: 2 } ]; + LeanAiurToplevel [ { num_obj: 3 } ]; LeanAiurFunction [ { num_obj: 2, num_8: 2 } ]; + LeanAiurCircuit [ { num_obj: 3 } ]; // Aiur FFI result structures (`Ix/Aiur/Semantics/BytecodeFfi.lean`, // `Ix/Aiur/Protocol.lean`). `IOBuffer` hashmaps cross the boundary as From 0489e127c16fd3b4df2cb003d571441a2b18c2d1 Mon Sep 17 00:00:00 2001 From: Gabriel Barreto Date: Mon, 3 Aug 2026 11:36:37 -0300 Subject: [PATCH 2/3] aiur: wire optional circuit-grouping application points (empty partitions) Route every site that compiles the IxVM kernel or the recursive-verifier toplevel for proving/verifying through `compileWithGroups` with a per-toplevel grouping datum (`IxVM.coldGroups`, `MultiStark.verifierColdGroups`) - CLI check/prove/verify, the ixvm test runner, the recursive-verifier tests, and the benches. `groupFunctions` resolves members by STRING name (the exact `toString` of the Global, the inverse of what statistics print, so measured groupings feed back verbatim). Both partitions start EMPTY, i.e. singleton circuits - behavior-identical to before; the data files are the single knob later commits turn. --- Benchmarks/RecursionDebug.lean | 6 +++--- Benchmarks/RecursiveVerifier.lean | 2 +- Benchmarks/Typecheck.lean | 4 ++-- Ix/Aiur/Compiler.lean | 15 +++++++++++++-- Ix/Cli/CheckCmd.lean | 6 +++--- Ix/Cli/ProveCmd.lean | 2 +- Ix/Cli/VerifyCmd.lean | 2 +- Ix/IxVM.lean | 1 + Ix/IxVM/ColdGroups.lean | 20 ++++++++++++++++++++ Ix/MultiStark.lean | 1 + Ix/MultiStark/VerifierColdGroups.lean | 20 ++++++++++++++++++++ Tests/Aiur/Aiur.lean | 4 ++-- Tests/Aiur/Common.lean | 2 +- Tests/Main.lean | 2 +- Tests/MultiStark.lean | 2 +- 15 files changed, 71 insertions(+), 18 deletions(-) create mode 100644 Ix/IxVM/ColdGroups.lean create mode 100644 Ix/MultiStark/VerifierColdGroups.lean diff --git a/Benchmarks/RecursionDebug.lean b/Benchmarks/RecursionDebug.lean index 8bc203020..847c91ab2 100644 --- a/Benchmarks/RecursionDebug.lean +++ b/Benchmarks/RecursionDebug.lean @@ -70,7 +70,7 @@ def proveConst (ixePath constName : String) (skipDeps : Bool) -- production toplevel no longer carries. let .ok toplevel := (if skipDeps then IxVM.ixVMFull else IxVM.ixVM) | IO.eprintln "IxVM toplevel merge failed"; return none - let .ok compiled := toplevel.compile + let .ok compiled := toplevel.compileWithGroups IxVM.coldGroups | IO.eprintln "IxVM compile failed"; return none let entrypoint := if skipDeps then `verify_const else `verify_claim let some funIdx := compiled.getFuncIdx entrypoint @@ -153,7 +153,7 @@ def main (args : List String) : IO UInt32 := do -- `--list-funcs`: dump the compiled verifier's funIdx → name table (for -- decoding fun_idx stacks printed by the Rust bytecode interpreter). if args.contains "--list-funcs" then - let .ok vCompiled := vTop.compile + let .ok vCompiled := vTop.compileWithGroups MultiStark.verifierColdGroups | IO.eprintln "multi-stark verifier compile failed"; return 1 let entries := vCompiled.nameMap.toArray.qsort (·.2 < ·.2) for (g, i) in entries do @@ -187,7 +187,7 @@ def main (args : List String) : IO UInt32 := do IO.println s!"ACCEPTED in {secs t0 t1} s: {Aiur.Value.ppDeref s.store depth v}" return 0 else - let .ok vCompiled := vTop.compile + let .ok vCompiled := vTop.compileWithGroups MultiStark.verifierColdGroups | IO.eprintln "multi-stark verifier compile failed"; return 1 let some vIdx := vCompiled.getFuncIdx `verify_multi_stark_proof | IO.eprintln "verify_multi_stark_proof entrypoint missing"; return 1 diff --git a/Benchmarks/RecursiveVerifier.lean b/Benchmarks/RecursiveVerifier.lean index e99d1f618..ed4ab5e3d 100644 --- a/Benchmarks/RecursiveVerifier.lean +++ b/Benchmarks/RecursiveVerifier.lean @@ -147,7 +147,7 @@ def main (args : List String) : IO UInt32 := do let vTop ← match MultiStark.multiStark with | .ok t => pure t | .error e => IO.eprintln s!"verifier merge failed: {e}"; return 1 - let vCompiled ← match vTop.compile with + let vCompiled ← match vTop.compileWithGroups MultiStark.verifierColdGroups with | .ok c => pure c | .error e => IO.eprintln s!"verifier compile failed: {e}"; return 1 let vIdx := vCompiled.getFuncIdx `verify_multi_stark_proof |>.get! diff --git a/Benchmarks/Typecheck.lean b/Benchmarks/Typecheck.lean index 5d3e6edee..bebde57a2 100644 --- a/Benchmarks/Typecheck.lean +++ b/Benchmarks/Typecheck.lean @@ -329,7 +329,7 @@ def runTypecheckCmd (p : Cli.Parsed) : IO UInt32 := do -- claim run, which is the honest reading of those numbers. let .ok toplevel := (if skipDeps then IxVM.ixVMFull else IxVM.ixVM) | throw (IO.userError "Merging IxVM kernel failed") - let .ok compiled := toplevel.compile + let .ok compiled := toplevel.compileWithGroups IxVM.coldGroups | throw (IO.userError "Compilation of IxVM kernel failed") let entrypoint := if skipDeps then `verify_const else `verify_claim let some funIdx := compiled.getFuncIdx entrypoint @@ -358,7 +358,7 @@ def runTypecheckCmd (p : Cli.Parsed) : IO UInt32 := do if !recursive then pure none else do let .ok vTop := MultiStark.multiStark | throw (IO.userError "Merging multi-stark verifier failed") - let .ok vCompiled := vTop.compile + let .ok vCompiled := vTop.compileWithGroups MultiStark.verifierColdGroups | throw (IO.userError "Compilation of multi-stark verifier failed") let some vIdx := vCompiled.getFuncIdx `verify_multi_stark_proof | throw (IO.userError "verify_multi_stark_proof entrypoint missing") diff --git a/Ix/Aiur/Compiler.lean b/Ix/Aiur/Compiler.lean index 752b2a878..0fe32a89b 100644 --- a/Ix/Aiur/Compiler.lean +++ b/Ix/Aiur/Compiler.lean @@ -59,9 +59,14 @@ per absorbed member. Errors if a name is unknown, unconstrained (it has no circuit to group), an entry function, listed twice, or if a group is empty. -/ def CompiledToplevel.groupFunctions (ct : CompiledToplevel) - (groups : Array (String × Array Lean.Name)) : + (groups : Array (String × Array String)) : Except String CompiledToplevel := do let t := ct.bytecode + -- Function names as printed (`toString` of the `Global`), the exact + -- inverse of what statistics reports — so measured groupings can be fed + -- back verbatim. + let byName : Std.HashMap String Bytecode.FunIdx := + ct.nameMap.fold (init := {}) fun acc g i => acc.insert (toString g) i -- Resolve and validate the groups into member-index arrays. let mut grouped : Std.HashMap Bytecode.FunIdx Nat := {} let mut resolved : Array (String × Array Bytecode.FunIdx) := #[] @@ -70,7 +75,7 @@ def CompiledToplevel.groupFunctions (ct : CompiledToplevel) throw s!"group {gname} is empty" let mut members := #[] for name in names do - let some i := ct.getFuncIdx name + let some i := byName[name]? | throw s!"group {gname}: unknown function {name}" let f := t.functions[i]! unless f.constrained do @@ -204,6 +209,12 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev circuits := bytecode.singletonCircuits fun i => reverseMap[i]?.getD s!"" } pure (CompiledToplevel.mk t bytecode nameMap) +/-- `compile`, then apply a function grouping (see +`CompiledToplevel.groupFunctions`). -/ +def Source.Toplevel.compileWithGroups (t : Source.Toplevel) + (groups : Array (String × Array String)) : Except String CompiledToplevel := do + (← t.compile).groupFunctions groups + /-- Progress helper: given success of the three `Except`-returning stages, `compile` as a whole returns `.ok` (the remaining stages — `deduplicate`, `needsCircuit`, the field-setter `mapIdx`, the name-map `fold`, and the diff --git a/Ix/Cli/CheckCmd.lean b/Ix/Cli/CheckCmd.lean index 565538782..ef6471162 100644 --- a/Ix/Cli/CheckCmd.lean +++ b/Ix/Cli/CheckCmd.lean @@ -656,7 +656,7 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do pure 1 pure go else do - let compiled ← match toplevel.compile with + let compiled ← match toplevel.compileWithGroups IxVM.coldGroups with | .error e => IO.eprintln s!"Compilation failed: {e}"; return 1 | .ok c => pure c let go (_ : Ix.Claim) (envHandle? : Option Aiur.EnvHandle) (target : Target) @@ -669,7 +669,7 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do return (← runShardCheckManifest manifest ixe k (fun c w l => runOne c none (.leanW w) l)) else do - let compiled ← match toplevel.compile with + let compiled ← match toplevel.compileWithGroups IxVM.coldGroups with | .error e => IO.eprintln s!"Compilation failed: {e}"; return 1 | .ok c => pure c return (← runShardCheckManifestNative manifest ixe k compiled printStats statsOut useBytecode) @@ -678,7 +678,7 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do return (← runShardCheckAll manifest ixe ((p.flag? "jobs").map (·.as! Nat)) (fun c w l => runOne c none (.leanW w) l)) else do - let compiled ← match toplevel.compile with + let compiled ← match toplevel.compileWithGroups IxVM.coldGroups with | .error e => IO.eprintln s!"Compilation failed: {e}"; return 1 | .ok c => pure c return (← runShardManifestAllNative manifest ixe diff --git a/Ix/Cli/ProveCmd.lean b/Ix/Cli/ProveCmd.lean index a35609d9a..787853b36 100644 --- a/Ix/Cli/ProveCmd.lean +++ b/Ix/Cli/ProveCmd.lean @@ -146,7 +146,7 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do let toplevel ← match IxVM.ixVM with | .error e => IO.eprintln s!"toplevel merging failed: {e}"; return 1 | .ok t => pure t - let compiled ← match toplevel.compile with + let compiled ← match toplevel.compileWithGroups IxVM.coldGroups with | .error e => IO.eprintln s!"compilation failed: {e}"; return 1 | .ok c => pure c let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters diff --git a/Ix/Cli/VerifyCmd.lean b/Ix/Cli/VerifyCmd.lean index 381352579..c2262b0d6 100644 --- a/Ix/Cli/VerifyCmd.lean +++ b/Ix/Cli/VerifyCmd.lean @@ -73,7 +73,7 @@ def verifyOneProof (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledTople def buildBackend : IO (Except String (Aiur.AiurSystem × Aiur.CompiledToplevel)) := do match IxVM.ixVM with | .error e => return .error s!"toplevel merging failed: {e}" - | .ok toplevel => match toplevel.compile with + | .ok toplevel => match toplevel.compileWithGroups IxVM.coldGroups with | .error e => return .error s!"compilation failed: {e}" | .ok compiled => return .ok (Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters, compiled) diff --git a/Ix/IxVM.lean b/Ix/IxVM.lean index 45c6b898d..0add22904 100644 --- a/Ix/IxVM.lean +++ b/Ix/IxVM.lean @@ -1,6 +1,7 @@ module public import Ix.Aiur.Meta public import Ix.IxVM.Core +public import Ix.IxVM.ColdGroups public import Ix.IxVM.ByteStream public import Ix.IxVM.Blake3 public import Ix.IxVM.RBTreeMap diff --git a/Ix/IxVM/ColdGroups.lean b/Ix/IxVM/ColdGroups.lean new file mode 100644 index 000000000..328731478 --- /dev/null +++ b/Ix/IxVM/ColdGroups.lean @@ -0,0 +1,20 @@ +module + +/-! +Circuit-grouping data for the IxVM kernel toplevel, applied wherever the +kernel is compiled for proving or verifying (see +`CompiledToplevel.groupFunctions`). Empty = no grouping: every constrained +function keeps its singleton circuit. Fill from measured workload +statistics; a stale grouping stays sound (grouping never affects +semantics), only less efficient. +-/ + +public section + +namespace IxVM + +def coldGroups : Array (String × Array String) := #[] + +end IxVM + +end diff --git a/Ix/MultiStark.lean b/Ix/MultiStark.lean index c99114428..6881ad134 100644 --- a/Ix/MultiStark.lean +++ b/Ix/MultiStark.lean @@ -11,6 +11,7 @@ public import Ix.MultiStark.Keccak public import Ix.MultiStark.Pcs public import Ix.MultiStark.SystemDeserialize public import Ix.MultiStark.Verifier +public import Ix.MultiStark.VerifierColdGroups public import Ix.MultiStark.Tests /-! diff --git a/Ix/MultiStark/VerifierColdGroups.lean b/Ix/MultiStark/VerifierColdGroups.lean new file mode 100644 index 000000000..58359883b --- /dev/null +++ b/Ix/MultiStark/VerifierColdGroups.lean @@ -0,0 +1,20 @@ +module + +/-! +Circuit-grouping data for the recursive-verifier toplevel, applied wherever +it is compiled for proving or verifying (see +`CompiledToplevel.groupFunctions`). Empty = no grouping: every constrained +function keeps its singleton circuit. Fill from measured workload +statistics; a stale grouping stays sound (grouping never affects +semantics), only less efficient. +-/ + +public section + +namespace MultiStark + +def verifierColdGroups : Array (String × Array String) := #[] + +end MultiStark + +end diff --git a/Tests/Aiur/Aiur.lean b/Tests/Aiur/Aiur.lean index 1e7b9fcc3..d62f0d047 100644 --- a/Tests/Aiur/Aiur.lean +++ b/Tests/Aiur/Aiur.lean @@ -886,8 +886,8 @@ def aiurTestCases : List AiurTestCase := [ ] /-- The grouping the `aiur` runner applies for the grouped environment. -/ -def testGroups : Array (String × Array Lean.Name) := - #[("test_group", #[`grouped_double, `grouped_pick, `grouped_sum_range])] +def testGroups : Array (String × Array String) := + #[("test_group", #["grouped_double", "grouped_pick", "grouped_sum_range"])] def groupedTestCases : List AiurTestCase := [ .prove `calls_grouped #[0, 5, 9] #[25] diff --git a/Tests/Aiur/Common.lean b/Tests/Aiur/Common.lean index 701462489..39ffb9305 100644 --- a/Tests/Aiur/Common.lean +++ b/Tests/Aiur/Common.lean @@ -73,7 +73,7 @@ structure AiurTestEnv where shapes : Array Aiur.CircuitShape def AiurTestEnv.build (toplevelFn : Except Aiur.Global Aiur.Source.Toplevel) - (groups : Array (String × Array Lean.Name) := #[]) : + (groups : Array (String × Array String) := #[]) : Except String AiurTestEnv := do let toplevel ← toplevelFn.mapError toString let compiled ← toplevel.compile diff --git a/Tests/Main.lean b/Tests/Main.lean index 5a4cc9264..32858d645 100644 --- a/Tests/Main.lean +++ b/Tests/Main.lean @@ -182,7 +182,7 @@ def ignoredRunners (env : Lean.Environment) : List (String × IO UInt32) := [ -- committed kernel system). let kernelUnitTests := .exec `kernel_unit_tests let serdeTest ← serdeNatAddComm env - match AiurTestEnv.build IxVM.ixVM, AiurTestEnv.build IxVM.ixVMFull with + match AiurTestEnv.build IxVM.ixVM IxVM.coldGroups, AiurTestEnv.build IxVM.ixVMFull with | .error e, _ | _, .error e => IO.eprintln s!"IxVM env build failed: {e}"; return 1 | .ok v2Env, .ok v2FullEnv => diff --git a/Tests/MultiStark.lean b/Tests/MultiStark.lean index c46b9477e..f59ded1ee 100644 --- a/Tests/MultiStark.lean +++ b/Tests/MultiStark.lean @@ -177,7 +177,7 @@ def endToEndSuite : IO UInt32 := do let vTop ← match MultiStark.multiStark with | .error e => IO.eprintln s!"verifier toplevel merge failed: {e}"; return 1 | .ok t => pure t - let vCompiled ← match vTop.compile with + let vCompiled ← match vTop.compileWithGroups MultiStark.verifierColdGroups with | .error e => IO.eprintln s!"verifier compilation failed: {e}"; return 1 | .ok c => pure c let vIdx ← match vCompiled.getFuncIdx `verify_multi_stark_proof with From c525d955bcff7f12f40f799e541d5a1eb717c4ee Mon Sep 17 00:00:00 2001 From: Gabriel Barreto Date: Wed, 12 Aug 2026 12:12:04 -0300 Subject: [PATCH 3/3] stats: print circuit layout shape columns (TEMP - revert before merge) Add Sel/Aux/Lkp columns (the circuit layout's selectors, auxiliaries, lookups; zero for memory/gadget rows) to the per-circuit statistics table. Grouping instrumentation only: merging is cheapest between circuits whose auxiliaries and lookups are CLOSE (both merge by max, selectors sum), so these columns are what a partition builder needs next to the width. Meant to be reverted once the partitions are chosen - this commit is self-contained in Ix/Aiur/Statistics.lean. --- Ix/Aiur/Statistics.lean | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/Ix/Aiur/Statistics.lean b/Ix/Aiur/Statistics.lean index 095a1d2d8..05603516c 100644 --- a/Ix/Aiur/Statistics.lean +++ b/Ix/Aiur/Statistics.lean @@ -45,6 +45,13 @@ structure CircuitStats where /-- FFT cost at the uncached height `height + cacheHits` (fixed-height gadget circuits keep their normal cost). Feeds `totalUncachedFftCost`. -/ uncachedFftCost : Float + -- TEMP (grouping instrumentation, revert with this commit): the circuit + -- layout shape, for picking group members — merging is cheapest between + -- circuits whose auxiliaries and lookups are CLOSE (both merge by max; + -- selectors sum). Zero for memory/gadget circuits (no function layout). + selectors : Nat := 0 + auxiliaries : Nat := 0 + lookups : Nat := 0 structure ExecutionStats where circuits : Array CircuitStats @@ -102,7 +109,9 @@ def computeStats (compiled : CompiledToplevel) (queryCounts : Array QueryCount) let (h, hits) := c.members.foldl (init := (0, 0)) fun (h, hits) i => let qc := queryCounts[i]! (h + qc.uniqueRows, hits + (qc.totalHits - qc.uniqueRows)) - mkStats c.name shape h hits + { mkStats c.name shape h hits with + selectors := c.layout.selectors, auxiliaries := c.layout.auxiliaries, + lookups := c.layout.lookups } let memoryCircuits := t.memorySizes.mapIdx fun i size => let shape := shapes[t.circuits.size + i]! let qc := queryCounts[nAllFuns + i]! @@ -162,9 +171,14 @@ def printStats (stats : ExecutionStats) : IO Unit := do let n := f.round.toUInt64.toNat toString n let wFftCost := stats.circuits.foldl (fun m cs => Nat.max m (formatSci cs.fftCost).length) 8 + -- TEMP (grouping instrumentation, revert with this commit) + let wSel := stats.circuits.foldl (fun m cs => Nat.max m (toString cs.selectors).length) 3 + let wAux := stats.circuits.foldl (fun m cs => Nat.max m (toString cs.auxiliaries).length) 3 + let wLkp := stats.circuits.foldl (fun m cs => Nat.max m (toString cs.lookups).length) 3 let wPct := 7 let wCum := 7 - let totalW := wName + 1 + wWidth + 1 + wHeight + 1 + wHits + 1 + wFftCost + 1 + wPct + 1 + wCum + let totalW := wName + 1 + wWidth + 1 + wSel + 1 + wAux + 1 + wLkp + 1 + + wHeight + 1 + wHits + 1 + wFftCost + 1 + wPct + 1 + wCum let totalWidth := stats.circuits.foldl (· + ·.width) 0 let savedPct := if stats.totalUncachedFftCost == 0.0 then "0.00%" @@ -177,14 +191,14 @@ def printStats (stats : ExecutionStats) : IO Unit := do IO.println s!"Total cache hits: {stats.totalCacheHits}" IO.println s!"Total saved cost: {savedPct}" IO.println sep - IO.println s!"{padRight "Name" wName} {padLeft "Width" wWidth} {padLeft "Height" wHeight} {padLeft "Hits" wHits} {padLeft "FFT cost" wFftCost} {padLeft "%" wPct} {padLeft "%++" wCum}" + IO.println s!"{padRight "Name" wName} {padLeft "Width" wWidth} {padLeft "Sel" wSel} {padLeft "Aux" wAux} {padLeft "Lkp" wLkp} {padLeft "Height" wHeight} {padLeft "Hits" wHits} {padLeft "FFT cost" wFftCost} {padLeft "%" wPct} {padLeft "%++" wCum}" IO.println sep let mut cumFftCost : Float := 0.0 for cs in stats.circuits do cumFftCost := cumFftCost + cs.fftCost let pct := formatPercent cs.fftCost stats.totalFftCost let cum := formatPercent cumFftCost stats.totalFftCost - IO.println s!"{padRight cs.name wName} {padLeft (toString cs.width) wWidth} {padLeft (toString cs.height) wHeight} {padLeft (toString cs.cacheHits) wHits} {padLeft (formatSci cs.fftCost) wFftCost} {padLeft pct wPct} {padLeft cum wCum}" + IO.println s!"{padRight cs.name wName} {padLeft (toString cs.width) wWidth} {padLeft (toString cs.selectors) wSel} {padLeft (toString cs.auxiliaries) wAux} {padLeft (toString cs.lookups) wLkp} {padLeft (toString cs.height) wHeight} {padLeft (toString cs.cacheHits) wHits} {padLeft (formatSci cs.fftCost) wFftCost} {padLeft pct wPct} {padLeft cum wCum}" end Aiur