diff --git a/docs/source/user_guide/optimization_passes.md b/docs/source/user_guide/optimization_passes.md index c0c90ba318..4481ab79c6 100644 --- a/docs/source/user_guide/optimization_passes.md +++ b/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, 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,9 +46,9 @@ 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. | @@ -65,8 +65,6 @@ A **control-flow graph** is a map of your kernel's basic blocks together with th Building and analyzing the CFG is the most expensive optimization in the pipeline, which is why it runs at most once per simplify stage rather than every round. -**One CFG per offloaded task.** The CFG optimization is built and run separately for each offloaded task, over that task's IR alone - never over the whole `qd.kernel` at once. This is both faster to analyze and safe: because each task is a separate device launch, a value held in a register in one task cannot survive into the next one, so there is never anything to forward across a task boundary anyway. Anything written to global memory is treated as potentially read by a later task, so no store another task might need is dropped. - ## Controlling the passes All of these are fields of `CompileConfig`, so you set them at `qd.init(...)` (or via the matching `QD_` environment variable). See [qd.init options](./init_options.md) for the full list and the environment-variable convention. @@ -79,14 +77,18 @@ 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 -These environment variables dump the IR so you can see the effect of each pass. Files are written to `debug_dump_path` (default `/tmp/ir/`): +These environment variables dump the IR so you can see the effect of each pass. Files are written to the directory set by the `debug_dump_path` option in `qd.init(...)` (default `/tmp/ir/`): - `QD_DUMP_IR=1` - writes an IR snapshot at each major pipeline stage (after lowering, before/after each simplify, after offload). - `QD_DUMP_SIMPLIFY=1` - writes an IR snapshot after every individual pass on every iteration of the simplify loop. Verbose, but it shows exactly which pass changed what. - `QD_DUMP_CFG=1` - writes the control-flow graph itself. (This also forces the CFG pass back onto the whole-kernel path so the complete graph can be dumped.) Setting `qd.init(print_ir=True)` prints the IR to the console at pipeline stages instead of writing files. + +## Under the hood: per-task scoping + +Once the kernel has been split into offloaded tasks, both CSE and the CFG optimization run over **one offloaded task's IR at a time**, never over the whole `qd.kernel` at once. This is both faster to analyze and safe: because each task is a separate device launch, a value held in a register in one task cannot survive into the next one, so there is never anything to deduplicate or forward across a task boundary. Anything written to global memory is treated as potentially read by a later task, so no store another task might need is dropped. diff --git a/quadrants/ir/transforms.h b/quadrants/ir/transforms.h index 30b13b2105..959754c115 100644 --- a/quadrants/ir/transforms.h +++ b/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); diff --git a/quadrants/program/compile_config.h b/quadrants/program/compile_config.h index 92c09a5fd8..85c7248f4f 100644 --- a/quadrants/program/compile_config.h +++ b/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; diff --git a/quadrants/transforms/compile_to_offloads.cpp b/quadrants/transforms/compile_to_offloads.cpp index f1a0918ea7..76c2f390a6 100644 --- a/quadrants/transforms/compile_to_offloads.cpp +++ b/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); diff --git a/quadrants/transforms/simplify.cpp b/quadrants/transforms/simplify.cpp index 60f3953dd3..4e50e01967 100644 --- a/quadrants/transforms/simplify.cpp +++ b/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 && diff --git a/quadrants/transforms/whole_kernel_cse.cpp b/quadrants/transforms/whole_kernel_cse.cpp index c7547d5f8b..8cfd60efa2 100644 --- a/quadrants/transforms/whole_kernel_cse.cpp +++ b/quadrants/transforms/whole_kernel_cse.cpp @@ -2,6 +2,7 @@ #include "quadrants/ir/analysis.h" #include "quadrants/ir/statements.h" #include "quadrants/ir/transforms.h" +#include "quadrants/ir/type_utils.h" #include "quadrants/ir/visitors.h" #include "quadrants/system/profiler.h" @@ -78,14 +79,38 @@ class WholeKernelCSE : public BasicStmtVisitor { std::vector>> scope_inserts_; DelayedIRModifier modifier_; + // When true, only address-computation statements (Global/External/MatrixPtr) are eliminated; all other statements + // are left untouched. Used pre-offload to merge same-address read/write pointers (the cheap, load-bearing part of + // whole-kernel CSE) without doing the expensive whole-kernel compute dedup, which is deferred to per-task CSE. + bool ptrs_only_ = false; + public: using BasicStmtVisitor::visit; - WholeKernelCSE() { + explicit WholeKernelCSE(bool ptrs_only = false) : ptrs_only_(ptrs_only) { allow_undefined_visitor = true; invoke_default_visitor = true; } + static bool is_ptr_stmt(Stmt *stmt) { + return stmt->is() || stmt->is() || stmt->is(); + } + + // What ptrs_only mode is allowed to eliminate: address-computation statements PLUS pure integer (addressing) + // arithmetic. The latter is required because two same-address pointers only prove equal (value_diff_ptr_index -> + // FindDirectValueBaseAndOffset) once their index arithmetic is itself merged to a common base statement. Global + // loads/stores are never eliminable (common_statement_eliminable()==false), so this stays sound; float solver + // compute is left untouched and deferred to per-task CSE. + static bool eligible_in_ptrs_only(Stmt *stmt) { + if (is_ptr_stmt(stmt)) { + return true; + } + if ((Type *)stmt->ret_type == nullptr) { + return false; + } + return is_integral(stmt->ret_type.get_element_type()); + } + bool is_done(Stmt *stmt) { return visited_.find(stmt->instance_id) != visited_.end(); } @@ -99,8 +124,12 @@ class WholeKernelCSE : public BasicStmtVisitor { // Use the dynamic type via `typeid(*stmt)` - `typeid(stmt)` operates on the pointer expression and returns the // `Stmt*` static type for every input, collapsing every statement class into the same hash component. auto hash_type = std::hash{}(std::type_index(typeid(*stmt))); - if (stmt->is() || stmt->is()) { - // special cases in common_statement_eliminable() + if (stmt->is() || stmt->is() || stmt->is()) { + // special cases in common_statement_eliminable(): these bucket by type alone so every same-typed candidate is + // compared via the value-based check below (definitely_same_address / same_value), rather than by operand-pointer + // identity. ExternalPtr needs this too: two accesses to the same ndarray address can have distinct index-compute + // statements (not yet merged), so operand-address hashing would split them into separate buckets and miss the + // merge -- leaving the read/write pointers split for cache_loop_invariant_global_vars (the ndarray break-flag). return hash_type; } auto op = stmt->get_operands(); @@ -160,6 +189,10 @@ class WholeKernelCSE : public BasicStmtVisitor { // container_statement does not need to be CSE-ed if (stmt->is_container_statement()) return; + // Pointers-only mode: skip every non-address / non-integer-addressing statement (leave float compute CSE to + // per-task CSE). + if (ptrs_only_ && !eligible_in_ptrs_only(stmt)) + return; // Generic visitor for all CSE-able statements. std::size_t hash_value = operand_hash(stmt); if (is_done(stmt)) { @@ -218,8 +251,8 @@ class WholeKernelCSE : public BasicStmtVisitor { } // Move common statements at the beginning or the end of both branches - // outside. - if (if_stmt->true_statements && if_stmt->false_statements) { + // outside. Skipped in pointers-only mode: we do not want to relocate arbitrary compute. + if (!ptrs_only_ && if_stmt->true_statements && if_stmt->false_statements) { auto &true_clause = if_stmt->true_statements; auto &false_clause = if_stmt->false_statements; if (irpass::analysis::same_statements(true_clause->statements[0].get(), false_clause->statements[0].get())) { @@ -246,8 +279,8 @@ class WholeKernelCSE : public BasicStmtVisitor { if_stmt->false_statements->accept(this); } - static bool run(IRNode *node) { - WholeKernelCSE eliminator; + static bool run(IRNode *node, bool ptrs_only = false) { + WholeKernelCSE eliminator(ptrs_only); bool modified = false; while (true) { node->accept(&eliminator); @@ -260,11 +293,111 @@ class WholeKernelCSE : public BasicStmtVisitor { } }; +namespace { + +// Collect the top-level offloaded tasks of |root| iff |root| is an already-offloaded kernel body (a Block whose +// statements are all OffloadedStmt). Empty otherwise. +std::vector collect_offloaded_tasks(IRNode *root) { + std::vector tasks; + auto *block = root->cast(); + if (block == nullptr || block->statements.empty()) { + return tasks; + } + for (auto &stmt : block->statements) { + if (!stmt->is()) { + return {}; + } + } + for (auto &stmt : block->statements) { + tasks.push_back(stmt->as()); + } + return tasks; +} + +// Run CSE on a single offloaded task, scoped to that task alone via a throwaway wrapper block. +bool cse_one_task(Block *parent, OffloadedStmt *off, bool ptrs_only = false) { + const int location = parent->locate(off); + QD_ASSERT(location != -1); + Block wrapper; + wrapper.insert(parent->extract(off)); + const bool modified = WholeKernelCSE::run(&wrapper, ptrs_only); + parent->insert(wrapper.extract(off), location); + return modified; +} + +} // namespace + namespace irpass { bool whole_kernel_cse(IRNode *root) { QD_AUTO_PROF; return WholeKernelCSE::run(root); } + +// Cheap whole-kernel merge of same-address pointer statements only (Global/External/MatrixPtr), leaving all compute +// alone. Run pre-offload (before the first flag_access) so a global's read and write pointers become one shared, +// activate=true pointer -- the precondition cache_loop_invariant_global_vars relies on to cache conditional/in-if +// stores. This is the load-bearing part of whole-kernel CSE; the expensive compute dedup stays per-task. +bool merge_global_ptrs(IRNode *root) { + QD_AUTO_PROF; + // Defensive: this is intended to run exactly once PRE-offload (single call in compile_to_offloads, before the first + // flag_access). It must never run post-offload -- per_task_cse handles pointers within each task there, and a + // whole-kernel pointer CSE over an offloaded kernel is both pointless and expensive. Skip if any top-level + // OffloadedStmt is present. + if (auto *block = root->cast()) { + for (auto &stmt : block->statements) { + if (stmt->is()) { + return false; + } + } + } + return WholeKernelCSE::run(root, /*ptrs_only=*/true); +} + +// Full per-task CSE run POST-offload, on the pre-split monolith, right after offload (before flag_access #2 and +// simplify_III's LICM). This is the post-offload analog of the pre-offload merge_global_ptrs: main runs +// whole_kernel_cse inside every full_simplify -- including the post-offload simplify_III that precedes +// cache_loop_invariant_global_vars -- which unifies each global's read and write pointers so the caching pass sees +// one shared, loop-invariant pointer and caches soundly. Per-task CSE (per_task_cse) skips the monolith and runs in +// the codegen workers, which is AFTER cache_loop, so without this pass cache_loop sees split read/write pointers and +// either caches a stale local (miscompile: the solver break flag never updates -> non-terminating loop / iteration +// cap) or must decline to cache (lost optimization, ~12% on contact-heavy GJK scenes). Fields are already unified by +// the pre-offload merge_global_ptrs, but ndarray accesses only become ExternalPtrStmts during offload, so they can +// only be unified here. Full CSE (not pointers-only) is required: an ExternalPtr merges with another only once their +// index-compute statements are themselves merged. Scoped per task (independent tasks are never cross-merged); still +// cheaper than main's cross-task whole_kernel_cse because each task's CSE buckets are smaller. +bool cse_offloaded_tasks(IRNode *root) { + QD_AUTO_PROF; + auto tasks = collect_offloaded_tasks(root); + if (tasks.empty()) { + return false; + } + auto *block = root->as(); + bool modified = false; + for (auto *off : tasks) { + if (cse_one_task(block, off)) { + modified = true; + } + } + return modified; +} + +// Per-offloaded-task CSE, parallelized across the codegen worker pool. The post-offload full_simplify passes +// (offload_to_executable: before_lower_access / simplify_IV / scalarize) run inside compile_task, which is enqueued +// per task to the codegen worker pool. At that point each worker's block holds exactly ONE offloaded task, so CSE +// runs here -> per task, in parallel, inside the simplify fixpoint (full quality). On the pre-split monolith +// (full_simplify in compile_to_offloads, where the block still holds every task) CSE is skipped: it is deferred to +// the parallel per-task pass above. Tasks are optimized independently, so doing all CSE in the workers is +// value-identical to doing it on the monolith -- just parallel. Before offload there are no offloaded tasks, so CSE +// is deferred likewise. +bool per_task_cse(IRNode *root) { + QD_AUTO_PROF; + auto tasks = collect_offloaded_tasks(root); + if (tasks.size() != 1) { + return false; + } + auto *block = root->as(); + return cse_one_task(block, tasks[0]); +} } // namespace irpass } // namespace quadrants::lang diff --git a/tests/python/test_cache_loop_invariant.py b/tests/python/test_cache_loop_invariant.py index 72146cd8e7..19fe2cbf56 100644 --- a/tests/python/test_cache_loop_invariant.py +++ b/tests/python/test_cache_loop_invariant.py @@ -61,6 +61,49 @@ def k(x: AnnotationType, result: AnnotationType): assert result[i] == m, f"result[{i}] = {result[i]}, expected {m}" +@pytest.mark.parametrize("use_ndarray", [False, True]) +@test_utils.test() +def test_conditional_store_to_loop_invariant_global(use_ndarray: bool) -> None: + """Regression: a loop-invariant global written *conditionally* inside an ``if`` must not read stale. + + ``flag[i]`` is invariant w.r.t. the inner ``j`` loop, so the cache_loop_invariant_global_vars pass may hoist its + load out of the loop. It is written conditionally (``if j >= threshold``) inside that loop, and the read must + observe the store. Caching is only sound when the read and write pointers to the same address are the same + statement; otherwise ``flag_access`` stamps the hoisted read ``activate=false``, the caching pass serves the + pre-loop value, and the store is lost -- ``acc`` then sums the stale ``0`` instead of the stored ``1`` (this broke + the rigid solver's convergence break-flag, an ~88% runtime regression / non-terminating loop). + + Whole-kernel CSE unifies those pointers on upstream. Per-task CSE restores the same precondition via the + ``merge_global_ptrs`` pass for the field path (pre-offload ``GlobalPtrStmt``s) and ``cse_offloaded_tasks`` for the + ndarray path (``ExternalPtrStmt``s, which only exist post-offload). Both are exercised here. + """ + n = 4 + m = 8 + threshold = 3 + + AnnotationType = qd.types.ndarray() if use_ndarray else qd.template() + TensorType = qd.ndarray if use_ndarray else qd.field + + @qd.kernel + def k(flag: AnnotationType, result: AnnotationType): + for i in range(n): # offloaded task + flag[i] = 0 + acc = 0 + for j in range(m): # inner loop; flag[i] is loop-invariant here + if j >= threshold: + flag[i] = 1 # conditional in-if store to the loop-invariant global + acc += flag[i] # must observe the store, not a stale cached load + result[i] = acc + + flag = TensorType(dtype=qd.i32, shape=(n,)) + result = TensorType(dtype=qd.i32, shape=(n,)) + + k(flag, result) + expected = m - threshold # flag == 1 for j in [threshold, m) + for i in range(n): + assert result[i] == expected, f"result[{i}] = {result[i]}, expected {expected}" + + @pytest.mark.parametrize("use_ndarray", [False, True]) @test_utils.test() def test_literal_index_load_not_cached_over_aliasing_store(use_ndarray: bool) -> None: