[Optimization] Make CSE per offload rather than whole kernel#790
Draft
hughperkins wants to merge 21 commits into
Draft
[Optimization] Make CSE per offload rather than whole kernel#790hughperkins wants to merge 21 commits into
hughperkins wants to merge 21 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: 
|
…desired defaults The PR now ships the intended behavior with no runtime toggles: - merge_global_ptrs always runs (drop QD_NO_PTR_MERGE gate) - ndarray two-phase exclusion always honored; field/SNode exclusion removed entirely (fields are handled upstream by merge_global_ptrs), dropping QD_LICM_EXCLUDE - drop QD_LICM_LOG debug logging (licm_log/describe_dest/[PTRMERGE]) Also drops now-unused snode-exclusion machinery and includes.
Collaborator
Author
|
Genesis unit tests on ndarray passing: 
|
Collaborator
Author
|
Genesis benchmarks ok-ish: 
|
…loop_invariant Recovers the runtime regression the ndarray-exclusion introduced on contact-heavy GJK scenes (duck_in_box_hard gjk=True: -10.47%). Root cause: cache_loop_invariant runs in offload_to_executable before any pointer-merging CSE (monolith simplify_III no-ops per-task CSE; offload_to_executable's own full_simplify runs after it). Pre-offload merge_global_ptrs can't reach ndarrays (not ExternalPtr yet pre-offload). Two changes: 1. operand_hash: special-case ExternalPtrStmt like GlobalPtrStmt so same-address ndarray pointers bucket together and merge via definitely_same_address even when their index compute statements are distinct (previously they only merged if operands were identical objects, which required full whole-kernel CSE -- what per-task CSE skips). 2. merge_offloaded_ptrs: run the cheap ptrs-only CSE per offloaded task right before cache_loop_invariant, unifying each address's read/write pointers (fields + ndarrays). cache_loop then caches soundly with one shared pointer, so the ndarray exclusion becomes inert (kept as a safety net) and the optimization is restored. Preserves per-task CSE.
…ag_access #2) Placing it before cache_loop_invariant was too late: flag_access #2 and simplify_III's LICM already split/hoisted the ndarray break-flag's read pointer by then, so cache_loop still couldn't cache it (duck_in_box_hard gjk=True stayed ~-12% vs main on same-node bench). Move the per-task pointer merge to immediately after offload -- the post-offload analog of the pre-offload merge_global_ptrs at flag_access #1 -- so the read/write ExternalPtrs are unified before any split/hoist and cache_loop can cache soundly.
…pstream The prior approach (ptrs-only merge + address-keyed cache_loop + two-phase ndarray exclusion) preserved correctness but left cache_loop unable to cache the ndarray break-flag, so duck_in_box_hard gjk=True stayed ~-12% vs main (same-node bench: main 2.81M / branch 2.47M FPS). Root cause: cache_loop needs a global's read and write pointers unified, which upstream gets from whole_kernel_cse running inside the post-offload full_simplify; per-task CSE defers to the codegen workers, which run after cache_loop. Fix: run full per-task CSE (cse_offloaded_tasks) on the monolith right after offload, before flag_access #2 and simplify_III's LICM -- the post-offload analog of the pre-offload merge_global_ptrs. This unifies each task's read/write pointers (crucially ndarray ExternalPtrs, which only exist post-offload and need full CSE to merge, since an ExternalPtr merges only once its index-compute is merged). With pointers unified, cache_loop_invariant_global_vars is sound exactly as upstream, so it is reverted verbatim to main -- dropping the address-keyed find_cache_entry AND the two-phase ndarray exclusion. Also special-case ExternalPtrStmt in operand_hash (like GlobalPtrStmt) so same-address ndarray pointers bucket together and merge via definitely_same_address. Same-node result: duck_in_box_hard gjk=True 2.85M FPS (>= main 2.81M), compile ~33s (< main 43s); ndarray detection hang test passes; field + ndarray regression test added.
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