Skip to content

perf: store per-module constant prefix trees in oleans#14362

Draft
Kha wants to merge 14 commits into
masterfrom
olean-const-trie
Draft

perf: store per-module constant prefix trees in oleans#14362
Kha wants to merge 14 commits into
masterfrom
olean-const-trie

Conversation

@Kha

@Kha Kha commented Jul 11, 2026

Copy link
Copy Markdown
Member

This PR replaces the eager construction of the imported-constants hash maps in finalizeImport (privateConstantMap, publicConstantMap, const2ModIdx; ~530k insertions for import Lean) with prefix trees over each module's constant names that are precomputed at olean write time and remain in the memory-mapped compacted regions. Import now only merges the per-module trees, allocating fresh nodes solely for name prefixes shared between modules (~3k nodes / ~81k edges for the stdlib), and lookups walk the merged view, disambiguating children by cached per-prefix Name hash. import Lean drops from ~527ms to ~393ms; olean sizes grow by ~6%.

Details:

  • New Lean.ConstTrie: generic per-module tree (ConstTrie α), merged import view (ImportedConsts α) with single-module subtrees borrowed from the regions and module-index tagging, and the new ConstMap structure (imported view + locals map).
  • ModuleData gains constTrie/extraConstTrie fields, filled centrally in saveModuleData(Parts). The extra-const tree replaces the 325k codegen-name entries of const2ModIdx (Kernel.Environment now holds allImportedConsts/importedExtraConsts instead, serving getModuleIdxFor?).
  • Duplicate imported constants are resolved (incl. subsumesInfo realizability subsumption and throwAlreadyImported) from the merge's reported collisions.
  • In saveModuleDataParts, all parts' data must stay alive across the chained compactor saves; the tries are therefore attached to all parts up front (the compactor tracks compacted objects by heap address, so per-part temporaries would corrupt later parts via address reuse).
  • stage0/src/stdlib_flags.h temporarily enables interpreter.prefer_native as this is an ABI-breaking change to Kernel.Environment/ConstMap (to be reverted by update-stage0). Stage 1 tests are meaningless on this commit since the stage 1 oleans lack the new ModuleData fields; validate against stage 2.

@Kha

Kha commented Jul 11, 2026

Copy link
Copy Markdown
Member Author

!bench

@leanprover-radar

leanprover-radar commented Jul 11, 2026

Copy link
Copy Markdown

Benchmark results for 7a68001 against 8006bb0 are in. There are significant results. @Kha

Warning

These warnings may indicate that the benchmark results are not directly comparable, for example due to changes in the runner configuration or hardware.

  • Bench repo commit hashes for run build differ between commits.
  • Bench repo commit hashes for run other differ between commits.
  • 🟥 build//instructions: +6.3T (+52.95%)

Large changes (55✅, 329🟥)

Too many entries to display here. View the full report on radar instead.

Medium changes (7✅, 1528🟥)

Too many entries to display here. View the full report on radar instead.

Small changes (5✅, 816🟥)

Too many entries to display here. View the full report on radar instead.

@github-actions github-actions Bot added toolchain-available A toolchain is available for this PR, at leanprover/lean4-pr-releases:pr-release-NNNN mathlib4-nightly-available A branch for this PR exists at leanprover-community/mathlib4-nightly-testing:lean-pr-testing-NNNN labels Jul 11, 2026
@leanprover-bot

leanprover-bot commented Jul 11, 2026

Copy link
Copy Markdown
Collaborator

Reference manual CI status:

@leanprover-bot leanprover-bot added the breaks-manual This is not necessarily a blocker for merging, but there needs to be a plan. label Jul 11, 2026
@Kha

Kha commented Jul 11, 2026

Copy link
Copy Markdown
Member Author

!bench mathlib

@leanprover-radar

leanprover-radar commented Jul 11, 2026

Copy link
Copy Markdown

