[Lang] whole_kernel_cse: 2.5x compile time speedup on large kernels

[duburcqa]

`whole_kernel_cse`: scoped MarkUndone, flat hash bucket, combined walk

Two commits, no behaviour change. Reduces the per-elimination IR-walking cost in WholeKernelCSE from O(N) full-IR sweeps to O(scope) walks (where `scope` is the SSA dominance frontier of the eliminated statement), replaces the per-scope hash-map stack with a flat hash table, swaps a string-allocating type comparison for an RTTI compare, and folds the MarkUndone walk into the replace-usages walk. Affects every kernel that runs through `irpass::full_simplify`; absolute savings are largest on large reverse-mode kernels where autodiff inflates the IR.

TL;DR

`irpass::whole_kernel_cse` runs three to four times per kernel inside `irpass::full_simplify`, with an outer fixpoint loop until no more eliminations are possible. The pass walks the kernel IR, hashes each statement by operand pointers, and replaces a statement with a previously-seen equivalent. Profiling shows it is the dominant compile-time cost on large autodiff kernels:

• `MarkUndone::run` walked `get_ir_root()` per elimination: O(M * V) total. SSA dominance bounds the user set to the eliminated stmt's parent block subtree plus ancestor top-level stmts (the same scope `replace_all_usages_with(nullptr, ...)` already covers).
• The visibility table was a stack of per-scope `unordered_map<size_t, unordered_set<Stmt*>>` requiring an O(D) map lookup per stmt visit, where D is nesting depth.
• Each elimination ran two scope sweeps - one for MarkUndone, one for `replace_all_usages_with` - both checking `has_operand`.
• `Stmt::type()` (used for the type-mismatch fast path in `common_statement_eliminable`) constructs a `StatementTypeNameVisitor`, dispatches `accept()`, and allocates a `std::string` per call.

Per-kernel measurements

Genesis Ant-scale reverse-mode test (`tests/test_grad.py::test_differentiable_rigid[gpu]`, macOS Metal, fresh compile cache via `QD_OFFLINE_CACHE=0`).

`kernel_forward_dynamics_without_qacc_c305_0` (V=9535, E=14331, 14 cyclic SCCs, 4241 AD-stacks):

Stage	`forward_dynamics` compile
Baseline	91 s
Scoped MarkUndone + flat bucket + RTTI	48 s
+ Combined `ReplaceAndMarkUndone` walker	36 s

Smaller kernels in the same test (forward velocity, update_cartesian_space, step_2, compute_qacc) see proportional improvements that are absolute-time invisible because their original CSE cost was already in the few-millisecond range.

Why

Sampling `pytest` during `kernel_forward_dynamics_without_qacc` compilation on macOS Metal showed 100% CPU in the `WholeKernelCSE::visit(Stmt*) -> MarkUndone::run -> BasicStmtVisitor::visit(Block*) -> IfStmt -> ...` recursion chain. The compile cost on this kernel scales with two compounding factors: many CSE candidates per pass, each triggering a full-IR sweep, and a deeply-nested `IfStmt` structure generated by the autodiff transform that inflates the visibility-stack lookup cost.

This pass runs on every kernel through `full_simplify`, so the optimisations apply equally to forward (`AutodiffMode::NONE` / `FORWARD`) and reverse compilations; reverse kernels were where the wallclock cost made the bottleneck visible because their IR is roughly two to three times the size of the forward equivalent.

Mechanism

1. Scope `MarkUndone` to the SSA dominance frontier

