Return groups#418
Draft
gabriel-barrett wants to merge 13 commits into
Draft
Conversation
gabriel-barrett
force-pushed
the
return-groups
branch
3 times, most recently
from
605a393 to
2e71e27
Compare
which collects all filtered sub-functions for each particular return group
gabriel-barrett
force-pushed
the
return-groups
branch
from
9059d94 to
afd477d
Compare
Lean now owns return-group splitting and layout recompute. Adds `Bytecode.Function.split` (with `returnGroups`/`filterGroup`/`fix`) and `Toplevel.computeFiltered`, populated in `Source.Toplevel.compile` after dedup + needsCircuit. Layout recompute reuses the existing `Compiler.Layout.blockLayout`. `Bytecode.Toplevel` gains `filteredFunctions : Array (Array (String × Function))`. The Rust FFI decoder reads it straight from Lean and interns group names into `Arc<str>` shared across every `Ctrl::Return` and every `filtered_functions` key. Deletes the Rust duplicates (`src/aiur/split.rs`, `src/aiur/layout.rs`) so there is one source of truth for bytecode shape. Tests: adds `shared_return_groups` exercising two arms sharing `even`, two sharing `odd`, and one unannotated (`""`) arm, aggregated into a single entry `shared_return_groups_all` so the STARK suite proves every group circuit in one run.
Source / Typed / Concrete / Simple keep `#[return_group(name)]` as a `String` (user-written labels). At lowering time `Concrete.Function.compile` allocates a `USize` per distinct group name in encounter order, stores the table in `Bytecode.Function.groupNames`, and stamps every `Ctrl.return` with the index. `Bytecode.Toplevel.filteredFunctions` becomes `Array (Array Function)` keyed positionally by group index. Rust mirrors: `Ctrl::Return(SelIdx, usize, Vec<ValIdx>)`, `Toplevel.filtered_functions: Vec<Vec<Function>>`, `QueryResult.return_group: usize`. The `Arc<str>` interner is gone — group identity is now a `Copy` integer. `LeanAiurCtrl` is declared with `num_usize: 1` on the `return` ctor so `get_usize(0)` is bounds-checked. `LeanAiurFunction` bumps `num_obj: 2 → 3` to carry `groupNames`; the Rust decoder reads it Lean-side only (display lookup in `Statistics`). Execute ships `(groupIdx, width, rows, hits)` quads instead of name strings; `Statistics.computeStats` resolves the display name via `t.functions[i].groupNames[groupIdx]`.
Replace the anonymous `Nat × Nat × Nat × Nat` quads / `Nat × Nat` pairs
shipped from the executor with `GroupStats { groupIdx, totalWidth,
uniqueRows, totalHits }` and `MemoryCount { uniqueRows, totalHits }`,
bundled into a `QueryCounts { functionStats, memoryCounts }` struct.
The FFI extern keeps the tuple shape (so the Rust side can build it
without matching a Lean structure ctor); `Bytecode.Toplevel.execute`
converts to the structured form at the public API boundary. Callers
(`Kernel.lean`, `Statistics.computeStats`) lose the `.1`/`.2.2.1`
indexing and use named field access.
gabriel-barrett
force-pushed
the
return-groups
branch
from
afd477d to
42c4ceb
Compare
Two sites dropped the per-function `groupNames` table:
1. `Dedup.deduplicate_newFunctions` rebuilt the `Function` struct as
`{ body, layout, entry, constrained }`, letting the field default
`:= #[""]` collapse the table to a singleton. Every dedup'd function
then split into one group at idx 0, so multi-group stats showed a
single row instead of one per path.
2. `Lower.addCase` snapshotted the compiler state per arm and restored
`initState` (patching only `selIdx`). Allocations made by
`allocCurrentGroup` inside the arm — pushes to `groupNames` /
`groupNameMap` — were discarded, so every arm reused index 0 and the
final table held only the surviving allocation.
Fix: thread `groupNames`/`groupNameMap` through both rollbacks. Verified
via `lake exe kernel Nat.add_comm` with three return-group annotations
on `blake3_compress_chunks` — stats now show one row per split.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before return groups, each Aiur function compiled to a single circuit covering all return paths. This means that if a circuit has cold paths, it contributes to the cost of the hot paths. With return groups you're able to annotate the paths that you want to be bundled in the same circuit, say have a dedicated circuit for each hot path and a single circuit for all cold paths