Benchmark results for leanprover-community/mathlib4-nightly-testing@48fbe78 against leanprover-community/mathlib4-nightly-testing@5edfd4c are in. There are significant results. @Kha

  • 🟥 main exited with code 1

No significant changes detected.

@Kha Kha force-pushed the olean-const-trie branch from 7a68001 to 736e8af Compare July 11, 2026 14:03
@Kha

Kha commented Jul 11, 2026

Copy link
Copy Markdown
Member Author

!bench

@leanprover-radar

leanprover-radar commented Jul 11, 2026

Copy link
Copy Markdown

Benchmark results for 736e8af against 8006bb0 are in. There are significant results. @Kha

Warning

These warnings may indicate that the benchmark results are not directly comparable, for example due to changes in the runner configuration or hardware.

  • Bench repo commit hashes for run build differ between commits.
  • Bench repo commit hashes for run other differ between commits.
  • 🟥 build//instructions: +4.5T (+38.02%)

Large changes (56✅, 232🟥)

Too many entries to display here. View the full report on radar instead.

Medium changes (5✅, 1101🟥)

Too many entries to display here. View the full report on radar instead.

Small changes (6✅, 1324🟥)

Too many entries to display here. View the full report on radar instead.

Kha added 2 commits July 11, 2026 14:09
This PR replaces the eager construction of the imported-constants hash maps in `finalizeImport` (`privateConstantMap`, `publicConstantMap`, `const2ModIdx`; ~530k insertions for `import Lean`) with prefix trees over each module's constant names that are precomputed at olean write time and remain in the memory-mapped compacted regions. Import now only merges the per-module trees, allocating fresh nodes solely for name prefixes shared between modules (~3k nodes / ~81k edges for the stdlib), and lookups walk the merged view, disambiguating children by cached per-prefix `Name` hash. `import Lean` drops from ~527ms to ~393ms; olean sizes grow by ~6%.

Details:

* New `Lean.ConstTrie`: generic per-module tree (`ConstTrie α`), merged import view (`ImportedConsts α`) with single-module subtrees borrowed from the regions and module-index tagging, and the new `ConstMap` structure (`imported` view + `locals` map).
* `ModuleData` gains `constTrie`/`extraConstTrie` fields, filled centrally in `saveModuleData(Parts)`. The extra-const tree replaces the 325k codegen-name entries of `const2ModIdx` (`Kernel.Environment` now holds `allImportedConsts`/`importedExtraConsts` instead, serving `getModuleIdxFor?`).
* Duplicate imported constants are resolved (incl. `subsumesInfo` realizability subsumption and `throwAlreadyImported`) from the merge's reported collisions.
* In `saveModuleDataParts`, all parts' data must stay alive across the chained compactor saves; the tries are therefore attached to all parts up front (the compactor tracks compacted objects by heap address, so per-part temporaries would corrupt later parts via address reuse).
* `stage0/src/stdlib_flags.h` temporarily enables `interpreter.prefer_native` as this is an ABI-breaking change to `Kernel.Environment`/`ConstMap` (to be reverted by `update-stage0`). Stage 1 tests are meaningless on this commit since the stage 1 oleans lack the new `ModuleData` fields; validate against stage 2.

Co-Authored-By: Claude
This PR speeds up constant lookups in the prefix-tree-backed environment by storing each trie node's child key hashes in a `ByteArray` alongside the children, so the per-level binary search runs over contiguous memory instead of chasing child and key pointers, in the memory-mapped per-module trees as well as the merged import spine. This supersedes the `indexed` hash-map wrappers on wide merged nodes, which could not cover wide nodes inside single-module region subtrees and made `Runtime.markPersistent` walk their entries on every import.