`MarkUndone::run` previously walked `modified_operand->get_ir_root()`. Replaced with the same scope `replace_all_usages_with(nullptr, ...)` already uses: walk `modified_operand->parent` (its containing block's full subtree) and then iterate ancestor blocks' top-level statements. Per SSA dominance every user of `modified_operand` lives in this scope, so we cannot miss a stale `visited_` entry. Cost per elimination drops from O(V) to O(scope size); on deeply-nested kernels the eliminated stmt usually lives well below the offloaded body, so the relevant subtree is much smaller than the whole IR.

2. Flat hash bucket plus per-scope insertion log

The visibility table was a `vector<unordered_map<size_t, unordered_set<Stmt*>>>` keyed by scope (one entry per active block). On each `visit(Stmt*)` we walked every scope's map. Replaced with:

• one global `unordered_map<size_t, vector<Stmt*>>` keyed by `operand_hash`;
• a `vector<vector<pair<size_t, Stmt*>>>` recording what was inserted at each scope depth.

`visit(Block*)` pushes a fresh insertion log entry, recurses, then on scope exit replays the entry to remove the corresponding stmts from the global table. Sibling subtrees see only ancestor inserts, identical to the per-scope-map semantics. Per-stmt visibility lookup is now O(1) hash plus an O(bucket) bucket walk instead of O(D) hash maps.

3. Combine MarkUndone and replace-usages into `ReplaceAndMarkUndone`

The MarkUndone walk and the `replace_all_usages_with` walk cover the same scope (parent block subtree plus ancestor top-level) and both check `has_operand(modified_operand)` per visited stmt. Folded into a single `ReplaceAndMarkUndone` pass that does both side-effects per visit. Halves the IR-walking cost per elimination.

4. RTTI type compare instead of `Stmt::type()`

`Stmt::type()` constructs a `StatementTypeNameVisitor`, dispatches `accept()` (virtual), and returns a freshly-allocated `std::string`. Called inside `common_statement_eliminable` for every (this_stmt, prev_stmt) pair that shares a hash bucket - which on large kernels was thousands of comparisons per pass.

Replaced with `typeid(*this_stmt) != typeid(*prev_stmt)`. Same correctness contract: the unchecked `prev_stmt->as<...>()` casts immediately below the early-return need this guard, otherwise a hash-bucket collision would reinterpret memory in `definitely_same_address`. After this PR the typeid compare is almost always true at runtime because `operand_hash` now also routes by `typeid(*stmt)` (see below); the guard is kept as the load-bearing safety net.

5. Bonus: `typeid(*stmt)` in `operand_hash`

`operand_hash` previously used `std::hashstd::type_index{}(std::type_index(typeid(stmt)))`. `typeid(stmt)` operates on the `Stmt*` static type for every input, collapsing every concrete statement class into the same hash component - bucket separation came only from operand pointer hashes. Switched to `typeid(*stmt)`, so different statement classes now land in different buckets when their operand sets happen to alias.

Side-effect audit

Concern	Where checked	Verdict
Visibility-table scope semantics	`visit(Block*)` pushes `scope_inserts_`; on exit, replay drops only this-scope entries from `visible_stmts_`	Sibling subtrees see only ancestor inserts; identical to the previous per-scope `unordered_map` stack
MarkUndone correctness on dominated users	New scope = `modified_operand->parent` subtree plus ancestor top-level stmts	SSA dominance guarantees every user of `modified_operand` lives in this scope; no stale `visited_` entry can survive
MarkUndone for stmts with `parent == nullptr`	Falls back to `get_ir_root()` walk	Defensive only; should not happen for IR built through normal compile paths
`ReplaceAndMarkUndone` matches `replace_all_usages_with` semantics	Inherits `BasicStmtVisitor` like the existing `StatementUsageReplace`; same scope traversal, same `has_operand` check, same call site `replace_operand_with`	Equivalent traversal; both side-effects fire on the same set of stmts
Type-mismatch guard after `operand_hash` typeid mix	`typeid(*this_stmt) != typeid(*prev_stmt)` early-return retained in `common_statement_eliminable`	Almost always true at runtime now (hash bucket already keyed by typeid) but load-bearing - the `prev_stmt->as<...>()` casts immediately below are unchecked static-cast wrappers; a hash collision without this guard would reinterpret memory
`typeid(*stmt)` on polymorphic Stmt	Stmt has virtual member functions (`quadrants/ir/ir.h`)	RTTI returns the dynamic type as required
Public API of `irpass::whole_kernel_cse`	Unchanged signature; same one-shot driver `WholeKernelCSE::run(IRNode*)`; same outer fixpoint loop	Callers unmodified

[claude]

⚠️ Code review skipped — your organization's overage spend limit has been reached.

Code review is billed via overage credits. To resume reviews, an organization admin can raise the monthly limit at claude.ai/admin-settings/claude-code.

Once credits are available, reopen this pull request to trigger a review.

[hughperkins]

@claude review

[hughperkins]

@claude review

[claude]

Verified the correctness arguments hold up (SSA-dominance scope matches StatementUsageReplace, the BasicStmtVisitor vs IRVisitor traversal divergence is benign because StructForStmt/MeshForStmt/OffloadedStmt have no Stmt* operands, and the flat-bucket scope cleanup is sound), but this bundles four subtle optimizations to a critical compiler pass — worth a human compiler-savvy reviewer.

Extended reasoning...

Overview

Single-file change to quadrants/transforms/whole_kernel_cse.cpp bundling four correctness-sensitive performance optimizations to the WholeKernelCSE pass: (1) scoping MarkUndone to the SSA dominance frontier instead of get_ir_root(), (2) replacing the per-scope vector<unordered_map<size_t, unordered_set<Stmt*>>> with a flat unordered_map<size_t, vector<Stmt*>> plus a per-scope insertion log, (3) fusing MarkUndone and replace_all_usages_with into a single ReplaceAndMarkUndone walker, and (4) swapping Stmt::type() string compares for typeid(*stmt) (and fixing operand_hash to use typeid(*stmt) instead of typeid(stmt), which was previously always the Stmt* static type). Public API of irpass::whole_kernel_cse is unchanged.

Security risks

None. This is a compiler IR transformation pass; no auth, crypto, network, or data-exposure surfaces. The risk model is correctness of generated kernel code, not security.

Level of scrutiny

High. whole_kernel_cse runs inside full_simplify for every kernel including reverse-mode autodiff kernels, and a miscompile here produces silently-wrong numerics. The new SSA-dominance scoping and the flat-bucket scope-cleanup logic are correctness-load-bearing — a stale visited_ entry or a leaked visibility entry across sibling scopes would be observable as either missed CSE or, worse, an incorrect cross-scope CSE replacement. I verified the scope claim matches StatementUsageReplace::run(nullptr, ...) (parent block subtree + ancestor top-level statements), the bucket-cleanup correctly preserves outer-scope entries via per-stmt std::find, and the BasicStmtVisitor-vs-IRVisitor traversal divergence is benign because the affected container stmts (StructForStmt, MeshForStmt, OffloadedStmt) have no Stmt* operands in their QD_STMT_DEF_FIELDS so has_operand is always false for them.

Other factors

The PR description includes a detailed side-effect audit table that matches what I see in the code, the per-kernel measurements give a concrete win (91s → 36s on forward_dynamics), and the bug-hunting pass did not flag anything. But the "no behavior change" claim has subtle caveats (the operand_hash change will re-bucket statements differently across runs, the traversal change is a different visitor base class, the visibility table is logically equivalent but a different data structure), and the correctness argument leans on SSA invariants that are easy to state and harder to verify against every IR-producer in the pipeline. A reviewer with compiler-pass context should double-check the SSA-dominance assumption holds for IR shapes produced by the autodiff transform — that's where the wallclock savings are largest and also where the IR is least conventional.

[duburcqa]

Review by local Claude Opus 4.7:

I've completed a careful review of PR #577. Here's my assessment.

Verdict

Looks correct and ships safely. The correctness rests on existing invariants in the codebase, the perf wins are real, and CI is green across all platforms. A few minor nits worth addressing.

What I verified

1. Scope narrowing of MarkUndone is sound (commit c7190dc02)
The new ReplaceAndMarkUndone::run walks the exact same scope irpass::replace_all_usages_with(nullptr, ...) was already using - confirmed by reading quadrants/transforms/statement_usage_replace.cpp:63-83:
- old_stmt->parent->accept(&walker) (full subtree of containing block)
- For ancestors: top-level statements only via parent_block() chain - no recursion into nested blocks of ancestor siblings

This is identical to the base-branch contract for replace_all_usages_with. By SSA dominance any user of old_stmt lives in this set, so the previous get_ir_root() walk was strictly redundant work. If base was correct, this is correct - and the existing
replace_all_usages_with already trusts the same scope.

The parent == nullptr defensive fallback to get_ir_root() matches the only edge case replace_all_usages_with doesn't handle (it would QD_ASSERT instead).

2. typeid(*stmt) is a genuine bug fix
typeid(stmt) operates on the pointer expression's static type (const Stmt *) and resolves at compile time to the same type_info for every input - the previous operand_hash had zero type discrimination in its hash_type component. typeid(*stmt) dereferences and uses RTTI
on the dynamic type. Stmt has virtual methods (accept, replace_operand_with, etc.) so RTTI works as expected. Main quadrants codebase doesn't compile with -fno-rtti (only external/SPIRV-Tools does).

3. Flat hash bucket + insertion log is semantically equivalent
- Within a single node->accept(&eliminator) traversal, visit(Block*) pushes a frame at entry and replays it at exit, so visible_stmts_ is back to empty at the end. Outer fixpoint iterations see a clean state.
- A stmt is registered at most once per visit(Stmt*) call (mutually exclusive branches), so no double-insert into a single scope's frame.
- Iteration order changes from unordered_set (undefined) to vector (insertion order). More deterministic; doesn't affect correctness because any matched prev_stmt is a valid replacement.

4. Combined ReplaceAndMarkUndone (commit c6f6431e2)
Both side-effects (visited_->erase + replace_operand_with) fire on exactly the stmts where has_operand(old_stmt_) is true. Original code did this in two separate scope traversals; combining them halves the walk cost without changing semantics.

5. Style + CI
- ASCII-only verified (grep -P "[^\x00-\x7F]" clean).
- Single-dash usage throughout, no em-dashes.
- All Build/Test jobs pass on Linux/Mac/Windows + AMD GPU + CUDA.

Nits / suggestions (non-blocking)

N1. Comment on common_statement_eliminable is slightly misleading.
whole_kernel_cse.cpp:120 claims as<...> are "unchecked static-cast convenience wrappers". They're actually dynamic_cast + QD_ASSERT(is<T>()) (ir.h:246-249). The hazard is still real - on type mismatch QD_ASSERT fires QD_ERROR which expands to __builtin_unreachable(), so
 it's UB - but "unchecked static-cast" describes the wrong mechanism. Suggest rewording to "the as<...> casts assert via QD_UNREACHABLE on type mismatch (UB in release)" or similar.

N2. Bucket erase on scope exit is O(B²).
whole_kernel_cse.cpp:196 does std::find(bucket.begin(), bucket.end(), stmt) for each entry in scope_inserts_.back(). For a large bucket where many inserts happen at the same depth, this is quadratic. In practice buckets are tiny because the new typeid(*stmt) mix gives
good separation, so this is unlikely to bite. But since insertion order is preserved and all inserts at this depth land at the end of the bucket, the cleanup could pop from the back in reverse order in O(1) per entry. Worth considering as a micro-follow-up.

N3. register_visible isn't a hot enough abstraction to be a method, but harmless.

What I'd flag for follow-up

- Adding a deterministic CSE-output fingerprint test would help future refactors of this pass; not blocking here.
- The if_stmt branch-merging path at whole_kernel_cse.cpp:222-241 re-parents stmts but doesn't MarkUndone for the moved stmt's operands' users. That's not introduced by this PR, but worth noting if anyone touches this file again.

LGTM as-is; the comment fix at N1 is the only thing I'd suggest before merge.

[github-actions]

Coverage Report (0eb454816)

Metric	Value
Diff coverage (changed lines only)	0%
Overall project coverage	73%

Total: 0 lines, 0 missing, 0% covered

[duburcqa]

env	batch_size	backend	gjk_collision	constraint_solver	compile_time_baseline_s	compile_time_new_s	compile_time_delta_pct	runtime_fps_baseline	runtime_fps_new	runtime_fps_delta_pct	realtime_factor_baseline	realtime_factor_new	realtime_factor_delta_pct
anymal_random	30000	cuda	-	-	82.9	61.9	-25.33	9423585	9634040	+2.23	94235.9	96340.4	+2.23
anymal_uniform	30000	cuda	-	-	89.5	62.2	-30.50	12235395	12941962	+5.77	122353.9	129419.6	+5.77
anymal_uniform_kinematic	0	cpu	-	-	51.4	34.3	-33.27	1992	2038	+2.31	19.9	20.4	+2.51
anymal_uniform_kinematic	30000	cuda	-	-	84.3	59.8	-29.06	10378314	10395688	+0.17	103783.1	103956.9	+0.17
anymal_zero	0	cpu	-	-	49.6	30.2	-39.11	6843	7466	+9.10	68.4	74.7	+9.21
anymal_zero	30000	cuda	-	-	87.9	61.7	-29.81	19072828	19591793	+2.72	190728.3	195917.9	+2.72
box_pyramid_3	4096	cuda	-	-	116.6	76.8	-34.13	980650	976615	-0.41	9806.5	9766.1	-0.41
box_pyramid_4	4096	cuda	-	-	112.9	78.4	-30.56	398200	397174	-0.26	3982.0	3971.7	-0.26
box_pyramid_5	4096	cuda	-	-	108.5	77.7	-28.39	141292	141645	+0.25	1412.9	1416.5	+0.25
box_pyramid_6	4096	cuda	False	-	111.5	76.1	-31.75	59399	59826	+0.72	594.0	598.3	+0.72
box_pyramid_6	4096	cuda	True	-	81.2	60.6	-25.37	61883	61462	-0.68	618.8	614.6	-0.68
dex_hand	4096	cuda	-	-	123.8	89.3	-27.87	17296	17225	-0.41	1081.0	1076.6	-0.41
duck_in_box_easy	30000	cuda	False	-	83.4	57.1	-31.53	26368796	26596389	+0.86	263688.0	265963.9	+0.86
duck_in_box_easy	30000	cuda	True	-	54.7	40.6	-25.78	9824790	9803261	-0.22	98247.9	98032.6	-0.22
duck_in_box_hard	0	cpu	-	-	52.6	33.1	-37.07	4982	5201	+4.40	49.8	52.0	+4.42
duck_in_box_hard	30000	cuda	False	-	82.8	57.2	-30.92	10040367	10278637	+2.37	100403.7	102786.4	+2.37
duck_in_box_hard	30000	cuda	True	-	53.5	40.8	-23.74	3557329	3531797	-0.72	35573.3	35318.0	-0.72
franka	30000	cuda	-	-	78.9	54.5	-30.93	21839164	21930939	+0.42	218391.6	219309.4	+0.42
franka_accessors	0	cpu	-	-	47.2	30.6	-35.17	1187	1221	+2.86	11.9	12.2	+2.52
franka_accessors	30000	cuda	-	-	84.9	58.8	-30.74	15371212	15694558	+2.10	153712.1	156945.6	+2.10
franka_free	30000	cuda	-	-	78.5	54.2	-30.96	32772988	32743296	-0.09	327729.9	327433.0	-0.09
franka_random	0	cpu	-	-	46.3	30.5	-34.13	5851	6489	+10.90	58.5	64.9	+10.94
franka_random	30000	cuda	-	CG	78.5	53.7	-31.59	16695511	16927661	+1.39	166955.1	169276.6	+1.39
franka_random	30000	cuda	-	Newton	76.8	52.5	-31.64	16721019	16848373	+0.76	167210.2	168483.7	+0.76
franka_random	30000	cuda	False	-	76.9	54.6	-29.00	16857886	16407363	-2.67	168578.9	164073.6	-2.67
franka_random	30000	cuda	True	-	50.5	36.8	-27.13	11378021	11497596	+1.05	113780.2	114976.0	+1.05
g1_fall	4096	cuda	-	Newton	97.7	77.8	-20.37	911275	910627	-0.07	4556.4	4553.1	-0.07
go2	4096	cuda	False	CG	80.8	55.7	-31.06	3610886	3610521	-0.01	36108.9	36105.2	-0.01
go2	4096	cuda	False	Newton	93.0	75.7	-18.60	4460541	4472146	+0.26	44605.4	44721.5	+0.26
go2	4096	cuda	True	-	83.4	56.2	-32.61	3240315	3253632	+0.41	32403.2	32536.3	+0.41
shadow_hand_cubes	0	cpu	-	-	48.8	32.6	-33.20	40	41	+2.50	1.3	1.4	+7.69
shadow_hand_cubes_sparse	0	cpu	-	-	47.7	31.4	-34.17	65	65	+0.00	2.2	2.2	+0.00

speed_comparison.csv

Comparing genesis/main + Quadrants 0.7.0 (baseline) vs Genesis-Embodied-AI/genesis-world#2743 + this PR

[hughperkins]

Fantastic! 🔥 🔥 🔥

ok to merge.