Skip to content

perf(aiur): reduce witness memory and proving FFT costs #635

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
arthurpaulino merged 7 commits into from
Sep 17, 2026
Merged

perf(aiur): reduce witness memory and proving FFT costs #635

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e2d7daf
perf(aiur): cache validated function claim shapes
arthurpaulino Sep 16, 2026
f1534a8
perf(aiur): tune lookup groups from compiled degree costs
arthurpaulino Sep 16, 2026
c00a791
perf(aiur): compact witness metadata and correct memory estimates
arthurpaulino Sep 16, 2026
1ea7a68
perf(aiur): consolidate byte lookup tables
arthurpaulino Sep 16, 2026
f29f77e
perf(aiur): use u16 limbs for u32 comparisons
arthurpaulino Sep 16, 2026
84b2396
perf(aiur): remove the intermediate u32 comparison carry constraint
arthurpaulino Sep 16, 2026
eaf84e0
test(aiur): refresh FFT pins for independent optimizations
arthurpaulino Sep 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Ix/Aiur/Compiler/Layout.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ def opLayout : Bytecode.Op → LayoutM Unit
| .u8Mul .. => do pushDegrees #[1, 1]; bumpAuxiliaries 2; bumpLookups
| .u8XorSplit7 .. | .u8XorSplit4 .. => do pushDegrees #[1, 1]; bumpAuxiliaries 2; bumpLookups
| .u8LessThan .. => do pushDegree 1; bumpAuxiliaries; bumpLookups
| .u32LessThan .. => do pushDegree 1; bumpAuxiliaries 12; bumpLookups 6
| .u32LessThan .. => do pushDegree 1; bumpAuxiliaries 6; bumpLookups 6
| .unconstrainedU32Add a b => do
let degrees ← (a ++ b).mapM getDegree
pushDegrees $ .replicate 4 1
Expand Down
10 changes: 6 additions & 4 deletions Ix/Aiur/Stages/Bytecode.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,12 +104,14 @@ structure FunctionLayout where
def FunctionLayout.width (l : FunctionLayout) : Nat :=
l.inputSize + l.selectors + l.auxiliaries

/-- Layout-only estimate of main and stage-2 width. This lacks the control
tree, compiled lookup degrees and PCS parameters; use the built system's
`circuitShapes` for actual widths after lookup tuning. -/
def FunctionLayout.totalWidth (l : FunctionLayout) : Nat :=
-- Stage 2 commits max(⌈L/k⌉, 1) chained partial accumulators (no message
-- inverses); see `multi_stark::lookup::stage2_width`. Mirrors the
-- synthesis grouping rule (`crates/aiur/src/synthesis.rs`): branchless
-- functions (one selector) have raw degree-1 lookup arguments, so their
-- lookups are grouped 2 per accumulator step.
-- inverses); see `multi_stark::lookup::stage2_width`. Retain the original
-- single-selector heuristic here; synthesis checks control flow and
-- retunes using compiled degrees and the FFT cost.
let slots := if l.selectors == 1 && l.lookups >= 2
then (l.lookups + 1) / 2
else max l.lookups 1
Expand Down
21 changes: 7 additions & 14 deletions Ix/Aiur/Stages/Codegen.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -611,9 +611,8 @@ private def emitU8Sub (out : Nat) (i j : ValIdx) : Array RustStmt :=
declVal (out + 1) (.field (.var "__b2_sub") "1")
]