Co-Authored-By: Claude
mathlib-nightly-testing Bot pushed a commit to leanprover-community/batteries that referenced this pull request Jul 11, 2026
mathlib-nightly-testing Bot pushed a commit to leanprover-community/mathlib4-nightly-testing that referenced this pull request Jul 11, 2026
leanprover-bot added a commit to leanprover/reference-manual that referenced this pull request Jul 11, 2026
@leanprover-bot leanprover-bot added builds-manual CI has verified that the Lean Language Reference builds against this PR and removed breaks-manual This is not necessarily a blocker for merging, but there needs to be a plan. labels Jul 11, 2026
@Kha Kha force-pushed the olean-const-trie branch from 736e8af to 2abb3e2 Compare July 11, 2026 15:26
@Kha

Kha commented Jul 11, 2026

Copy link
Copy Markdown
Member Author

!bench

@leanprover-radar

leanprover-radar commented Jul 11, 2026

Copy link
Copy Markdown

Benchmark results for 2abb3e2 against 8006bb0 are in. There are significant results. @Kha

Warning

These warnings may indicate that the benchmark results are not directly comparable, for example due to changes in the runner configuration or hardware.

  • Bench repo commit hashes for run build differ between commits.
  • Bench repo commit hashes for run other differ between commits.
  • 🟥 build//instructions: +8.2T (+69.20%)

Large changes (49✅, 433🟥)

Too many entries to display here. View the full report on radar instead.

Medium changes (9✅, 1889🟥)

Too many entries to display here. View the full report on radar instead.

Small changes (4✅, 369🟥)

Too many entries to display here. View the full report on radar instead.

mathlib-nightly-testing Bot pushed a commit to leanprover-community/batteries that referenced this pull request Jul 11, 2026
mathlib-nightly-testing Bot pushed a commit to leanprover-community/mathlib4-nightly-testing that referenced this pull request Jul 11, 2026
leanprover-bot added a commit to leanprover/reference-manual that referenced this pull request Jul 11, 2026
This PR restores the intended final form of the trie child index, which the previous commit accidentally carried in an interim form due to a workspace snapshot race: instead of a binary search that composes each probed hash from eight byte loads, the child index stores `[k][2^k + 1 u16 offsets][one fingerprint byte per child]`; the query hash's top `k` bits select a window of on average at most one child, a fingerprint byte filters it, and the final name component comparison (sound on its own, as siblings always differ in their final component) confirms the match.

Co-Authored-By: Claude
@Kha

Kha commented Jul 11, 2026

Copy link
Copy Markdown
Member Author

!bench

@leanprover-radar

leanprover-radar commented Jul 11, 2026

Copy link
Copy Markdown

Benchmark results for e957cab against 8006bb0 are in. There are significant results. @Kha

Warning

These warnings may indicate that the benchmark results are not directly comparable, for example due to changes in the runner configuration or hardware.

  • Bench repo commit hashes for run build differ between commits.
  • Bench repo commit hashes for run other differ between commits.
  • 🟥 build//instructions: +2.6T (+21.58%)

Large changes (52✅, 126🟥)

Too many entries to display here. View the full report on radar instead.

Medium changes (7✅, 560🟥)

Too many entries to display here. View the full report on radar instead.

Small changes (6✅, 1925🟥)

Too many entries to display here. View the full report on radar instead.

mathlib-nightly-testing Bot pushed a commit to leanprover-community/batteries that referenced this pull request Jul 11, 2026
mathlib-nightly-testing Bot pushed a commit to leanprover-community/mathlib4-nightly-testing that referenced this pull request Jul 11, 2026
leanprover-bot added a commit to leanprover/reference-manual that referenced this pull request Jul 11, 2026
The process-global lookup cache marked its Name keys and view roots persistent on insertion. `lean_mark_persistent` plain-writes reference counts while recursing through the object graph, racing with atomic RC updates from concurrent elaboration threads; a lost update can leave a supposedly persistent object with a live count that later drops to zero, freeing memory the cache still references. This caused flaky SIGSEGVs and `loose bvar` panics in clean stdlib builds.

Pin roots and keys with plain leaked references instead (cache hits only compare them, without reference counting), and mark cached values multi-threaded via `lean_mark_mt`, which only mutates single-thread-owned objects and is therefore race-free at this call site.

