[Optimization] Make CSE per offload rather than whole kernel

docs/source/user_guide/optimization_passes.md

@@ -12,29 +12,29 @@ Let's start by defining terms that will be used throughout the page and are nece
 - **IR (intermediate representation)** - the compiler's internal version of your kernel: a flat list of small, explicitly-typed instructions, sitting between your Python source and the final machine code. Every pass reads and rewrites the IR; none of it is something you write by hand.
 - **Pass** - one transformation step over the IR. An *optimization* pass rewrites the IR into a form that produces the **same results** but runs faster or uses less memory. (Some passes are not optimizations but *lowering* steps - they translate high-level constructs into lower-level ones; this page focuses on the optimizations.)
 - **Basic block** - a straight-line run of instructions with no branches into or out of the middle. Control flow (`if`, loops) connects blocks together.
-- **Offloaded task** - after the *offload* step, your kernel is split into one or more tasks, and each task becomes a single device launch: one GPU grid launch on a GPU backend, or one parallel loop on CPU. A simple kernel is usually one task; a kernel with, say, a short serial preamble followed by a big parallel loop becomes several tasks that run back to back.
+- **Offloaded task** - after the *offload* step, your kernel is split into one or more tasks, and each task becomes a single device launch: one GPU launch on a GPU backend, or one parallel loop on CPU. A simple kernel is usually one task; a kernel with, say, a short serial preamble followed by a big parallel loop becomes several tasks that run back to back.
 
 ## The compile pipeline at a glance
 
 Compilation runs as a fixed sequence of stages. Optimization passes are interleaved with the lowering steps that gradually turn high-level IR into device code:
 
 ```
-Python (AST)
-   │  lower to IR, type-check
-   ▼
-high-level IR  ──► simplify ──► (autodiff, if requested) ──► simplify
-   │
-   ▼
+Python (AST = abstract syntax tree)
+   |  lower to IR, type-check
+   v
+high-level IR  --> simplify --> (autodiff = automatic differentiation, if requested) --> simplify
+   |
+   v
 offload  (split the kernel into offloaded tasks)
-   │
-   ▼
-per-task IR  ──► simplify ──► lower memory access ──► simplify
-   │
-   ▼
-backend codegen  (LLVM → PTX/SASS, or SPIR-V, …)
+   |
+   v
+per-task IR  --> simplify --> lower memory access --> simplify
+   |
+   v
+backend codegen  (translate IR into the device machine code your GPU runs)
 ```
 
-The "simplify" boxes are all the same routine (internally `full_simplify`), invoked at several points. Most of the interesting optimization work happens inside it.
+The "simplify" boxes are all the same routine (internally `full_simplify`), invoked at several points. Most of the interesting optimization work happens inside it. The optional autodiff step, run only when you ask Quadrants for gradients, is covered in [Automatic differentiation](./autodiff.md).
 
 ## The simplify loop
 
@@ -46,16 +46,18 @@ In the order they run each round:
 |------|--------------|
 | Extract constant | Lifts constant values out of larger expressions into standalone constant instructions, so the passes below can recognize and reuse them. |
 | Unreachable-code elimination | Removes branches that can never be taken (e.g. the body of an `if` whose condition is always false). |
-| Binary-op / algebraic simplification | Applies arithmetic identities: `x * 1 → x`, `x + 0 → x`, `x * 2 → x + x`, and similar peephole rewrites. |
+| Binary-op / algebraic simplification | Applies arithmetic identities: `x * 1 -> x`, `x + 0 -> x`, `x * 2 -> x + x`, and similar local rewrites over a short window of instructions ("peephole" optimizations). |
 | Constant folding | Pre-computes expressions whose inputs are all known at compile time: `2 * 3 → 6`. |
-| Dead-code elimination (**DIE**) | Drops instructions whose results are never used. Runs several times per round, after passes that tend to create newly-dead instructions. |
+| Dead-instruction elimination (**DIE**) | Drops instructions whose results are never used. Runs several times per round, after passes that tend to create newly-dead instructions. |
 | Loop-invariant code motion (**LICM**) | Hoists a computation that produces the same value on every iteration out of the loop, so it runs once instead of N times. |
 | Local simplify | Peephole cleanups within a block. |
 | Common-subexpression elimination (**CSE**) | Finds an identical expression computed more than once and computes it a single time, reusing the result. |
 | Control-flow-graph (**CFG**) optimization | Memory-focused optimizations that need a whole-task view; see the next section. Runs only once per stage (it is the most expensive pass). |
 
 Two of these - CSE and CFG optimization - run only when `opt_level > 0` (the default is `1`).
 
+**CSE is scoped per offloaded task.** Once the kernel has been split, common-subexpression elimination runs over one offloaded task's IR at a time rather than the whole kernel at once - the same per-task scoping described for [CFG optimization](#control-flow-graph-cfg-optimization) below, and for the same reason: each task is a separate device launch, so there is nothing to deduplicate across a task boundary.
+
 ## Control-flow-graph (CFG) optimization
 
 A **control-flow graph** is a map of your kernel's basic blocks together with the branches connecting them. It lets the compiler answer questions of the form "if execution reaches *here*, what must already have happened?" - which is exactly what is needed to optimize reads and writes to memory. Two such optimizations run on the CFG:
@@ -79,7 +81,7 @@ All of these are fields of `CompileConfig`, so you set them at `qd.init(...)` (o
 | `constant_folding` | `True` | Enables the constant-folding pass. |
 | `fast_math` | `True` | Allows IEEE-relaxed floating-point rewrites (e.g. fusing a multiply and add). Covered in [qd.init options](./init_options.md#fast_math). |
 
-For everyday use, leave them at their defaults. The most common deliberate change is `cfg_optimization=False` when iterating on a kernel whose compile time is in your way. Note that, in general, changing these options is relatively fragile since the Quadrants tests run assuming the default values.
+For everyday use, leave them at their defaults - they are the best-supported and most reliable configuration. The most common deliberate change is `cfg_optimization=False` when iterating on a kernel whose compile time is in your way.
 
 ## Inspecting what the compiler did
 

quadrants/ir/transforms.h

@@ -46,6 +46,9 @@ bool alg_simp(IRNode *root, const CompileConfig &config);
 bool demote_operations(IRNode *root, const CompileConfig &config);
 bool binary_op_simplify(IRNode *root, const CompileConfig &config);
 bool whole_kernel_cse(IRNode *root);
+bool per_task_cse(IRNode *root);
+bool merge_global_ptrs(IRNode *root);
+bool cse_offloaded_tasks(IRNode *root);
 bool extract_constant(IRNode *root, const CompileConfig &config);
 bool unreachable_code_elimination(IRNode *root);
 bool loop_invariant_code_motion(IRNode *root, const CompileConfig &config);

quadrants/program/compile_config.h

@@ -25,6 +25,10 @@ struct CompileConfig {
   bool lower_access;
   bool simplify_after_lower_access;
   bool move_loop_invariant_outside_if;
+  // Load-bearing optimization on contact-heavy solves (e.g. duck_in_box). The pass caches a loop-invariant global
+  // load into a local, which is only sound when a global's read and write pointers are the same statement. Under
+  // per-task CSE that unification is restored by merge_global_ptrs (pre-offload, fields) and cse_offloaded_tasks
+  // (post-offload, ndarrays) -- see compile_to_offloads.cpp -- so this pass itself is unchanged from upstream.
   bool cache_loop_invariant_global_vars{true};
   bool demote_dense_struct_fors;
   bool advanced_optimization;

quadrants/transforms/compile_to_offloads.cpp

@@ -131,6 +131,17 @@ void compile_to_offloads(IRNode *ir,
     irpass::analysis::verify_if_debug(ir, config);
   }
 
+  // Merge a global's separate read/write GlobalPtrStmts (same address) into one shared, activate=true pointer BEFORE
+  // this first flag_access, so flag_access cannot stamp a read-only (activate=false) copy that the CSE eliminability
+  // rule then refuses to re-merge with the in-loop write. Without it, cache_loop_invariant_global_vars sees a split
+  // read/write and cannot cache conditional/in-if stores -> the -88% solver break-flag bug + the lost duck_in_box
+  // optimization. On main this fell out of whole_kernel_cse running inside every full_simplify fixpoint; per-task CSE
+  // does no pre-offload whole-kernel CSE, so we do this one cheap, pointers-only pass here instead (arithmetic is
+  // already canonical after simplify_I, so a single call is enough; running it in the fixpoint was a +12-22s
+  // compile regression for no extra benefit).
+  irpass::merge_global_ptrs(ir);
+  irpass::analysis::verify_if_debug(ir, config);
+
   irpass::flag_access(ir);
   irpass::analysis::verify_if_debug(ir, config);
 
@@ -141,6 +152,18 @@ void compile_to_offloads(IRNode *ir,
   irpass::analysis::verify_if_debug(ir, config);
 
   dump_ir("after_offload");
+
+  // Full per-task CSE now, before flag_access #2 splits a global's read/write pointers by access flag and before
+  // simplify_III's LICM hoists the read pointer out of the loop. This restores the pointer-unification that main
+  // 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.
+  // Gated on opt_level like all other CSE (per_task_cse / upstream whole_kernel_cse): at opt_level 0 there is no CSE
+  // to require pointer unification, matching upstream behaviour.
+  if (config.opt_level > 0) {
+    irpass::cse_offloaded_tasks(ir);
+  }
+
   // NOTE: There was an additional CFG pass here, removed in
   // https://github.com/taichi-dev/taichi/pull/8691
   irpass::flag_access(ir);

quadrants/transforms/simplify.cpp

@@ -570,10 +570,15 @@ void full_simplify(IRNode *root, const CompileConfig &config, const FullSimplify
         modified = true;
       if (should_dump)
         dump_step("10_die", iteration);
-      if (config.opt_level > 0 && whole_kernel_cse(root))
+      if (config.opt_level > 0 && per_task_cse(root))
         modified = true;
       if (should_dump)
-        dump_step("11_whole_kernel_cse", iteration);
+        dump_step("11_per_task_cse", iteration);
+      // NOTE: the pre-offload same-address pointer merge (merge_global_ptrs) used to run here, in the fixpoint. It 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 was a +12-22s compile regression (measured on
+      // franka/duck/box_pyramid) for zero extra benefit, so it is now a single call in compile_to_offloads.
       // Don't do this time-consuming optimization pass again if the IR is
       // not modified.
       if (config.opt_level > 0 && first_iteration && config.cfg_optimization &&