/-- `Op::U32LessThan`: mirror execute.rs lines 477-505. Pure
compare + 6-byte-pair range-check via
`bytes2_queries.bump_range_check` (constrained mode only). -/
/-- `Op::U32LessThan`: mirror the checked comparison in `execute.rs`,
recording six scalar u16 range queries in constrained mode. -/
private def emitU32LessThan (out : Nat) (x y : ValIdx) : Array RustStmt :=
let blockExpr : String :=
s!"\{ let __a_val = __v_{x}.as_canonical_u64();" ++
Expand All @@ -622,17 +621,11 @@ private def emitU32LessThan (out : Nat) (x y : ValIdx) : Array RustStmt :=
s!" let __b_u32 = u32::try_from(__b_val).ok().ok_or(ExecError::U32OutOfRange(__b_val))?;" ++
s!" let __result = G::from_bool(__a_u32 < __b_u32);" ++
s!" if !unconstrained \{" ++
s!" let __x_bytes = __a_u32.to_le_bytes();" ++
s!" let __z_bytes = __b_u32.to_le_bytes();" ++
s!" let __c_u32 = __b_u32.wrapping_sub(__a_u32).wrapping_sub(1);" ++
s!" let __y_bytes = __c_u32.to_le_bytes();" ++
s!" record.bytes2_queries.bump_range_check(&G::from_u8(__x_bytes[0]), &G::from_u8(__x_bytes[1]));" ++
s!" record.bytes2_queries.bump_range_check(&G::from_u8(__x_bytes[2]), &G::from_u8(__x_bytes[3]));" ++
s!" record.bytes2_queries.bump_range_check(&G::from_u8(__y_bytes[0]), &G::from_u8(__y_bytes[1]));" ++
s!" record.bytes2_queries.bump_range_check(&G::from_u8(__y_bytes[2]), &G::from_u8(__y_bytes[3]));" ++
s!" record.bytes2_queries.bump_range_check(&G::from_u8(__z_bytes[0]), &G::from_u8(__z_bytes[1]));" ++
s!" record.bytes2_queries.bump_range_check(&G::from_u8(__z_bytes[2]), &G::from_u8(__z_bytes[3]));" ++
s!" } __result }"
s!" let __c_u32 = __a_u32.wrapping_sub(__b_u32);" ++
s!" for __word in [__a_u32, __c_u32, __b_u32] \{" ++
s!" record.bytes2_queries.bump_u16_range_check((__word & 0xffff) as u16);" ++
s!" record.bytes2_queries.bump_u16_range_check((__word >> 16) as u16);" ++
s!" } } __result }"
#[.letStmt false s!"__v_{out}" (some "G") (.lit blockExpr)]