Co-Authored-By: Claude
@Kha

Kha commented Jul 12, 2026

Copy link
Copy Markdown
Member Author

!bench

@leanprover-radar

leanprover-radar commented Jul 12, 2026

Copy link
Copy Markdown

Benchmark results for 4aa7198 against 8006bb0 are in. There are significant results. @Kha

Warning

These warnings may indicate that the benchmark results are not directly comparable, for example due to changes in the runner configuration or hardware.

  • Bench repo commit hashes for run build differ between commits.
  • Bench repo commit hashes for run other differ between commits.
  • 🟥 build//instructions: +475.2G (+4.01%)

Large changes (51✅, 42🟥)

Too many entries to display here. View the full report on radar instead.

Medium changes (6✅, 24🟥)

Too many entries to display here. View the full report on radar instead.

Small changes (8✅, 2466🟥)

Too many entries to display here. View the full report on radar instead.

The cache replaced entries on slot collision and leaked the replaced entry, as concurrent readers may still hold it. Colliding hot keys then alternate clobbering each other, leaking one entry per alternation; the `misc/import Std.Data.DHashMap.Internal.RawLemmas` benchmark peaked at 2.5x baseline memory this way.

Probe a short window (8 slots) instead and never replace: a miss inserts into the window's first empty slot, or not at all if the window is full. This bounds the cache's memory by the table size while keeping hits lock-free; entries are still never freed. Lookup instruction counts on `Init.Data.Sum`, `Init.Data.List.Lemmas`, and the riscv elaboration benchmark are unchanged within 0.3%.

Co-Authored-By: Claude
@Kha

Kha commented Jul 12, 2026

Copy link
Copy Markdown
Member Author

!bench

@leanprover-radar

leanprover-radar commented Jul 12, 2026

Copy link
Copy Markdown

Benchmark results for 58c655e against 8006bb0 are in. There are significant results. @Kha

Warning

These warnings may indicate that the benchmark results are not directly comparable, for example due to changes in the runner configuration or hardware.

  • Bench repo commit hashes for run build differ between commits.
  • Bench repo commit hashes for run other differ between commits.
  • 🟥 build//instructions: +540.0G (+4.56%)

Large changes (51✅, 44🟥)

Too many entries to display here. View the full report on radar instead.

Medium changes (7✅, 24🟥)

Too many entries to display here. View the full report on radar instead.

Small changes (7✅, 2466🟥)

Too many entries to display here. View the full report on radar instead.

mathlib-nightly-testing Bot pushed a commit to leanprover-community/batteries that referenced this pull request Jul 12, 2026
mathlib-nightly-testing Bot pushed a commit to leanprover-community/mathlib4-nightly-testing that referenced this pull request Jul 12, 2026
leanprover-bot added a commit to leanprover/reference-manual that referenced this pull request Jul 12, 2026
The cache accessors re-box the shared cached entry into a fresh `Option` on every call and all callers immediately match on the result, so inlining fuses away the intermediate allocation. Saves ~0.2% instructions on elaboration-heavy modules.

Co-Authored-By: Claude
@mathlib-lean-pr-testing mathlib-lean-pr-testing Bot added the breaks-mathlib This is not necessarily a blocker for merging: but there needs to be a plan label Jul 12, 2026
@mathlib-lean-pr-testing

mathlib-lean-pr-testing Bot commented Jul 12, 2026

Copy link
Copy Markdown

Mathlib CI status (docs):

The import-time spine merge re-read each child key's cached hash through two pointer indirections per sort comparison, grew the flattened candidate array incrementally, and allocated a singleton group plus a recursive call for every child declared by only one module, which is the overwhelmingly common case. Cache the hash in the flattened elements, preallocate the array exactly, and borrow single-module children directly. Saves 10-13% instructions on import-heavy modules (`Lean` umbrella 1.98G -> 1.77G locally).

Co-Authored-By: Claude
@Kha

Kha commented Jul 13, 2026

