[Optimization] Make CSE per offload rather than whole kernel#790
Open
hughperkins wants to merge 27 commits into
Open
[Optimization] Make CSE per offload rather than whole kernel#790hughperkins wants to merge 27 commits into
hughperkins wants to merge 27 commits into
Conversation
…M cache Run CSE per offloaded task, but only once each task is isolated in the codegen worker pool: the post-offload full_simplify passes (offload_to_executable's before_lower_access / simplify_IV / scalarize) execute inside compile_task, which is enqueued per task to the worker pool, so per_task_cse fires there (one task per worker, inside the simplify fixpoint = full quality, in parallel). CSE is skipped on the pre-split monolith (the block still holds every task) and deferred to those parallel per-task passes. Because tasks are optimized independently, doing all CSE in the workers is value-identical to the serial variant -- just parallel. cache_loop_invariant_global_vars defaulted off for the same soundness reason as the serial variant.
…stale break-flag) cache_loop_invariant_global_vars previously keyed its per-loop cache by GlobalPtrStmt* identity. With per-task CSE (post-offload) the read and write GlobalPtrStmts to the same address (e.g. the solver's improved[i_b] break flag) are distinct statements, so pointer-identity keying allocated separate locals -> in-loop reads never saw in-loop writes -> stale break flag -> solver ran to the iteration cap (the -88% per-offload-CSE regression). find_cache_entry now also matches an existing entry by definitely_same_address, so read+write share one local. Re-enables the pass (compile_config default true) which is load-bearing on contact-heavy solves (duck_in_box). QD_LICM_LOG=1 prints per-access cache decisions for diagnosis.
…able (conditional) accesses QD_LICM_LOG on anymal_zero showed address-keying alone is insufficient: the pass is unsound whenever a snode has an access that is cache-eligible (offload-unique, static index, non-atomic) yet not cacheable at its own site - typically a store inside an if-block that LICM did not hoist (move_loop_invariant_outside_if off). That store bypasses the cached loop-invariant local while reads still read the local -> stale (the solver break-flag / -88% regression). fix_addr could never help because such stores never reach cache_global_to_local. Phase 1 traverses each task and marks any such snode UNSAFE; phase 2 caches as before but skips UNSAFE snodes, leaving all their accesses on global (always correct). Same-depth read/write pairs still unify via the address-aware find_cache_entry, so the pass keeps its optimization on the clean (contact-heavy/duck) snodes.
…lobal_ptrs) WholeKernelCSE gains a ptrs_only mode that eliminates only Global/External/MatrixPtr statements (no compute CSE, no if-branch hoisting). Exposed as irpass::merge_global_ptrs and called in compile_to_offloads right before the first flag_access. This restores the one precondition cache_loop_invariant_global_vars needs on contact-heavy scenes: each global's read and write pointers become a single shared activate=true pointer, so conditional/in-if stores become hoistable/cacheable (recovering the duck_in_box optimization) - while the expensive whole-kernel compute dedup stays deferred to per-task CSE. QD_NO_PTR_MERGE=1 disables for A/B.
…E log Pointers-only mode now also eliminates pure integer (index) compute, not just Global/External/MatrixPtr. Two same-address pointers only satisfy definitely_same_address once their index arithmetic collapses to a common base statement (value_diff_ptr_index/FindDirectValueBaseAndOffset), so merging pointers alone was a near no-op. Loads and stores remain non-eliminable so soundness is unchanged; float compute stays deferred to per-task CSE. QD_LICM_LOG=1 prints [PTRMERGE] round/modified counts.
Pre-offload the pass found nothing (PTRMERGE modified=0): offload is what clones each global's address computation into per-task read/write GlobalPtrStmts. Run the merge on the offloaded IR right before flag_access #2 so the still-activate=true duplicates unify into one shared pointer before flag_access can stamp the read-only copy activate=false. This is the split that blocks cache_loop_invariant_global_vars.
…up), drop standalone Evidence (solve_monolith IR): main's whole_kernel_cse runs in full_simplify's fixpoint at simplify_I, before the first flag_access, deduping read/write pointers so flag_access never splits them (after_offload: 1 split). The per-task branch did no CSE pre-offload, so flag_access split ~16 addresses; a single standalone post-offload merge could not undo it (activate=false read dominates; CSE won't merge a later activate=true write into it) - only 16->13. Running the cheap pointer+integer merge in the fixpoint (like main) fires at simplify_I before the split. Removed the ineffective standalone post-offload call.
Now that merge_global_ptrs unifies read/write pointers in the full_simplify fixpoint (splits 16->1 like main), the two-phase UNSAFE exclusion is a redundant workaround that also disables caching for contact-heavy (duck) fields -> suspected cause of the residual duck regression despite splits being fixed. Toggle to A/B.
…the root-cause fix) Bench confirms merge_global_ptrs (fixpoint) alone fixes both problems: anymal_zero at full baseline (was -88%) AND duck_in_box within noise of main (+0.5% / -1.0% / +1.6% / -2.0%, vs -8..-13% with the exclusion on). The two-phase UNSAFE exclusion recovered anymal but kept duck at pass-off, so it is now off by default. QD_LICM_EXCLUDE=1 re-enables it as a safety valve.
…ssion) The pass is only load-bearing pre-offload (unify read/write pointers before the first flag_access). Post-offload, per_task_cse already merges pointers within each task, so running merge_global_ptrs in every post-offload full_simplify fixpoint iteration (per task, in the codegen workers) was pure redundant work -> ~3.5% compile regression across the suite. Now skips as soon as any top-level OffloadedStmt is present. Pre-offload behavior (the correctness fix) is unchanged.
…plify fixpoint The pre-offload pointer merge only needs to happen ONCE before the first flag_access (compile_to_offloads), to unify a global's read/write pointers before flag_access stamps the read-only copy activate=false. Running a whole-kernel pointer CSE in every fixpoint iteration of every pre-offload phase (simplify_I/II + autodiff) was a +12-22s compile regression on collision-heavy scenes (franka/duck/box_pyramid; anymal ~0) for zero extra benefit -- same-build A/B: franka_random default 120.8s vs no_merge 108.5s. Now a single irpass::merge_global_ptrs(ir) right before flag_access #1; arithmetic is already canonical after simplify_I so one pass suffices. Correctness path (splits->1, -88% bug + duck optimization) unchanged.
Collaborator
Author
|
Genesis benchmarks: 
|
Collaborator
Author
|
Benchmarks on actual genesis CI ran ok:
|
hughperkins
marked this pull request as ready for review
July 22, 2026 15:20
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9d60308817
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| // gets from whole_kernel_cse running inside the post-offload full_simplify (per-task CSE otherwise defers to the | ||
| // codegen workers, which run after cache_loop_invariant_global_vars). Needed for ndarrays, which only become | ||
| // ExternalPtrStmts during offload and so cannot be reached by the pre-offload merge_global_ptrs. See the pass. | ||
| irpass::cse_offloaded_tasks(ir); |
There was a problem hiding this comment.
When users initialize with opt_level=0, this new post-offload call still runs full per-task CSE over every offloaded task, even though the simplify loop below continues to gate CSE on config.opt_level > 0 and docs/source/user_guide/optimization_passes.md documents opt_level=0 as disabling CSE for faster compilation. This makes the public opt_level knob no longer reliably disable the expensive CSE pass; either guard this standalone call (and the related pointer merge if it remains part of CSE semantics) with the same option or update the public docs/semantics.
AGENTS.md reference: AGENTS.md:L15-L22
Useful? React with 👍 / 👎.
…er_task_cse cse_offloaded_tasks is the post-offload analog of whole_kernel_cse, which upstream runs inside full_simplify gated on opt_level>0; per_task_cse is gated the same way. Running it unconditionally would do CSE at opt_level 0 (where upstream does none), contradicting the documented behaviour (optimization_passes.md: CSE runs only when opt_level>0). At opt_level 0 there is no CSE requiring pointer unification anyway.
Reflect the per-task CSE migration: CSE runs over one offloaded task's IR at a time (in the per-task simplify stages, in parallel across worker threads) rather than the whole kernel, mirroring the existing per-task CFG note; and a one-shot per-task CSE runs right after offload so downstream memory optimizations see each global's read/write as a single shared access.
Run the CI check_doc_quality rubric against optimization_passes.md: - gloss autodiff inline + link autodiff.md (Rule 1 undefined term) - gloss backend codegen targets (PTX/SASS/SPIR-V) as device machine code (Rule 1) - drop the internal 'tests assume defaults' maintenance note (Rule 2 scope) - fix DIE expansion: dead-instruction (matches irpass::die), not dead-code
The non-ASCII CI check flags box-drawing/arrow glyphs in changed doc lines. Converting the pipeline diagram to ASCII (-->/|/v) keeps autodiff glossed inline (= automatic differentiation) for the doc-quality check while satisfying the non-ASCII rule.
Collaborator
Author
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d732af9135
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| bool cse_one_task(Block *parent, OffloadedStmt *off, bool ptrs_only = false) { | ||
| const int location = parent->locate(off); | ||
| QD_ASSERT(location != -1); | ||
| Block wrapper; |
There was a problem hiding this comment.
When a task contains a MatrixPtrStmt (for example dynamic vector/matrix element access, especially with real_matrix_scalarize=False), this temporary wrapper severs the offload from the original root block's callable. WholeKernelCSE then visits that MatrixPtrStmt, whose common_statement_eliminable() calls get_callable() through the parent chain; with this wrapper the chain ends at a null callable, so the assertion fires during cse_offloaded_tasks/per_task_cse instead of compiling the supported kernel. Preserve the original parent on the wrapper or avoid extracting the offload into a parentless block.
Useful? React with 👍 / 👎.
Doc quality (CI Cursor agent is stricter than a local check): drop undefined jargon it flagged/would flag - PTX/SASS/LLVM/SPIR-V (line 34), AST (glossed), 'grid' (removed), 'peephole' (defined inline); and remove internal pass-scheduling/worker-thread prose (line 59) that the end user cannot control or observe. Line wrapping: reflow the test_cache_loop_invariant docstring to fill to ~120 cols (was ~94-105, flagged as under-wrapped).
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
1 participant
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.
Issue: #
Brief Summary
copilot:summary
Walkthrough
copilot:walkthrough