private def u32PackExpr (xs : Array ValIdx) : String :=
Expand Down
166 changes: 83 additions & 83 deletions Tests/Ix/IxVM.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -439,89 +439,89 @@ private def nameOfString (str : String) : Lean.Name :=
and an improvement has to be acknowledged by re-pinning. These pins
include v3 contract decoding, canonicality checks, and scoped claim headers. -/
private def kernelCheckEntries : List (String × Nat) := [
("HEq", 129_580_132),
("HEq.rec", 133_703_388),
("Eq.rec", 133_072_615),
("Nat", 129_633_587),
("Nat.add", 170_479_894),
("Nat.add_comm", 322_598_714),
("Nat.decEq", 377_100_753),
("Nat.decLe", 814_998_141),
("Nat.sub_le_of_le_add", 1_957_472_572),
("Nat.shiftRight_succ", 1_455_068_309),
("Trans.mk", 137_107_353),
("Array.append_assoc", 9_081_433_824),
("Vector.append", 9_287_002_309),
("IxVMPrim.nat_add_lit", 213_488_292),
("IxVMPrim.nat_sub_lit", 228_506_721),
("IxVMPrim.nat_mul_lit", 204_344_721),
("IxVMPrim.nat_mul_big", 202_896_358),
("IxVMPrim.nat_div_lit", 1_420_411_368),
("IxVMPrim.nat_mod_lit", 1_447_938_067),
("IxVMPrim.nat_succ_lit", 144_339_930),
("IxVMPrim.nat_pred_lit", 166_168_846),
("IxVMPrim.nat_gcd_lit", 2_232_501_380),
("IxVMPrim.nat_land_lit", 3_652_583_530),
("IxVMPrim.nat_lor_lit", 3_654_583_905),
("IxVMPrim.nat_xor_lit", 3_674_202_062),
("IxVMPrim.nat_shl_lit", 234_243_500),
("IxVMPrim.nat_shr_lit", 1_435_151_445),
("IxVMPrim.nat_pow_big", 395_904_832),
("IxVMPrim.nat_beq_lit", 201_412_466),
("IxVMPrim.nat_ble_lit", 196_817_405),
("IxVMPrim.nat_cases_big", 166_401_353),
("IxVMPrim.nat_dec_le", 833_212_679),
("IxVMPrim.nat_dec_lt", 844_424_937),
("IxVMPrim.nat_dec_eq", 416_937_906),
("IxVMPrim.str_size_lit", 2_563_571_976),
("IxVMPrim.bv_to_nat_lit", 2_135_250_014),
("IxVMInd.Even", 206_886_515),
("IxVMInd.Odd", 206_892_115),
("IxVMInd.Even.rec", 226_129_215),
("IxVMInd.Odd.rec", 226_131_072),
("IxVMInd.IdxTeleN.rec", 161_370_130),
("IxVMInd.IdxTeleB.rec", 161_368_435),
("IxVMInd.SoloA.rec", 156_211_923),
("IxVMInd.SoloB.rec", 156_211_088),
("IxVMInd.UnsafeSquash", 131_473_733),
("IxVMInd.Tree", 131_267_145),
("IxVMInd.Tree.rec", 142_162_255),
("IxVMInd.DedupM", 134_625_951),
("IxVMInd.DedupM.rec", 149_282_091),
("IxVMInd.DepthM", 133_016_231),
("IxVMInd.DepthM.rec", 144_730_870),
("String.Internal.append", 2_533_658_506),
("_private.Init.Prelude.0.Lean.extractMainModule._unsafe_rec", 3_750_129_764),
("Lean.Syntax.rec", 2_597_021_992),
("IxVMInd.AuxTie", 326_741_066),
("IxVMInd.AuxTie.rec", 373_345_395),
("IxVMInd.HiddenIdx", 130_558_032),
("IxVMInd.HiddenIdx.rec", 133_948_354),
("IxVMInd.thmMajorUse", 517_920_579),
("IxVMInd.partialKRec", 150_421_901),
("IxVMInd.deepRebase", 221_891_341),
("String.Slice.Pattern.Model.NoPrefixPatternModel.rec", 3_515_782_465),
("Lean.Widget.TaggedText.rec", 2_566_008_609),
("Lean.Doc.Part.rec", 2_611_709_383),
("Lean.Doc.Block.rec", 2_847_867_384),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup1.A", 132_585_241),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup1.A.rec", 136_255_212),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup1.A.rec_1", 135_184_337),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup1.A.rec_2", 135_184_337),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup2.A.rec_1", 135_184_337),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedupMixed.M", 132_875_994),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedupMixed.M.rec", 145_089_906),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedupMixed.M.rec_1", 145_088_324),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedupMixed.M.rec_2", 135_184_337),
("strOfListFoldSize", 2_848_829_120),
("strOfListFoldSizeAscii", 2_849_606_441),
("IxVMPrim.lazy_ble_offset", 234_904_821),
("IxVMPrim.lazy_unit_cast", 558_233_712),
("IxVMPrim.sizeof_unit", 161_110_680),
("IxVMPerf.let_continuations", 165_727_825),
("IxVMPerf.mul_row", 424_822_409),
("IxVMPerf.context_tower", 518_991_146),
("IxVMPerf.mul_wide", 268_439_307),
("HEq", 108_495_574),
("HEq.rec", 112_384_237),
("Eq.rec", 111_757_088),
("Nat", 108_546_656),
("Nat.add", 146_389_791),
("Nat.add_comm", 286_989_161),
("Nat.decEq", 339_867_991),
("Nat.decLe", 744_859_627),
("Nat.sub_le_of_le_add", 1_808_798_827),
("Nat.shiftRight_succ", 1_339_891_681),
("Trans.mk", 115_473_215),
("Array.append_assoc", 8_482_232_650),
("Vector.append", 8_671_958_125),
("IxVMPrim.nat_add_lit", 185_732_193),
("IxVMPrim.nat_sub_lit", 199_325_550),
("IxVMPrim.nat_mul_lit", 177_066_589),
("IxVMPrim.nat_mul_big", 175_788_118),
("IxVMPrim.nat_div_lit", 1_307_242_268),
("IxVMPrim.nat_mod_lit", 1_332_536_900),
("IxVMPrim.nat_succ_lit", 121_887_619),
("IxVMPrim.nat_pred_lit", 141_744_465),
("IxVMPrim.nat_gcd_lit", 2_061_457_700),
("IxVMPrim.nat_land_lit", 3_376_640_812),
("IxVMPrim.nat_lor_lit", 3_378_398_517),
("IxVMPrim.nat_xor_lit", 3_395_994_274),
("IxVMPrim.nat_shl_lit", 204_310_958),
("IxVMPrim.nat_shr_lit", 1_320_506_095),
("IxVMPrim.nat_pow_big", 361_514_602),
("IxVMPrim.nat_beq_lit", 174_672_222),
("IxVMPrim.nat_ble_lit", 170_365_817),
("IxVMPrim.nat_cases_big", 142_102_154),
("IxVMPrim.nat_dec_le", 761_507_789),
("IxVMPrim.nat_dec_lt", 771_720_110),
("IxVMPrim.nat_dec_eq", 376_436_724),
("IxVMPrim.str_size_lit", 2_360_108_000),
("IxVMPrim.bv_to_nat_lit", 1_969_340_217),
("IxVMInd.Even", 179_459_070),
("IxVMInd.Odd", 179_464_671),
("IxVMInd.Even.rec", 197_286_374),
("IxVMInd.Odd.rec", 197_288_230),
("IxVMInd.IdxTeleN.rec", 137_630_968),
("IxVMInd.IdxTeleB.rec", 137_629_273),
("IxVMInd.SoloA.rec", 132_815_226),
("IxVMInd.SoloB.rec", 132_814_391),
("IxVMInd.UnsafeSquash", 110_174_981),
("IxVMInd.Tree", 110_068_718),
("IxVMInd.Tree.rec", 120_158_594),
("IxVMInd.DedupM", 113_066_695),
("IxVMInd.DedupM.rec", 126_681_373),
("IxVMInd.DepthM", 111_645_324),
("IxVMInd.DepthM.rec", 122_589_004),
("String.Internal.append", 2_333_030_898),
("_private.Init.Prelude.0.Lean.extractMainModule._unsafe_rec", 3_460_502_868),
("Lean.Syntax.rec", 2_391_207_067),
("IxVMInd.AuxTie", 287_499_396),
("IxVMInd.AuxTie.rec", 330_644_761),
("IxVMInd.HiddenIdx", 109_373_096),
("IxVMInd.HiddenIdx.rec", 112_518_422),
("IxVMInd.thmMajorUse", 469_426_294),
("IxVMInd.partialKRec", 127_699_359),
("IxVMInd.deepRebase", 193_715_227),
("String.Slice.Pattern.Model.NoPrefixPatternModel.rec", 3_250_343_797),
("Lean.Widget.TaggedText.rec", 2_363_428_548),
("Lean.Doc.Part.rec", 2_405_755_939),
("Lean.Doc.Block.rec", 2_628_895_151),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup1.A", 111_219_505),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup1.A.rec", 114_568_146),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup1.A.rec_1", 113_736_777),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup1.A.rec_2", 113_736_777),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedup2.A.rec_1", 113_736_777),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedupMixed.M", 111_504_500),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedupMixed.M.rec", 122_803_609),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedupMixed.M.rec_1", 122_802_027),
("_private.Tests.Ix.Compile.Mutual.0.Tests.Ix.Compile.Mutual.AuxDedupMixed.M.rec_2", 113_736_777),
("strOfListFoldSize", 2_621_459_865),
("strOfListFoldSizeAscii", 2_622_117_638),
("IxVMPrim.lazy_ble_offset", 205_262_728),
("IxVMPrim.lazy_unit_cast", 506_194_228),
("IxVMPrim.sizeof_unit", 137_095_643),
("IxVMPerf.let_continuations", 143_072_956),
("IxVMPerf.mul_row", 384_322_113),
("IxVMPerf.context_tower", 467_936_574),
("IxVMPerf.mul_wide", 236_260_962),
]

/-- Variant of `kernelChecks`, pinned to the baseline
Expand Down
4 changes: 2 additions & 2 deletions Tests/Main.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,8 +281,8 @@ def ignoredRunners (env : Lean.Environment) : List (String × IO UInt32) := [
let actual :=
(Aiur.computeStats vmEnv.compiled qc vmEnv.shapes).totalFftCost.round.toUInt64.toNat
pure (LSpec.test
s!"Shard pipeline FFT matches: expected 7_072_190_269, got {actual}"
(actual = 7_072_190_269))
s!"Shard pipeline FFT matches: expected 6_434_780_974, got {actual}"
(actual = 6_434_780_974))
LSpec.lspecIO
(.ofList [("ixvm",
[fullSeq, aiurSeq, arenaSeq, exploitSeq, paritySeq, shardSeq])]) []),
Expand Down
Loading