Copy link
Copy Markdown
Member Author

!bench

@leanprover-radar

leanprover-radar commented Jul 13, 2026

Copy link
Copy Markdown

Benchmark results for 255e10e against 8006bb0 are in. There are significant results. @Kha

Warning

These warnings may indicate that the benchmark results are not directly comparable, for example due to changes in the runner configuration or hardware.

  • Bench repo commit hashes for run build differ between commits.
  • Bench repo commit hashes for run other differ between commits.
  • 🟥 build//instructions: +390.7G (+3.30%)

Large changes (51✅, 31🟥)

Too many entries to display here. View the full report on radar instead.

Medium changes (6✅, 23🟥)

Too many entries to display here. View the full report on radar instead.

Small changes (10✅, 2476🟥)

Too many entries to display here. View the full report on radar instead.

Batteries folds over `env.constants.map₁` via the `HashMap.foldM` name (`EnvSearch`, `Util.Cache`), which the accessor set so far only covered as `foldlM`.

Co-Authored-By: Claude
mathlib-nightly-testing Bot pushed a commit to leanprover-community/batteries that referenced this pull request Jul 13, 2026
mathlib-nightly-testing Bot pushed a commit to leanprover-community/mathlib4-nightly-testing that referenced this pull request Jul 13, 2026
leanprover-bot added a commit to leanprover/reference-manual that referenced this pull request Jul 13, 2026
@Kha

Kha commented Jul 13, 2026

Copy link
Copy Markdown
Member Author

!bench mathlib

@leanprover-radar

leanprover-radar commented Jul 13, 2026

Copy link
Copy Markdown

Benchmark results for leanprover-community/mathlib4-nightly-testing@ce9e9e4 against leanprover-community/mathlib4-nightly-testing@5edfd4c are in. There are significant results. @Kha

  • 🟥 main exited with code 137

No significant changes detected.

Kha added 2 commits July 13, 2026 14:34
The region fix walk assumed well-formed input: a mis-sized object desynchronizes the walk, after which fixups land at wrong offsets and readers consume untranslated saved pointers. With corrupted artifacts in a build cache this manifested as flaky import-time segfaults under highly parallel mathlib builds, silent self-propagation of the corruption into freshly written oleans, and non-terminating imports (a misparsed zero-size object stops the walk from advancing).

Detect the desync instead: the walk must advance on every object and end exactly at the region end, and saved pointers outside every known region range now throw instead of hitting release-mode undefined behavior. On the write side, refuse to compact scalar arrays with element size 0, which no legitimate constructor produces and which is exactly the corruption shape that desynchronizes readers (`lean_sarray_byte_size` degenerates to the header size).

Reproducer: occupying a module's preferred `mmap` base address (forcing the relocation walk) while importing an affected file previously segfaulted or hung; it now fails with `compacted region: object walk failed to advance, corrupt olean file?` naming the file.

Co-Authored-By: Claude
The compat accessor recomputed the map by folding over all imported constants on every call. `Batteries.Tactic.Lint.getDeclsInPackage` reads it inside a per-declaration fold, turning linting into a quadratic walk over millions of constants on mathlib — `lake lint` effectively never terminated (the bench and mathlib CI lint timeouts).

Store the mapping as a `Thunk` on `Kernel.Environment`, filled in `finalizeImport` from the merged views: the runtime forces thunks exactly once, safely across threads (`lean_thunk_get_core` claims the closure atomically and `mark_mt`s the value), so downstream per-declaration reads are constant time again after a single fold, matching the eagerly-built field this accessor replaced. Code that manually reconstructs a `Kernel.Environment` with different views carries the old thunk, which matches how the precomputed field behaved under the same surgery.

Verified: `lake lint` on a from-scratch batteries build previously hung indefinitely in `getDeclsInPackage` and now completes.

Co-Authored-By: Claude
@Kha Kha removed the mathlib4-nightly-available A branch for this PR exists at leanprover-community/mathlib4-nightly-testing:lean-pr-testing-NNNN label Jul 13, 2026
@Kha

