Function grouping

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

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!

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")

Ix/Aiur/Compiler.lean

@@ -42,6 +42,73 @@ 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 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) := #[]
+  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 := byName[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 +157,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,13 +197,24 @@ 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!"<fn {i}>" }
   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

Ix/Aiur/Compiler/Lower.lean

@@ -619,7 +619,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
 

Ix/Aiur/Stages/Bytecode.lean

@@ -115,9 +115,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

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
@@ -81,44 +88,39 @@ 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!"<fn {i}>"
-        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 with
+      selectors := c.layout.selectors, auxiliaries := c.layout.auxiliaries,
+      lookups := c.layout.lookups }
   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)
@@ -169,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%"
@@ -184,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
 

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

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

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)