Kha commented Jul 13, 2026

Copy link
Copy Markdown
Member Author

!bench mathlib

@leanprover-radar

leanprover-radar commented Jul 13, 2026

Copy link
Copy Markdown

Benchmark results for leanprover-community/mathlib4-nightly-testing@a04b6d1 against leanprover-community/mathlib4-nightly-testing@5edfd4c are in. There are significant results. @Kha

  • 🟥 build//instructions: +7.0T (+4.62%)

Large changes (1🟥)

  • 1 hidden

Small changes (3978🟥)

  • 🟥 build/module/Aesop.BuiltinRules.Assumption//instructions: +294.3M (+10.71%)
  • 🟥 build/module/Aesop.BuiltinRules.DestructProducts//instructions: +283.6M (+9.13%)
  • 🟥 build/module/Aesop.BuiltinRules.Split//instructions: +273.7M (+13.21%)
  • 🟥 build/module/Aesop.BuiltinRules//instructions: +270.4M (+16.81%)
  • 🟥 build/module/Aesop.Frontend.Attribute//instructions: +282.0M (+10.58%)
  • 🟥 build/module/Aesop.Frontend.Command//instructions: +290.8M (+5.86%)
  • 🟥 build/module/Aesop.Frontend.Saturate//instructions: +266.2M (+6.43%)
  • 🟥 build/module/Aesop.Main//instructions: +267.1M (+7.03%)
  • 🟥 build/module/Aesop//instructions: +257.8M (+17.03%)
  • 🟥 build/module/Batteries.Data.String.Lemmas//instructions: +1.1G (+3.21%)
  • 🟥 build/module/Mathlib.Algebra.AddConstMap.Basic//instructions: +964.1M (+3.86%)
  • 🟥 build/module/Mathlib.Algebra.AddConstMap.Equiv//instructions: +788.0M (+10.19%)
  • 🟥 build/module/Mathlib.Algebra.AddTorsor.Basic//instructions: +529.1M (+18.89%)
  • 🟥 build/module/Mathlib.Algebra.AddTorsor.Defs//instructions: +316.5M (+16.07%)
  • 🟥 build/module/Mathlib.Algebra.AffineMonoid.Basic//instructions: +611.4M (+15.66%)
  • 🟥 build/module/Mathlib.Algebra.AffineMonoid.Embedding//instructions: +931.5M (+17.62%)
  • 🟥 build/module/Mathlib.Algebra.AffineMonoid.Irreducible//instructions: +638.9M (+7.25%)
  • 🟥 build/module/Mathlib.Algebra.AffineMonoid.UniqueSums//instructions: +833.7M (+15.07%)
  • 🟥 build/module/Mathlib.Algebra.Algebra.Bilinear//instructions: +697.6M (+4.26%)
  • 🟥 build/module/Mathlib.Algebra.Algebra.Equiv//instructions: +820.1M (+1.98%)
  • and 3958 more

mathlib-nightly-testing Bot pushed a commit to leanprover-community/batteries that referenced this pull request Jul 13, 2026
@github-actions github-actions Bot added the mathlib4-nightly-available A branch for this PR exists at leanprover-community/mathlib4-nightly-testing:lean-pr-testing-NNNN label Jul 13, 2026
mathlib-nightly-testing Bot pushed a commit to leanprover-community/mathlib4-nightly-testing that referenced this pull request Jul 13, 2026
leanprover-bot added a commit to leanprover/reference-manual that referenced this pull request Jul 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaks-mathlib This is not necessarily a blocker for merging: but there needs to be a plan builds-manual CI has verified that the Lean Language Reference builds against this PR mathlib4-nightly-available A branch for this PR exists at leanprover-community/mathlib4-nightly-testing:lean-pr-testing-NNNN toolchain-available A toolchain is available for this PR, at leanprover/lean4-pr-releases:pr-release-NNNN

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants