From 41b3d94c81b679ef49f7816cc48b29c1436a77ed Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 4 May 2026 12:48:16 +0200 Subject: [PATCH 1/5] [AutoDiff] Adstack load+store eliminations: EliminateRecomputableAdStackPushes pass + leaf extensions --- quadrants/transforms/auto_diff.cpp | 671 ++++++++++++++++++++++++++++- 1 file changed, 668 insertions(+), 3 deletions(-) diff --git a/quadrants/transforms/auto_diff.cpp b/quadrants/transforms/auto_diff.cpp index c0ac2b6579..df75343517 100644 --- a/quadrants/transforms/auto_diff.cpp +++ b/quadrants/transforms/auto_diff.cpp @@ -52,6 +52,148 @@ class NonLinearOps { BinaryOpType::min, BinaryOpType::max}; }; +// Returns true iff `stmt`'s transitive operand DAG terminates at recomputable leaves via side-effect-free +// interior ops only. Used by `EliminateRecomputableAdStackPushes` and `BackupSSA::generic_visit` to decide +// whether a forward SSA value can be reconstructed in the reverse-pass scope from already-stack-backed allocas +// + kernel-args + constants + loop indices, instead of being spilled to a per-iteration adstack or to +// `BackupSSA::load`'s last-iteration plain alloca. +// +// Recomputable leaves: AdStackLoadTopStmt (re-readable via cloned load), AdStackAllocaStmt (the stack itself, +// shared not cloned), ArgLoadStmt (kernel-arg, immutable within the launch), ConstStmt, LoopIndexStmt (clonable +// to read the reverse-direction loop's index, which matches the forward iteration the reverse is currently +// processing). Side-effect-free interior ops: UnaryOp, BinaryOp, TernaryOp, MatrixPtr, GlobalPtr, ExternalPtr. +// +// `GlobalLoadStmt` and `ExternalPtrAccessStmt`-style reads are intentionally not recomputable: the underlying +// global memory may be mutated mid-kernel by a sibling task, so reading at a different IR position can yield a +// different value. `LocalLoadStmt` similarly aliases mutable allocas; the reverse pass must read its forward +// value through the dedicated spill machinery. +// +// Caller passes a `cache` to share memoization across multiple queries on the same DAG (diamond shapes). +class RecomputableChainAnalyzer { + public: + static bool is_recomputable(Stmt *stmt, std::unordered_map &cache) { + auto it = cache.find(stmt); + if (it != cache.end()) + return it->second; + // Tentatively false to break cycles in pathological IR (real SSA DAGs are acyclic, but the cache also + // serves as a visited set during recursion). + cache[stmt] = false; + bool result = check(stmt, cache); + cache[stmt] = result; + return result; + } + + private: + static bool check(Stmt *stmt, std::unordered_map &cache) { + // Recomputable leaves: ConstStmt, ArgLoadStmt, AdStackLoadTopStmt, AdStackAllocaStmt. LoopIndexStmt is + // intentionally excluded - cloning a LoopIndexStmt copies the reference to the forward RangeForStmt, + // but the cloned consumer lives inside the reverse RangeForStmt (a separate stmt with its own loop + // index), so the cloned read points at undefined state and silently double-accumulates gradients + // (`test_adstack_sum_linear`). A future MakeAdjoint coordination pass that tracks + // forward-RangeFor-to-reverse-RangeFor mapping could lift this restriction. + // + // AdStackLoadTopStmt as a leaf is correct under the dominance + control-flow-consumer + self-load + // guards in `EliminateRecomputableAdStackPushes::run_one_pass`. The reasoning: + // + // - In FORWARD: the eliminated stack S's body push dominates each `AdStackLoadTopStmt(S)` (the + // dominance guard), and the chain leaves' pushes dominate the chain's evaluation point. Each + // stack referenced by a chain leaf has a stable "top" value within one forward iteration body + // (no pops in forward), so substituting load_top(S) with the chain re-evaluates to the same + // value (no iteration shift). + // + // - In REVERSE: `MakeAdjoint` visits forward stmts in reverse order, emitting reverse code at a + // cursor that advances one position per visit. For a forward stmt F at position P_F, the + // emitted reverse code lands at cursor position roughly inversely correlated with P_F. The + // dominance guard ensures every chain-leaf stack T has its body push at a position P_T with + // P_T < P (where P is load_top(S)'s forward position). Therefore T's pop in reverse (emitted + // when MakeAdjoint visits T's body push) lands AFTER the chain consumer's reverse emission + // (the consumer's forward position is P > P_T). At the consumer's reverse cursor, T has not + // been popped yet, so load_top(T) returns T's iter-k-push value - matching what the original + // load_top(S) returned in forward iter k. + // + // The control-flow-consumer guard in `run_one_pass` covers a separate issue: stacks whose load_tops + // are direct operands of IfStmt cond / RangeFor begin/end. `MakeAdjoint::visit(IfStmt)` runs a + // dedicated snap-stack fixup that ONLY triggers when the cond is a bare `AdStackLoadTopStmt` (line + // 2168-2202). Eliminating the cond stack converts the cond into a compound stmt, the snap-stack + // does not trigger, and the reverse cond falls back to BackupSSA's load(op) which is single-slot + // last-iter only - silent gradient corruption on multi-iter loops. + if (stmt->is() || stmt->is() || stmt->is() || + stmt->is()) { + return true; + } + // GlobalLoadStmt as recomputable: the load reads a SNode value via GlobalPtrStmt. The cloned chain in + // the reverse pass re-issues the same load - safe iff the global is not mutated between the forward + // chain evaluation and the cloned re-read. Within a single kernel execution, forward writes complete + // before reverse runs, so the reverse re-read sees the kernel's final post-write state. If a global is + // mutated by the forward and then re-read by the reverse clone, the values can differ. + // + // Most rigid-step kernel chain leaves are reads of input parameters (mass, inertia, joint params, + // morphology) that the kernel does NOT write. For those, the re-read returns the same value. + // Conservative analysis: a future safety check could prove the SNode is read-only in the kernel; until + // then this path relies on the test suite's gradient-correctness asserts to surface any mutated-global + // miscompilation. + bool is_interior = stmt->is() || stmt->is() || stmt->is() || + stmt->is() || stmt->is() || stmt->is() || + stmt->is(); + if (!is_interior) { + return false; + } + auto operands = stmt->get_operands(); + for (auto *op : operands) { + if (op == nullptr) + continue; + if (!is_recomputable(op, cache)) + return false; + } + return true; + } +}; + +// Clones the SSA chain rooted at `src` into the IR, inserting cloned stmts before `insert_point`. Returns the +// cloned root. Per-stmt cache shared across one resolution materializes each SSA value at most once: diamond +// DAGs see two consumers but get one shared clone. `AdStackAllocaStmt` is treated as a leaf and shared (not +// cloned) - the stack itself is a unique storage handle that must not be duplicated. +// +// Pop-ordering safety: cloned `AdStackLoadTopStmt`s read the live top at the cloned position. `MakeAdjoint` +// emits `AdStackPopStmt` for each surviving `AdStackPushStmt`, and the existing reverse-pass scheme places +// the pop AFTER all uses of that stack within the reverse iteration (uses include both the original +// `AdStackLoadTopStmt`s emitted by `ReplaceLocalVarWithStacks` and the consumers' clones). For loop-carried +// allocas the pop fires early to expose the iteration's INPUT primal as the new top, which is exactly the +// value the recomputed chain needs - the existing per-consumer clone path at `BackupSSA::generic_visit` line +// ~2697 relies on this same property and has been correct in production. +class RecomputableChainCloner { + public: + static Stmt *clone_at(Stmt *src, Stmt *insert_point, std::unordered_map &cache) { + auto it = cache.find(src); + if (it != cache.end()) + return it->second; + Stmt *cloned = nullptr; + if (src->is()) { + // The alloca is shared, not cloned: every load reads the same physical stack. + cloned = src; + } else if (src->is() || src->is() || src->is()) { + auto cloned_unique = src->clone(); + cloned = insert_point->insert_before_me(std::move(cloned_unique)); + // For AdStackLoadTopStmt clones, the cloned stmt's `stack` operand still points at the original + // AdStackAllocaStmt - that's the desired sharing. + } else { + // Compound op: clone first, then walk operands and rewire each to a recursive clone. + auto cloned_unique = src->clone(); + cloned = insert_point->insert_before_me(std::move(cloned_unique)); + int n = src->num_operands(); + for (int i = 0; i < n; i++) { + auto *op = src->operand(i); + if (op != nullptr) { + Stmt *new_op = clone_at(op, cloned, cache); + cloned->set_operand(i, new_op); + } + } + } + cache[src] = cloned; + return cloned; + } +}; + class IndependentBlocksJudger : public BasicStmtVisitor { public: using BasicStmtVisitor::visit; @@ -504,7 +646,23 @@ class AdStackAllocaJudger : public BasicStmtVisitor { void visit(LocalStoreStmt *stmt) override { if (stmt->dest == target_alloca_backup_) { load_only_ = false; - if (local_loaded_) { + // Gate the load+store-implies-stack-needed rule on actually being inside a dynamic RangeForStmt at the + // point this evidence accumulates. The rule's purpose is to preserve cross-iteration RAW dependencies + // (`for j: p, q = q, p + q` Fibonacci-style) that BackupSSA's single overwrite-each-iteration alloca + // cannot back. With no enclosing dynamic for-loop the IB body executes once: there is no cross- + // iteration RAW to preserve, and the "load+store" pattern is just an in-block accumulator that the + // reverse pass handles via plain SSA cloning. Promoting such allocas under a static-unrolled loop body + // wastes one AdStack per accumulator (one push per unrolled-iter store + one load_top per unrolled- + // iter load) without any reverse-pass consumer needing per-iter replay - that is the unrolled-overhead + // bug Plan B targets. + // + // `dynamic_for_depth_` is incremented in `visit(RangeForStmt)` and decremented on exit. The judger + // walks the IB tree from the alloca's enclosing block, so depth here reflects exactly the nesting of + // *dynamic* for-loops between the alloca and the current load/store. StructFor / WhileStmt do not + // increment because their bodies still execute per-iter and need the same RAW protection (StructFor + // is the kernel-level offload-loop in some cases, but its body is a per-thread independent block; + // load+store there is the same shape as for a top-level alloca). + if (local_loaded_ && dynamic_for_depth_ > 0) { is_stack_needed_ = true; } } @@ -641,13 +799,16 @@ class AdStackAllocaJudger : public BasicStmtVisitor { return; } + dynamic_for_depth_++; stmt->body->accept(this); + dynamic_for_depth_--; } static bool run(AllocaStmt *target_alloca) { AdStackAllocaJudger judger; judger.target_alloca_ = target_alloca; judger.target_alloca_backup_ = target_alloca; + judger.dynamic_for_depth_ = 0; target_alloca->parent->accept(&judger); return (!judger.load_only_) && judger.is_stack_needed_; } @@ -680,6 +841,11 @@ class AdStackAllocaJudger : public BasicStmtVisitor { bool is_stack_needed_ = false; bool local_loaded_ = false; bool load_only_ = true; + // Nesting depth of dynamic `RangeForStmt` containers between the alloca's enclosing block and the current + // visit cursor. Static-unrolled `qd.static(range(...))` loops are removed by the AST transformer before the + // judger sees the IR, so they do not contribute to depth. The load+store-implies-stack-needed rule fires + // only when this depth is positive; see the rationale in `visit(LocalStoreStmt)`. + int dynamic_for_depth_ = 0; }; class RegulateTensorTypedStatements : public BasicStmtVisitor { @@ -1092,6 +1258,472 @@ class ReplaceLocalVarWithStacks : public BasicStmtVisitor { } }; +// Eliminate AdStackAllocaStmts whose pushed value is recomputable from already-stack-backed allocas, kernel +// args, constants, and loop indices. Runs between `ReplaceLocalVarWithStacks` and `MakeAdjoint`, so the reverse +// pass is generated against the cleaned IR (no spurious AdStackPushStmt / AdStackLoadTopStmt scaffolding for +// values the reverse can reconstruct on the fly via cloned forward DAGs). +// +// Eligibility per AdStackAllocaStmt S in an independent block: +// 1. S is written by exactly one AdStackPushStmt (single-push pattern). Multi-push allocas hold loop-carried +// state where each iteration's push depends on the previous - the reverse pass cannot reconstruct +// iteration k's value from iteration (k-1)'s value, so the stack is genuinely needed. +// 2. The pushed value's transitive operand DAG is recomputable per RecomputableChainAnalyzer (leaves at +// AdStackLoadTop / AdStackAlloca / ArgLoad / Const / LoopIndex; interior side-effect-free ops only). +// 3. S has no AdStackLoadTopStmt with `return_ptr=true` consumers (those return a pointer aliasing the slot +// and a follow-up store would not be modeled by an SSA replacement). +// +// Action on eligibility: replace every `AdStackLoadTopStmt(S)` with the original pushed SSA value (the +// `AdStackPushStmt::v`), then erase the push and the alloca. The pushed SSA value's chain stays in the forward +// IR; downstream consumers in the forward pass now reference it directly as SSA, and `MakeAdjoint` plus +// `BackupSSA` reconstruct the chain on demand in the reverse pass via the same DAG-clone path that +// `BackupSSA::generic_visit` exercises for cross-block SSA references. +// +// Iterates to fixed point: eliminating one stack can newly expose another stack's chain as recomputable +// (S2's pushed value chained through `AdStackLoadTopStmt(S1)`; once S1 is gone, the chain may collapse into +// pure SSA chained through S1's pushed-value chain instead). Each pass eliminates at least one stack or +// terminates. +// +// Cost model: this pass trades forward-pass adstack pushes (memory write + top-pointer bump per push) for +// extra arithmetic in the reverse pass (the cloned DAG re-executes the forward chain). For pure-arithmetic +// chains rooted at a single loop-carried alloca - the dominant shape on Genesis-style rigid-step kernels - +// the win is typically 5-10x: each push is a memory op crossing the L1 boundary on every iteration, while +// the recomputed arithmetic stays in registers and reuses warmed-up sin/cos/exp pipeline state. +class EliminateRecomputableAdStackPushes { + public: + static void run(Block *ib) { + // Iterate to fixed point. Each pass either eliminates at least one stack or terminates. The hard cap + // protects against analysis bugs (e.g. an eligibility check that returns true on the same stmt twice + // in a row); under a correct implementation each pass strictly reduces the AdStackAllocaStmt count, so + // the actual iteration bound is the number of stacks at entry. 1024 is well above any realistic kernel. + for (int i = 0; i < 1024; i++) { + if (!run_one_pass(ib)) + return; + } + } + + private: + // True iff `s` is a literal-zero ConstStmt or a MatrixInitStmt whose every element is a literal-zero + // ConstStmt. Matches the init-zero push pattern that ReplaceLocalVarWithStacks emits right after each + // AdStackAllocaStmt: scalar allocas get a ConstStmt(0), tensor-typed allocas get a MatrixInitStmt of + // ConstStmt(0)s. Both are erased along with the alloca during elimination - they only matter if a load + // could observe them before the body push fires, which the standard PromoteSSA2LocalVar + + // ReplaceLocalVarWithStacks codegen never produces (loads always follow the body push within the loop + // iter that emits them). + static bool is_zero_init_value(Stmt *s) { + if (auto *c = s->cast()) { + return c->val.equal_value(0); + } + if (auto *mi = s->cast()) { + for (auto *elem : mi->values) { + auto *c = elem->cast(); + if (c == nullptr || !c->val.equal_value(0)) + return false; + } + return true; + } + return false; + } + + // Returns true if any stack was eliminated this pass. + static bool run_one_pass(Block *ib) { + // Collect every AdStackAllocaStmt anywhere within the IB. The IB is the root of a contiguous AD scope, so + // a single walk is sufficient. + std::vector stacks; + auto collected = irpass::analysis::gather_statements(ib, [&](Stmt *s) { return s->is(); }); + for (auto *s : collected) { + stacks.push_back(s->as()); + } + + bool modified = false; + std::unordered_map recomputable_cache; + + for (auto *stack : stacks) { + // Re-classify users of `stack` each iteration: a previous elimination on a downstream stack may have + // removed users that previously disqualified `stack`. + std::vector pushes; + std::vector load_tops; + bool disqualified = false; + + auto users = irpass::analysis::gather_statements(ib, [&](Stmt *s) { + if (auto *p = s->cast()) + return p->stack == stack; + if (auto *lt = s->cast()) + return lt->stack == stack; + if (auto *po = s->cast()) + return po->stack == stack; + if (auto *aa = s->cast()) + return aa->stack == stack; + if (auto *la = s->cast()) + return la->stack == stack; + return false; + }); + for (auto *user : users) { + if (auto *p = user->cast()) { + pushes.push_back(p); + } else if (auto *lt = user->cast()) { + if (lt->return_ptr) { + // return_ptr=true returns a pointer into the slot; replacing with an SSA value loses the + // pointer-identity contract. Keep the stack. + disqualified = true; + break; + } + load_tops.push_back(lt); + } else if (user->is() || user->is() || + user->is()) { + // Pre-MakeAdjoint, pop / adj-acc / load-top-adj should not appear yet for this stack. If they do, + // some upstream pass has already touched it - keep as-is to avoid double-rewrites. + disqualified = true; + break; + } + } + if (disqualified || pushes.empty()) { + continue; + } + + // ReplaceLocalVarWithStacks emits a const-zero "init" AdStackPushStmt immediately after the + // AdStackAllocaStmt as the stack's initial value (auto_diff.cpp:867-872 in pre-fix tree). Real + // user-level allocas in a loop body have at most ONE additional "body" push per iteration. Treat + // const-zero pushes whose value is a `ConstStmt` with all-zero `TypedConstant` as inits and require at + // most one non-init "body" push for elimination eligibility. Multi-body-push allocas hold loop-carried + // state (each iteration's push depends on the previous), so the reverse pass cannot reconstruct + // iteration k from iteration k-1 and the stack must stay. + AdStackPushStmt *body_push = nullptr; + std::vector init_pushes; + for (auto *p : pushes) { + if (is_zero_init_value(p->v)) { + init_pushes.push_back(p); + } else { + if (body_push != nullptr) { + disqualified = true; + break; + } + body_push = p; + } + } + if (disqualified || body_push == nullptr) { + continue; + } + + Stmt *pushed_val = body_push->v; + if (!RecomputableChainAnalyzer::is_recomputable(pushed_val, recomputable_cache)) { + continue; + } + // Loop-carried self-reference: when the pushed value's transitive operand DAG includes an + // AdStackLoadTopStmt of the SAME stack we are about to eliminate, the value is `read prev iter from + // stack -> compute -> push next iter`. Rewriting load_top($S) -> pushed_value_SSA leaves a self-cycle + // in the SSA graph (`acc_new = acc_new + sin(a)`) which both corrupts the IR and loses the iteration + // recurrence the stack was carrying. The recomputable-chain analyzer happily accepts such chains + // because AdStackLoadTopStmt is a recomputable leaf in general; the additional check below specifically + // disqualifies self-loaded stacks. + // + // The classic shape this protects: `acc = 0.0; for j in range(N): acc += sin(...)`. After + // PromoteSSA2LocalVar + ReplaceLocalVarWithStacks the IR has `$acc = stack_alloc; init push; + // for j: $tmp = stack_load_top $acc; $new = $tmp + sin(...); push $acc, val=$new`. The chain `$tmp + + // sin(...)` is recomputable per the leaf rules (LoadTop is a leaf), but rewriting load_top($acc) to + // ($tmp + sin(...)) makes the SSA graph self-reference $new and the reverse pass loses every + // iteration's accumulator state. + // Read-before-write protection: substituting `AdStackLoadTopStmt(S)` with the body push's value SSA + // chain only preserves semantics when the body push DOMINATES every load_top in the forward IR - + // that is, every load_top reads what the body push just wrote in the same iteration. The + // PromoteSSA2LocalVar + ReplaceLocalVarWithStacks pipeline emits exactly this shape for required-def + // spills: `alloca; push val=def; load_top` in document order, with the push immediately following + // the def and loads occurring later in the same iteration. + // + // Loop-carried recurrences violate the dominance rule. Take the canonical Fibonacci shape + // `p, q = q, p + q` lowered to IR: + // + // $tmp_p = load_top($p) # reads PREVIOUS iter's push into $p + // $tmp_q = load_top($p) + load_top($q) + // push $p, val=$tmp_p # writes NEXT iter's value into $p + // push $q, val=$tmp_q + // + // Here load_top($p) reads BEFORE the body push of $p within the same iter, so it observes iter + // (k-1)'s value, not iter k's. Substituting `load_top($p) -> $tmp_p`'s SSA chain (which reads + // load_top($q) at iter k) gives iter k's q-value instead of iter (k-1)'s p-value, off by one + // iteration and producing zero gradients on the Fibonacci-style regression test pinned by + // `test_ad_fibonacci_index`. + // + // The check below asserts that for every load_top in `load_tops`, the body push is its predecessor + // within the same containing block (or, equivalently, load_top comes AFTER the body push in + // document order at the body push's block level). Loads that live inside nested control flow + // (`IfStmt`, `RangeForStmt`) under the body push's block are also fine because the body push + // dominates them. We approximate dominance with the lexical "push's block contains load_top's + // ancestor block AND push position < load_top's ancestor's position in that block" check. + auto block_position = [](Stmt *s) -> int { + Block *b = s->parent; + if (b == nullptr) + return -1; + for (size_t i = 0; i < b->statements.size(); i++) { + if (b->statements[i].get() == s) + return static_cast(i); + } + return -1; + }; + Block *push_block = body_push->parent; + int push_pos = block_position(body_push); + bool dominates_all_loads = true; + for (auto *lt : load_tops) { + Stmt *cursor = lt; + Block *cursor_block = cursor->parent; + while (cursor_block != nullptr && cursor_block != push_block) { + cursor = cursor_block->parent_stmt(); + if (cursor == nullptr) + break; + cursor_block = cursor->parent; + } + if (cursor_block != push_block) { + // load_top is outside the push's block scope: cannot establish dominance. + dominates_all_loads = false; + break; + } + int cursor_pos = block_position(cursor); + if (cursor_pos <= push_pos) { + // load_top precedes the body push (or its enclosing container does): it reads the previous + // iteration's value, not iter k's. Substitution would shift iterations by one. + dominates_all_loads = false; + break; + } + } + if (!dominates_all_loads) { + continue; + } + + // Reverse-position correctness for chain leaves pushed in the same block as S. + // + // After elimination, every reverse stmt that consumed `load_top(S)` instead consumes the chain + // (cloned by `BackupSSA::generic_visit` at the consumer's reverse position). The cloned chain reads + // `load_top(T)` for each chain leaf T, where the read happens at the consumer's reverse cursor. + // + // `MakeAdjoint` visits forward stmts in REVERSE order, emitting at a cursor that advances per visit. + // For a forward stmt at position P, its reverse emissions land "later" in reverse block iff P is + // smaller. T's pop is emitted when `MakeAdjoint` visits T's body push at position P_T. The cloned + // chain at the consumer's reverse position reads `load_top(T)` POST-pop iff T's pop has fired by then, + // i.e. iff visit(P_T) precedes visit(C_pos), i.e. iff P_T > C_pos for every consumer C of S's + // load_tops. + // + // Forward chain evaluates at S's push position P_S, before any of T's same-block body pushes (since + // we already required P_T > P_S via the dominance check above for S's loads). So forward chain reads + // T's pre-iter-k-push value. Reverse clone post-pop also returns T's pre-iter-k-push value (= iter k + // INPUT for loop-carried T). Match. + // + // Without this check, the canonical Fibonacci-style shape (`p, q = q, p + q; b[q] += a[q]`) where S + // stores `p + q` and chain leaves p_stack / q_stack are pushed AFTER S but BEFORE the GlobalPtr + // index consumer would silently corrupt gradients: the reverse clone at the GlobalPtr's adjoint + // emission reads p_stack and q_stack PRE-pop (iter k POST-push values), summing to a different value + // than the forward chain at P_S used. + // + // Chain leaves T whose stacks have NO body push in S's containing block are stable within the iter + // (their tops do not change during the iter), so neither forward nor reverse evaluation order shifts + // their values. Excluded from this check. + auto find_consumers_of_load_tops = [&](std::vector &out_positions) { + std::function walk = [&](Block *b, int container_pos_in_push_block) { + for (size_t i = 0; i < b->statements.size(); i++) { + Stmt *s = b->statements[i].get(); + int effective_pos = (b == push_block) ? static_cast(i) : container_pos_in_push_block; + for (auto *op : s->get_operands()) { + if (op == nullptr) + continue; + for (auto *lt : load_tops) { + if (op == lt) { + out_positions.push_back(effective_pos); + break; + } + } + } + if (auto *if_s = s->cast()) { + if (if_s->true_statements) + walk(if_s->true_statements.get(), effective_pos); + if (if_s->false_statements) + walk(if_s->false_statements.get(), effective_pos); + } else if (auto *rf = s->cast()) { + if (rf->body) + walk(rf->body.get(), effective_pos); + } else if (auto *sf = s->cast()) { + if (sf->body) + walk(sf->body.get(), effective_pos); + } + } + }; + walk(push_block, -1); + }; + std::vector consumer_positions; + find_consumers_of_load_tops(consumer_positions); + int max_consumer_pos = -1; + for (int p : consumer_positions) { + if (p > max_consumer_pos) + max_consumer_pos = p; + } + + // Collect all distinct AdStackAllocaStmt referenced via AdStackLoadTopStmt leaves in pushed_val's chain. + std::unordered_set chain_leaf_stacks; + std::unordered_set walked_for_leaves; + std::function collect_leaves = [&](Stmt *s) { + if (walked_for_leaves.count(s)) + return; + walked_for_leaves.insert(s); + if (auto *lt = s->cast()) { + if (auto *as = lt->stack->cast()) + chain_leaf_stacks.insert(as); + return; + } + if (s->is() || s->is() || s->is()) + return; + for (auto *op : s->get_operands()) { + if (op != nullptr) + collect_leaves(op); + } + }; + collect_leaves(pushed_val); + + bool reverse_safe = true; + for (auto *T : chain_leaf_stacks) { + // Find T's body pushes (non-init) in push_block. Only same-block pushes can interfere: pushes in + // outer blocks happen once per outer iteration, so their tops are stable across the inner iter. + for (size_t i = 0; i < push_block->statements.size(); i++) { + Stmt *s = push_block->statements[i].get(); + if (auto *p = s->cast()) { + if (p->stack == T && !is_zero_init_value(p->v)) { + if (static_cast(i) <= max_consumer_pos) { + reverse_safe = false; + break; + } + } + } + } + if (!reverse_safe) + break; + } + if (!reverse_safe) { + continue; + } + + bool is_self_loaded = false; + std::unordered_set visited; + std::function walk = [&](Stmt *s) { + if (is_self_loaded || visited.count(s)) + return; + visited.insert(s); + if (auto *lt = s->cast()) { + if (lt->stack == stack) { + is_self_loaded = true; + return; + } + // load_top of a different stack: the operand chain stops here at this leaf for the self-check. + return; + } + if (s->is() || s->is() || s->is()) { + return; // leaf, not a self-load + } + for (auto *op : s->get_operands()) { + if (op != nullptr) + walk(op); + } + }; + walk(pushed_val); + if (is_self_loaded) { + continue; + } + + // Control-flow-cond consumers: `MakeAdjoint::visit(IfStmt)` (auto_diff.cpp:2131-) detects bare + // `AdStackLoadTopStmt` conds and emits a dedicated 1-push-per-execution snap-stack so the reverse + // IfStmt's cond reads the forward-time cond value rather than re-reading the stack at the consumer + // position (which sits BEFORE the per-iter pops in the reverse scope and would see a stale top). + // Compound conds rely on `BackupSSA`'s `load(op)` spill, which inserts a `LocalStore` immediately + // after the forward cond - per-iter, single alloca, last-write-wins. + // + // If we eliminate a stack whose load_top is the bare cond of an IfStmt (or feeds the begin/end of an + // inner RangeForStmt, by the same argument applied to MakeAdjoint::visit(RangeForStmt)), the rewrite + // turns the cond / loop-bound into an inlined recomputed SSA chain. The snap-stack guard at line + // 2157 stops firing (the cond is no longer a bare AdStackLoadTopStmt), and `BackupSSA`'s + // recomputable-clone fallback positions a fresh `AdStackLoadTopStmt` of an enclosing loop-carried + // stack at the consumer's IR location - which sits BEFORE the per-iter pops, returning the post- + // iteration value of the loop-carried stack instead of the iteration-k value the cond was originally + // computed against. Net effect: silent gradient corruption in any if/loop nested inside a dynamic + // for-loop with a loop-carried alloca. + // + // Keep stacks whose load_tops have such consumers. Note that this is structurally narrower than + // "don't eliminate stacks with cross-IR control-flow uses": `cmp_lt(load_top, threshold) -> IfStmt` + // is fine when the cond is THIS stack's pushed value (the stack itself is the cond stack), because + // that's the snap-stack-eligible shape. The consumer check below trips only on a load_top of THIS + // stack feeding directly into a control-flow-shaping operand of another stmt. + // `irpass::analysis::gather_statements` walks via `BasicStmtVisitor`, whose visit overrides for + // IfStmt / RangeForStmt / StructForStmt do not invoke the per-stmt test predicate on the container + // itself - only on stmts inside their bodies. Walk container stmts manually here. + bool has_control_flow_consumer = false; + std::function walk_block = [&](Block *b) { + if (has_control_flow_consumer) + return; + for (auto &owned : b->statements) { + Stmt *s = owned.get(); + if (auto *if_s = s->cast()) { + for (auto *lt : load_tops) { + if (if_s->cond == lt) { + has_control_flow_consumer = true; + return; + } + } + if (if_s->true_statements) + walk_block(if_s->true_statements.get()); + if (has_control_flow_consumer) + return; + if (if_s->false_statements) + walk_block(if_s->false_statements.get()); + } else if (auto *rf = s->cast()) { + for (auto *lt : load_tops) { + if (rf->begin == lt || rf->end == lt) { + has_control_flow_consumer = true; + return; + } + } + if (rf->body) + walk_block(rf->body.get()); + } else if (auto *sf = s->cast()) { + if (sf->body) + walk_block(sf->body.get()); + } else if (auto *off = s->cast()) { + if (off->body) + walk_block(off->body.get()); + if (off->tls_prologue) + walk_block(off->tls_prologue.get()); + if (off->bls_prologue) + walk_block(off->bls_prologue.get()); + if (off->bls_epilogue) + walk_block(off->bls_epilogue.get()); + if (off->tls_epilogue) + walk_block(off->tls_epilogue.get()); + } + } + }; + walk_block(ib); + if (has_control_flow_consumer) { + continue; + } + + // Eligible: rewrite each load_top to use the pushed value directly. The pushed value lives in the + // forward IR and dominates each load_top by SSA construction (the load_top reads what the push wrote; + // both are inside the same loop body in the dynamic-loop case, with the push preceding the load_top). + for (auto *lt : load_tops) { + irpass::replace_all_usages_with(ib, lt, pushed_val); + lt->parent->erase(lt); + } + // Erase init-zero pushes (they only matter if a load could observe them, but the rewriting above just + // routed every load to the body-pushed SSA value), the body push, and the alloca itself. + for (auto *p : init_pushes) { + p->parent->erase(p); + } + body_push->parent->erase(body_push); + stack->parent->erase(stack); + + // Invalidate the cache: the eliminated stack might have been referenced by other recomputable chains + // we evaluated earlier in this pass, but those evaluations only matter for stacks we eliminate THIS + // pass; re-running from scratch on the next iteration recomputes them. + recomputable_cache.clear(); + modified = true; + } + return modified; + } +}; + class ReverseOuterLoops : public BasicStmtVisitor { using BasicStmtVisitor::visit; @@ -2715,13 +3347,36 @@ class BackupSSA : public BasicStmtVisitor { } else if (op->is()) { stmt->set_operand(i, stmt->insert_before_me(op->clone())); } else { - auto alloca = load(op); - stmt->set_operand(i, stmt->insert_before_me(Stmt::make(alloca))); + // Recomputable-chain fallback before the last-iter `load(op)` spill: when the cross-block SSA op + // is rooted in a DAG of side-effect-free arithmetic over already-stack-backed allocas, kernel-args, + // constants, and loop indices, clone the chain into the reverse scope at the consumer site. The + // cloned chain reads stack tops live (matching `MakeAdjoint`'s pop ordering, which fires pops + // AFTER all uses of a stack within one reverse iter) and reads kernel-args / constants / loop + // indices via direct clones - exactly the path that was already correct for AdStackLoadTopStmt and + // ArgLoadStmt operands. + // + // The pre-existing `load(op)` fallback below remains correct for genuinely non-recomputable + // cross-block ops (a forward `GlobalLoadStmt` of a needs_grad SNode whose value the reverse must + // read at last-write rather than recompute, for instance) and for shapes outside one of our + // independent blocks. Adding the recomputable path above the fallback strictly subsets the + // previous behaviour: a chain that fails the predicate falls through to the old `load(op)` line. + if (RecomputableChainAnalyzer::is_recomputable(op, recomputable_cache_)) { + std::unordered_map clone_cache; + Stmt *cloned = RecomputableChainCloner::clone_at(op, stmt, clone_cache); + stmt->set_operand(i, cloned); + } else { + auto alloca = load(op); + stmt->set_operand(i, stmt->insert_before_me(Stmt::make(alloca))); + } } } } } + // Memoization cache for `RecomputableChainAnalyzer::is_recomputable` queries within one BackupSSA run. + // Re-used across all generic_visit calls; invariant during the visit because forward IR is read-only here. + std::unordered_map recomputable_cache_; + void visit(Stmt *stmt) override { generic_visit(stmt); } @@ -3110,6 +3765,16 @@ void auto_diff(IRNode *root, const CompileConfig &config, AutodiffMode autodiff_ ib->accept(&replace); type_check(root, config); + // Drop AdStackAllocas whose pushed value is recomputable from already-stack-backed allocas + args + + // const + loop-index. Trades forward-pass adstack memory traffic (one push + one load per iter per + // intermediate spill) for cloned arithmetic in the reverse scope, which `BackupSSA::generic_visit` + // generates on demand via the same RecomputableChainCloner path. Must run after + // `ReplaceLocalVarWithStacks` (so the analyzer sees the AdStackAlloca shape, not the alloca-with-store + // shape PromoteSSA2LocalVar emits) and before `MakeAdjoint` (so the reverse pass is generated against + // the cleaned forward IR with no spurious push/load scaffolding). + EliminateRecomputableAdStackPushes::run(ib); + type_check(root, config); + MakeAdjoint::run(ib); type_check(root, config); BackupSSA::run(ib); From 161505bf10b3a729ce409d2a0ac48b7016aec37a Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 4 May 2026 17:54:16 +0200 Subject: [PATCH 2/5] [AutoDiff] Reflow EliminateRecomputableAdStackPushes / MakeAdjoint::visit(IfStmt) comments to 120 cols --- quadrants/transforms/auto_diff.cpp | 72 +++++++++++++++++------------- 1 file changed, 42 insertions(+), 30 deletions(-) diff --git a/quadrants/transforms/auto_diff.cpp b/quadrants/transforms/auto_diff.cpp index df75343517..3beb7930fd 100644 --- a/quadrants/transforms/auto_diff.cpp +++ b/quadrants/transforms/auto_diff.cpp @@ -1624,28 +1624,27 @@ class EliminateRecomputableAdStackPushes { continue; } - // Control-flow-cond consumers: `MakeAdjoint::visit(IfStmt)` (auto_diff.cpp:2131-) detects bare - // `AdStackLoadTopStmt` conds and emits a dedicated 1-push-per-execution snap-stack so the reverse - // IfStmt's cond reads the forward-time cond value rather than re-reading the stack at the consumer - // position (which sits BEFORE the per-iter pops in the reverse scope and would see a stale top). - // Compound conds rely on `BackupSSA`'s `load(op)` spill, which inserts a `LocalStore` immediately - // after the forward cond - per-iter, single alloca, last-write-wins. + // Control-flow-cond consumers: by the time the IR reaches this pass, every IfStmt cond and RangeFor begin/end + // is either a bare `AdStackLoadTopStmt` (loop-carried local promoted via `PromoteSSA2LocalVar` -> + // `AdStackAllocaJudger::visit(IfStmt|RangeForStmt)` -> `ReplaceLocalVarWithStacks`) or a stmt outside any + // adstack-bearing alloca (no load_top in the chain at all). `MakeAdjoint::visit(IfStmt)` (auto_diff.cpp around + // line 2452) caps the bare-load_top case with a 1-push-per-execution snap-stack when the if body itself pushes + // to the cond's backing stack so the reverse cond reads the forward-time value, not a stack top mutated by the + // body. // - // If we eliminate a stack whose load_top is the bare cond of an IfStmt (or feeds the begin/end of an - // inner RangeForStmt, by the same argument applied to MakeAdjoint::visit(RangeForStmt)), the rewrite - // turns the cond / loop-bound into an inlined recomputed SSA chain. The snap-stack guard at line - // 2157 stops firing (the cond is no longer a bare AdStackLoadTopStmt), and `BackupSSA`'s - // recomputable-clone fallback positions a fresh `AdStackLoadTopStmt` of an enclosing loop-carried - // stack at the consumer's IR location - which sits BEFORE the per-iter pops, returning the post- - // iteration value of the loop-carried stack instead of the iteration-k value the cond was originally - // computed against. Net effect: silent gradient corruption in any if/loop nested inside a dynamic - // for-loop with a loop-carried alloca. + // If we eliminate a stack whose load_top IS the bare cond of an IfStmt / inner RangeFor, the rewrite turns the + // cond / loop-bound into an inlined recomputed SSA chain. The snap-stack guard at the `AdStackLoadTopStmt`-cond + // check in `MakeAdjoint::visit(IfStmt)` stops firing (the cond is no longer bare), and `BackupSSA`'s clone path + // positions a fresh `AdStackLoadTopStmt` of an enclosing loop-carried stack at the consumer's IR location - + // which sits BEFORE the per-iter pops, returning the post-iteration value instead of the iteration-k value the + // cond was originally computed against. Net effect: silent gradient corruption in any if/loop nested inside a + // dynamic for-loop with a loop-carried alloca. // - // Keep stacks whose load_tops have such consumers. Note that this is structurally narrower than - // "don't eliminate stacks with cross-IR control-flow uses": `cmp_lt(load_top, threshold) -> IfStmt` - // is fine when the cond is THIS stack's pushed value (the stack itself is the cond stack), because - // that's the snap-stack-eligible shape. The consumer check below trips only on a load_top of THIS - // stack feeding directly into a control-flow-shaping operand of another stmt. + // Keep stacks whose load_tops have such consumers. The consumer check below trips only on a load_top of THIS + // stack feeding DIRECTLY into a control-flow-shaping operand of another stmt; that is also the exhaustive set + // of unsafe-elim shapes here, because compound-cond cases (arithmetic over a load_top feeding a cond) reach + // this pass with the cond rewritten into a separate adstack-promoted value via the alloca-promotion pipeline + // above and so look like the bare case from this guard's POV. // `irpass::analysis::gather_statements` walks via `BasicStmtVisitor`, whose visit overrides for // IfStmt / RangeForStmt / StructForStmt do not invoke the per-stmt test predicate on the container // itself - only on stmts inside their bodies. Walk container stmts manually here. @@ -2465,16 +2464,29 @@ class MakeAdjoint : public ADTransform { Stmt *reverse_cond = if_stmt->cond; AdStackAllocaStmt *snap_stack_ptr = nullptr; // Narrow guard: only the bare `AdStackLoadTopStmt` shape needs the explicit snapshot below. A compound cond (e.g. - // `BinaryOp(cmp_lt, AdStackLoadTopStmt(x_stack), threshold)` from `if x < threshold` when `x` has been promoted to - // an adstack by `ReplaceLocalVarWithStacks`) is already handled correctly by `BackupSSA::generic_visit`'s - // else-branch (`load(op)` path at the end of that function): it spills the forward-time value of the whole cond - // stmt - including the embedded `AdStackLoadTopStmt` read - into a dedicated alloca via a `LocalStoreStmt` emitted - // immediately after the forward cond, then the reverse IfStmt's operand becomes a `LocalLoadStmt` of that alloca. - // That captures the forward-time cond exactly. The bare-`AdStackLoadTopStmt` case is special because - // `generic_visit` takes a different branch for that shape (clone-branch): it emits a fresh `AdStackLoadTopStmt` at - // reverse time, which re-reads the stack top AFTER the body's pushes and therefore sees the wrong cond value. The - // snap-stack below is the dedicated fix for that single shape - no recursive walk needed for compound conds - // because the spill branch already covers them. + // `cmp_lt(load_top(x_stack) + 0.1, threshold)` from `if x + 0.1 < threshold` when `x` has been promoted to an + // adstack by `ReplaceLocalVarWithStacks`) reaches this visitor as a BARE `AdStackLoadTopStmt` cond anyway - the + // cmp / arithmetic value goes through `PromoteSSA2LocalVar`'s required-defs set (the IfStmt cond path adds the + // cond's value-producing op), then `AdStackAllocaJudger::visit(IfStmt)` (around line 763) marks its alloca + // stack-needed because it feeds the cond, then `ReplaceLocalVarWithStacks` promotes the alloca to an adstack. By + // the time control reaches here, `if_stmt->cond` is `AdStackLoadTopStmt` of that snap-promoted adstack and the + // body of the if does NOT push to that stack (the cmp value is pushed once just before the IfStmt, never inside), + // so the `body_pushes_to_stack` guard below is false and we correctly skip the additional snap-stack. Per-iter + // cond values are preserved by the alloca-promotion pipeline. + // + // The bare-`AdStackLoadTopStmt` case the snap-stack below handles is the OTHER shape: a load_top whose backing + // stack IS pushed to inside the if body (e.g. short-circuit lowering of `&&` pushes the rhs onto the same stack + // that holds the cond). Without the snap-stack, `BackupSSA::generic_visit`'s clone-branch for `AdStackLoadTopStmt` + // emits a fresh `AdStackLoadTopStmt` at the reverse cursor, which then reads the post-body-push top and sees the + // wrong cond value. The snap-stack here decouples the cond value from the body's pushes by capturing it once just + // before the IfStmt and reading it back at the matching reverse cursor. + // + // Earlier comments here claimed that compound conds rely on `BackupSSA::generic_visit`'s `load(op)` else-branch + // ("single alloca, last-write-wins" spill) for correctness. That description was incomplete: the alloca-promotion- + // to-adstack pipeline catches compound conds before they reach BackupSSA, so the spill is a FALLBACK for shapes + // the alloca-promotion missed, not the load-bearing path for compound conds in general. Verified empirically on + // 2026-05-04 with a loop-carried local v + compound cond `v + 0.1 > threshold` + nonlinear use of v in the if + // body: gradients match analytic on origin/main and on C3. if (if_stmt->cond->is()) { auto *cond_stack = if_stmt->cond->as()->stack->as(); if (body_pushes_to_stack(if_stmt, cond_stack)) { From 2f4575329bbbdb5de2b48ffd58df5be15c3a5d3e Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 4 May 2026 21:47:29 +0200 Subject: [PATCH 3/5] [AutoDiff] Split auto_diff.cpp into stage-grouped files under quadrants/transforms/auto_diff/ --- cmake/QuadrantsCore.cmake | 1 + quadrants/transforms/auto_diff.cpp | 3905 ----------------- quadrants/transforms/auto_diff/auto_diff.cpp | 84 + .../transforms/auto_diff/auto_diff_common.h | 439 ++ .../eliminate_recomputable_pushes.cpp | 479 ++ .../auto_diff/forward_state_spill.cpp | 667 +++ .../auto_diff/forward_state_spill.h | 30 + quadrants/transforms/auto_diff/ir_shaping.cpp | 646 +++ quadrants/transforms/auto_diff/ir_shaping.h | 28 + .../transforms/auto_diff/make_adjoint.cpp | 857 ++++ quadrants/transforms/auto_diff/make_adjoint.h | 12 + quadrants/transforms/auto_diff/make_dual.cpp | 345 ++ quadrants/transforms/auto_diff/make_dual.h | 11 + .../auto_diff/post_adjoint_cleanup.cpp | 324 ++ .../auto_diff/post_adjoint_cleanup.h | 21 + quadrants/transforms/auto_diff/validation.cpp | 264 ++ quadrants/transforms/auto_diff/validation.h | 23 + 17 files changed, 4231 insertions(+), 3905 deletions(-) delete mode 100644 quadrants/transforms/auto_diff.cpp create mode 100644 quadrants/transforms/auto_diff/auto_diff.cpp create mode 100644 quadrants/transforms/auto_diff/auto_diff_common.h create mode 100644 quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp create mode 100644 quadrants/transforms/auto_diff/forward_state_spill.cpp create mode 100644 quadrants/transforms/auto_diff/forward_state_spill.h create mode 100644 quadrants/transforms/auto_diff/ir_shaping.cpp create mode 100644 quadrants/transforms/auto_diff/ir_shaping.h create mode 100644 quadrants/transforms/auto_diff/make_adjoint.cpp create mode 100644 quadrants/transforms/auto_diff/make_adjoint.h create mode 100644 quadrants/transforms/auto_diff/make_dual.cpp create mode 100644 quadrants/transforms/auto_diff/make_dual.h create mode 100644 quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp create mode 100644 quadrants/transforms/auto_diff/post_adjoint_cleanup.h create mode 100644 quadrants/transforms/auto_diff/validation.cpp create mode 100644 quadrants/transforms/auto_diff/validation.h diff --git a/cmake/QuadrantsCore.cmake b/cmake/QuadrantsCore.cmake index fdd1b166c6..4e1bd4c606 100644 --- a/cmake/QuadrantsCore.cmake +++ b/cmake/QuadrantsCore.cmake @@ -69,6 +69,7 @@ file(GLOB QUADRANTS_CORE_SOURCE "quadrants/struct/*" "quadrants/system/*" "quadrants/transforms/*" + "quadrants/transforms/auto_diff/*" "quadrants/platform/cuda/*" "quadrants/platform/amdgpu/*" "quadrants/platform/mac/*" "quadrants/platform/windows/*" "quadrants/codegen/*.cpp" "quadrants/codegen/*.h" diff --git a/quadrants/transforms/auto_diff.cpp b/quadrants/transforms/auto_diff.cpp deleted file mode 100644 index 3beb7930fd..0000000000 --- a/quadrants/transforms/auto_diff.cpp +++ /dev/null @@ -1,3905 +0,0 @@ -#include "quadrants/ir/analysis.h" -#include "quadrants/ir/ir.h" -#include "quadrants/ir/statements.h" -#include "quadrants/ir/transforms.h" -#include "quadrants/ir/visitors.h" -#include "quadrants/transforms/utils.h" - -#include -#include -#include -#include -#include - -namespace quadrants::lang { - -template -Stmt *insert_const(const DataType &dtype, Stmt *stmt, const T &value, bool insert_before_me = false) { - auto type = dtype.ptr_removed(); - Stmt *zero = nullptr; - if (insert_before_me) - zero = stmt->insert_before_me(Stmt::make(TypedConstant(type.get_element_type(), value))); - else - zero = stmt->insert_after_me(Stmt::make(TypedConstant(type.get_element_type(), value))); - - if (type->is()) { - auto t_dtype = type->as(); - std::vector values(t_dtype->get_num_elements(), zero); - if (insert_before_me) { - zero = zero->insert_before_me(Stmt::make(values)); - } else { - zero = zero->insert_after_me(Stmt::make(values)); - } - zero->ret_type = type; - } - return zero; -} - -class IndependentBlockMetaData { - public: - bool is_ib = true; - bool is_smallest_ib = true; -}; - -class NonLinearOps { - public: - inline static const std::set ternary_collections{TernaryOpType::select}; - inline static const std::set unary_collections{ - UnaryOpType::abs, UnaryOpType::sin, UnaryOpType::cos, UnaryOpType::tan, UnaryOpType::tanh, UnaryOpType::asin, - UnaryOpType::acos, UnaryOpType::exp, UnaryOpType::log, UnaryOpType::sqrt, UnaryOpType::rsqrt}; - inline static const std::set binary_collections{BinaryOpType::mul, BinaryOpType::div, - BinaryOpType::atan2, BinaryOpType::pow, - BinaryOpType::min, BinaryOpType::max}; -}; - -// Returns true iff `stmt`'s transitive operand DAG terminates at recomputable leaves via side-effect-free -// interior ops only. Used by `EliminateRecomputableAdStackPushes` and `BackupSSA::generic_visit` to decide -// whether a forward SSA value can be reconstructed in the reverse-pass scope from already-stack-backed allocas -// + kernel-args + constants + loop indices, instead of being spilled to a per-iteration adstack or to -// `BackupSSA::load`'s last-iteration plain alloca. -// -// Recomputable leaves: AdStackLoadTopStmt (re-readable via cloned load), AdStackAllocaStmt (the stack itself, -// shared not cloned), ArgLoadStmt (kernel-arg, immutable within the launch), ConstStmt, LoopIndexStmt (clonable -// to read the reverse-direction loop's index, which matches the forward iteration the reverse is currently -// processing). Side-effect-free interior ops: UnaryOp, BinaryOp, TernaryOp, MatrixPtr, GlobalPtr, ExternalPtr. -// -// `GlobalLoadStmt` and `ExternalPtrAccessStmt`-style reads are intentionally not recomputable: the underlying -// global memory may be mutated mid-kernel by a sibling task, so reading at a different IR position can yield a -// different value. `LocalLoadStmt` similarly aliases mutable allocas; the reverse pass must read its forward -// value through the dedicated spill machinery. -// -// Caller passes a `cache` to share memoization across multiple queries on the same DAG (diamond shapes). -class RecomputableChainAnalyzer { - public: - static bool is_recomputable(Stmt *stmt, std::unordered_map &cache) { - auto it = cache.find(stmt); - if (it != cache.end()) - return it->second; - // Tentatively false to break cycles in pathological IR (real SSA DAGs are acyclic, but the cache also - // serves as a visited set during recursion). - cache[stmt] = false; - bool result = check(stmt, cache); - cache[stmt] = result; - return result; - } - - private: - static bool check(Stmt *stmt, std::unordered_map &cache) { - // Recomputable leaves: ConstStmt, ArgLoadStmt, AdStackLoadTopStmt, AdStackAllocaStmt. LoopIndexStmt is - // intentionally excluded - cloning a LoopIndexStmt copies the reference to the forward RangeForStmt, - // but the cloned consumer lives inside the reverse RangeForStmt (a separate stmt with its own loop - // index), so the cloned read points at undefined state and silently double-accumulates gradients - // (`test_adstack_sum_linear`). A future MakeAdjoint coordination pass that tracks - // forward-RangeFor-to-reverse-RangeFor mapping could lift this restriction. - // - // AdStackLoadTopStmt as a leaf is correct under the dominance + control-flow-consumer + self-load - // guards in `EliminateRecomputableAdStackPushes::run_one_pass`. The reasoning: - // - // - In FORWARD: the eliminated stack S's body push dominates each `AdStackLoadTopStmt(S)` (the - // dominance guard), and the chain leaves' pushes dominate the chain's evaluation point. Each - // stack referenced by a chain leaf has a stable "top" value within one forward iteration body - // (no pops in forward), so substituting load_top(S) with the chain re-evaluates to the same - // value (no iteration shift). - // - // - In REVERSE: `MakeAdjoint` visits forward stmts in reverse order, emitting reverse code at a - // cursor that advances one position per visit. For a forward stmt F at position P_F, the - // emitted reverse code lands at cursor position roughly inversely correlated with P_F. The - // dominance guard ensures every chain-leaf stack T has its body push at a position P_T with - // P_T < P (where P is load_top(S)'s forward position). Therefore T's pop in reverse (emitted - // when MakeAdjoint visits T's body push) lands AFTER the chain consumer's reverse emission - // (the consumer's forward position is P > P_T). At the consumer's reverse cursor, T has not - // been popped yet, so load_top(T) returns T's iter-k-push value - matching what the original - // load_top(S) returned in forward iter k. - // - // The control-flow-consumer guard in `run_one_pass` covers a separate issue: stacks whose load_tops - // are direct operands of IfStmt cond / RangeFor begin/end. `MakeAdjoint::visit(IfStmt)` runs a - // dedicated snap-stack fixup that ONLY triggers when the cond is a bare `AdStackLoadTopStmt` (line - // 2168-2202). Eliminating the cond stack converts the cond into a compound stmt, the snap-stack - // does not trigger, and the reverse cond falls back to BackupSSA's load(op) which is single-slot - // last-iter only - silent gradient corruption on multi-iter loops. - if (stmt->is() || stmt->is() || stmt->is() || - stmt->is()) { - return true; - } - // GlobalLoadStmt as recomputable: the load reads a SNode value via GlobalPtrStmt. The cloned chain in - // the reverse pass re-issues the same load - safe iff the global is not mutated between the forward - // chain evaluation and the cloned re-read. Within a single kernel execution, forward writes complete - // before reverse runs, so the reverse re-read sees the kernel's final post-write state. If a global is - // mutated by the forward and then re-read by the reverse clone, the values can differ. - // - // Most rigid-step kernel chain leaves are reads of input parameters (mass, inertia, joint params, - // morphology) that the kernel does NOT write. For those, the re-read returns the same value. - // Conservative analysis: a future safety check could prove the SNode is read-only in the kernel; until - // then this path relies on the test suite's gradient-correctness asserts to surface any mutated-global - // miscompilation. - bool is_interior = stmt->is() || stmt->is() || stmt->is() || - stmt->is() || stmt->is() || stmt->is() || - stmt->is(); - if (!is_interior) { - return false; - } - auto operands = stmt->get_operands(); - for (auto *op : operands) { - if (op == nullptr) - continue; - if (!is_recomputable(op, cache)) - return false; - } - return true; - } -}; - -// Clones the SSA chain rooted at `src` into the IR, inserting cloned stmts before `insert_point`. Returns the -// cloned root. Per-stmt cache shared across one resolution materializes each SSA value at most once: diamond -// DAGs see two consumers but get one shared clone. `AdStackAllocaStmt` is treated as a leaf and shared (not -// cloned) - the stack itself is a unique storage handle that must not be duplicated. -// -// Pop-ordering safety: cloned `AdStackLoadTopStmt`s read the live top at the cloned position. `MakeAdjoint` -// emits `AdStackPopStmt` for each surviving `AdStackPushStmt`, and the existing reverse-pass scheme places -// the pop AFTER all uses of that stack within the reverse iteration (uses include both the original -// `AdStackLoadTopStmt`s emitted by `ReplaceLocalVarWithStacks` and the consumers' clones). For loop-carried -// allocas the pop fires early to expose the iteration's INPUT primal as the new top, which is exactly the -// value the recomputed chain needs - the existing per-consumer clone path at `BackupSSA::generic_visit` line -// ~2697 relies on this same property and has been correct in production. -class RecomputableChainCloner { - public: - static Stmt *clone_at(Stmt *src, Stmt *insert_point, std::unordered_map &cache) { - auto it = cache.find(src); - if (it != cache.end()) - return it->second; - Stmt *cloned = nullptr; - if (src->is()) { - // The alloca is shared, not cloned: every load reads the same physical stack. - cloned = src; - } else if (src->is() || src->is() || src->is()) { - auto cloned_unique = src->clone(); - cloned = insert_point->insert_before_me(std::move(cloned_unique)); - // For AdStackLoadTopStmt clones, the cloned stmt's `stack` operand still points at the original - // AdStackAllocaStmt - that's the desired sharing. - } else { - // Compound op: clone first, then walk operands and rewire each to a recursive clone. - auto cloned_unique = src->clone(); - cloned = insert_point->insert_before_me(std::move(cloned_unique)); - int n = src->num_operands(); - for (int i = 0; i < n; i++) { - auto *op = src->operand(i); - if (op != nullptr) { - Stmt *new_op = clone_at(op, cloned, cache); - cloned->set_operand(i, new_op); - } - } - } - cache[src] = cloned; - return cloned; - } -}; - -class IndependentBlocksJudger : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - - void visit(LocalLoadStmt *stmt) override { - QD_ASSERT(stmt->src->is() || stmt->src->is() || stmt->src->is()); - touched_allocas_.insert(stmt->src); - } - - void visit(LocalStoreStmt *stmt) override { - QD_ASSERT(stmt->dest->is() || stmt->dest->is() || - stmt->dest->is()); - touched_allocas_.insert(stmt->dest); - } - - void visit(AtomicOpStmt *stmt) override { - // We don't need to check the global atomics inside the range for-loops - // because - // 1. If the range for-loop is innermost, they will be captured by - // MakeAdjoint anyway - // 2. If the range for-loop is not innermost, they will be processed by - // another IndependentBlocksJudger - if (is_inside_loop_) - return; - - Stmt *dest = stmt->dest; - if (dest->is()) { - dest = dest->as()->origin; - } - - if (dest->is()) { - if (dest->as() - ->base_ptr->as() - ->ret_type.ptr_removed() - ->as() - ->elements() - .size() > TypeFactory::GRAD_PTR_POS_IN_NDARRAY) { - qualified_glb_operations_ = true; - } - } else { - QD_ASSERT(dest->is()); - if (dest->as()->snode->has_adjoint()) { - qualified_glb_operations_ = true; - } - } - } - - void visit(GlobalLoadStmt *stmt) override { - // We don't need to check the global load inside the range for-loops - // because - // 1. If the range for-loop is innermost, they will be captured by - // MakeAdjoint anyway - // 2. If the range for-loop is not innermost, they will be processed by - // another IndependentBlocksJudger - if (is_inside_loop_) - return; - - Stmt *src = stmt->src; - if (src->is()) { - src = src->as()->origin; - } - - if ((src->is() && src->as() - ->base_ptr->as() - ->ret_type.ptr_removed() - ->as() - ->elements() - .size() > TypeFactory::GRAD_PTR_POS_IN_NDARRAY) || - (src->is() && src->as()->snode->has_adjoint())) { - qualified_glb_operations_ = true; - } - } - - void visit(RangeForStmt *stmt) override { - inner_most_loop_ = false; - is_inside_loop_ = true; - stmt->body->accept(this); - is_inside_loop_ = false; - } - - static void run(IRNode *root, IndependentBlockMetaData &ib_meta_data) { - IndependentBlocksJudger Judger; - Block *block = root->as(); - root->accept(&Judger); - std::set outside_blocks; - // Collect all parent blocks (i.e. outside blocks) of the current block for - // local load/store stmt checks - for (auto b = block->parent_block(); b; b = b->parent_block()) { - if (b) - outside_blocks.insert(b); - } - for (const auto &alloca : Judger.touched_allocas_) { - // Test if the alloca belongs to the current block - if (outside_blocks.find(alloca->parent) != outside_blocks.end()) { - // This block is not an IB since it loads/modifies outside variables - ib_meta_data.is_ib = false; - } - } - - // To judge whether a block is an IB - // - No local load/store to allocas *outside* itself has been strictly - // enforced - - // To judge whether a block is a smallest IB - // - If the #1 is satisfied, either an inner most loop or a block without - // global atomics / global load is an IB - ib_meta_data.is_smallest_ib = ib_meta_data.is_ib && (Judger.qualified_glb_operations_ || Judger.inner_most_loop_); - } - - private: - std::set touched_allocas_; - bool qualified_glb_operations_ = false; - bool inner_most_loop_ = true; - bool is_inside_loop_ = false; -}; - -// Remove the duplicated IBs, remove blocks who are others' children because -// each block should only be processed once -class DuplicateIndependentBlocksCleaner : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - - void check_children_ib(Block *target_block) { - // Remove the block if it is the child of the block being visiting - if (independent_blocks_cleaned_.find(target_block) != independent_blocks_cleaned_.end()) { - independent_blocks_cleaned_.erase(target_block); - } - } - - void visit(StructForStmt *stmt) override { - check_children_ib(stmt->body.get()); - stmt->body->accept(this); - } - void visit(RangeForStmt *stmt) override { - check_children_ib(stmt->body.get()); - stmt->body->accept(this); - } - - static std::set run(const std::vector> &raw_IBs) { - DuplicateIndependentBlocksCleaner cleaner; - // Remove duplicate IBs - for (auto const &item : raw_IBs) { - cleaner.independent_blocks_cleaned_.insert(item.second); - } - // No clean is needed if only one IB exists - if (cleaner.independent_blocks_cleaned_.size() > 1) { - // Check from the block with smallest depth, ensure no duplicate visit - // happens - for (const auto &block : cleaner.independent_blocks_cleaned_) { - block->accept(&cleaner); - } - } - return cleaner.independent_blocks_cleaned_; - } - - private: - std::set independent_blocks_cleaned_; -}; - -// Do automatic differentiation pass in the reverse order (reverse-mode AD) - -// Independent Block (IB): blocks (i.e. loop bodies) whose iterations are -// independent of previous iterations and outer scopes. IBs are where the -// MakeAdjoint pass happens. IBs may contain if's and for-loops. - -// IBs are not always the inner-most for loop body. If the inner-most for-loop -// has iterations that carry iteration-dependent variables, it's not an IB. - -// Clearly the outermost level is always an IB, but we want IBs to be as small -// as possible. Outside IBs, we just need to reverse the for-loop orders. - -// Figure out the IBs. -class IdentifyIndependentBlocks : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - - void visit(WhileStmt *stmt) override { - QD_ERROR("WhileStmt is not supported in AutoDiff."); - } - - void visit(ContinueStmt *stmt) override { - QD_ERROR("ContinueStmt is not supported in AutoDiff."); - } - - void visit(WhileControlStmt *stmt) override { - QD_ERROR("WhileControlStmt (break) is not supported in AutoDiff."); - } - - void visit_loop_body(Block *block) { - auto ib_meta_data = IndependentBlockMetaData(); - // An IB has no local load/store to allocas *outside* itself - // Note: - // - Local atomics should have been demoted before this pass. - // - It is OK for an IB to have more than two for loops. - // - No global load/atomics operations to the global variables which - // require gradient - if (block->statements.empty()) { - // A empty block shoud be a smallest IB - ib_meta_data.is_ib = true; - ib_meta_data.is_smallest_ib = true; - } else { - IndependentBlocksJudger::run(block, ib_meta_data); - } - - if (ib_meta_data.is_smallest_ib) { - independent_blocks_.push_back({depth_, block}); - } else if (ib_meta_data.is_ib) { - current_ib_ = block; - block->accept(this); - } else { - if (depth_ <= 1) { - QD_ASSERT(depth_ == 1); - // The top level block is already not an IB, store it - independent_blocks_.push_back({depth_ - 1, block}); - } else { - independent_blocks_.push_back({depth_ - 1, block->parent_block()}); - } - } - } - - void visit(StructForStmt *stmt) override { - QD_ASSERT(depth_ == 0); - depth_++; - current_ib_ = stmt->body.get(); - visit_loop_body(stmt->body.get()); - depth_--; - } - - void visit(RangeForStmt *stmt) override { - if (depth_ == 0) { - current_ib_ = stmt->body.get(); - } - depth_++; - visit_loop_body(stmt->body.get()); - depth_--; - } - - static std::set run(IRNode *root) { - IdentifyIndependentBlocks pass; - Block *block = root->as(); - bool has_for = false; - for (auto &s : block->statements) { - if (s->is() || s->is()) { - has_for = true; - } - } - if (!has_for) { - // The whole block is an IB - pass.independent_blocks_.push_back({0, block}); - } else { - root->accept(&pass); - } - // Sort the IBs by their depth from shallow to deep - std::sort( - pass.independent_blocks_.begin(), pass.independent_blocks_.end(), - [](const std::pair &a, const std::pair &b) -> bool { return a.first < b.first; }); - - QD_ASSERT(!pass.independent_blocks_.empty()); - return DuplicateIndependentBlocksCleaner::run(pass.independent_blocks_); - } - - private: - std::vector> independent_blocks_; - int depth_{0}; - Block *current_ib_{nullptr}; -}; - -// Note that SSA does not mean the instruction will be executed at most once. -// For instructions that may be executed multiple times, we treat them as a -// mutable local variables. -class PromoteSSA2LocalVar : public BasicStmtVisitor { - using BasicStmtVisitor::visit; - - explicit PromoteSSA2LocalVar(Block *block) { - alloca_block_ = block; - invoke_default_visitor = true; - execute_once_ = true; - } - - // Demand-driven `required_defs_` set: the SSA defining stmts that downstream consumers actually require to be - // available at every iteration. A consumer requires its operand iff its adjoint formula reads it - the precise - // set is the operands of any non-linear unary / binary / ternary op (per `NonLinearOps::*_collections`), the - // indices of any `GlobalPtrStmt` / `ExternalPtrStmt` (the reverse pass replays the load), and the `cond` of any - // `IfStmt` / the `begin`/`end` of any `RangeForStmt` (the reverse pass clones the control flow). Stmts outside - // this set are left in pure SSA form: their values stay register-resident inside the forward pass, and - // `MakeAdjoint`'s reverse-pass formulas (which never read those values directly) generate correct adjoint - // accumulations against the adstack-backed defs that ARE in the set. Skipping promotion for the rest of the - // body collapses the alloca + LocalStore + LocalLoad triple-multiplier on unrolled IR that emitted a spill + - // reload pair per non-required arithmetic op with no reverse-pass consumer for the spilled value. - // - // Operands of `LocalStoreStmt` are not added because `MakeAdjoint::visit(LocalStoreStmt)` reads only - // `adjoint(stmt->dest)`, not the forward `stmt->val`. Linear ops (add / sub / mod / cmp / neg / floor / ceil / - // cast / logic_not / bit-ops) are likewise excluded because their adjoint formulas read only `adjoint(stmt)`. - static void compute_required_defs(Block *block, std::unordered_set &out) { - std::function walk = [&](Block *b) { - for (auto &owned : b->statements) { - Stmt *stmt = owned.get(); - if (auto *u = stmt->cast()) { - if (NonLinearOps::unary_collections.find(u->op_type) != NonLinearOps::unary_collections.end()) { - out.insert(u->operand); - } - } else if (auto *bin = stmt->cast()) { - if (NonLinearOps::binary_collections.find(bin->op_type) != NonLinearOps::binary_collections.end()) { - out.insert(bin->lhs); - out.insert(bin->rhs); - } - } else if (auto *tern = stmt->cast()) { - if (NonLinearOps::ternary_collections.find(tern->op_type) != NonLinearOps::ternary_collections.end()) { - out.insert(tern->op1); - out.insert(tern->op2); - out.insert(tern->op3); - } - } else if (auto *gp = stmt->cast()) { - for (auto *idx : gp->indices) { - out.insert(idx); - } - } else if (auto *ep = stmt->cast()) { - for (auto *idx : ep->indices) { - out.insert(idx); - } - } else if (auto *mp = stmt->cast()) { - // Reverse-mode `MakeAdjoint::visit(MatrixPtrStmt)` reads `stmt->offset` for dynamic-index adjoint - // routing, so a per-iteration-varying offset producer left in pure SSA would be backed by BackupSSA's - // single overwrite-each-iteration alloca and the reverse pass would read the last forward offset for - // every iteration. Promote it to alloca so `AdStackAllocaJudger::visit(MatrixPtrStmt)` can - // adstack-promote loop-varying offsets. Speculative fix: not exhibited by any failing test in the AD - // suite today, but the analysis is sound and the cost of leaving it unfixed is silent gradient - // corruption on `tensor[i + j]`-style local-tensor indexing inside a serial range-for. - out.insert(mp->offset); - } else if (auto *if_s = stmt->cast()) { - out.insert(if_s->cond); - if (if_s->true_statements) { - walk(if_s->true_statements.get()); - } - if (if_s->false_statements) { - walk(if_s->false_statements.get()); - } - } else if (auto *rf = stmt->cast()) { - out.insert(rf->begin); - out.insert(rf->end); - walk(rf->body.get()); - } else if (auto *sf = stmt->cast()) { - if (sf->body) { - walk(sf->body.get()); - } - } else if (auto *while_s = stmt->cast()) { - if (while_s->body) { - walk(while_s->body.get()); - } - } - } - }; - walk(block); - } - - void visit(Stmt *stmt) override { - if (execute_once_) - return; - if (!(stmt->is() || stmt->is() || stmt->is() || - stmt->is() || stmt->is() || stmt->is())) { - // TODO: this list may be incomplete - return; - } - - // `AllocaStmt`s always need to be hoisted to the top of the IB regardless of consumer analysis: a user-level - // `var = ...` construct inside a loop body must own a fixed slot at the IB's entry so every iteration shares it - // (cross-iteration accumulators are exactly the shape that drives the hoist). The demand-driven gate only - // applies to value-producing stmts (UnaryOp / BinaryOp / TernaryOp / GlobalLoad / LoopIndex) where the - // alloca + LocalStore + LocalLoad triple is purely a reverse-pass-readable spill - those skip when no consumer - // requires the value. The `LocalStoreStmt`s emitted here are placeholders that `ReplaceLocalVarWithStacks` - // rewrites into `AdStackPushStmt`s downstream. - if (stmt->is()) { - auto dtype = stmt->ret_type.ptr_removed(); - auto alloc = Stmt::make(dtype); - auto alloc_ptr = alloc.get(); - QD_ASSERT(alloca_block_); - alloca_block_->insert(std::move(alloc), 0); - immediate_modifier_->replace_usages_with(stmt, alloc_ptr); - - auto zero = insert_const(dtype, stmt, 0); - zero->insert_after_me(Stmt::make(alloc_ptr, zero)); - stmt->parent->erase(stmt); - return; - } - - if (required_defs_.find(stmt) == required_defs_.end()) { - return; - } - - auto alloc = Stmt::make(stmt->ret_type.ptr_removed()); - auto alloc_ptr = alloc.get(); - QD_ASSERT(alloca_block_); - alloca_block_->insert(std::move(alloc), 0); - auto load = stmt->insert_after_me(Stmt::make(alloc_ptr)); - immediate_modifier_->replace_usages_with(stmt, load); - // Create the load first so that the operand of the store does not get rewritten to point at the load (the - // SSA value `stmt` is still the right thing to spill; only the downstream consumers see the load). - stmt->insert_after_me(Stmt::make(alloc_ptr, stmt)); - } - - void visit(RangeForStmt *stmt) override { - auto old_execute_once = execute_once_; - execute_once_ = false; // loop body may be executed many times - stmt->body->accept(this); - execute_once_ = old_execute_once; - } - - private: - Block *alloca_block_{nullptr}; - bool execute_once_; - std::unordered_set required_defs_; - // ImmediateIRModifier collapses each `replace_usages_with` from a whole-tree walk (O(N)) to a constant-time - // operand-pointer rewrite: the modifier gathers every (consumer, operand_index) pair feeding any existing stmt - // once at construction (one O(N) pass), and per-replacement just looks up the table for `old_stmt` and rewrites - // each consumer's operand pointer. Without it, `PromoteSSA2LocalVar` ran K (= number of promoted defs) full IR - // walks - O(K*N), the dominant Quadrants-IR-side compile cost on unrolled reverse-mode bodies. - std::unique_ptr immediate_modifier_; - - public: - static void run(Block *block) { - PromoteSSA2LocalVar pass(block); - compute_required_defs(block, pass.required_defs_); - pass.immediate_modifier_ = std::make_unique(block); - block->accept(&pass); - } -}; - -class AdStackAllocaJudger : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - // Find the usage of the stmt recursively along the LocalLoadStmt - void visit(LocalLoadStmt *stmt) override { - if (stmt->src == target_alloca_) { - local_loaded_ = true; - target_alloca_ = stmt; - } - } - - // Track whether the alloca has any store at all so a load-only alloca (no adstack-relevant data flow either - // direction) can short-circuit `run()` regardless of what the per-op visitors below find. The decision of - // whether the alloca needs adstack promotion is made entirely by the precise visitors: non-linear unary / - // binary / ternary, GlobalPtr / ExternalPtr index, and IfStmt / RangeForStmt bound. Plus an alloca that is - // both loaded and stored anywhere in the IB is treated as loop-carried, which is needed for kernels like - // `for j: p, q = q, p + q` where the reverse pass routes the gradient through the cross-iteration - // recurrence and BackupSSA's single overwrite-each-iteration alloca cannot back the read-after-write across - // iterations. The visit-order-dependent load+store evidence here is conservative: any alloca with both a - // load and a store inside the IB triggers it, including pure accumulators whose adjoint formulas don't - // actually need per-iteration values - the slight over-promotion cost is the price of correctness on - // Fibonacci-style recurrences (silent gradient corruption otherwise). - void visit(LocalStoreStmt *stmt) override { - if (stmt->dest == target_alloca_backup_) { - load_only_ = false; - // Gate the load+store-implies-stack-needed rule on actually being inside a dynamic RangeForStmt at the - // point this evidence accumulates. The rule's purpose is to preserve cross-iteration RAW dependencies - // (`for j: p, q = q, p + q` Fibonacci-style) that BackupSSA's single overwrite-each-iteration alloca - // cannot back. With no enclosing dynamic for-loop the IB body executes once: there is no cross- - // iteration RAW to preserve, and the "load+store" pattern is just an in-block accumulator that the - // reverse pass handles via plain SSA cloning. Promoting such allocas under a static-unrolled loop body - // wastes one AdStack per accumulator (one push per unrolled-iter store + one load_top per unrolled- - // iter load) without any reverse-pass consumer needing per-iter replay - that is the unrolled-overhead - // bug Plan B targets. - // - // `dynamic_for_depth_` is incremented in `visit(RangeForStmt)` and decremented on exit. The judger - // walks the IB tree from the alloca's enclosing block, so depth here reflects exactly the nesting of - // *dynamic* for-loops between the alloca and the current load/store. StructFor / WhileStmt do not - // increment because their bodies still execute per-iter and need the same RAW protection (StructFor - // is the kernel-level offload-loop in some cases, but its body is a per-thread independent block; - // load+store there is the same shape as for a top-level alloca). - if (local_loaded_ && dynamic_for_depth_ > 0) { - is_stack_needed_ = true; - } - } - } - - // Check if the alloca is load only - void visit(AtomicOpStmt *stmt) override { - if (stmt->dest == target_alloca_backup_) - load_only_ = false; - } - - // The stack is needed if the alloca serves as the index of any global variables. Same cursor-vs-backup - // pattern as visit(IfStmt)/visit(RangeForStmt) below: `index` is always a value-producing stmt (typically a - // `LocalLoadStmt` reading the alloca, or a `ConstStmt`), never the alloca itself. The raw `index == - // target_alloca_` comparison only matches the first load's instance the `visit(LocalLoadStmt)` cursor - // advanced to - any subsequent load of the same alloca used as a different GlobalPtr index slips through. - // Resolve the LocalLoad chain and compare `ll->src` against `target_alloca_backup_` to catch every load. - void visit(GlobalPtrStmt *stmt) override { - if (is_stack_needed_) - return; - for (const auto &index : stmt->indices) { - auto *index_ll = index->cast(); - if (index_ll && index_ll->src == target_alloca_backup_) - is_stack_needed_ = true; - } - } - - void visit(ExternalPtrStmt *stmt) override { - if (is_stack_needed_) - return; - for (const auto &index : stmt->indices) { - auto *index_ll = index->cast(); - if (index_ll && index_ll->src == target_alloca_backup_) - is_stack_needed_ = true; - } - } - - // Reverse-mode `MakeAdjoint::visit(MatrixPtrStmt)` reads `stmt->offset` for dynamic-index adjoint routing, so - // a runtime-varying offset whose value comes from this alloca needs adstack promotion - otherwise BackupSSA - // backs the offset with a single overwrite-each-iteration slot and the reverse pass routes every iteration's - // adjoint into the last forward offset's slot. Same cursor-vs-backup pattern as the index visitors above. - void visit(MatrixPtrStmt *stmt) override { - if (is_stack_needed_) - return; - auto *offset_ll = stmt->offset->cast(); - if (offset_ll && offset_ll->src == target_alloca_backup_) - is_stack_needed_ = true; - } - - // Check whether the target alloca is fed into a non-linear unary op. Same cursor-vs-backup pattern as - // visit(GlobalPtrStmt) above: `stmt->operand` is a value-producing stmt (typically LocalLoad), never the - // alloca itself, so resolve the LocalLoad chain and compare against the backup. - void visit(UnaryOpStmt *stmt) override { - if (is_stack_needed_) - return; - if (NonLinearOps::unary_collections.find(stmt->op_type) != NonLinearOps::unary_collections.end()) { - auto *operand_ll = stmt->operand->cast(); - if (operand_ll && operand_ll->src == target_alloca_backup_) - is_stack_needed_ = true; - } - } - - // Check whether the target alloca is fed into a non-linear binary op. Same cursor-vs-backup pattern. - void visit(BinaryOpStmt *stmt) override { - if (is_stack_needed_) - return; - if (NonLinearOps::binary_collections.find(stmt->op_type) != NonLinearOps::binary_collections.end()) { - auto *lhs_ll = stmt->lhs->cast(); - auto *rhs_ll = stmt->rhs->cast(); - if ((lhs_ll && lhs_ll->src == target_alloca_backup_) || (rhs_ll && rhs_ll->src == target_alloca_backup_)) - is_stack_needed_ = true; - } - } - - // Check whether the target alloca is fed into a non-linear ternary op. Same cursor-vs-backup pattern. - void visit(TernaryOpStmt *stmt) override { - if (is_stack_needed_) - return; - if (NonLinearOps::ternary_collections.find(stmt->op_type) != NonLinearOps::ternary_collections.end()) { - auto *op1_ll = stmt->op1->cast(); - auto *op2_ll = stmt->op2->cast(); - auto *op3_ll = stmt->op3->cast(); - if ((op1_ll && op1_ll->src == target_alloca_backup_) || (op2_ll && op2_ll->src == target_alloca_backup_) || - (op3_ll && op3_ll->src == target_alloca_backup_)) - is_stack_needed_ = true; - } - } - - // Check whether the target alloca feeds the condition of an if stmt. `stmt->cond` is always a - // value-producing stmt - typically a direct `LocalLoadStmt` reading the alloca, but also commonly a - // `BinaryOpStmt` wrapping such a load (e.g. `j < i+1`). Walk the expression chain to catch every load of - // the target alloca: the raw `stmt->cond == target_alloca_` comparison the old code used only matched the - // first-visited load's instance, and a direct `cast` still misses the BinaryOp case that - // `visit(BinaryOpStmt)` cannot catch (comparison ops are linear and so not in `NonLinearOps`). Covers the - // shape defensively: IR simplification currently collapses most BinaryOp-wrapped conds before the judger - // sees them, so no failing regression test pins it today, but the fix is structurally correct for future - // IR changes that preserve the BinaryOp wrapping. - void visit(IfStmt *stmt) override { - if (is_stack_needed_) - return; - - if (feeds_target_alloca(stmt->cond)) { - is_stack_needed_ = true; - return; - } - - if (stmt->true_statements) - stmt->true_statements->accept(this); - if (stmt->false_statements) - stmt->false_statements->accept(this); - } - - // Check whether the target alloca feeds the begin or end of a range-for bound. Under reverse-mode AD, if an - // inner for-loop's bound is an enclosing loop-carried counter (the canonical triangular-nested - // `for k in range(j)` shape, or the `range(j+1)` / `range(n-i)` shapes where the bound is a linear arithmetic - // expression of a loop-carried alloca), its reverse clone must read the bound from the per-iteration forward - // value; without an adstack the reverse pass sees only the last forward value and the inner loop over- or - // under-runs, silently corrupting gradients for the earliest inner indices (those visited most often across - // outer iterations). This check is the only thing that promotes such a loop-counter alloca - - // `visit(LocalStoreStmt)`'s `local_loaded_` short-circuit does not fire because the counter is only LOAD-ed - // inside the inner-loop bound, not LOAD-then-STORE-ed. Walk the expression chain through - // `feeds_target_alloca` so both direct LocalLoads (`range(j)`) and LocalLoads nested under linear ops - // (`range(j+1)`, `range(n-i)`, ...) trigger promotion. The BinaryOp-wrapped case is defensively covered: IR - // simplification currently collapses most such bounds before the judger sees them, so no failing regression - // test pins it today, but the walker is structurally correct for future IR changes that preserve the - // wrapping. The raw-cast direct `LocalLoadStmt` variant pinned by `test_adstack_inner_for_bound_is_enclosing - // _loop_index` remains covered - that shape takes the first branch of the walker trivially. - void visit(RangeForStmt *stmt) override { - if (is_stack_needed_) - return; - - if (feeds_target_alloca(stmt->begin) || feeds_target_alloca(stmt->end)) { - is_stack_needed_ = true; - return; - } - - dynamic_for_depth_++; - stmt->body->accept(this); - dynamic_for_depth_--; - } - - static bool run(AllocaStmt *target_alloca) { - AdStackAllocaJudger judger; - judger.target_alloca_ = target_alloca; - judger.target_alloca_backup_ = target_alloca; - judger.dynamic_for_depth_ = 0; - target_alloca->parent->accept(&judger); - return (!judger.load_only_) && judger.is_stack_needed_; - } - - private: - // Recursively walk a value expression to decide whether it transitively reads `target_alloca_backup_` via a - // `LocalLoadStmt`. Used by `visit(IfStmt)` and `visit(RangeForStmt)` to detect the target alloca feeding a - // bound or condition even when wrapped by linear ops (e.g. `range(j+1)`, `j < i+1`). Linear binary/unary - // ops are traversed because `visit(BinaryOpStmt)`/`visit(UnaryOpStmt)` only flag *non-linear* ops - their - // linear-op path does not otherwise promote the alloca. `ConstStmt`s and unrelated values return false and - // terminate the recursion; the walker is always finite because SSA IR guarantees acyclic operand graphs. - bool feeds_target_alloca(Stmt *expr) const { - if (auto *ll = expr->cast()) { - return ll->src == target_alloca_backup_; - } - if (auto *bop = expr->cast()) { - return feeds_target_alloca(bop->lhs) || feeds_target_alloca(bop->rhs); - } - if (auto *uop = expr->cast()) { - return feeds_target_alloca(uop->operand); - } - if (auto *top = expr->cast()) { - return feeds_target_alloca(top->op1) || feeds_target_alloca(top->op2) || feeds_target_alloca(top->op3); - } - return false; - } - - Stmt *target_alloca_; - Stmt *target_alloca_backup_; - bool is_stack_needed_ = false; - bool local_loaded_ = false; - bool load_only_ = true; - // Nesting depth of dynamic `RangeForStmt` containers between the alloca's enclosing block and the current - // visit cursor. Static-unrolled `qd.static(range(...))` loops are removed by the AST transformer before the - // judger sees the IR, so they do not contribute to depth. The load+store-implies-stack-needed rule fires - // only when this depth is positive; see the rationale in `visit(LocalStoreStmt)`. - int dynamic_for_depth_ = 0; -}; - -class RegulateTensorTypedStatements : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - DelayedIRModifier delayed_modifier_; - - explicit RegulateTensorTypedStatements() { - } - - template - void process_store_stmt(Store *stmt) { - QD_ASSERT(stmt->template is() || stmt->template is()); - - if (stmt->dest->template is()) { - auto matrix_ptr_stmt = stmt->dest->template as(); - auto orig_stmt = matrix_ptr_stmt->origin; - - if (!orig_stmt->ret_type.ptr_removed()->template is()) { - return; - } - - auto tensor_type = orig_stmt->ret_type.ptr_removed()->template as(); - auto num_elements = tensor_type->get_num_elements(); - - if (matrix_ptr_stmt->offset->template is()) { - /* - [Static index] - Fwd: - $0 = alloca <4 x i32> - $1 = load $0 - $2 = matrix ptr $1, 2 // offset = 2 - $3 : local store $2, $val - - Replaced: - $0 = alloca <4 x i32> - $1 = load $0 - $2 = matrix ptr $1, 2 // --> erase - - $3 = matrix ptr $1, 0 - $4 = load $3 - - $5 = matrix ptr $1, 1 - $6 = load $5 - - $7 = matrix ptr $1, 3 - $8 = load $7 - - $9 = matrix init [$4, $6, $val, $8] - - $10 : store $0, $9 - */ - int offset = matrix_ptr_stmt->offset->template as()->val.val_int32(); - - QD_ASSERT(offset < num_elements); - - std::vector values; - for (int i = 0; i < num_elements; i++) { - if (i == offset) { - values.push_back(stmt->val); - continue; - } - - auto const_i = insert_const(PrimitiveType::i32, stmt, i, true); - auto matrix_ptr_stmt_i = Stmt::make(orig_stmt, const_i); - matrix_ptr_stmt_i->ret_type = tensor_type->get_element_type(); - - auto local_load_stmt_i = Stmt::make(matrix_ptr_stmt_i.get()); - local_load_stmt_i->ret_type = tensor_type->get_element_type(); - - values.push_back(local_load_stmt_i.get()); - - stmt->insert_before_me(std::move(matrix_ptr_stmt_i)); - stmt->insert_before_me(std::move(local_load_stmt_i)); - } - - auto matrix_init_stmt = Stmt::make(values); - matrix_init_stmt->ret_type = tensor_type; - - auto store_stmt = Stmt::make(orig_stmt, matrix_init_stmt.get()); - stmt->insert_before_me(std::move(matrix_init_stmt)); - stmt->replace_with(std::move(store_stmt)); - - return; - - } else { - /* - [Dynamic index] - Fwd: - $0 = alloca <4 x i32> - $1 = load $0 - $2 = matrix ptr $1, $offset // offset = 2 - $3 : local store $2, $val - - Replaced: - $0 = alloca <4 x i32> - - $1 = load $0 - $2 = matrix init [$val, $val, $val, $val] - - $3 = matrix init [$offset, $offset, $offset, $offset] - $4 = matrix init [0, 1, 2, 3] - - $5 = bin_eq $3, $4 - $6 = select $5, $2, $1 - - $7 : store $0, $6 - */ - auto tensor_type = orig_stmt->ret_type.ptr_removed()->template as(); - auto num_elements = tensor_type->get_num_elements(); - - auto tensor_shape = tensor_type->get_shape(); - auto index_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::i32); - - std::vector val_values(num_elements, stmt->val); - std::vector offset_values(num_elements, matrix_ptr_stmt->offset); - std::vector index_values(num_elements); - for (int i = 0; i < num_elements; i++) { - index_values[i] = insert_const(PrimitiveType::i32, stmt, i, true); - } - - auto matrix_val = Stmt::make(val_values); - matrix_val->ret_type = tensor_type; - - auto matrix_offset = Stmt::make(offset_values); - matrix_offset->ret_type = index_tensor_type; - - auto matrix_index = Stmt::make(index_values); - matrix_index->ret_type = index_tensor_type; - auto cmp_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::u1); - auto matrix_eq = Stmt::make(BinaryOpType::cmp_eq, matrix_offset.get(), matrix_index.get()); - matrix_eq->ret_type = cmp_tensor_type; - - auto orig_value = Stmt::make(orig_stmt); - orig_value->ret_type = tensor_type; - - auto matrix_select = - Stmt::make(TernaryOpType::select, matrix_eq.get(), matrix_val.get(), orig_value.get()); - matrix_select->ret_type = tensor_type; - - auto store_stmt = Stmt::make(orig_stmt, matrix_select.get()); - - stmt->insert_before_me(std::move(matrix_val)); - stmt->insert_before_me(std::move(matrix_offset)); - stmt->insert_before_me(std::move(matrix_index)); - stmt->insert_before_me(std::move(matrix_eq)); - stmt->insert_before_me(std::move(orig_value)); - stmt->insert_before_me(std::move(matrix_select)); - stmt->replace_with(std::move(store_stmt)); - return; - } - } - } - - void visit(LocalStoreStmt *stmt) override { - process_store_stmt(stmt); - } - - void visit(GlobalStoreStmt *stmt) override { - process_store_stmt(stmt); - } - - static void run(IRNode *root) { - RegulateTensorTypedStatements pass; - root->accept(&pass); - } -}; - -class ReplaceLocalVarWithStacks : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - int ad_stack_size; - DelayedIRModifier delayed_modifier_; - - explicit ReplaceLocalVarWithStacks(int ad_stack_size) : ad_stack_size(ad_stack_size) { - } - - void visit(AllocaStmt *alloc) override { - bool is_stack_needed = AdStackAllocaJudger::run(alloc); - if (is_stack_needed) { - auto dtype = alloc->ret_type.ptr_removed(); - auto stack_alloca = Stmt::make(dtype, ad_stack_size); - auto stack_alloca_ptr = stack_alloca.get(); - - alloc->replace_with(VecStatement(std::move(stack_alloca))); - - // Note that unlike AllocaStmt, AdStackAllocaStmt does NOT have an 0 as - // initial value. Therefore here we push an initial 0 value. - auto zero = insert_const(dtype, stack_alloca_ptr, 0); - zero->insert_after_me(Stmt::make(stack_alloca_ptr, zero)); - } - } - - void visit(LocalLoadStmt *stmt) override { - if (stmt->src->is()) { - auto stack_load = Stmt::make(stmt->src); - stack_load->ret_type = stmt->ret_type; - - stmt->replace_with(std::move(stack_load)); - return; - } - - // Slot load from a stack-backed tensor. After `visit(MatrixPtrStmt)`, `stmt->src` is of the form - // `MatrixPtrStmt(AdStackLoadTopStmt(stack, return_ptr=true), offset)`. A direct load through that pointer - // leaves the store-to-load forwarding walker in `ir/control_flow_graph.cpp` with no reaching definition, - // because the only producer for the stack's top slots is an `AdStackPushStmt` (tagged `ir_traits::Load`, - // invisible to `get_store_destination`). Replace the load with a full-tensor `AdStackLoadTopStmt` - // materialized into a fresh regular `AllocaStmt`, then re-subscript it - a plain alloca + LocalStore - // sequence is a shape the reach-in walker can trace end-to-end. - if (stmt->src->is()) { - auto matrix_ptr = stmt->src->as(); - if (matrix_ptr->origin->is() && matrix_ptr->origin->as()->return_ptr) { - auto stack = matrix_ptr->origin->as()->stack; - QD_ASSERT(stack->is()); - auto tensor_type = stack->ret_type.ptr_removed(); - - auto full_load = Stmt::make(stack); - full_load->ret_type = tensor_type; - auto full_load_ptr = full_load.get(); - - auto fresh_alloca = Stmt::make(tensor_type); - auto fresh_alloca_ptr = fresh_alloca.get(); - fresh_alloca->ret_type = tensor_type; - fresh_alloca->ret_type.set_is_pointer(true); - - auto fresh_store = Stmt::make(fresh_alloca_ptr, full_load_ptr); - - auto new_matrix_ptr = Stmt::make(fresh_alloca_ptr, matrix_ptr->offset); - new_matrix_ptr->ret_type = stmt->ret_type; - - auto new_load = Stmt::make(new_matrix_ptr.get()); - new_load->ret_type = stmt->ret_type; - - stmt->insert_before_me(std::move(full_load)); - stmt->insert_before_me(std::move(fresh_alloca)); - stmt->insert_before_me(std::move(fresh_store)); - stmt->insert_before_me(std::move(new_matrix_ptr)); - stmt->replace_with(std::move(new_load)); - } - } - } - - void visit(LocalStoreStmt *stmt) override { - if (stmt->dest->is()) { - auto matrix_ptr_stmt = stmt->dest->as(); - if (matrix_ptr_stmt->origin->is()) { - auto stack_top_stmt = matrix_ptr_stmt->origin->as(); - QD_ASSERT(stack_top_stmt->return_ptr == true); - - if (!stack_top_stmt->ret_type.ptr_removed()->is()) { - return; - } - - auto tensor_type = stack_top_stmt->ret_type.ptr_removed()->as(); - auto num_elements = tensor_type->get_num_elements(); - - if (matrix_ptr_stmt->offset->is()) { - /* - [Static index] - Load the full current top as a tensor via `AdStackLoadTopStmt` and merge the new value at `offset` - using a boolean mask + `select`. Mirrors the dynamic-index lowering below so that every slot of the - new pushed tensor derives from either `stmt->val` or the loaded top tensor and the IR contains no - per-slot `LocalLoadStmt` on a stack-backed `MatrixPtrStmt`. - - Why that invariant matters: the store-to-load forwarding walker in `ir/control_flow_graph.cpp` does - not treat `AdStackPushStmt` as a reaching definition (it is tagged `ir_traits::Load`, so - `get_store_destination` returns nothing for it), so a `LocalLoadStmt(MatrixPtrStmt(stack_top_ptr, - i))` inserted here has no reaching def and ends up reading an uninitialized adjoint slot in the - reverse kernel. Keep the `AdStackLoadTopStmt(stack)` + mask-select shape when touching this path. - - Fwd: - $1 = alloca <4 x i32> - $2 = matrix ptr $1, 2 // offset = 2 - $3 : local store $2, $val - - Replaced: - $1 = alloca <4 x i32> - - $2 = matrix init [$val, $val, $val, $val] - $3 = matrix init [false, false, true, false] // mask with `offset == i` - - $4 = ad stack load top (full tensor) $1 - $5 = select $3, $2, $4 - - $6 : stack push $1, $5 - */ - int offset = matrix_ptr_stmt->offset->as()->val.val_int32(); - - QD_ASSERT(offset < num_elements); - - auto tensor_shape = tensor_type->get_shape(); - auto cmp_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::u1); - - std::vector val_values(num_elements, stmt->val); - std::vector mask_values(num_elements); - for (int i = 0; i < num_elements; i++) { - mask_values[i] = insert_const(PrimitiveType::u1, stmt, i == offset ? 1 : 0, true); - } - - auto matrix_val = Stmt::make(val_values); - matrix_val->ret_type = tensor_type; - - auto matrix_mask = Stmt::make(mask_values); - matrix_mask->ret_type = cmp_tensor_type; - - auto matrix_alloca_value = Stmt::make(stack_top_stmt->stack); - matrix_alloca_value->ret_type = tensor_type; - - auto matrix_select = Stmt::make(TernaryOpType::select, matrix_mask.get(), matrix_val.get(), - matrix_alloca_value.get()); - matrix_select->ret_type = tensor_type; - - auto stack_push = Stmt::make(stack_top_stmt->stack, matrix_select.get()); - - stmt->insert_before_me(std::move(matrix_val)); - stmt->insert_before_me(std::move(matrix_mask)); - stmt->insert_before_me(std::move(matrix_alloca_value)); - stmt->insert_before_me(std::move(matrix_select)); - stmt->replace_with(std::move(stack_push)); - - return; - - } else { - /* - [Dynamic index] - Fwd: - $1 = alloca <4 x i32> - $2 = matrix ptr $1, $offset // offset = 2 - $3 : local store $2, $val - - Replaced: - $1 = alloca <4 x i32> - - $2 = matrix init [$val, $val, $val, $val] - - $3 = matrix init [$offset, $offset, $offset, $offset] - $4 = matrix init [0, 1, 2, 3] - - $5 = bin_eq $3, $4 - $6 = select $5, $2, $1 - - $7 : store $1, $6 - */ - auto tensor_type = stack_top_stmt->ret_type.ptr_removed()->as(); - auto num_elements = tensor_type->get_num_elements(); - - auto tensor_shape = tensor_type->get_shape(); - auto index_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::i32); - - std::vector val_values(num_elements, stmt->val); - std::vector offset_values(num_elements, matrix_ptr_stmt->offset); - std::vector index_values(num_elements); - for (int i = 0; i < num_elements; i++) { - index_values[i] = insert_const(PrimitiveType::i32, stmt, i, true); - } - - auto matrix_val = Stmt::make(val_values); - matrix_val->ret_type = tensor_type; - - auto matrix_offset = Stmt::make(offset_values); - matrix_offset->ret_type = index_tensor_type; - - auto matrix_index = Stmt::make(index_values); - matrix_index->ret_type = index_tensor_type; - - auto cmp_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::u1); - auto matrix_eq = Stmt::make(BinaryOpType::cmp_eq, matrix_offset.get(), matrix_index.get()); - matrix_eq->ret_type = cmp_tensor_type; - - auto matrix_alloca_value = Stmt::make(stack_top_stmt->stack); - matrix_alloca_value->ret_type = tensor_type; - - auto matrix_select = Stmt::make(TernaryOpType::select, matrix_eq.get(), matrix_val.get(), - matrix_alloca_value.get()); - matrix_select->ret_type = tensor_type; - - auto stack_push = Stmt::make(stack_top_stmt->stack, matrix_select.get()); - - stmt->insert_before_me(std::move(matrix_val)); - stmt->insert_before_me(std::move(matrix_offset)); - stmt->insert_before_me(std::move(matrix_index)); - stmt->insert_before_me(std::move(matrix_eq)); - stmt->insert_before_me(std::move(matrix_alloca_value)); - stmt->insert_before_me(std::move(matrix_select)); - stmt->replace_with(std::move(stack_push)); - - return; - } - } - } - - // Non Tensor-type - if (stmt->dest->is()) - stmt->replace_with(Stmt::make(stmt->dest, stmt->val)); - } - - void visit(MatrixPtrStmt *stmt) override { - if (stmt->origin->is()) { - auto stack_top = Stmt::make(stmt->origin, true /*is_ptr*/); - stack_top->ret_type = stmt->origin->ret_type; - stack_top->ret_type.set_is_pointer(true); - - Stmt *stack_top_stmt = stack_top.get(); - stmt->insert_before_me(std::move(stack_top)); - - auto new_matrix_ptr_stmt = Stmt::make(stack_top_stmt, stmt->offset); - new_matrix_ptr_stmt->ret_type = stmt->ret_type; - stmt->replace_with(std::move(new_matrix_ptr_stmt)); - } - } -}; - -// Eliminate AdStackAllocaStmts whose pushed value is recomputable from already-stack-backed allocas, kernel -// args, constants, and loop indices. Runs between `ReplaceLocalVarWithStacks` and `MakeAdjoint`, so the reverse -// pass is generated against the cleaned IR (no spurious AdStackPushStmt / AdStackLoadTopStmt scaffolding for -// values the reverse can reconstruct on the fly via cloned forward DAGs). -// -// Eligibility per AdStackAllocaStmt S in an independent block: -// 1. S is written by exactly one AdStackPushStmt (single-push pattern). Multi-push allocas hold loop-carried -// state where each iteration's push depends on the previous - the reverse pass cannot reconstruct -// iteration k's value from iteration (k-1)'s value, so the stack is genuinely needed. -// 2. The pushed value's transitive operand DAG is recomputable per RecomputableChainAnalyzer (leaves at -// AdStackLoadTop / AdStackAlloca / ArgLoad / Const / LoopIndex; interior side-effect-free ops only). -// 3. S has no AdStackLoadTopStmt with `return_ptr=true` consumers (those return a pointer aliasing the slot -// and a follow-up store would not be modeled by an SSA replacement). -// -// Action on eligibility: replace every `AdStackLoadTopStmt(S)` with the original pushed SSA value (the -// `AdStackPushStmt::v`), then erase the push and the alloca. The pushed SSA value's chain stays in the forward -// IR; downstream consumers in the forward pass now reference it directly as SSA, and `MakeAdjoint` plus -// `BackupSSA` reconstruct the chain on demand in the reverse pass via the same DAG-clone path that -// `BackupSSA::generic_visit` exercises for cross-block SSA references. -// -// Iterates to fixed point: eliminating one stack can newly expose another stack's chain as recomputable -// (S2's pushed value chained through `AdStackLoadTopStmt(S1)`; once S1 is gone, the chain may collapse into -// pure SSA chained through S1's pushed-value chain instead). Each pass eliminates at least one stack or -// terminates. -// -// Cost model: this pass trades forward-pass adstack pushes (memory write + top-pointer bump per push) for -// extra arithmetic in the reverse pass (the cloned DAG re-executes the forward chain). For pure-arithmetic -// chains rooted at a single loop-carried alloca - the dominant shape on Genesis-style rigid-step kernels - -// the win is typically 5-10x: each push is a memory op crossing the L1 boundary on every iteration, while -// the recomputed arithmetic stays in registers and reuses warmed-up sin/cos/exp pipeline state. -class EliminateRecomputableAdStackPushes { - public: - static void run(Block *ib) { - // Iterate to fixed point. Each pass either eliminates at least one stack or terminates. The hard cap - // protects against analysis bugs (e.g. an eligibility check that returns true on the same stmt twice - // in a row); under a correct implementation each pass strictly reduces the AdStackAllocaStmt count, so - // the actual iteration bound is the number of stacks at entry. 1024 is well above any realistic kernel. - for (int i = 0; i < 1024; i++) { - if (!run_one_pass(ib)) - return; - } - } - - private: - // True iff `s` is a literal-zero ConstStmt or a MatrixInitStmt whose every element is a literal-zero - // ConstStmt. Matches the init-zero push pattern that ReplaceLocalVarWithStacks emits right after each - // AdStackAllocaStmt: scalar allocas get a ConstStmt(0), tensor-typed allocas get a MatrixInitStmt of - // ConstStmt(0)s. Both are erased along with the alloca during elimination - they only matter if a load - // could observe them before the body push fires, which the standard PromoteSSA2LocalVar + - // ReplaceLocalVarWithStacks codegen never produces (loads always follow the body push within the loop - // iter that emits them). - static bool is_zero_init_value(Stmt *s) { - if (auto *c = s->cast()) { - return c->val.equal_value(0); - } - if (auto *mi = s->cast()) { - for (auto *elem : mi->values) { - auto *c = elem->cast(); - if (c == nullptr || !c->val.equal_value(0)) - return false; - } - return true; - } - return false; - } - - // Returns true if any stack was eliminated this pass. - static bool run_one_pass(Block *ib) { - // Collect every AdStackAllocaStmt anywhere within the IB. The IB is the root of a contiguous AD scope, so - // a single walk is sufficient. - std::vector stacks; - auto collected = irpass::analysis::gather_statements(ib, [&](Stmt *s) { return s->is(); }); - for (auto *s : collected) { - stacks.push_back(s->as()); - } - - bool modified = false; - std::unordered_map recomputable_cache; - - for (auto *stack : stacks) { - // Re-classify users of `stack` each iteration: a previous elimination on a downstream stack may have - // removed users that previously disqualified `stack`. - std::vector pushes; - std::vector load_tops; - bool disqualified = false; - - auto users = irpass::analysis::gather_statements(ib, [&](Stmt *s) { - if (auto *p = s->cast()) - return p->stack == stack; - if (auto *lt = s->cast()) - return lt->stack == stack; - if (auto *po = s->cast()) - return po->stack == stack; - if (auto *aa = s->cast()) - return aa->stack == stack; - if (auto *la = s->cast()) - return la->stack == stack; - return false; - }); - for (auto *user : users) { - if (auto *p = user->cast()) { - pushes.push_back(p); - } else if (auto *lt = user->cast()) { - if (lt->return_ptr) { - // return_ptr=true returns a pointer into the slot; replacing with an SSA value loses the - // pointer-identity contract. Keep the stack. - disqualified = true; - break; - } - load_tops.push_back(lt); - } else if (user->is() || user->is() || - user->is()) { - // Pre-MakeAdjoint, pop / adj-acc / load-top-adj should not appear yet for this stack. If they do, - // some upstream pass has already touched it - keep as-is to avoid double-rewrites. - disqualified = true; - break; - } - } - if (disqualified || pushes.empty()) { - continue; - } - - // ReplaceLocalVarWithStacks emits a const-zero "init" AdStackPushStmt immediately after the - // AdStackAllocaStmt as the stack's initial value (auto_diff.cpp:867-872 in pre-fix tree). Real - // user-level allocas in a loop body have at most ONE additional "body" push per iteration. Treat - // const-zero pushes whose value is a `ConstStmt` with all-zero `TypedConstant` as inits and require at - // most one non-init "body" push for elimination eligibility. Multi-body-push allocas hold loop-carried - // state (each iteration's push depends on the previous), so the reverse pass cannot reconstruct - // iteration k from iteration k-1 and the stack must stay. - AdStackPushStmt *body_push = nullptr; - std::vector init_pushes; - for (auto *p : pushes) { - if (is_zero_init_value(p->v)) { - init_pushes.push_back(p); - } else { - if (body_push != nullptr) { - disqualified = true; - break; - } - body_push = p; - } - } - if (disqualified || body_push == nullptr) { - continue; - } - - Stmt *pushed_val = body_push->v; - if (!RecomputableChainAnalyzer::is_recomputable(pushed_val, recomputable_cache)) { - continue; - } - // Loop-carried self-reference: when the pushed value's transitive operand DAG includes an - // AdStackLoadTopStmt of the SAME stack we are about to eliminate, the value is `read prev iter from - // stack -> compute -> push next iter`. Rewriting load_top($S) -> pushed_value_SSA leaves a self-cycle - // in the SSA graph (`acc_new = acc_new + sin(a)`) which both corrupts the IR and loses the iteration - // recurrence the stack was carrying. The recomputable-chain analyzer happily accepts such chains - // because AdStackLoadTopStmt is a recomputable leaf in general; the additional check below specifically - // disqualifies self-loaded stacks. - // - // The classic shape this protects: `acc = 0.0; for j in range(N): acc += sin(...)`. After - // PromoteSSA2LocalVar + ReplaceLocalVarWithStacks the IR has `$acc = stack_alloc; init push; - // for j: $tmp = stack_load_top $acc; $new = $tmp + sin(...); push $acc, val=$new`. The chain `$tmp + - // sin(...)` is recomputable per the leaf rules (LoadTop is a leaf), but rewriting load_top($acc) to - // ($tmp + sin(...)) makes the SSA graph self-reference $new and the reverse pass loses every - // iteration's accumulator state. - // Read-before-write protection: substituting `AdStackLoadTopStmt(S)` with the body push's value SSA - // chain only preserves semantics when the body push DOMINATES every load_top in the forward IR - - // that is, every load_top reads what the body push just wrote in the same iteration. The - // PromoteSSA2LocalVar + ReplaceLocalVarWithStacks pipeline emits exactly this shape for required-def - // spills: `alloca; push val=def; load_top` in document order, with the push immediately following - // the def and loads occurring later in the same iteration. - // - // Loop-carried recurrences violate the dominance rule. Take the canonical Fibonacci shape - // `p, q = q, p + q` lowered to IR: - // - // $tmp_p = load_top($p) # reads PREVIOUS iter's push into $p - // $tmp_q = load_top($p) + load_top($q) - // push $p, val=$tmp_p # writes NEXT iter's value into $p - // push $q, val=$tmp_q - // - // Here load_top($p) reads BEFORE the body push of $p within the same iter, so it observes iter - // (k-1)'s value, not iter k's. Substituting `load_top($p) -> $tmp_p`'s SSA chain (which reads - // load_top($q) at iter k) gives iter k's q-value instead of iter (k-1)'s p-value, off by one - // iteration and producing zero gradients on the Fibonacci-style regression test pinned by - // `test_ad_fibonacci_index`. - // - // The check below asserts that for every load_top in `load_tops`, the body push is its predecessor - // within the same containing block (or, equivalently, load_top comes AFTER the body push in - // document order at the body push's block level). Loads that live inside nested control flow - // (`IfStmt`, `RangeForStmt`) under the body push's block are also fine because the body push - // dominates them. We approximate dominance with the lexical "push's block contains load_top's - // ancestor block AND push position < load_top's ancestor's position in that block" check. - auto block_position = [](Stmt *s) -> int { - Block *b = s->parent; - if (b == nullptr) - return -1; - for (size_t i = 0; i < b->statements.size(); i++) { - if (b->statements[i].get() == s) - return static_cast(i); - } - return -1; - }; - Block *push_block = body_push->parent; - int push_pos = block_position(body_push); - bool dominates_all_loads = true; - for (auto *lt : load_tops) { - Stmt *cursor = lt; - Block *cursor_block = cursor->parent; - while (cursor_block != nullptr && cursor_block != push_block) { - cursor = cursor_block->parent_stmt(); - if (cursor == nullptr) - break; - cursor_block = cursor->parent; - } - if (cursor_block != push_block) { - // load_top is outside the push's block scope: cannot establish dominance. - dominates_all_loads = false; - break; - } - int cursor_pos = block_position(cursor); - if (cursor_pos <= push_pos) { - // load_top precedes the body push (or its enclosing container does): it reads the previous - // iteration's value, not iter k's. Substitution would shift iterations by one. - dominates_all_loads = false; - break; - } - } - if (!dominates_all_loads) { - continue; - } - - // Reverse-position correctness for chain leaves pushed in the same block as S. - // - // After elimination, every reverse stmt that consumed `load_top(S)` instead consumes the chain - // (cloned by `BackupSSA::generic_visit` at the consumer's reverse position). The cloned chain reads - // `load_top(T)` for each chain leaf T, where the read happens at the consumer's reverse cursor. - // - // `MakeAdjoint` visits forward stmts in REVERSE order, emitting at a cursor that advances per visit. - // For a forward stmt at position P, its reverse emissions land "later" in reverse block iff P is - // smaller. T's pop is emitted when `MakeAdjoint` visits T's body push at position P_T. The cloned - // chain at the consumer's reverse position reads `load_top(T)` POST-pop iff T's pop has fired by then, - // i.e. iff visit(P_T) precedes visit(C_pos), i.e. iff P_T > C_pos for every consumer C of S's - // load_tops. - // - // Forward chain evaluates at S's push position P_S, before any of T's same-block body pushes (since - // we already required P_T > P_S via the dominance check above for S's loads). So forward chain reads - // T's pre-iter-k-push value. Reverse clone post-pop also returns T's pre-iter-k-push value (= iter k - // INPUT for loop-carried T). Match. - // - // Without this check, the canonical Fibonacci-style shape (`p, q = q, p + q; b[q] += a[q]`) where S - // stores `p + q` and chain leaves p_stack / q_stack are pushed AFTER S but BEFORE the GlobalPtr - // index consumer would silently corrupt gradients: the reverse clone at the GlobalPtr's adjoint - // emission reads p_stack and q_stack PRE-pop (iter k POST-push values), summing to a different value - // than the forward chain at P_S used. - // - // Chain leaves T whose stacks have NO body push in S's containing block are stable within the iter - // (their tops do not change during the iter), so neither forward nor reverse evaluation order shifts - // their values. Excluded from this check. - auto find_consumers_of_load_tops = [&](std::vector &out_positions) { - std::function walk = [&](Block *b, int container_pos_in_push_block) { - for (size_t i = 0; i < b->statements.size(); i++) { - Stmt *s = b->statements[i].get(); - int effective_pos = (b == push_block) ? static_cast(i) : container_pos_in_push_block; - for (auto *op : s->get_operands()) { - if (op == nullptr) - continue; - for (auto *lt : load_tops) { - if (op == lt) { - out_positions.push_back(effective_pos); - break; - } - } - } - if (auto *if_s = s->cast()) { - if (if_s->true_statements) - walk(if_s->true_statements.get(), effective_pos); - if (if_s->false_statements) - walk(if_s->false_statements.get(), effective_pos); - } else if (auto *rf = s->cast()) { - if (rf->body) - walk(rf->body.get(), effective_pos); - } else if (auto *sf = s->cast()) { - if (sf->body) - walk(sf->body.get(), effective_pos); - } - } - }; - walk(push_block, -1); - }; - std::vector consumer_positions; - find_consumers_of_load_tops(consumer_positions); - int max_consumer_pos = -1; - for (int p : consumer_positions) { - if (p > max_consumer_pos) - max_consumer_pos = p; - } - - // Collect all distinct AdStackAllocaStmt referenced via AdStackLoadTopStmt leaves in pushed_val's chain. - std::unordered_set chain_leaf_stacks; - std::unordered_set walked_for_leaves; - std::function collect_leaves = [&](Stmt *s) { - if (walked_for_leaves.count(s)) - return; - walked_for_leaves.insert(s); - if (auto *lt = s->cast()) { - if (auto *as = lt->stack->cast()) - chain_leaf_stacks.insert(as); - return; - } - if (s->is() || s->is() || s->is()) - return; - for (auto *op : s->get_operands()) { - if (op != nullptr) - collect_leaves(op); - } - }; - collect_leaves(pushed_val); - - bool reverse_safe = true; - for (auto *T : chain_leaf_stacks) { - // Find T's body pushes (non-init) in push_block. Only same-block pushes can interfere: pushes in - // outer blocks happen once per outer iteration, so their tops are stable across the inner iter. - for (size_t i = 0; i < push_block->statements.size(); i++) { - Stmt *s = push_block->statements[i].get(); - if (auto *p = s->cast()) { - if (p->stack == T && !is_zero_init_value(p->v)) { - if (static_cast(i) <= max_consumer_pos) { - reverse_safe = false; - break; - } - } - } - } - if (!reverse_safe) - break; - } - if (!reverse_safe) { - continue; - } - - bool is_self_loaded = false; - std::unordered_set visited; - std::function walk = [&](Stmt *s) { - if (is_self_loaded || visited.count(s)) - return; - visited.insert(s); - if (auto *lt = s->cast()) { - if (lt->stack == stack) { - is_self_loaded = true; - return; - } - // load_top of a different stack: the operand chain stops here at this leaf for the self-check. - return; - } - if (s->is() || s->is() || s->is()) { - return; // leaf, not a self-load - } - for (auto *op : s->get_operands()) { - if (op != nullptr) - walk(op); - } - }; - walk(pushed_val); - if (is_self_loaded) { - continue; - } - - // Control-flow-cond consumers: by the time the IR reaches this pass, every IfStmt cond and RangeFor begin/end - // is either a bare `AdStackLoadTopStmt` (loop-carried local promoted via `PromoteSSA2LocalVar` -> - // `AdStackAllocaJudger::visit(IfStmt|RangeForStmt)` -> `ReplaceLocalVarWithStacks`) or a stmt outside any - // adstack-bearing alloca (no load_top in the chain at all). `MakeAdjoint::visit(IfStmt)` (auto_diff.cpp around - // line 2452) caps the bare-load_top case with a 1-push-per-execution snap-stack when the if body itself pushes - // to the cond's backing stack so the reverse cond reads the forward-time value, not a stack top mutated by the - // body. - // - // If we eliminate a stack whose load_top IS the bare cond of an IfStmt / inner RangeFor, the rewrite turns the - // cond / loop-bound into an inlined recomputed SSA chain. The snap-stack guard at the `AdStackLoadTopStmt`-cond - // check in `MakeAdjoint::visit(IfStmt)` stops firing (the cond is no longer bare), and `BackupSSA`'s clone path - // positions a fresh `AdStackLoadTopStmt` of an enclosing loop-carried stack at the consumer's IR location - - // which sits BEFORE the per-iter pops, returning the post-iteration value instead of the iteration-k value the - // cond was originally computed against. Net effect: silent gradient corruption in any if/loop nested inside a - // dynamic for-loop with a loop-carried alloca. - // - // Keep stacks whose load_tops have such consumers. The consumer check below trips only on a load_top of THIS - // stack feeding DIRECTLY into a control-flow-shaping operand of another stmt; that is also the exhaustive set - // of unsafe-elim shapes here, because compound-cond cases (arithmetic over a load_top feeding a cond) reach - // this pass with the cond rewritten into a separate adstack-promoted value via the alloca-promotion pipeline - // above and so look like the bare case from this guard's POV. - // `irpass::analysis::gather_statements` walks via `BasicStmtVisitor`, whose visit overrides for - // IfStmt / RangeForStmt / StructForStmt do not invoke the per-stmt test predicate on the container - // itself - only on stmts inside their bodies. Walk container stmts manually here. - bool has_control_flow_consumer = false; - std::function walk_block = [&](Block *b) { - if (has_control_flow_consumer) - return; - for (auto &owned : b->statements) { - Stmt *s = owned.get(); - if (auto *if_s = s->cast()) { - for (auto *lt : load_tops) { - if (if_s->cond == lt) { - has_control_flow_consumer = true; - return; - } - } - if (if_s->true_statements) - walk_block(if_s->true_statements.get()); - if (has_control_flow_consumer) - return; - if (if_s->false_statements) - walk_block(if_s->false_statements.get()); - } else if (auto *rf = s->cast()) { - for (auto *lt : load_tops) { - if (rf->begin == lt || rf->end == lt) { - has_control_flow_consumer = true; - return; - } - } - if (rf->body) - walk_block(rf->body.get()); - } else if (auto *sf = s->cast()) { - if (sf->body) - walk_block(sf->body.get()); - } else if (auto *off = s->cast()) { - if (off->body) - walk_block(off->body.get()); - if (off->tls_prologue) - walk_block(off->tls_prologue.get()); - if (off->bls_prologue) - walk_block(off->bls_prologue.get()); - if (off->bls_epilogue) - walk_block(off->bls_epilogue.get()); - if (off->tls_epilogue) - walk_block(off->tls_epilogue.get()); - } - } - }; - walk_block(ib); - if (has_control_flow_consumer) { - continue; - } - - // Eligible: rewrite each load_top to use the pushed value directly. The pushed value lives in the - // forward IR and dominates each load_top by SSA construction (the load_top reads what the push wrote; - // both are inside the same loop body in the dynamic-loop case, with the push preceding the load_top). - for (auto *lt : load_tops) { - irpass::replace_all_usages_with(ib, lt, pushed_val); - lt->parent->erase(lt); - } - // Erase init-zero pushes (they only matter if a load could observe them, but the rewriting above just - // routed every load to the body-pushed SSA value), the body push, and the alloca itself. - for (auto *p : init_pushes) { - p->parent->erase(p); - } - body_push->parent->erase(body_push); - stack->parent->erase(stack); - - // Invalidate the cache: the eliminated stack might have been referenced by other recomputable chains - // we evaluated earlier in this pass, but those evaluations only matter for stacks we eliminate THIS - // pass; re-running from scratch on the next iteration recomputes them. - recomputable_cache.clear(); - modified = true; - } - return modified; - } -}; - -class ReverseOuterLoops : public BasicStmtVisitor { - using BasicStmtVisitor::visit; - - private: - explicit ReverseOuterLoops(const std::set &IB) : loop_depth_(0), ib_(IB) { - } - - bool is_ib(Block *block) const { - return std::find(ib_.begin(), ib_.end(), block) != ib_.end(); - } - - // Sibling for-loops inside a non-IB container block execute their reverse-mode companions - // in the container's forward order by default, because MakeAdjoint only touches IB-level bodies - // and nothing else permutes the enclosing order. Reverse-mode AD requires the opposite: if the - // forward body runs `for_A; for_B` and for_B's reverse depends on reads produced by for_A's - // forward run, the reverse pass must execute `rev-for_B; rev-for_A` so for_A's reverse sees the - // adjoints for_B has populated (e.g. `cdof[i]=x[i]; cdofvel[i]=cdof[i]*vel[i]` silently returns - // x.grad=0 otherwise: rev-for_A clears cdof.grad before rev-for_B has populated it). - // - // Naive pairwise swap of for-loop positions is unsafe whenever a non-loop stmt between two - // for-loops feeds the later sibling's SSA operand chain (e.g. a GlobalLoad that supplies a - // dynamic trip count): after the swap, the consumer for-loop ends up before its producer and - // the IR verifier rejects the block. Before swapping, hoist any such producer (and its - // transitive in-block dependencies) to the slot just before the first sibling for-loop. Non-loop - // stmts unrelated to for-loop operands stay at their original indices; memory ordering between - // non-loop stmts is preserved because the hoist keeps them in their original relative order and - // only moves them upward over for-loops (which produce no SSA value and cannot be the source of - // a missed memory read for a non-loop that gets hoisted above them). - // - // The top-level kernel block is handled by `reverse_segments` before this pass, so we only - // reorder inside nested non-IB blocks here. - static void reverse_for_loop_order_in_place(Block *block) { - const int n = (int)block->statements.size(); - std::vector for_indices; - for (int i = 0; i < n; ++i) { - Stmt *s = block->statements[i].get(); - if (s->is() || s->is()) { - for_indices.push_back(i); - } - } - if (for_indices.size() < 2) { - return; - } - const int first_for = for_indices.front(); - - std::unordered_map pos_of; - pos_of.reserve(n); - for (int i = 0; i < n; ++i) { - pos_of[block->statements[i].get()] = i; - } - - // Walk the SSA operand graph of every for-loop (restricted to this block). Any in-block stmt - // that (a) the operand closure reaches and (b) sits at or after `first_for` gets flagged for - // hoisting: after swap, that stmt must precede every for-loop, not just the ones it feeds. - std::unordered_set must_hoist; - std::vector stack; - auto push_if_internal = [&](Stmt *s) { - if (s == nullptr) { - return; - } - auto it = pos_of.find(s); - if (it == pos_of.end() || it->second < first_for) { - return; - } - if (must_hoist.insert(s).second) { - stack.push_back(s); - } - }; - // Seed the hoist frontier from both the for-loop's direct SSA operands (`begin`, `end`) and - // from every stmt nested inside the for-loop's body that references an outer-block stmt as a - // free variable. The body-use gather is what catches the case where the later sibling - // for-loop consumes a non-loop outer-block stmt `S` inside its body (e.g. `for_B: body reads - // S`) rather than through `for_B`'s range bound: `RangeForStmt::get_operands()` returns only - // `{begin, end}`, so without walking the body `S` would miss `must_hoist`, the pairwise swap - // would place `for_B` ahead of `S`, and the IR verifier would reject the SSA violation. - for (int fi : for_indices) { - for (Stmt *op : block->statements[fi]->get_operands()) { - push_if_internal(op); - } - Stmt *for_stmt = block->statements[fi].get(); - irpass::analysis::gather_statements(for_stmt, [&](Stmt *body_stmt) { - for (Stmt *op : body_stmt->get_operands()) { - push_if_internal(op); - } - return false; - }); - } - while (!stack.empty()) { - Stmt *s = stack.back(); - stack.pop_back(); - for (Stmt *op : s->get_operands()) { - push_if_internal(op); - } - } - // For-loops themselves end up in `must_hoist` only because their own operand-closure reached - // them; they do not get hoisted as non-loop producers - strip them here to keep `must_hoist` - // to "non-loop stmts that need to move above all for-loops". - for (int fi : for_indices) { - must_hoist.erase(block->statements[fi].get()); - } - - std::vector> new_stmts; - new_stmts.reserve(n); - // Stmts strictly before `first_for` keep their original slot. - for (int i = 0; i < first_for; ++i) { - new_stmts.push_back(std::move(block->statements[i])); - } - // Hoisted non-loop stmts slot in here, in their original relative order. - for (int i = first_for; i < n; ++i) { - if (must_hoist.count(block->statements[i].get()) != 0) { - new_stmts.push_back(std::move(block->statements[i])); - } - } - // Remainder (for-loops and non-hoisted non-loops) in original order, with for-loops swapped - // pairwise inside this suffix. - std::vector> suffix; - std::vector suffix_for_positions; - for (int i = first_for; i < n; ++i) { - auto &sp = block->statements[i]; - if (!sp) { - continue; - } - bool is_for = sp->is() || sp->is(); - if (is_for) { - suffix_for_positions.push_back((int)suffix.size()); - } - suffix.push_back(std::move(sp)); - } - for (int lo = 0, hi = (int)suffix_for_positions.size() - 1; lo < hi; ++lo, --hi) { - std::swap(suffix[suffix_for_positions[lo]], suffix[suffix_for_positions[hi]]); - } - for (auto &s : suffix) { - new_stmts.push_back(std::move(s)); - } - - QD_ASSERT((int)new_stmts.size() == n); - block->statements.clear(); - for (auto &s : new_stmts) { - block->statements.push_back(std::move(s)); - } - } - - void visit(StructForStmt *stmt) override { - loop_depth_ += 1; - if (!is_ib(stmt->body.get())) { - stmt->body->accept(this); - reverse_for_loop_order_in_place(stmt->body.get()); - } - loop_depth_ -= 1; - } - - void visit(RangeForStmt *stmt) override { - if (loop_depth_ >= 1) { - stmt->reversed = !stmt->reversed; - } - loop_depth_ += 1; - if (!is_ib(stmt->body.get())) { - stmt->body->accept(this); - reverse_for_loop_order_in_place(stmt->body.get()); - } - loop_depth_ -= 1; - } - - // Deliberately no `visit(IfStmt *)` override, although sibling for-loops can live directly inside an if-branch - // block (`true_statements` / `false_statements`) the same way they live inside a for-body. The default - // `BasicStmtVisitor::visit(IfStmt *)` recurses into both branches so inner `RangeForStmt::body`s still get the - // sibling-reorder treatment via the range-for visitor above, but `reverse_for_loop_order_in_place` is never - // invoked on the branch block itself. That is intentional: `MakeAdjoint::visit(IfStmt *)` below emits the adjoint - // if-stmt by iterating each branch's statements in reverse order (`for (int i = size - 1; i >= 0; --i)` in its - // `true_statements` / `false_statements` loops), which achieves the same sibling-for reordering effect that the - // missing override here would provide. Overriding `visit(IfStmt)` in this class is therefore a no-op on the - // generated adjoint code. Keep the comment rather than the override so the visitor-coverage gap is documented. - - int loop_depth_; - std::set ib_; - - public: - static void run(IRNode *root, const std::set &IB) { - ReverseOuterLoops pass(IB); - root->accept(&pass); - } -}; - -// Base class for both reverse (make adjoint) and forward (make dual) mode -class ADTransform : public IRVisitor { - protected: - Stmt *constant(float32 x, DataType dtype = PrimitiveType::unknown) { - dtype.set_is_pointer(false); - if (!dtype->is()) - return insert(TypedConstant(x)); - - auto tensor_type = dtype->as(); - auto num_elements = tensor_type->get_num_elements(); - std::vector values; - for (int i = 0; i < num_elements; i++) { - values.push_back(insert(TypedConstant(x))); - } - auto matrix_init_stmt = insert(values); - matrix_init_stmt->ret_type = tensor_type; - return matrix_init_stmt; - } - - // utils - Stmt *sgn(Stmt *inp) { - return insert(UnaryOpType::sgn, load(inp)); - } - - // utils - Stmt *negate(Stmt *inp) { - return insert(UnaryOpType::neg, load(inp)); - } - - Stmt *sqrt(Stmt *inp) { - return insert(UnaryOpType::sqrt, load(inp)); - } - - Stmt *rsqrt(Stmt *inp) { - return insert(UnaryOpType::rsqrt, load(inp)); - } - - Stmt *mul(Stmt *op1, Stmt *op2) { - return insert(BinaryOpType::mul, load(op1), load(op2)); - } - - Stmt *sqr(Stmt *op1) { - return mul(op1, op1); - } - - Stmt *add(Stmt *op1, Stmt *op2) { - return insert(BinaryOpType::add, load(op1), load(op2)); - } - - Stmt *cmp_lt(Stmt *op1, Stmt *op2) { - return insert(BinaryOpType::cmp_lt, load(op1), load(op2)); - } - - Stmt *sub(Stmt *op1, Stmt *op2) { - return insert(BinaryOpType::sub, load(op1), load(op2)); - } - - Stmt *div(Stmt *op1, Stmt *op2) { - return insert(BinaryOpType::div, load(op1), load(op2)); - } - - Stmt *sel(Stmt *op1, Stmt *op2, Stmt *op3) { - return insert(TernaryOpType::select, load(op1), load(op2), load(op3)); - } - - Stmt *cos(Stmt *op1) { - return insert(UnaryOpType::cos, load(op1)); - } - - Stmt *sin(Stmt *op1) { - return insert(UnaryOpType::sin, load(op1)); - } - - Stmt *log(Stmt *op1) { - return insert(UnaryOpType::log, load(op1)); - } - - Stmt *tan(Stmt *op1) { - return insert(UnaryOpType::tan, load(op1)); - } - - Stmt *tanh(Stmt *op1) { - return insert(UnaryOpType::tanh, load(op1)); - } - - Stmt *exp(Stmt *op1) { - return insert(UnaryOpType::exp, load(op1)); - } - - Stmt *pow(Stmt *op1, Stmt *op2) { - return insert(BinaryOpType::pow, load(op1), load(op2)); - } - - public: - virtual Stmt *insert_grad_stmt(std::unique_ptr &&stmt) = 0; - - template - Stmt *insert(Args &&...args) { - return insert_grad_stmt(Stmt::make(args...)); - } - - template - Stmt *insert_const_for_grad(const DataType &dtype, Stmt *stmt, const T &val) { - auto zero = insert(TypedConstant(dtype.ptr_removed().get_element_type(), val)); - if (dtype.ptr_removed()->is()) { - auto t_dtype = dtype.ptr_removed()->as(); - std::vector values(t_dtype->get_num_elements(), zero); - zero = insert(values); - zero->ret_type = dtype.ptr_removed(); - } - return zero; - } - - void visit(AllocaStmt *alloca) override { - // do nothing. - } - - void visit(AdStackAllocaStmt *alloca) override { - // do nothing. - } - - void visit(ArgLoadStmt *stmt) override { - // do nothing. - } - - void visit(GetElementStmt *stmt) override { - // do nothing - } - - void visit(LoopIndexStmt *stmt) override { - // do nothing. - } - - void visit(MatrixPtrStmt *stmt) override { - // do nothing. - } - - void visit(PrintStmt *print_stmt) override { - // do nothing - } - - void visit(ConstStmt *const_stmt) override { - // do nothing - } - - void visit(ReturnStmt *stmt) override { - // do nothing - } - - void visit(WhileControlStmt *stmt) override { - QD_NOT_IMPLEMENTED - } - - void visit(ContinueStmt *stmt) override { - QD_NOT_IMPLEMENTED; - } - - void visit(WhileStmt *stmt) override { - QD_NOT_IMPLEMENTED - } - - void visit(GlobalPtrStmt *stmt) override { - // do nothing - } - - void visit(ExternalPtrStmt *stmt) override { - // do nothing - } - - void visit(ExternalTensorShapeAlongAxisStmt *stmt) override { - // do nothing - } - - void visit(DecorationStmt *stmt) override { - // do nothing - } - - Stmt *load(Stmt *alloc) { - QD_ASSERT(alloc != nullptr); - if (alloc->is() || alloc->is()) { - return insert(alloc); - } else { - // non alloca - return alloc; - } - } - - bool gradients_stopped(GlobalLoadStmt *stmt, SNode *snode) { - for (auto block = stmt->parent; block; block = block->parent_block()) { - for (auto s : block->stop_gradients) { - if (s == snode) { - return true; - } - } - } - return false; - } - - void visit(AssertStmt *stmt) override { - // do nothing - } - - void visit(RangeAssumptionStmt *stmt) override { - // do nothing - } - - void visit(LinearizeStmt *stmt) override { - // do nothing - } - - void visit(IntegerOffsetStmt *stmt) override { - // do nothing - } - - void visit(RandStmt *stmt) override { - QD_ERROR("RandStmt not supported in AutoDiff for now."); - } -}; - -// Generate the adjoint version of an independent block -class MakeAdjoint : public ADTransform { - public: - using ADTransform::visit; - Block *current_block; - Block *alloca_block; - // Backup the forward pass (the forward pass might be modified during the - // MakeAdjoint) for search whether a GlobalLoadStmt is inside a for-loop when - // allocating adjoint (see the function `adjoint`) Should be stored - // 1. Before entering a for-loop body - // 2. Before entering a if-stmt - // Should be restored after processing every statement in the two cases above - Block *forward_backup; - // IB root: stays constant across visitor recursion. Used when we need to allocate - // persistent storage that must survive enclosing for-loop iterations (e.g. the - // dedicated ad-stacks that snapshot IfStmt conds in visit(IfStmt)). - Block *ib_root; - std::map adjoint_stmt; - - explicit MakeAdjoint(Block *block) { - current_block = nullptr; - alloca_block = block; - forward_backup = block; - ib_root = block; - } - - static void run(Block *block) { - auto p = MakeAdjoint(block); - block->accept(&p); - } - - // Does `if_stmt`'s true/false body contain any AdStackPushStmt targeting `stack`? Recursive to - // catch pushes nested inside further control flow (if-in-if, if-in-for). Used by visit(IfStmt) - // to gate cond-snapshotting. Must be narrow: snapshotting every if-stmt would add an - // AdStackAllocaStmt per if, and determine_ad_stack_size cannot size stacks whose push/pop pair - // is only reachable through branches its Bellman-Ford walk considers "unreached" -- codegen then - // aborts with "Adaptive autodiff stack's size should have been determined" and the extras also - // spam "Unused autodiff stack should have been eliminated" for every untouched snap stack. Only - // when the body actually pushes onto the cond's backing stack does BackupSSA's reverse-time - // clone of load_top read a post-body value rather than the forward cond (the real bug); in every - // other case the clone is already correct and a snapshot would be dead weight. - static bool block_pushes_to_stack(Block *block, Stmt *stack) { - if (!block) - return false; - for (auto &stmt : block->statements) { - if (auto *push = stmt->cast()) { - if (push->stack == stack) - return true; - } - if (auto *inner_if = stmt->cast()) { - if (block_pushes_to_stack(inner_if->true_statements.get(), stack)) - return true; - if (block_pushes_to_stack(inner_if->false_statements.get(), stack)) - return true; - } - if (auto *inner_for = stmt->cast()) { - if (block_pushes_to_stack(inner_for->body.get(), stack)) - return true; - } - if (auto *inner_for = stmt->cast()) { - if (block_pushes_to_stack(inner_for->body.get(), stack)) - return true; - } - } - return false; - } - - static bool body_pushes_to_stack(IfStmt *if_stmt, Stmt *stack) { - return block_pushes_to_stack(if_stmt->true_statements.get(), stack) || - block_pushes_to_stack(if_stmt->false_statements.get(), stack); - } - - // TODO: current block might not be the right block to insert adjoint - // instructions! - void visit(Block *block) override { - std::vector statements; - // always make a copy since the list can be modified. - for (auto &stmt : block->statements) { - statements.push_back(stmt.get()); - } - std::reverse(statements.begin(), statements.end()); // reverse-mode AD... - for (auto stmt : statements) { - current_block = block; - stmt->accept(this); - } - } - - Stmt *insert_grad_stmt(std::unique_ptr &&stmt) override { - auto ptr = stmt.get(); - current_block->insert(std::move(stmt), -1); - return ptr; - } - - // Accumulate [value] to the adjoint of [primal] - void accumulate(Stmt *primal, Stmt *value) { - auto alloca_ = adjoint(primal); - if (!alloca_ || alloca_->is()) { - return; // primal may be int variable - } - if (alloca_->is()) { - auto alloca = alloca_->cast(); - if (is_real(alloca->ret_type.get_element_type())) { - insert(alloca, load(value)); - } - } else { - QD_ASSERT(alloca_->is()); - auto alloca = alloca_->as(); - auto local_load = insert(alloca); - local_load->ret_type = alloca->ret_type.ptr_removed(); - insert(alloca, add(local_load, value)); - } - } - - Stmt *adjoint(Stmt *stmt) { - DataType adjoint_dtype = stmt->ret_type.ptr_removed(); - if (adjoint_dtype->is()) { - DataType prim_dtype = PrimitiveType::f32; - if (is_real(adjoint_dtype.get_element_type())) { - prim_dtype = adjoint_dtype.get_element_type(); - } - adjoint_dtype = - TypeFactory::get_instance().get_tensor_type(adjoint_dtype->as()->get_shape(), prim_dtype); - } else if (stmt->is()) { - // pass - } else if (!is_real(adjoint_dtype) || stmt->is()) { - return constant(0); - } - - if (adjoint_stmt.find(stmt) == adjoint_stmt.end()) { - // normal SSA cases - - // create the alloca - // auto alloca = - // Stmt::make(get_current_program().config.gradient_dt); - // maybe it's better to use the statement data type than the default type - auto alloca = Stmt::make(adjoint_dtype); - adjoint_stmt[stmt] = alloca.get(); - - // We need to insert the alloca in the block of GlobalLoadStmt when the - // GlobalLoadStmt is not inside the currently-processed range-for. - // Code sample: - // a and b require grad - // Case 1 (GlobalLoadStmt is outside the for-loop, compute 5 times and - // accumulate once, alloca history value is needed): - // for i in range(5): - // p = a[i] - // q = b[i] - // for _ in range(5) - // q += p - - // Case 2 (GlobalLoadStmt is inside the for-loop, compute once and - // accumulate immediately, alloca history value can be discarded): - // for i in range(5): - // q = b[i] - // for _ in range(5) - // q += a[i] - if (stmt->is() && forward_backup->locate(stmt->as()) == -1) { - // Case 1: the GlobalLoadStmt lives in a block outside the currently-processed range-for iteration. Its - // adjoint must persist across all iterations of the inner reversed loop, so the alloca cannot live in the - // current alloca_block (which would be the inner reversed loop body). Walk up from the primal's enclosing - // block until we hit one whose owning statement unconditionally dominates both the forward and the reverse - // code (a loop / offloaded / kernel body, not an if / while body): visit(IfStmt) emits the reverse code - // into a brand new sibling IfStmt, not back into the forward if-body, so an alloca placed inside the - // forward branch is SSA-invalid from the reverse branch's point of view and gets DCE'd into silently-zero - // gradients. - Block *target = stmt->as()->parent; - while (target != nullptr) { - Stmt *parent_stmt = target->parent_stmt(); - if (parent_stmt == nullptr || parent_stmt->is() || parent_stmt->is() || - parent_stmt->is() || parent_stmt->is()) { - break; - } - target = parent_stmt->parent; - } - // Reaching a null target means the primal's enclosing-block chain is broken (an unparented block). Falling - // back to alloca_block here would silently restore the pre-fix bug (adjoint eliminated as DCE on the - // reverse branch); hard-assert instead so malformed IR surfaces loudly. - QD_ASSERT(target != nullptr); - target->insert(std::move(alloca), 0); - } else { - alloca_block->insert(std::move(alloca), 0); - } - } - return adjoint_stmt[stmt]; - } - - // For ops in `NonLinearOps::unary_collections` the reverse formula must NOT read the forward `stmt` - // directly - only `stmt->operand` (adstack-backed), `adjoint(stmt)`, and constants. BackupSSA spills the - // forward stmt to a single plain alloca overwritten each iteration, so reading `stmt` from a reversed - // dynamic loop would use the last-iteration value regardless of which reverse iteration is running. This - // helper walks the value-tree at IR-transform time and asserts the invariant; paired with the Python - // meta-test `test_unary_collections_audit` (which catches "forgot to add to unary_collections"), it - // covers both halves of the class of bugs the per-op audit used to miss. - void accumulate_unary_operand_checked(UnaryOpStmt *stmt, Stmt *value) { - if (NonLinearOps::unary_collections.find(stmt->op_type) != NonLinearOps::unary_collections.end()) { - std::unordered_set visited; - std::function walk = [&](const Stmt *s) { - if (!s || visited.count(s)) - return; - visited.insert(s); - QD_ASSERT_INFO(s != stmt, - "MakeAdjoint adjoint formula for UnaryOpType::{} reads the forward stmt directly. It " - "must read `stmt->operand` (adstack-backed) instead - see the tan/tanh/exp branches " - "for the recompute pattern.", - unary_op_type_name(stmt->op_type)); - for (auto *op : s->get_operands()) - walk(op); - }; - walk(value); - } - accumulate(stmt->operand, value); - } - - void visit(UnaryOpStmt *stmt) override { - auto acc = [&](Stmt *value) { accumulate_unary_operand_checked(stmt, value); }; - if (stmt->op_type == UnaryOpType::floor || stmt->op_type == UnaryOpType::ceil) { - // do nothing - } else if (stmt->op_type == UnaryOpType::neg) { - accumulate(stmt->operand, negate(adjoint(stmt))); - } else if (stmt->op_type == UnaryOpType::abs) { - acc(mul(adjoint(stmt), sgn(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::sin) { - acc(mul(adjoint(stmt), cos(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::cos) { - acc(negate(mul(adjoint(stmt), sin(stmt->operand)))); - } else if (stmt->op_type == UnaryOpType::tan) { - // d/dx tan(x) = 1 + tan(x)^2. Recompute tan(operand) rather than reusing the forward value: the primal is - // per-iteration inside dynamic loops but BackupSSA only spills forward values to a single plain alloca, so - // reading the forward tan would use the last-iteration value in the reversed loop. The operand, in contrast, - // rides the adstack through its LocalLoad, so a fresh tan on it is per-iteration correct. - acc(mul(adjoint(stmt), add(constant(1, stmt->ret_type), sqr(tan(stmt->operand))))); - } else if (stmt->op_type == UnaryOpType::tanh) { - // Recompute tanh(operand) in the reverse pass instead of reusing the forward stmt value. In dynamic loops - // BackupSSA spills the forward stmt to a single plain alloca overwritten each iteration, so the reversed - // loop would read the last-iteration tanh for every backward step. The operand rides the adstack through - // LocalLoad, so a fresh tanh on it is per-iteration correct. Trade-off: tanh is evaluated twice per - // iteration (once forward, once backward); caching the forward value on the adstack is a future optimization. - acc(mul(adjoint(stmt), sub(constant(1, stmt->ret_type), sqr(tanh(stmt->operand))))); - } else if (stmt->op_type == UnaryOpType::asin) { - acc(mul(adjoint(stmt), - div(constant(1, stmt->ret_type), sqrt(sub(constant(1, stmt->ret_type), sqr(stmt->operand)))))); - } else if (stmt->op_type == UnaryOpType::acos) { - acc(mul(adjoint(stmt), - negate(div(constant(1, stmt->ret_type), sqrt(sub(constant(1, stmt->ret_type), sqr(stmt->operand))))))); - } else if (stmt->op_type == UnaryOpType::exp) { - // See the tanh case above: recompute exp on the adstack-backed operand so the reversed loop sees the - // per-iteration value rather than the last-forward value spilled by BackupSSA. Same trade-off as tanh: exp - // is evaluated twice per iteration (once forward, once backward); caching the forward value on the adstack - // is a future optimization. - acc(mul(adjoint(stmt), exp(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::log) { - // No recompute workaround needed: the reverse formula `1 / operand` reads `stmt->operand` directly (which - // is adstack-backed via LocalLoad inside dynamic loops), not the forward `log(operand)` stmt value. - acc(div(adjoint(stmt), stmt->operand)); - } else if (stmt->op_type == UnaryOpType::sqrt) { - // No recompute workaround needed: the reverse formula reads `stmt->operand` (adstack-backed via LocalLoad - // inside dynamic loops, gated on `unary_collections` membership) and recomputes `sqrt(operand)` from it, - // not the forward `sqrt(operand)` stmt value. Unlike tanh/exp this case was already operand-based before - // the recompute fix landed; the structure mirrors log above. - acc(mul(adjoint(stmt), div(constant(0.5f, stmt->ret_type), sqrt(stmt->operand)))); - } else if (stmt->op_type == UnaryOpType::rsqrt) { - acc(mul(adjoint(stmt), - mul(constant(-0.5f, stmt->ret_type), pow(rsqrt(stmt->operand), constant(3, stmt->ret_type))))); - } else if (stmt->op_type == UnaryOpType::cast_value) { - if (is_real(stmt->cast_type.get_element_type()) && is_real(stmt->operand->ret_type.get_element_type())) { - accumulate(stmt->operand, adjoint(stmt)); - } - } else if (stmt->op_type == UnaryOpType::logic_not) { - // do nothing - } else { - QD_P(unary_op_type_name(stmt->op_type)); - QD_NOT_IMPLEMENTED; - } - } - - void visit(BinaryOpStmt *bin) override { - if (bin->op_type == BinaryOpType::add) { - accumulate(bin->lhs, adjoint(bin)); - accumulate(bin->rhs, adjoint(bin)); - } else if (bin->op_type == BinaryOpType::sub) { - accumulate(bin->lhs, adjoint(bin)); - accumulate(bin->rhs, negate(adjoint(bin))); - } else if (bin->op_type == BinaryOpType::mul) { - // d (x * y) = y * dx + x * dy - accumulate(bin->lhs, mul(adjoint(bin), bin->rhs)); - accumulate(bin->rhs, mul(adjoint(bin), bin->lhs)); - } else if (bin->op_type == BinaryOpType::mod) { - // Do nothing - } else if (bin->op_type == BinaryOpType::div) { - accumulate(bin->lhs, div(adjoint(bin), bin->rhs)); - accumulate(bin->rhs, negate(div(mul(adjoint(bin), bin->lhs), mul(bin->rhs, bin->rhs)))); - } else if (bin->op_type == BinaryOpType::atan2) { - auto numerator = add(sqr(bin->lhs), sqr(bin->rhs)); - accumulate(bin->lhs, div(mul(adjoint(bin), bin->rhs), numerator)); - accumulate(bin->rhs, negate(div(mul(adjoint(bin), bin->lhs), numerator))); - } else if (bin->op_type == BinaryOpType::pow) { - // d (x ^ y) = x ^ (y-1) * (y * dx + log(x) * x * dy) - auto common_coeff = pow(bin->lhs, sub(bin->rhs, constant(1, bin->ret_type))); // x ^ (y-1) - accumulate(bin->lhs, mul(adjoint(bin), mul(bin->rhs, common_coeff))); - accumulate(bin->rhs, mul(adjoint(bin), mul(log(bin->lhs), mul(bin->lhs, common_coeff)))); - } else if (bin->op_type == BinaryOpType::min || bin->op_type == BinaryOpType::max) { - auto cmp = bin->op_type == BinaryOpType::min ? cmp_lt(bin->lhs, bin->rhs) : cmp_lt(bin->rhs, bin->lhs); - auto zero = insert_const_for_grad(bin->ret_type, bin, 0); - accumulate(bin->lhs, sel(cmp, adjoint(bin), zero)); - accumulate(bin->rhs, sel(cmp, zero, adjoint(bin))); - } else if (bin->op_type == BinaryOpType::floordiv) { - // do nothing - } else if (is_comparison(bin->op_type) || is_bit_op(bin->op_type) || binary_is_logical(bin->op_type)) { - // do nothing - - } else { - QD_WARN("gradient of binary op {}\n{}", binary_op_type_name(bin->op_type), bin->get_tb()); - QD_NOT_IMPLEMENTED; - } - } - - void visit(TernaryOpStmt *stmt) override { - QD_ASSERT(stmt->op_type == TernaryOpType::select); - auto zero = insert_const_for_grad(stmt->ret_type, stmt, 0); - accumulate(stmt->op2, insert(TernaryOpType::select, stmt->op1, load(adjoint(stmt)), zero)); - accumulate(stmt->op3, insert(TernaryOpType::select, stmt->op1, zero, load(adjoint(stmt)))); - } - - void visit(IfStmt *if_stmt) override { - // Snapshot a stack-backed forward cond into a dedicated 1-push-per-if-execution ad-stack, - // but only when the cond's backing stack is also pushed inside the if body (e.g. short-circuit - // lowering pushes the rhs of `&&` onto the same stack that holds the cond). Without this, - // BackupSSA's clone of `if_stmt->cond` in the reverse block reads the cond stack AFTER the - // body-pushes rather than the forward-time cond value - the reverse IfStmt flips, pop counts - // drift, and gradients come out silently zero. A dedicated stack has exactly one push per - // forward if-execution, so the reverse load_top matches the forward cond. - // - // Guarded by the body-push check because snapshotting indiscriminately adds AdStackAllocaStmts - // that go through `determine_ad_stack_size` unused on every other if-stmt in the kernel - the - // adaptive-size pass emits "Unused autodiff stack should have been eliminated" warnings and - // the codegen step then fails with "Adaptive autodiff stack's size should have been determined". - Stmt *reverse_cond = if_stmt->cond; - AdStackAllocaStmt *snap_stack_ptr = nullptr; - // Narrow guard: only the bare `AdStackLoadTopStmt` shape needs the explicit snapshot below. A compound cond (e.g. - // `cmp_lt(load_top(x_stack) + 0.1, threshold)` from `if x + 0.1 < threshold` when `x` has been promoted to an - // adstack by `ReplaceLocalVarWithStacks`) reaches this visitor as a BARE `AdStackLoadTopStmt` cond anyway - the - // cmp / arithmetic value goes through `PromoteSSA2LocalVar`'s required-defs set (the IfStmt cond path adds the - // cond's value-producing op), then `AdStackAllocaJudger::visit(IfStmt)` (around line 763) marks its alloca - // stack-needed because it feeds the cond, then `ReplaceLocalVarWithStacks` promotes the alloca to an adstack. By - // the time control reaches here, `if_stmt->cond` is `AdStackLoadTopStmt` of that snap-promoted adstack and the - // body of the if does NOT push to that stack (the cmp value is pushed once just before the IfStmt, never inside), - // so the `body_pushes_to_stack` guard below is false and we correctly skip the additional snap-stack. Per-iter - // cond values are preserved by the alloca-promotion pipeline. - // - // The bare-`AdStackLoadTopStmt` case the snap-stack below handles is the OTHER shape: a load_top whose backing - // stack IS pushed to inside the if body (e.g. short-circuit lowering of `&&` pushes the rhs onto the same stack - // that holds the cond). Without the snap-stack, `BackupSSA::generic_visit`'s clone-branch for `AdStackLoadTopStmt` - // emits a fresh `AdStackLoadTopStmt` at the reverse cursor, which then reads the post-body-push top and sees the - // wrong cond value. The snap-stack here decouples the cond value from the body's pushes by capturing it once just - // before the IfStmt and reading it back at the matching reverse cursor. - // - // Earlier comments here claimed that compound conds rely on `BackupSSA::generic_visit`'s `load(op)` else-branch - // ("single alloca, last-write-wins" spill) for correctness. That description was incomplete: the alloca-promotion- - // to-adstack pipeline catches compound conds before they reach BackupSSA, so the spill is a FALLBACK for shapes - // the alloca-promotion missed, not the load-bearing path for compound conds in general. Verified empirically on - // 2026-05-04 with a loop-carried local v + compound cond `v + 0.1 > threshold` + nonlinear use of v in the if - // body: gradients match analytic on origin/main and on C3. - if (if_stmt->cond->is()) { - auto *cond_stack = if_stmt->cond->as()->stack->as(); - if (body_pushes_to_stack(if_stmt, cond_stack)) { - auto cond_type = if_stmt->cond->ret_type.ptr_removed(); - // Size the snap stack the same way as the cond stack it mirrors: one forward push per - // if-execution matched by one reverse pop. Reusing cond_stack->max_size keeps the snap - // stack exempt from `determine_ad_stack_size` when the cond stack itself was built with a - // fixed size, which is always true when ReplaceLocalVarWithStacks ran with a non-zero - // `ad_stack_size` (the only configuration we currently support for stack-based reverse AD). - auto snap_stack = Stmt::make(cond_type, cond_stack->max_size); - snap_stack_ptr = snap_stack->as(); - // Allocate at the IB root so the stack persists across enclosing for-loop iterations. - ib_root->insert(std::move(snap_stack), 0); - // Per-execution forward push of the cond value, just before the forward if-stmt. No - // initial zero push: the reverse load_top always runs after a matching forward push, so - // leaving the stack empty at entry is both correct and avoids a dead store that DSE would - // otherwise drop (and that `determine_ad_stack_size` would then miscount). - if_stmt->insert_before_me(Stmt::make(snap_stack_ptr, if_stmt->cond)); - // Per-execution reverse load of the snapshotted cond, emitted in the current reverse block. - reverse_cond = insert(snap_stack_ptr); - reverse_cond->ret_type = cond_type; - } - } - - auto new_if = Stmt::make_typed(reverse_cond); - if (if_stmt->true_statements) { - new_if->set_true_statements(std::make_unique()); - auto old_current_block = current_block; - // Backup forward pass - forward_backup = if_stmt->true_statements.get(); - - current_block = new_if->true_statements.get(); - for (int i = if_stmt->true_statements->statements.size() - 1; i >= 0; i--) { - if_stmt->true_statements->statements[i]->accept(this); - // Restore forward pass - forward_backup = if_stmt->true_statements.get(); - } - - current_block = old_current_block; - } - if (if_stmt->false_statements) { - new_if->set_false_statements(std::make_unique()); - auto old_current_block = current_block; - - // Backup forward pass - forward_backup = if_stmt->false_statements.get(); - - current_block = new_if->false_statements.get(); - for (int i = if_stmt->false_statements->statements.size() - 1; i >= 0; i--) { - if_stmt->false_statements->statements[i]->accept(this); - // Restore forward pass - forward_backup = if_stmt->false_statements.get(); - } - current_block = old_current_block; - } - insert_grad_stmt(std::move(new_if)); - if (snap_stack_ptr) { - // One pop per reverse if-execution, paired with the forward push above. - insert(snap_stack_ptr); - } - } - - void visit(RangeForStmt *for_stmt) override { - auto new_for = for_stmt->clone(); - auto new_for_ptr = new_for->as(); - new_for_ptr->reversed = !new_for_ptr->reversed; - insert_grad_stmt(std::move(new_for)); - const int len = new_for_ptr->body->size(); - - for (int i = 0; i < len; i++) { - new_for_ptr->body->erase(0); - } - - std::vector statements; - // always make a copy since the list can be modified. - for (auto &stmt : for_stmt->body->statements) { - statements.push_back(stmt.get()); - } - std::reverse(statements.begin(), statements.end()); // reverse-mode AD... - auto old_alloca_block = alloca_block; - auto old_current_block = current_block; - auto old_forward_backup = forward_backup; // store the block which is not inside the current IB, - // such as outer most loop - // Backup the forward pass - forward_backup = for_stmt->body.get(); - for (auto stmt : statements) { - alloca_block = new_for_ptr->body.get(); - current_block = new_for_ptr->body.get(); - stmt->accept(this); - // Restore the forward pass - forward_backup = for_stmt->body.get(); - } - // Restore current_block. Missing here before: if this RangeForStmt is visited from within another compound - // stmt (notably visit(IfStmt)), that outer visitor will continue iterating its own body in reverse after we - // return and emit further reverse stmts. Without this restore those emissions land in the reversed-for's - // body instead of the outer block, producing silently-wrong gradients whenever a runtime-guarded if wraps a - // for-loop with loop-carried variables (the reverse loop body ends up over-popping the adstack and emitting - // the x.grad accumulation on every iteration instead of once). - current_block = old_current_block; - forward_backup = old_forward_backup; - alloca_block = old_alloca_block; - } - - void visit(StructForStmt *for_stmt) override { - // Save/restore mirrors visit(RangeForStmt) above. Rationale: visit(Block) inside `body->accept(this)` - // sets current_block = for_stmt->body at the start of every iteration, so on return current_block - // points at the struct-for's body. An enclosing compound visitor (e.g. visit(IfStmt)) that resumes - // iterating its children in reverse after this StructForStmt needs current_block and alloca_block to - // still be its own, not this for's; otherwise subsequent reverse emissions land inside the struct-for - // body and any adjoint alloca lives in a block the enclosing if-branch cannot reach. forward_backup - // must be saved too because `visit(IfStmt)` mutates it without restoring, so a nested if inside the - // struct-for body leaves it pointing at the if-branch block, which then survives past this visitor and - // mis-routes `adjoint()` on GlobalLoadStmts for later siblings at the enclosing scope. - auto old_alloca_block = alloca_block; - auto old_current_block = current_block; - auto old_forward_backup = forward_backup; - alloca_block = for_stmt->body.get(); - for_stmt->body->accept(this); - current_block = old_current_block; - alloca_block = old_alloca_block; - forward_backup = old_forward_backup; - } - - // Equivalent to AdStackLoadTopStmt when no stack is needed - void visit(LocalLoadStmt *stmt) override { - // QD_ASSERT(!needs_grad(stmt->ret_type)); - if (is_real(stmt->ret_type.get_element_type())) - accumulate(stmt->src, load(adjoint(stmt))); - } - - // Equivalent to AdStackPushStmt when no stack is needed - void visit(LocalStoreStmt *stmt) override { - accumulate(stmt->val, load(adjoint(stmt->dest))); - - // Clear the adjoint of the dest after local store, - // Because LocalStoreStmt overwrites the dest, - // 1. If the alloca is inside a loop, the adjoint of this alloca of this - // iteration should be cleared after this iteration has been done - // 2. If the alloca serves as the dest of multiple LocalStoreStmt, only the - // last LocalStoreStmt should be taken account of - auto dest_type = stmt->dest->ret_type.ptr_removed(); - if (is_real(dest_type.get_element_type())) { - auto dtype = dest_type; - auto zero = insert_const_for_grad(dtype, stmt, 0); - insert(adjoint(stmt->dest), zero); - } - } - - void visit(AdStackLoadTopStmt *stmt) override { - if (is_real(stmt->ret_type.get_element_type())) - insert(stmt->stack, load(adjoint(stmt))); - } - - void visit(AdStackPushStmt *stmt) override { - accumulate(stmt->v, insert(stmt->stack)); - insert(stmt->stack); - } - - void visit(GlobalLoadStmt *stmt) override { - // issue global store to adjoint - - if (stmt->src->is() || - (stmt->src->is() && stmt->src->as()->origin->is())) { - ExternalPtrStmt *src = nullptr; - bool is_ptr_offset = false; - if (stmt->src->is()) { - src = stmt->src->as()->origin->as(); - is_ptr_offset = true; - } else { - src = stmt->src->as(); - } - auto arg = src->base_ptr->as(); - if (arg->ret_type.ptr_removed()->as()->elements().size() > TypeFactory::GRAD_PTR_POS_IN_NDARRAY) { - QD_ASSERT_INFO(!src->is_grad, - "Cannot automatically differentiate through a grad " - "tensor, if you really want to do that, pass the grad " - "tensor into the kernel directly"); - auto adj_ptr = - insert(src->base_ptr, src->indices, src->ndim, src->element_shape, /*is_grad=*/true); - adj_ptr->ret_type = src->ret_type; - - if (is_ptr_offset) { - adj_ptr = insert(adj_ptr, stmt->src->as()->offset); - adj_ptr->ret_type = stmt->src->ret_type; - adj_ptr->ret_type.set_is_pointer(true); - } - insert(AtomicOpType::add, adj_ptr, load(adjoint(stmt))); - } - return; - } - - if (stmt->src->is() || - (stmt->src->is() && stmt->src->as()->origin->is())) { - GlobalPtrStmt *src = nullptr; - bool is_ptr_offset = false; - if (stmt->src->is()) { - is_ptr_offset = true; - src = stmt->src->as()->origin->as(); - } else { - src = stmt->src->as(); - } - - auto snode = src->snode; - if (!snode->has_adjoint()) { - // No adjoint SNode. Do nothing - return; - } - if (gradients_stopped(stmt, snode)) { - // gradients stopped, do nothing. - return; - } - QD_ASSERT(snode->get_adjoint() != nullptr); - snode = snode->get_adjoint(); - auto adj_ptr = insert(snode, src->indices); - adj_ptr->ret_type = src->ret_type; - if (is_ptr_offset) { - adj_ptr = insert(adj_ptr, stmt->src->as()->offset); - } - insert(AtomicOpType::add, adj_ptr, load(adjoint(stmt))); - return; - } - } - - void visit(GlobalStoreStmt *stmt) override { - // erase and replace with global load adjoint - - Stmt *adjoint_ptr{nullptr}; - if (stmt->dest->is() || - (stmt->dest->is() && stmt->dest->as()->origin->is())) { - ExternalPtrStmt *dest = nullptr; - bool is_ptr_offset = false; - if (stmt->dest->is()) { - is_ptr_offset = true; - dest = stmt->dest->as()->origin->as(); - } else { - dest = stmt->dest->as(); - } - - auto arg = dest->base_ptr->as(); - if (arg->ret_type.ptr_removed()->as()->elements().size() <= TypeFactory::GRAD_PTR_POS_IN_NDARRAY) { - return; - } - QD_ASSERT_INFO(!dest->is_grad, - "Cannot automatically differentiate through a grad " - "tensor, if you really want to do that, pass the grad " - "tensor into the kernel directly"); - adjoint_ptr = insert(dest->base_ptr, dest->indices, dest->ndim, dest->element_shape, - /*is_grad=*/true); - adjoint_ptr->ret_type = dest->ret_type; - - if (is_ptr_offset) { - adjoint_ptr = insert(adjoint_ptr, stmt->dest->as()->offset); - adjoint_ptr->ret_type = stmt->dest->ret_type; - adjoint_ptr->ret_type.set_is_pointer(true); - } - - accumulate(stmt->val, insert(adjoint_ptr)); - } - - if (stmt->dest->is() || - (stmt->dest->is() && stmt->dest->as()->origin->is())) { - GlobalPtrStmt *dest = nullptr; - bool is_ptr_offset = false; - if (stmt->dest->is()) { - is_ptr_offset = true; - dest = stmt->dest->as()->origin->as(); - } else { - dest = stmt->dest->as(); - } - - auto snode = dest->snode; - if (!snode->has_adjoint()) { - // no gradient (likely integer types) - return; - } - QD_ASSERT(snode->get_adjoint() != nullptr); - snode = snode->get_adjoint(); - adjoint_ptr = insert(snode, dest->indices); - adjoint_ptr->ret_type = dest->ret_type; - if (is_ptr_offset) { - adjoint_ptr = insert(adjoint_ptr, stmt->dest->as()->offset); - } - accumulate(stmt->val, insert(adjoint_ptr)); - } - - // Clear the gradient after accumulation finished. - auto zero = insert_const_for_grad(adjoint_ptr->ret_type.ptr_removed(), stmt, 0); - insert(adjoint_ptr, zero); - - stmt->parent->erase(stmt); - } - - void visit(AtomicOpStmt *stmt) override { - if (stmt->dest->is() || - (stmt->dest->is() && stmt->dest->as()->origin->is())) { - ExternalPtrStmt *dest = nullptr; - bool is_ptr_offset = false; - if (stmt->dest->is()) { - is_ptr_offset = true; - dest = stmt->dest->as()->origin->as(); - } else { - dest = stmt->dest->as(); - } - - auto arg = dest->base_ptr->as(); - if (arg->ret_type.ptr_removed()->as()->elements().size() > TypeFactory::GRAD_PTR_POS_IN_NDARRAY) { - QD_ASSERT_INFO(!dest->is_grad, - "Cannot automatically differentiate through a grad " - "tensor, if you really want to do that, pass the grad " - "tensor into the kernel directly"); - auto adjoint_ptr = - insert(dest->base_ptr, dest->indices, dest->ndim, dest->element_shape, /*is_grad=*/true); - adjoint_ptr->ret_type = dest->ret_type; - - if (is_ptr_offset) { - adjoint_ptr = insert(adjoint_ptr, stmt->dest->as()->offset); - - adjoint_ptr->ret_type = stmt->dest->ret_type; - adjoint_ptr->ret_type.set_is_pointer(true); - } - adjoint_ptr->ret_type = dest->ret_type; - accumulate(stmt->val, insert(adjoint_ptr)); - stmt->parent->erase(stmt); - } - return; - } - - if (stmt->dest->is() || - (stmt->dest->is() && stmt->dest->as()->origin->is())) { - GlobalPtrStmt *dest = nullptr; - bool is_ptr_offset = false; - if (stmt->dest->is()) { - is_ptr_offset = true; - dest = stmt->dest->as()->origin->as(); - } else { - dest = stmt->dest->as(); - } - - auto snode = dest->snode; - if (!snode->has_adjoint()) { - // no gradient (likely integer types) - return; - } - - QD_ASSERT(snode->get_adjoint() != nullptr); - snode = snode->get_adjoint(); - auto adjoint_ptr = insert(snode, dest->indices); - adjoint_ptr->ret_type = dest->ret_type; - if (is_ptr_offset) { - adjoint_ptr = insert(adjoint_ptr, stmt->dest->as()->offset); - } - accumulate(stmt->val, insert(adjoint_ptr)); - stmt->parent->erase(stmt); - return; - } - } - - void visit(MatrixPtrStmt *stmt) override { - if (stmt->origin->is() || stmt->origin->is()) { - /* - The case of MatrixPtrStmt(GlobalPtrStmt, ...) is already handled in - GlobalPtrStmt, GlobalStoreStmt and AtomicStmt - - TODO(zhanlue): Try to separate out the chain rule for MatrixPtrStmt from - GlobalPtrStmt, GlobalStoreStmt and AtomicStmt and migrate the logics - here. - */ - return; - } - - DataType prim_dtype = PrimitiveType::f32; - if (is_real(stmt->ret_type.ptr_removed().get_element_type())) { - prim_dtype = stmt->ret_type.ptr_removed().get_element_type(); - } - - Stmt *adjoint_value = nullptr; - if (stmt->offset->is()) { - /* - [Static index] - Fwd: - $0 = alloca <4 x i32> - $1 = matrix ptr $0, 2 // offset = 2 - - Adjoint: - $3 = matrix init [0, 0, $1_adj, 0] // adjoint_value - - accumulate($0_adj, $3) - */ - int offset = stmt->offset->as()->val.val_int32(); - - auto tensor_type = stmt->origin->ret_type.ptr_removed()->as(); - int num_elements = tensor_type->get_num_elements(); - - auto zero = insert_const_for_grad(prim_dtype, stmt, 0); - std::vector values; - for (int i = 0; i < num_elements; i++) { - if (i == offset) { - values.push_back(load(adjoint(stmt))); - } else { - values.push_back(zero); - } - } - auto matrix_init_stmt = insert(values); - matrix_init_stmt->ret_type = tensor_type; - - adjoint_value = matrix_init_stmt; - - } else { - /* - [Dynamic index] - Fwd: - $0 = alloca <4 x i32> - $1 = matrix ptr $0, $offset - - Adjoint: - $3 = matrix init [0.0, 0.0, 0.0, 0.0] - $4 = matrix init [$1_adj, $1_adj, $1_adj, $1_adj] - - $5 = matrix init [0, 1, 2, 3] - $6 = matrix init [offset, offset, offset, offset] - $7 = bin_eq $6, $5 - $8 = select $7, $4, $3 // adjoint_value - - accumulate($0_adj, $7) - */ - auto tensor_type = stmt->origin->ret_type.ptr_removed()->as(); - auto tensor_shape = tensor_type->get_shape(); - int num_elements = tensor_type->get_num_elements(); - - auto zero = insert_const_for_grad(prim_dtype, stmt, 0); - auto stmt_adj = load(adjoint(stmt)); - - std::vector zero_values(num_elements, zero); - std::vector stmt_adj_values(num_elements, stmt_adj); - std::vector offset_values(num_elements, stmt->offset); - std::vector indices_values(num_elements); - for (size_t i = 0; i < num_elements; i++) { - indices_values[i] = insert(TypedConstant((int32)i)); - } - - auto zero_matrix_init_stmt = insert(zero_values); - zero_matrix_init_stmt->ret_type = tensor_type; - auto stmt_adj_matrix_init_stmt = insert(stmt_adj_values); - stmt_adj_matrix_init_stmt->ret_type = tensor_type; - - auto index_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::i32); - auto indices_matrix_init_stmt = insert(indices_values); - indices_matrix_init_stmt->ret_type = index_tensor_type; - - auto offset_matrix_init_stmt = insert(offset_values); - offset_matrix_init_stmt->ret_type = index_tensor_type; - auto cmp_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::u1); - auto bin_eq_stmt = insert(BinaryOpType::cmp_eq, offset_matrix_init_stmt, indices_matrix_init_stmt); - bin_eq_stmt->ret_type = cmp_tensor_type; - - auto select_stmt = - insert(TernaryOpType::select, bin_eq_stmt, stmt_adj_matrix_init_stmt, zero_matrix_init_stmt); - adjoint_value = select_stmt; - } - - accumulate(stmt->origin, adjoint_value); - } - - void visit(MatrixInitStmt *stmt) override { - auto adjoint_ptr = adjoint(stmt); - - auto tensor_type = stmt->ret_type->as(); - int num_elements = tensor_type->get_num_elements(); - - for (size_t i = 0; i < num_elements; i++) { - auto const_i = insert_const_for_grad(PrimitiveType::i32, stmt, i); - - auto matrix_ptr_stmt_i = insert(adjoint_ptr, const_i); - matrix_ptr_stmt_i->ret_type = tensor_type->get_element_type(); - matrix_ptr_stmt_i->ret_type.set_is_pointer(true); - - accumulate(stmt->values[i], load(matrix_ptr_stmt_i)); - } - } -}; - -// Forward mode autodiff -class MakeDual : public ADTransform { - public: - using ADTransform::visit; - Stmt *current_stmt; - Block *current_block; - Block *alloca_block; - std::map dual_stmt; - - explicit MakeDual(Block *block) { - current_stmt = nullptr; - alloca_block = block; - current_block = block; - } - - static void run(Block *block) { - auto p = MakeDual(block); - block->accept(&p); - } - - Stmt *insert_grad_stmt(std::unique_ptr &&stmt) override { - auto ptr = stmt.get(); - current_stmt = current_stmt->insert_after_me(std::move(stmt)); - return ptr; - } - - void visit(Block *block) override { - std::vector statements; - // always make a copy since the list can be modified. - for (auto &stmt : block->statements) { - statements.push_back(stmt.get()); - } - for (auto stmt : statements) { - current_stmt = stmt; - stmt->accept(this); - } - } - - // Accumulate [value] to the dual of [primal] - void accumulate(Stmt *primal, Stmt *value) { - auto alloca_ = dual(primal); - if (!alloca_ || alloca_->is()) - return; // primal may be int variable - - QD_ASSERT(alloca_->is()); - auto alloca = alloca_->as(); - auto local_load = insert(alloca); - insert(alloca, add(local_load, value)); - } - - Stmt *dual(Stmt *stmt) { - auto dual_type = stmt->ret_type.ptr_removed(); - if (!is_real(dual_type.get_element_type()) || stmt->is()) { - return constant(0); - } - if (dual_stmt.find(stmt) == dual_stmt.end()) { - // normal SSA cases - - // create the alloca - // auto alloca = - // Stmt::make(get_current_program().config.gradient_dt); - // maybe it's better to use the statement data type than the default type - auto alloca = Stmt::make(dual_type); - dual_stmt[stmt] = alloca.get(); - - // TODO: check whether there are any edge cases for the alloca_block - alloca_block->insert(std::move(alloca), 0); - } - return dual_stmt[stmt]; - } - - void visit(UnaryOpStmt *stmt) override { - if (stmt->op_type == UnaryOpType::neg) { - accumulate(stmt, negate(dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::abs) { - accumulate(stmt, mul(sgn(stmt->operand), dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::sin) { - accumulate(stmt, mul(cos(stmt->operand), dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::cos) { - accumulate(stmt, negate(mul(sin(stmt->operand), dual(stmt->operand)))); - } else if (stmt->op_type == UnaryOpType::tan) { - // d/dx tan(x) = 1 + tan(x)^2. Forward mode executes in primal order, so `stmt` is the - // current-iteration tan value - no BackupSSA stale-value concern; reusing it is per-iteration - // correct. - accumulate(stmt, mul(add(constant(1), sqr(stmt)), dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::tanh) { - accumulate(stmt, mul(sub(constant(1), sqr(stmt)), dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::asin) { - accumulate(stmt, mul(div(constant(1), sqrt(sub(constant(1), sqr(stmt->operand)))), dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::acos) { - accumulate(stmt, mul(negate(div(constant(1), sqrt(sub(constant(1), sqr(stmt->operand))))), dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::exp) { - accumulate(stmt, mul(stmt, dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::log) { - accumulate(stmt, div(dual(stmt->operand), stmt->operand)); - } else if (stmt->op_type == UnaryOpType::sqrt) { - accumulate(stmt, mul(div(constant(0.5f), sqrt(stmt->operand)), dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::rsqrt) { - accumulate(stmt, mul(mul(constant(-0.5f), pow(rsqrt(stmt->operand), constant(3))), dual(stmt->operand))); - } else if (stmt->op_type == UnaryOpType::cast_value) { - if (is_real(stmt->cast_type.get_element_type()) && is_real(stmt->operand->ret_type.get_element_type())) { - accumulate(stmt, dual(stmt->operand)); - } - } else if (stmt->op_type == UnaryOpType::logic_not) { - // do nothing - } else { - QD_P(unary_op_type_name(stmt->op_type)); - QD_NOT_IMPLEMENTED - } - } - - void visit(BinaryOpStmt *bin) override { - if (bin->op_type == BinaryOpType::add) { - accumulate(bin, dual(bin->lhs)); - accumulate(bin, dual(bin->rhs)); - } else if (bin->op_type == BinaryOpType::sub) { - accumulate(bin, dual(bin->lhs)); - accumulate(bin, negate(dual(bin->rhs))); - } else if (bin->op_type == BinaryOpType::mul) { - // d (x * y) = y * dx + x * dy - accumulate(bin, mul(bin->lhs, dual(bin->rhs))); - accumulate(bin, mul(bin->rhs, dual(bin->lhs))); - } else if (bin->op_type == BinaryOpType::mod) { - // Do nothing - } else if (bin->op_type == BinaryOpType::div) { - accumulate(bin, div(dual(bin->lhs), bin->rhs)); - accumulate(bin, negate(div(mul(dual(bin->rhs), bin->lhs), mul(bin->rhs, bin->rhs)))); - } else if (bin->op_type == BinaryOpType::atan2) { - auto numerator = add(sqr(bin->lhs), sqr(bin->rhs)); - accumulate(bin, div(mul(bin->rhs, dual(bin->lhs)), numerator)); - accumulate(bin, negate(div(mul(bin->lhs, dual(bin->rhs)), numerator))); - } else if (bin->op_type == BinaryOpType::pow) { - // d (x ^ y) = x ^ (y-1) * (y * dx + log(x) * x * dy) - auto common_coeff = pow(bin->lhs, sub(bin->rhs, constant(1))); // x ^ (y-1) - accumulate(bin, mul(dual(bin->lhs), mul(bin->rhs, common_coeff))); - accumulate(bin, mul(dual(bin->rhs), mul(log(bin->lhs), mul(bin->lhs, common_coeff)))); - } else if (bin->op_type == BinaryOpType::min || bin->op_type == BinaryOpType::max) { - auto cmp = bin->op_type == BinaryOpType::min ? cmp_lt(bin->lhs, bin->rhs) : cmp_lt(bin->rhs, bin->lhs); - auto zero = insert_const_for_grad(bin->ret_type, bin, 0); - accumulate(bin, sel(cmp, dual(bin->lhs), zero)); - accumulate(bin, sel(cmp, zero, dual(bin->rhs))); - } else if (bin->op_type == BinaryOpType::floordiv) { - // do nothing - } else if (is_comparison(bin->op_type) || is_bit_op(bin->op_type)) { - // do nothing - } else { - QD_WARN("gradient of binary op {}\n{}", binary_op_type_name(bin->op_type), bin->get_tb()); - QD_NOT_IMPLEMENTED - } - } - - void visit(TernaryOpStmt *stmt) override { - QD_ASSERT(stmt->op_type == TernaryOpType::select); - auto zero = insert_const_for_grad(stmt->ret_type, stmt, 0); - accumulate(stmt, insert(TernaryOpType::select, stmt->op1, load(dual(stmt->op2)), zero)); - accumulate(stmt, insert(TernaryOpType::select, stmt->op1, zero, load(dual(stmt->op3)))); - } - - void visit(IfStmt *if_stmt) override { - if (if_stmt->true_statements) { - std::vector true_statements; - for (auto &stmt : if_stmt->true_statements->statements) { - true_statements.push_back(stmt.get()); - } - - for (auto stmt : true_statements) { - current_stmt = stmt; - stmt->accept(this); - } - } - if (if_stmt->false_statements) { - std::vector false_statements; - for (auto &stmt : if_stmt->false_statements->statements) { - false_statements.push_back(stmt.get()); - } - - for (auto stmt : false_statements) { - current_stmt = stmt; - stmt->accept(this); - } - } - } - - void visit(RangeForStmt *for_stmt) override { - std::vector statements; - // always make a copy since the list can be modified. - for (auto &stmt : for_stmt->body->statements) { - statements.push_back(stmt.get()); - } - auto previous_alloca_block = alloca_block; - alloca_block = for_stmt->body.get(); - for (auto stmt : statements) { - current_stmt = stmt; - stmt->accept(this); - } - alloca_block = previous_alloca_block; - } - - void visit(StructForStmt *for_stmt) override { - // Save/restore mirrors visit(RangeForStmt) above and MakeAdjoint::visit(StructForStmt). An enclosing - // compound visitor that resumes iterating its body after this StructForStmt needs alloca_block to still - // point at its own block, not the sparse-for body, so new dual allocas land where the enclosing reverse - // code can reach them. - auto previous_alloca_block = alloca_block; - alloca_block = for_stmt->body.get(); - for_stmt->body->accept(this); - alloca_block = previous_alloca_block; - } - - void visit(LocalLoadStmt *stmt) override { - // QD_ASSERT(!needs_grad(stmt->ret_type)); - accumulate(stmt, dual(stmt->src)); - } - - void visit(LocalStoreStmt *stmt) override { - // Clear the dual of the dest before local store, - // Because LocalStoreStmt overwrites the dest, - // If the alloca serves as the dest of multiple LocalStoreStmt, only the - // last LocalStoreStmt should be taken account of, i.e, its history should - // be cleared - auto dtype = stmt->dest->ret_type.ptr_removed(); - if (is_real(dtype.get_element_type())) { - auto zero = insert_const_for_grad(dtype, stmt, 0); - insert(dual(stmt->dest), zero); - } - - accumulate(stmt->dest, dual(stmt->val)); - } - - void visit(GlobalLoadStmt *stmt) override { - // issue global store to dual - GlobalPtrStmt *src = nullptr; - bool is_ptr_offset = false; - if (stmt->src->is()) { - is_ptr_offset = true; - src = stmt->src->as()->origin->as(); - } else { - src = stmt->src->as(); - } - auto snode = src->snode; - if (!snode->has_dual()) { - // No dual SNode. Do nothing - return; - } - if (gradients_stopped(stmt, snode)) { - // gradients stopped, do nothing. - return; - } - QD_ASSERT(snode->get_dual() != nullptr); - snode = snode->get_dual(); - auto dual_ptr = insert(snode, src->indices); - dual_ptr->ret_type = src->ret_type; - if (is_ptr_offset) { - dual_ptr = insert(dual_ptr, stmt->src->as()->offset); - } - accumulate(stmt, insert(dual_ptr)); - } - - void visit(GlobalStoreStmt *stmt) override { - GlobalPtrStmt *dest = nullptr; - bool is_ptr_offset = false; - if (stmt->dest->is()) { - is_ptr_offset = true; - dest = stmt->dest->as()->origin->as(); - } else { - dest = stmt->dest->as(); - } - auto snode = dest->snode; - if (!snode->has_dual()) { - // no gradient (likely integer types) - return; - } - QD_ASSERT(snode->get_dual() != nullptr); - snode = snode->get_dual(); - auto dual_ptr = insert(snode, dest->indices); - dual_ptr->ret_type = dest->ret_type; - if (is_ptr_offset) { - dual_ptr = insert(dual_ptr, stmt->dest->as()->offset); - } - insert(AtomicOpType::add, dual_ptr, load(dual(stmt->val))); - } - - void visit(AtomicOpStmt *stmt) override { - GlobalPtrStmt *dest = nullptr; - bool is_ptr_offset = false; - if (stmt->dest->is()) { - is_ptr_offset = true; - dest = stmt->dest->as()->origin->as(); - } else { - dest = stmt->dest->as(); - } - auto snode = dest->snode; - if (!snode->has_dual()) { - // no gradient (likely integer types) - return; - } - QD_ASSERT(snode->get_dual() != nullptr); - snode = snode->get_dual(); - auto dual_ptr = insert(snode, dest->indices); - dual_ptr->ret_type = dest->ret_type; - if (is_ptr_offset) { - dual_ptr = insert(dual_ptr, stmt->dest->as()->offset); - } - insert(AtomicOpType::add, dual_ptr, load(dual(stmt->val))); - } - - void visit(MatrixInitStmt *stmt) override { - std::vector duals; - for (auto &s : stmt->values) { - duals.push_back(dual(s)); - } - auto dual_stmt = insert(duals); - dual_stmt->ret_type = stmt->ret_type; - - accumulate(stmt, dual_stmt); - } - - void visit(MatrixPtrStmt *stmt) override { - if (stmt->origin->is()) { - // Handled in GlobalLoadStmt and GlobalStoreStmt - return; - } - - auto origin_dual = dual(stmt->origin); - auto origin_dual_ptr = insert(origin_dual, stmt->offset); - origin_dual_ptr->ret_type = stmt->ret_type; - - accumulate(stmt, origin_dual_ptr); - } -}; - -class BackupSSA : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - - Block *independent_block; - std::map backup_alloca; - - explicit BackupSSA(Block *independent_block) : independent_block(independent_block) { - allow_undefined_visitor = true; - invoke_default_visitor = true; - } - - Stmt *load(Stmt *stmt) { - if (backup_alloca.find(stmt) == backup_alloca.end()) { - auto alloca = Stmt::make(stmt->ret_type.ptr_removed()); - auto alloca_ptr = alloca.get(); - independent_block->insert(std::move(alloca), 0); - auto local_store = Stmt::make(alloca_ptr, stmt); - stmt->insert_after_me(std::move(local_store)); - backup_alloca[stmt] = alloca_ptr; - } - return backup_alloca[stmt]; - } - - void generic_visit(Stmt *stmt) { - std::vector leaf_to_root; - auto t = stmt->parent; - while (t != nullptr) { - leaf_to_root.push_back(t); - t = t->parent_block(); - } - int num_operands = stmt->get_operands().size(); - for (int i = 0; i < num_operands; i++) { - auto op = stmt->operand(i); - if (op == nullptr) { - continue; - } - if (std::find(leaf_to_root.begin(), leaf_to_root.end(), op->parent) == leaf_to_root.end() && - !op->is()) { - if (op->is()) { - // Just create another AdStackLoadTopStmt - stmt->set_operand(i, stmt->insert_before_me(op->clone())); - } else if (op->is()) { - // Backup AdStackAllocaStmt because it should not be local stored and - // local loaded - auto stack_alloca = op->as(); - if (backup_alloca.find(op) == backup_alloca.end()) { - auto backup_stack_alloca = Stmt::make(stack_alloca->dt, stack_alloca->max_size); - auto backup_stack_alloca_ptr = backup_stack_alloca.get(); - independent_block->insert(std::move(backup_stack_alloca), 0); - backup_alloca[op] = backup_stack_alloca_ptr; - // Replace usages of all blocks i.e., the entry point for the - // replace is the top level block - irpass::replace_all_usages_with(leaf_to_root.back(), op, backup_stack_alloca_ptr); - // Erase the outdated AdStackAllocaStmt - op->parent->erase(op); - } - } else if (op->is()) { - stmt->set_operand(i, stmt->insert_before_me(op->clone())); - } else { - // Recomputable-chain fallback before the last-iter `load(op)` spill: when the cross-block SSA op - // is rooted in a DAG of side-effect-free arithmetic over already-stack-backed allocas, kernel-args, - // constants, and loop indices, clone the chain into the reverse scope at the consumer site. The - // cloned chain reads stack tops live (matching `MakeAdjoint`'s pop ordering, which fires pops - // AFTER all uses of a stack within one reverse iter) and reads kernel-args / constants / loop - // indices via direct clones - exactly the path that was already correct for AdStackLoadTopStmt and - // ArgLoadStmt operands. - // - // The pre-existing `load(op)` fallback below remains correct for genuinely non-recomputable - // cross-block ops (a forward `GlobalLoadStmt` of a needs_grad SNode whose value the reverse must - // read at last-write rather than recompute, for instance) and for shapes outside one of our - // independent blocks. Adding the recomputable path above the fallback strictly subsets the - // previous behaviour: a chain that fails the predicate falls through to the old `load(op)` line. - if (RecomputableChainAnalyzer::is_recomputable(op, recomputable_cache_)) { - std::unordered_map clone_cache; - Stmt *cloned = RecomputableChainCloner::clone_at(op, stmt, clone_cache); - stmt->set_operand(i, cloned); - } else { - auto alloca = load(op); - stmt->set_operand(i, stmt->insert_before_me(Stmt::make(alloca))); - } - } - } - } - } - - // Memoization cache for `RecomputableChainAnalyzer::is_recomputable` queries within one BackupSSA run. - // Re-used across all generic_visit calls; invariant during the visit because forward IR is read-only here. - std::unordered_map recomputable_cache_; - - void visit(Stmt *stmt) override { - generic_visit(stmt); - } - - void visit(IfStmt *stmt) override { - generic_visit(stmt); - BasicStmtVisitor::visit(stmt); - } - - // generic_visit spills cross-block operands (the for-loop's `begin` and `end`) the same way it does for an - // IfStmt's cond. MakeAdjoint clones a forward for-loop into the reverse scope and shares the clone's - // `begin`/`end` pointers with the forward stmt; when those operands live inside the forward for's body (e.g. - // inner `for k in range(j)` where `j` is an enclosing loop's index promoted to a per-iter adstack), the reverse - // clone's operand no longer dominates its use. generic_visit's AdStackLoadTopStmt branch handles this by - // inserting a fresh AdStackLoadTop in the reverse scope, which reads the correct per-iteration value. - void visit(RangeForStmt *stmt) override { - generic_visit(stmt); - stmt->body->accept(this); - } - - void visit(StructForStmt *stmt) override { - generic_visit(stmt); - stmt->body->accept(this); - } - - void visit(WhileStmt *stmt) override { - QD_ERROR("WhileStmt not supported in AutoDiff for now."); - } - - void visit(Block *block) override { - std::vector statements; - // always make a copy since the list can be modified. - for (auto &stmt : block->statements) { - statements.push_back(stmt.get()); - } - for (auto stmt : statements) { - QD_ASSERT(!stmt->erased); - stmt->accept(this); - } - } - - public: - static void run(Block *block) { - BackupSSA pass(block); - block->accept(&pass); - } -}; - -namespace irpass { - -// clang-format off -/* -Support for TensorType: How to handle MatrixPtrStmt & MatrixInitStmt - -[Original Quadrants Code] - -@ti.kernel -def test(...): - b = ti.Vector([0, 1, 2, 3]) - b[2] = 100 - y = b[3] * b[2] * x - - -[Forward] [Forward-Replaced] [Backward] -$b = alloca Tensor<4 x i32> --> $b = adstack alloca <4 x i32> -$1 = matrix init [0, 1, 2, 3] --> $1 = matrix init [0, 1, 2, 3] - adstack pop -$2: local store $b, $1 --> adstack push $1 --> acc($1_adj, adstack top adj()) - -$3 = matrix ptr $b, 2 --> $2 = adstack top(is_ptr=True) --> adstack acc adj($2_adj) - - acc($2_adj, $14) - $3 = matrix ptr $2, 0 --> $14 = matrix_init({$3_adj, 0, 0, 0}) - - acc($2_adj, $13) - $4 = matrix ptr $2, 2 --> $13 = matrix_init({0, 0, $4_adj, 0}) - - acc($2_adj, $12) - $5 = matrix ptr $2, 3 --> $12 = matrix_init({0, 0, 0, $5_adj}) - - $6 = load($3) --> acc($3_adj, $6_adj) - $7 = load($4) --> acc($4_adj, $7_adj) - $8 = load($5) --> acc($5_adj, $8_adj) - - acc($8_adj, matrix ptr($9_adj, 3)) - acc($7_adj, matrix ptr($9_adj, 2)) - acc(100_adj, matrix ptr($9_adj, 1)) - $9 = matrix_init($6,100,$7,$8) --> acc($6_adj, matrix ptr($9_adj, 0)) - - adstack pop -$4 = local store $3, 100 --> adstack push $9 --> acc($9_adj, adstack top adj()) - - $10 = adstack top(is_ptr=True) --> adstack acc adj($10_adj) - - acc($10_adj, $18) -$5 = matrix ptr $b, 3 --> $11 = matrix ptr $10, 3 --> $18 = matrix_init({0, 0, 0, $11_adj}) -$b3 = local load $5 --> $b3 = local load $11 --> acc($11_adj, $b3_adj) - - $12 = adstack top(is_ptr=True) --> adstack acc adj($12_adj) - - acc($12_adj, $17) -$6 = matrix ptr $b, 2 --> $13 = matrix ptr $12, 2 --> $17 = matrix_init({0, 0, $13_adj, 0}) -$b2 = local load $6 --> $b2 = local load $13 --> acc($13_adj, $b2_adj) - - acc($b3_adj, $15) - acc($b2_adj, $16) - $16 = mul($tmp_adj, $b3) -$tmp = mul b3, b2 --> $tmp = mul $b3, $b2 --> $15 = mul($tmp_adj, $b2) - - acc($tmp_adj, $14) -$y = mul $tmp, $x --> $y = mul $tmp, $x --> $14 = mul($y_adj, $x) -*/ -// clang-format on -// Within a single straight-line block, multiple `AdStackLoadTopStmt` reads of the same stack with no -// intervening `AdStackPushStmt` / `AdStackPopStmt` for that stack return the same value, and multiple -// `AdStackLoadTopAdjStmt` reads are equivalent under the same conditions plus no intervening -// `AdStackAccAdjointStmt`. After Python-side static unrolling collapses an inner loop into straight-line IR, -// every iteration's read of an outer-loop-invariant adstack value emits a separate LoadTop in the same block; -// each individual load is cheap (read u64 count + GEP) but unrolled-loop counts of hundreds to thousands -// inflate PTX size and ptxas register-allocator cost. This pass walks each block, caches the most recent -// LoadTop / LoadTopAdj per stack, replaces subsequent same-stack reads with the cached SSA value until a -// Push / Pop / AccAdjoint invalidates the cache for that stack, and conservatively clears the cache when -// crossing into nested control flow (IfStmt / RangeForStmt / StructForStmt / WhileStmt) where unseen -// push/pop/AccAdjoint may appear. -class CoalesceAdStackLoads : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - - CoalesceAdStackLoads() { - invoke_default_visitor = true; - allow_undefined_visitor = true; - } - - void visit(Block *block) override { - std::unordered_map primal_cache; - std::unordered_map adjoint_cache; - // Iterate over a snapshot so erase()s during the walk do not invalidate the iteration. The - // DelayedIRModifier still applies erases at the end via `modify_ir()` so SSA users get rewritten in one - // pass; the snapshot only protects the visitor's own walk. - std::vector stmts; - stmts.reserve(block->statements.size()); - for (auto &s : block->statements) { - stmts.push_back(s.get()); - } - for (Stmt *s : stmts) { - if (auto *lt = s->cast()) { - // `return_ptr=true` returns a pointer into the stack slot rather than loading a value; coalescing - // those is unsound because subsequent stores via the pointer would alias with each other through - // the cached SSA value. The non-pointer variant is the common case driven by reverse-pass formulas. - if (lt->return_ptr) { - continue; - } - auto it = primal_cache.find(lt->stack); - if (it != primal_cache.end()) { - irpass::replace_all_usages_with(lt->parent, lt, it->second); - modifier_.erase(lt); - } else { - primal_cache[lt->stack] = lt; - } - } else if (auto *la = s->cast()) { - auto it = adjoint_cache.find(la->stack); - if (it != adjoint_cache.end()) { - irpass::replace_all_usages_with(la->parent, la, it->second); - modifier_.erase(la); - } else { - adjoint_cache[la->stack] = la; - } - } else if (auto *p = s->cast()) { - primal_cache.erase(p->stack); - adjoint_cache.erase(p->stack); - } else if (auto *po = s->cast()) { - primal_cache.erase(po->stack); - adjoint_cache.erase(po->stack); - } else if (auto *aa = s->cast()) { - // Accumulating into the top adjoint slot does not alter the primal half, so the primal cache for - // the same stack stays valid; only the adjoint cache must be invalidated. - adjoint_cache.erase(aa->stack); - } else if (s->is() || s->is() || s->is() || s->is()) { - // The pass is intentionally intra-block: any push / pop hidden inside a nested block would - // invalidate caches asymmetrically across branches and make a sound merge complex. Conservatively - // drop both caches before descending so reads after the nested block see no stale entries. - primal_cache.clear(); - adjoint_cache.clear(); - s->accept(this); - } else { - // Any other statement is assumed inert with respect to the AdStack count headers; recurse into - // potential nested blocks (defensively) but leave the caches intact. - s->accept(this); - } - } - } - - static bool run(IRNode *root) { - CoalesceAdStackLoads pass; - root->accept(&pass); - return pass.modifier_.modify_ir(); - } - - private: - DelayedIRModifier modifier_; -}; - -// Detect cross-iteration read-after-write through a `needs_grad` global field at the body of an offload-level -// range-for. `MakeAdjoint` treats each iteration of an offload-level loop as independent and only chases -// loop-carried dependencies through `AdStackAllocaJudger`, which inspects local `AllocaStmt`s - cross-iteration -// dependencies through `GlobalStoreStmt` / `GlobalLoadStmt` on a `has_adjoint()` SNode are silently dropped. -// A kernel like `out[i] = out[i-1] + x[None]` therefore returns wrong gradients with no diagnostic, even with -// `ad_stack_experimental_enabled`. Pair every same-SNode (store, load) and raise when the index Stmt*s differ. -// Same-Stmt indices (`out[i] = out[i] + x[None]`) are in-place accumulation and the per-iteration backward -// pass handles those correctly; pointer equality on the SSA index Stmts after the `pre_autodiff` simplify -// (`compile_to_offloads.cpp:121`) is enough to distinguish the two cases. Walks the body of every direct-child -// `RangeForStmt` of `root`; nested for-loops' bodies are skipped via `inside_nested_` because their cross-iter -// behavior is governed by the existing IB / adstack machinery, not by this offload-level guard. -class OffloadLevelGlobalCrossIterRAWChecker : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - - void visit(GlobalStoreStmt *stmt) override { - if (inside_nested_) - return; - if (auto *gp = extract_global_ptr(stmt->dest)) { - if (gp->snode->has_adjoint()) - stores_[gp->snode].push_back(gp); - } - } - - void visit(GlobalLoadStmt *stmt) override { - if (inside_nested_) - return; - if (auto *gp = extract_global_ptr(stmt->src)) { - if (gp->snode->has_adjoint()) - loads_[gp->snode].push_back(gp); - } - } - - void visit(RangeForStmt *stmt) override { - walk_nested_loop_body(stmt->body.get()); - } - - void visit(StructForStmt *stmt) override { - walk_nested_loop_body(stmt->body.get()); - } - - void visit(MeshForStmt *stmt) override { - walk_nested_loop_body(stmt->body.get()); - } - - void visit(WhileStmt *stmt) override { - walk_nested_loop_body(stmt->body.get()); - } - - static void run(Block *offload_body) { - OffloadLevelGlobalCrossIterRAWChecker checker; - offload_body->accept(&checker); - for (const auto &kv : checker.stores_) { - auto *snode = kv.first; - auto load_it = checker.loads_.find(snode); - if (load_it == checker.loads_.end()) - continue; - for (auto *store : kv.second) { - // The store's index must reference the offload's `LoopIndexStmt` on at least one axis - otherwise it is - // iteration-independent (writes the same slot every iteration) and there is no cross-iter chain to drop. - if (!any_axis_references_loop_index(store)) - continue; - for (auto *load : load_it->second) { - // Skip loads whose index is iteration-independent on every axis (constants, captures from outer - // scopes, ...). Such loads read from a slice of the SNode that the loop never writes to in this - // iteration, so they are not a cross-iter RAW hazard. - if (!any_axis_references_loop_index(load)) - continue; - if (!indices_have_no_cross_iter_dependency(store, load)) { - QD_ERROR( - "Cross-iteration read-after-write on the `needs_grad` global field `{}` at the offload-level " - "loop body is not supported in reverse-mode autodiff: the cross-iteration adjoint chain is " - "not generated and the backward pass returns wrong gradients with no diagnostic. Wrap the " - "loop in an outer for-loop in the kernel so the inner range becomes adstack-eligible, or " - "restructure the kernel so the offload-level loop has no global field cross-iteration " - "read-after-write. Sibling-component access on a disjoint axis like `out[i, 0] = out[i, 1]` " - "is allowed; in-place accumulation `out[i] = out[i] + x` (same index Stmt on every axis) is " - "also fine.", - snode->get_node_type_name_hinted()); - } - } - } - } - } - - private: - void walk_nested_loop_body(Block *body) { - bool prev = inside_nested_; - inside_nested_ = true; - body->accept(this); - inside_nested_ = prev; - } - - static GlobalPtrStmt *extract_global_ptr(Stmt *s) { - if (auto *gp = s->cast()) - return gp; - if (auto *mp = s->cast()) { - if (auto *gp = mp->origin->cast()) - return gp; - } - return nullptr; - } - - // Decide whether two `GlobalPtrStmt`s on the same SNode index iter-independent slices of that SNode (no - // cross-iter dependency). The rules are per-axis: - // - axis where neither index references `LoopIndexStmt`: constants or outer-scope captures, may differ - // (sibling-component access on a disjoint axis like `state[i, 0]` vs `state[i, 1]`); never a cross-iter - // dependency. - // - axis where both indices reference `LoopIndexStmt`: must be the same SSA `Stmt*` (e.g. both bare - // `LoopIndexStmt(i)`). If one is `LoopIndex` and the other is `LoopIndex - 1` or similar, the load reads - // a slot the store wrote in a different iteration: this is cross-iter. - // - axis where exactly one references `LoopIndexStmt`: load and store have iter-dependent vs iter- - // independent indexing on the same axis; conservatively flag as cross-iter, since whether the constant - // index aliases an iteration of the loop is not provable here. - static bool indices_have_no_cross_iter_dependency(GlobalPtrStmt *a, GlobalPtrStmt *b) { - if (a->indices.size() != b->indices.size()) - return false; - for (std::size_t i = 0; i < a->indices.size(); ++i) { - bool a_refs = references_loop_index(a->indices[i]); - bool b_refs = references_loop_index(b->indices[i]); - if (a_refs != b_refs) - return false; - if (a_refs && b_refs && a->indices[i] != b->indices[i]) - return false; - } - return true; - } - - // Walk an index `Stmt*` and decide whether it (transitively) references a `LoopIndexStmt` of any enclosing - // loop. Recurses through linear arithmetic ops because an index like `i - 1` lowers to - // `BinaryOpStmt(LoopIndexStmt(i), Sub, ConstStmt(1))` and we want to recognise the `LoopIndexStmt` underneath. - // Conservatively returns false for anything else (constants, ndarray loads from outer scope, ...) - those - // express iteration-independent indexing that the cross-iter guard intentionally exempts. - static bool references_loop_index(Stmt *s) { - if (s == nullptr) - return false; - if (s->is()) - return true; - if (auto *bo = s->cast()) - return references_loop_index(bo->lhs) || references_loop_index(bo->rhs); - if (auto *uo = s->cast()) - return references_loop_index(uo->operand); - if (auto *to = s->cast()) - return references_loop_index(to->op1) || references_loop_index(to->op2) || references_loop_index(to->op3); - return false; - } - - static bool any_axis_references_loop_index(GlobalPtrStmt *gp) { - for (auto *idx : gp->indices) { - if (references_loop_index(idx)) - return true; - } - return false; - } - - std::unordered_map> stores_; - std::unordered_map> loads_; - bool inside_nested_ = false; -}; - -void auto_diff(IRNode *root, const CompileConfig &config, AutodiffMode autodiff_mode, bool use_stack) { - QD_AUTO_PROF; - if (autodiff_mode == AutodiffMode::kReverse) { - if (auto *root_block = root->cast()) { - // Top-level range-fors at the kernel root become the offload-level loop. Walk each direct child once - // and refuse cross-iteration global RAW on a needs_grad SNode before the AD transformation runs - any - // diagnostic from inside `MakeAdjoint` would already point at the partially-rewritten reverse IR and - // be much harder to read. StructFor offloads are out of scope: their loop variable is an SNode index - // produced by the runtime, not a user-controlled integer arithmeticed against itself, so the - // `out[i-1]` shape this guard targets does not arise. - for (auto &s : root_block->statements) { - if (auto *rf = s->cast()) - OffloadLevelGlobalCrossIterRAWChecker::run(rf->body.get()); - } - } - RegulateTensorTypedStatements::run(root); - if (use_stack) { - auto IB = IdentifyIndependentBlocks::run(root); - ReverseOuterLoops::run(root, IB); - - for (auto ib : IB) { - PromoteSSA2LocalVar::run(ib); - ReplaceLocalVarWithStacks replace(config.ad_stack_size); - ib->accept(&replace); - type_check(root, config); - - // Drop AdStackAllocas whose pushed value is recomputable from already-stack-backed allocas + args + - // const + loop-index. Trades forward-pass adstack memory traffic (one push + one load per iter per - // intermediate spill) for cloned arithmetic in the reverse scope, which `BackupSSA::generic_visit` - // generates on demand via the same RecomputableChainCloner path. Must run after - // `ReplaceLocalVarWithStacks` (so the analyzer sees the AdStackAlloca shape, not the alloca-with-store - // shape PromoteSSA2LocalVar emits) and before `MakeAdjoint` (so the reverse pass is generated against - // the cleaned forward IR with no spurious push/load scaffolding). - EliminateRecomputableAdStackPushes::run(ib); - type_check(root, config); - - MakeAdjoint::run(ib); - type_check(root, config); - BackupSSA::run(ib); - // After MakeAdjoint emits the reverse-pass body, an outer-loop-invariant value pulled into the - // reverse direction by `accumulate_unary_operand_checked` becomes a fresh `AdStackLoadTopStmt` per - // user-side use; in straight-line unrolled IR those reads coalesce to one load per stack per block, - // which the dedicated pass handles before the IR reaches `irpass::analysis::verify_if_debug`. - CoalesceAdStackLoads::run(ib); - irpass::analysis::verify_if_debug(root, config); - } - } else { - auto IB = IdentifyIndependentBlocks::run(root); - ReverseOuterLoops::run(root, IB); - type_check(root, config); - for (auto ib : IB) { - MakeAdjoint::run(ib); - type_check(root, config); - BackupSSA::run(ib); - CoalesceAdStackLoads::run(ib); - irpass::analysis::verify_if_debug(root, config); - } - } - } else if (autodiff_mode == AutodiffMode::kForward) { - // Forward mode autodiff - Block *block = root->as(); - MakeDual::run(block); - } - type_check(root, config); - irpass::analysis::verify_if_debug(root, config); -} - -class GloablDataAccessRuleChecker : public BasicStmtVisitor { - public: - using BasicStmtVisitor::visit; - - void visit(GlobalLoadStmt *stmt) override { - GlobalPtrStmt *src = nullptr; - if (stmt->src->is()) { - src = stmt->src->as(); - } else { - QD_ASSERT(stmt->src->is()); - src = stmt->src->as()->origin->as(); - } - auto snode = src->snode; - if (!snode->has_adjoint_checkbit()) { - return; - } - QD_ASSERT(snode->get_adjoint_checkbit() != nullptr); - snode = snode->get_adjoint_checkbit(); - auto global_ptr = stmt->insert_after_me(Stmt::make(snode, src->indices)); - auto dtype = global_ptr->ret_type.ptr_removed(); - - auto one = insert_const(dtype, global_ptr, 1, false /*insert_before_me*/); - one->insert_after_me(Stmt::make(global_ptr, one)); - } - - void visit_gloabl_store_stmt_and_atomic_add(Stmt *stmt, GlobalPtrStmt *dest) { - auto snode = dest->snode; - if (!snode->has_adjoint_checkbit()) { - return; - } - QD_ASSERT(snode->get_adjoint_checkbit() != nullptr); - snode = snode->get_adjoint_checkbit(); - auto global_ptr = stmt->insert_before_me(Stmt::make(snode, dest->indices)); - auto global_load = stmt->insert_before_me(Stmt::make(global_ptr)); - auto dtype = global_ptr->ret_type.ptr_removed(); - auto zero = insert_const(dtype, stmt, 0, /*insert_before_me=*/true); - auto check_equal = stmt->insert_before_me(Stmt::make(BinaryOpType::cmp_eq, global_load, zero)); - std::string msg = fmt::format( - "(kernel={}) Breaks the global data access rule. Snode {} is " - "overwritten unexpectedly.", - kernel_name_, dest->snode->get_node_type_name()); - msg += "\n" + stmt->get_tb(); - - stmt->insert_before_me(Stmt::make(check_equal, msg, std::vector())); - } - - void visit(GlobalStoreStmt *stmt) override { - GlobalPtrStmt *dest = nullptr; - if (stmt->dest->is()) { - dest = stmt->dest->as(); - } else { - QD_ASSERT(stmt->dest->is()); - dest = stmt->dest->as()->origin->as(); - } - visit_gloabl_store_stmt_and_atomic_add(stmt, dest); - } - - void visit(AtomicOpStmt *stmt) override { - GlobalPtrStmt *dest = nullptr; - if (stmt->dest->is()) { - dest = stmt->dest->as(); - } else { - QD_ASSERT(stmt->dest->is()); - dest = stmt->dest->as()->origin->as(); - } - visit_gloabl_store_stmt_and_atomic_add(stmt, dest); - } - - static void run(IRNode *root, const std::string &kernel_name) { - GloablDataAccessRuleChecker checker; - checker.kernel_name_ = kernel_name; - root->accept(&checker); - } - - private: - std::string kernel_name_; -}; - -void differentiation_validation_check(IRNode *root, const CompileConfig &config, const std::string &kernel_name) { - return irpass::GloablDataAccessRuleChecker::run(root, kernel_name); -} - -} // namespace irpass - -} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/auto_diff.cpp b/quadrants/transforms/auto_diff/auto_diff.cpp new file mode 100644 index 0000000000..177a42a345 --- /dev/null +++ b/quadrants/transforms/auto_diff/auto_diff.cpp @@ -0,0 +1,84 @@ +#include "quadrants/transforms/auto_diff/forward_state_spill.h" +#include "quadrants/transforms/auto_diff/ir_shaping.h" +#include "quadrants/transforms/auto_diff/make_adjoint.h" +#include "quadrants/transforms/auto_diff/make_dual.h" +#include "quadrants/transforms/auto_diff/post_adjoint_cleanup.h" +#include "quadrants/transforms/auto_diff/validation.h" + +#include "quadrants/ir/analysis.h" +#include "quadrants/ir/ir.h" +#include "quadrants/ir/transforms.h" + +namespace quadrants::lang { + +namespace irpass { + +void auto_diff(IRNode *root, const CompileConfig &config, AutodiffMode autodiff_mode, bool use_stack) { + QD_AUTO_PROF; + if (autodiff_mode == AutodiffMode::kReverse) { + if (auto *root_block = root->cast()) { + // Top-level range-fors at the kernel root become the offload-level loop. Walk each direct child once + // and refuse cross-iteration global RAW on a needs_grad SNode before the AD transformation runs - any + // diagnostic from inside `MakeAdjoint` would already point at the partially-rewritten reverse IR and + // be much harder to read. StructFor offloads are out of scope: their loop variable is an SNode index + // produced by the runtime, not a user-controlled integer arithmeticed against itself, so the + // `out[i-1]` shape this guard targets does not arise. + for (auto &s : root_block->statements) { + if (auto *rf = s->cast()) + offload_level_global_cross_iter_raw_check(rf->body.get()); + } + } + regulate_tensor_typed_statements(root); + if (use_stack) { + auto IB = identify_independent_blocks(root); + reverse_outer_loops(root, IB); + + for (auto ib : IB) { + promote_ssa_to_local_var(ib); + replace_local_var_with_stacks(ib, config.ad_stack_size); + type_check(root, config); + + // Drop AdStackAllocas whose pushed value is recomputable from already-stack-backed allocas + args + + // const + loop-index. Trades forward-pass adstack memory traffic (one push + one load per iter per + // intermediate spill) for cloned arithmetic in the reverse scope, which `BackupSSA::generic_visit` + // generates on demand via the same RecomputableChainCloner path. Must run after + // `replace_local_var_with_stacks` (so the analyzer sees the AdStackAlloca shape, not the alloca-with-store + // shape promote_ssa_to_local_var emits) and before `make_adjoint` (so the reverse pass is generated + // against the cleaned forward IR with no spurious push/load scaffolding). + eliminate_recomputable_ad_stack_pushes(ib); + type_check(root, config); + + make_adjoint(ib); + type_check(root, config); + backup_ssa(ib); + // After MakeAdjoint emits the reverse-pass body, an outer-loop-invariant value pulled into the + // reverse direction by `accumulate_unary_operand_checked` becomes a fresh `AdStackLoadTopStmt` per + // user-side use; in straight-line unrolled IR those reads coalesce to one load per stack per block, + // which the dedicated pass handles before the IR reaches `irpass::analysis::verify_if_debug`. + coalesce_ad_stack_loads(ib); + irpass::analysis::verify_if_debug(root, config); + } + } else { + auto IB = identify_independent_blocks(root); + reverse_outer_loops(root, IB); + type_check(root, config); + for (auto ib : IB) { + make_adjoint(ib); + type_check(root, config); + backup_ssa(ib); + coalesce_ad_stack_loads(ib); + irpass::analysis::verify_if_debug(root, config); + } + } + } else if (autodiff_mode == AutodiffMode::kForward) { + // Forward mode autodiff + Block *block = root->as(); + make_dual(block); + } + type_check(root, config); + irpass::analysis::verify_if_debug(root, config); +} + +} // namespace irpass + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/auto_diff_common.h b/quadrants/transforms/auto_diff/auto_diff_common.h new file mode 100644 index 0000000000..ca067045b7 --- /dev/null +++ b/quadrants/transforms/auto_diff/auto_diff_common.h @@ -0,0 +1,439 @@ +#pragma once + +#include "quadrants/ir/analysis.h" +#include "quadrants/ir/ir.h" +#include "quadrants/ir/statements.h" +#include "quadrants/ir/transforms.h" +#include "quadrants/ir/visitors.h" +#include "quadrants/transforms/utils.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace quadrants::lang { + +// ---------------------------------------------------------------------------- +// Shared helpers used across multiple autodiff translation units. +// ---------------------------------------------------------------------------- + +template +Stmt *insert_const(const DataType &dtype, Stmt *stmt, const T &value, bool insert_before_me = false) { + auto type = dtype.ptr_removed(); + Stmt *zero = nullptr; + if (insert_before_me) + zero = stmt->insert_before_me(Stmt::make(TypedConstant(type.get_element_type(), value))); + else + zero = stmt->insert_after_me(Stmt::make(TypedConstant(type.get_element_type(), value))); + + if (type->is()) { + auto t_dtype = type->as(); + std::vector values(t_dtype->get_num_elements(), zero); + if (insert_before_me) { + zero = zero->insert_before_me(Stmt::make(values)); + } else { + zero = zero->insert_after_me(Stmt::make(values)); + } + zero->ret_type = type; + } + return zero; +} + +class IndependentBlockMetaData { + public: + bool is_ib = true; + bool is_smallest_ib = true; +}; + +class NonLinearOps { + public: + inline static const std::set ternary_collections{TernaryOpType::select}; + inline static const std::set unary_collections{ + UnaryOpType::abs, UnaryOpType::sin, UnaryOpType::cos, UnaryOpType::tan, UnaryOpType::tanh, UnaryOpType::asin, + UnaryOpType::acos, UnaryOpType::exp, UnaryOpType::log, UnaryOpType::sqrt, UnaryOpType::rsqrt}; + inline static const std::set binary_collections{BinaryOpType::mul, BinaryOpType::div, + BinaryOpType::atan2, BinaryOpType::pow, + BinaryOpType::min, BinaryOpType::max}; +}; + +// ---------------------------------------------------------------------------- +// Recomputable chain analyzer + cloner: decide whether a forward SSA value can +// be reconstructed in the reverse-pass scope from already-stack-backed allocas, +// kernel args, constants, and loop indices, and clone the DAG at a target +// insertion point. Cross-stage shared infrastructure: used by +// `EliminateRecomputableAdStackPushes` (forward_state_spill stage) to drop +// pushes whose values are recomputable, and by `BackupSSA::generic_visit` +// (post_adjoint_cleanup stage) to clone such chains in place of cross-block +// SSA reads. +// ---------------------------------------------------------------------------- + +// Returns true iff `stmt`'s transitive operand DAG terminates at recomputable leaves via side-effect-free +// interior ops only. Used by `EliminateRecomputableAdStackPushes` and `BackupSSA::generic_visit` to decide +// whether a forward SSA value can be reconstructed in the reverse-pass scope from already-stack-backed allocas +// + kernel-args + constants + loop indices, instead of being spilled to a per-iteration adstack or to +// `BackupSSA::load`'s last-iteration plain alloca. +// +// Recomputable leaves: AdStackLoadTopStmt (re-readable via cloned load), AdStackAllocaStmt (the stack itself, +// shared not cloned), ArgLoadStmt (kernel-arg, immutable within the launch), ConstStmt, LoopIndexStmt (clonable +// to read the reverse-direction loop's index, which matches the forward iteration the reverse is currently +// processing). Side-effect-free interior ops: UnaryOp, BinaryOp, TernaryOp, MatrixPtr, GlobalPtr, ExternalPtr. +// +// `GlobalLoadStmt` and `ExternalPtrAccessStmt`-style reads are intentionally not recomputable: the underlying +// global memory may be mutated mid-kernel by a sibling task, so reading at a different IR position can yield a +// different value. `LocalLoadStmt` similarly aliases mutable allocas; the reverse pass must read its forward +// value through the dedicated spill machinery. +// +// Caller passes a `cache` to share memoization across multiple queries on the same DAG (diamond shapes). +class RecomputableChainAnalyzer { + public: + static bool is_recomputable(Stmt *stmt, std::unordered_map &cache) { + auto it = cache.find(stmt); + if (it != cache.end()) + return it->second; + // Tentatively false to break cycles in pathological IR (real SSA DAGs are acyclic, but the cache also + // serves as a visited set during recursion). + cache[stmt] = false; + bool result = check(stmt, cache); + cache[stmt] = result; + return result; + } + + private: + static bool check(Stmt *stmt, std::unordered_map &cache) { + // Recomputable leaves: ConstStmt, ArgLoadStmt, AdStackLoadTopStmt, AdStackAllocaStmt. LoopIndexStmt is + // intentionally excluded - cloning a LoopIndexStmt copies the reference to the forward RangeForStmt, + // but the cloned consumer lives inside the reverse RangeForStmt (a separate stmt with its own loop + // index), so the cloned read points at undefined state and silently double-accumulates gradients + // (`test_adstack_sum_linear`). A future MakeAdjoint coordination pass that tracks + // forward-RangeFor-to-reverse-RangeFor mapping could lift this restriction. + // + // AdStackLoadTopStmt as a leaf is correct under the dominance + control-flow-consumer + self-load + // guards in `EliminateRecomputableAdStackPushes::run_one_pass`. The reasoning: + // + // - In FORWARD: the eliminated stack S's body push dominates each `AdStackLoadTopStmt(S)` (the + // dominance guard), and the chain leaves' pushes dominate the chain's evaluation point. Each + // stack referenced by a chain leaf has a stable "top" value within one forward iteration body + // (no pops in forward), so substituting load_top(S) with the chain re-evaluates to the same + // value (no iteration shift). + // + // - In REVERSE: `MakeAdjoint` visits forward stmts in reverse order, emitting reverse code at a + // cursor that advances one position per visit. For a forward stmt F at position P_F, the + // emitted reverse code lands at cursor position roughly inversely correlated with P_F. The + // dominance guard ensures every chain-leaf stack T has its body push at a position P_T with + // P_T < P (where P is load_top(S)'s forward position). Therefore T's pop in reverse (emitted + // when MakeAdjoint visits T's body push) lands AFTER the chain consumer's reverse emission + // (the consumer's forward position is P > P_T). At the consumer's reverse cursor, T has not + // been popped yet, so load_top(T) returns T's iter-k-push value - matching what the original + // load_top(S) returned in forward iter k. + // + // The control-flow-consumer guard in `run_one_pass` covers a separate issue: stacks whose load_tops + // are direct operands of IfStmt cond / RangeFor begin/end. `MakeAdjoint::visit(IfStmt)` runs a + // dedicated snap-stack fixup that ONLY triggers when the cond is a bare `AdStackLoadTopStmt` (line + // 2168-2202). Eliminating the cond stack converts the cond into a compound stmt, the snap-stack + // does not trigger, and the reverse cond falls back to BackupSSA's load(op) which is single-slot + // last-iter only - silent gradient corruption on multi-iter loops. + if (stmt->is() || stmt->is() || stmt->is() || + stmt->is()) { + return true; + } + // GlobalLoadStmt as recomputable: the load reads a SNode value via GlobalPtrStmt. The cloned chain in + // the reverse pass re-issues the same load - safe iff the global is not mutated between the forward + // chain evaluation and the cloned re-read. Within a single kernel execution, forward writes complete + // before reverse runs, so the reverse re-read sees the kernel's final post-write state. If a global is + // mutated by the forward and then re-read by the reverse clone, the values can differ. + // + // Most rigid-step kernel chain leaves are reads of input parameters (mass, inertia, joint params, + // morphology) that the kernel does NOT write. For those, the re-read returns the same value. + // Conservative analysis: a future safety check could prove the SNode is read-only in the kernel; until + // then this path relies on the test suite's gradient-correctness asserts to surface any mutated-global + // miscompilation. + bool is_interior = stmt->is() || stmt->is() || stmt->is() || + stmt->is() || stmt->is() || stmt->is() || + stmt->is(); + if (!is_interior) { + return false; + } + auto operands = stmt->get_operands(); + for (auto *op : operands) { + if (op == nullptr) + continue; + if (!is_recomputable(op, cache)) + return false; + } + return true; + } +}; + +// Clones the SSA chain rooted at `src` into the IR, inserting cloned stmts before `insert_point`. Returns the +// cloned root. Per-stmt cache shared across one resolution materializes each SSA value at most once: diamond +// DAGs see two consumers but get one shared clone. `AdStackAllocaStmt` is treated as a leaf and shared (not +// cloned) - the stack itself is a unique storage handle that must not be duplicated. +// +// Pop-ordering safety: cloned `AdStackLoadTopStmt`s read the live top at the cloned position. `MakeAdjoint` +// emits `AdStackPopStmt` for each surviving `AdStackPushStmt`, and the existing reverse-pass scheme places +// the pop AFTER all uses of that stack within the reverse iteration (uses include both the original +// `AdStackLoadTopStmt`s emitted by `ReplaceLocalVarWithStacks` and the consumers' clones). For loop-carried +// allocas the pop fires early to expose the iteration's INPUT primal as the new top, which is exactly the +// value the recomputed chain needs - the existing per-consumer clone path at `BackupSSA::generic_visit` line +// ~2697 relies on this same property and has been correct in production. +class RecomputableChainCloner { + public: + static Stmt *clone_at(Stmt *src, Stmt *insert_point, std::unordered_map &cache) { + auto it = cache.find(src); + if (it != cache.end()) + return it->second; + Stmt *cloned = nullptr; + if (src->is()) { + // The alloca is shared, not cloned: every load reads the same physical stack. + cloned = src; + } else if (src->is() || src->is() || src->is()) { + auto cloned_unique = src->clone(); + cloned = insert_point->insert_before_me(std::move(cloned_unique)); + // For AdStackLoadTopStmt clones, the cloned stmt's `stack` operand still points at the original + // AdStackAllocaStmt - that's the desired sharing. + } else { + // Compound op: clone first, then walk operands and rewire each to a recursive clone. + auto cloned_unique = src->clone(); + cloned = insert_point->insert_before_me(std::move(cloned_unique)); + int n = src->num_operands(); + for (int i = 0; i < n; i++) { + auto *op = src->operand(i); + if (op != nullptr) { + Stmt *new_op = clone_at(op, cloned, cache); + cloned->set_operand(i, new_op); + } + } + } + cache[src] = cloned; + return cloned; + } +}; + +// ---------------------------------------------------------------------------- +// ADTransform: shared base for reverse-mode (MakeAdjoint) and forward-mode +// (MakeDual) IR builders. All methods are inline so derived classes in +// separate translation units can use them without ODR concerns. +// ---------------------------------------------------------------------------- +class ADTransform : public IRVisitor { + protected: + Stmt *constant(float32 x, DataType dtype = PrimitiveType::unknown) { + dtype.set_is_pointer(false); + if (!dtype->is()) + return insert(TypedConstant(x)); + + auto tensor_type = dtype->as(); + auto num_elements = tensor_type->get_num_elements(); + std::vector values; + for (int i = 0; i < num_elements; i++) { + values.push_back(insert(TypedConstant(x))); + } + auto matrix_init_stmt = insert(values); + matrix_init_stmt->ret_type = tensor_type; + return matrix_init_stmt; + } + + // utils + Stmt *sgn(Stmt *inp) { + return insert(UnaryOpType::sgn, load(inp)); + } + + // utils + Stmt *negate(Stmt *inp) { + return insert(UnaryOpType::neg, load(inp)); + } + + Stmt *sqrt(Stmt *inp) { + return insert(UnaryOpType::sqrt, load(inp)); + } + + Stmt *rsqrt(Stmt *inp) { + return insert(UnaryOpType::rsqrt, load(inp)); + } + + Stmt *mul(Stmt *op1, Stmt *op2) { + return insert(BinaryOpType::mul, load(op1), load(op2)); + } + + Stmt *sqr(Stmt *op1) { + return mul(op1, op1); + } + + Stmt *add(Stmt *op1, Stmt *op2) { + return insert(BinaryOpType::add, load(op1), load(op2)); + } + + Stmt *cmp_lt(Stmt *op1, Stmt *op2) { + return insert(BinaryOpType::cmp_lt, load(op1), load(op2)); + } + + Stmt *sub(Stmt *op1, Stmt *op2) { + return insert(BinaryOpType::sub, load(op1), load(op2)); + } + + Stmt *div(Stmt *op1, Stmt *op2) { + return insert(BinaryOpType::div, load(op1), load(op2)); + } + + Stmt *sel(Stmt *op1, Stmt *op2, Stmt *op3) { + return insert(TernaryOpType::select, load(op1), load(op2), load(op3)); + } + + Stmt *cos(Stmt *op1) { + return insert(UnaryOpType::cos, load(op1)); + } + + Stmt *sin(Stmt *op1) { + return insert(UnaryOpType::sin, load(op1)); + } + + Stmt *log(Stmt *op1) { + return insert(UnaryOpType::log, load(op1)); + } + + Stmt *tan(Stmt *op1) { + return insert(UnaryOpType::tan, load(op1)); + } + + Stmt *tanh(Stmt *op1) { + return insert(UnaryOpType::tanh, load(op1)); + } + + Stmt *exp(Stmt *op1) { + return insert(UnaryOpType::exp, load(op1)); + } + + Stmt *pow(Stmt *op1, Stmt *op2) { + return insert(BinaryOpType::pow, load(op1), load(op2)); + } + + public: + virtual Stmt *insert_grad_stmt(std::unique_ptr &&stmt) = 0; + + template + Stmt *insert(Args &&...args) { + return insert_grad_stmt(Stmt::make(args...)); + } + + template + Stmt *insert_const_for_grad(const DataType &dtype, Stmt *stmt, const T &val) { + auto zero = insert(TypedConstant(dtype.ptr_removed().get_element_type(), val)); + if (dtype.ptr_removed()->is()) { + auto t_dtype = dtype.ptr_removed()->as(); + std::vector values(t_dtype->get_num_elements(), zero); + zero = insert(values); + zero->ret_type = dtype.ptr_removed(); + } + return zero; + } + + void visit(AllocaStmt *alloca) override { + // do nothing. + } + + void visit(AdStackAllocaStmt *alloca) override { + // do nothing. + } + + void visit(ArgLoadStmt *stmt) override { + // do nothing. + } + + void visit(GetElementStmt *stmt) override { + // do nothing + } + + void visit(LoopIndexStmt *stmt) override { + // do nothing. + } + + void visit(MatrixPtrStmt *stmt) override { + // do nothing. + } + + void visit(PrintStmt *print_stmt) override { + // do nothing + } + + void visit(ConstStmt *const_stmt) override { + // do nothing + } + + void visit(ReturnStmt *stmt) override { + // do nothing + } + + void visit(WhileControlStmt *stmt) override { + QD_NOT_IMPLEMENTED + } + + void visit(ContinueStmt *stmt) override { + QD_NOT_IMPLEMENTED; + } + + void visit(WhileStmt *stmt) override { + QD_NOT_IMPLEMENTED + } + + void visit(GlobalPtrStmt *stmt) override { + // do nothing + } + + void visit(ExternalPtrStmt *stmt) override { + // do nothing + } + + void visit(ExternalTensorShapeAlongAxisStmt *stmt) override { + // do nothing + } + + void visit(DecorationStmt *stmt) override { + // do nothing + } + + Stmt *load(Stmt *alloc) { + QD_ASSERT(alloc != nullptr); + if (alloc->is() || alloc->is()) { + return insert(alloc); + } else { + // non alloca + return alloc; + } + } + + bool gradients_stopped(GlobalLoadStmt *stmt, SNode *snode) { + for (auto block = stmt->parent; block; block = block->parent_block()) { + for (auto s : block->stop_gradients) { + if (s == snode) { + return true; + } + } + } + return false; + } + + void visit(AssertStmt *stmt) override { + // do nothing + } + + void visit(RangeAssumptionStmt *stmt) override { + // do nothing + } + + void visit(LinearizeStmt *stmt) override { + // do nothing + } + + void visit(IntegerOffsetStmt *stmt) override { + // do nothing + } + + void visit(RandStmt *stmt) override { + QD_ERROR("RandStmt not supported in AutoDiff for now."); + } +}; + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp b/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp new file mode 100644 index 0000000000..6230174365 --- /dev/null +++ b/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp @@ -0,0 +1,479 @@ +#include "quadrants/transforms/auto_diff/auto_diff_common.h" +#include "quadrants/transforms/auto_diff/forward_state_spill.h" + +namespace quadrants::lang { + +namespace { + +// Eliminate AdStackAllocaStmts whose pushed value is recomputable from already-stack-backed allocas, kernel +// args, constants, and loop indices. Runs between `ReplaceLocalVarWithStacks` and `MakeAdjoint`, so the reverse +// pass is generated against the cleaned IR (no spurious AdStackPushStmt / AdStackLoadTopStmt scaffolding for +// values the reverse can reconstruct on the fly via cloned forward DAGs). +// +// Eligibility per AdStackAllocaStmt S in an independent block: +// 1. S is written by exactly one AdStackPushStmt (single-push pattern). Multi-push allocas hold loop-carried +// state where each iteration's push depends on the previous - the reverse pass cannot reconstruct +// iteration k's value from iteration (k-1)'s value, so the stack is genuinely needed. +// 2. The pushed value's transitive operand DAG is recomputable per RecomputableChainAnalyzer (leaves at +// AdStackLoadTop / AdStackAlloca / ArgLoad / Const / LoopIndex; interior side-effect-free ops only). +// 3. S has no AdStackLoadTopStmt with `return_ptr=true` consumers (those return a pointer aliasing the slot +// and a follow-up store would not be modeled by an SSA replacement). +// +// Action on eligibility: replace every `AdStackLoadTopStmt(S)` with the original pushed SSA value (the +// `AdStackPushStmt::v`), then erase the push and the alloca. The pushed SSA value's chain stays in the forward +// IR; downstream consumers in the forward pass now reference it directly as SSA, and `MakeAdjoint` plus +// `BackupSSA` reconstruct the chain on demand in the reverse pass via the same DAG-clone path that +// `BackupSSA::generic_visit` exercises for cross-block SSA references. +// +// Iterates to fixed point: eliminating one stack can newly expose another stack's chain as recomputable +// (S2's pushed value chained through `AdStackLoadTopStmt(S1)`; once S1 is gone, the chain may collapse into +// pure SSA chained through S1's pushed-value chain instead). Each pass eliminates at least one stack or +// terminates. +// +// Cost model: this pass trades forward-pass adstack pushes (memory write + top-pointer bump per push) for +// extra arithmetic in the reverse pass (the cloned DAG re-executes the forward chain). For pure-arithmetic +// chains rooted at a single loop-carried alloca - the dominant shape on Genesis-style rigid-step kernels - +// the win is typically 5-10x: each push is a memory op crossing the L1 boundary on every iteration, while +// the recomputed arithmetic stays in registers and reuses warmed-up sin/cos/exp pipeline state. +class EliminateRecomputableAdStackPushes { + public: + static void run(Block *ib) { + // Iterate to fixed point. Each pass either eliminates at least one stack or terminates. The hard cap + // protects against analysis bugs (e.g. an eligibility check that returns true on the same stmt twice + // in a row); under a correct implementation each pass strictly reduces the AdStackAllocaStmt count, so + // the actual iteration bound is the number of stacks at entry. 1024 is well above any realistic kernel. + for (int i = 0; i < 1024; i++) { + if (!run_one_pass(ib)) + return; + } + } + + private: + // True iff `s` is a literal-zero ConstStmt or a MatrixInitStmt whose every element is a literal-zero + // ConstStmt. Matches the init-zero push pattern that ReplaceLocalVarWithStacks emits right after each + // AdStackAllocaStmt: scalar allocas get a ConstStmt(0), tensor-typed allocas get a MatrixInitStmt of + // ConstStmt(0)s. Both are erased along with the alloca during elimination - they only matter if a load + // could observe them before the body push fires, which the standard PromoteSSA2LocalVar + + // ReplaceLocalVarWithStacks codegen never produces (loads always follow the body push within the loop + // iter that emits them). + static bool is_zero_init_value(Stmt *s) { + if (auto *c = s->cast()) { + return c->val.equal_value(0); + } + if (auto *mi = s->cast()) { + for (auto *elem : mi->values) { + auto *c = elem->cast(); + if (c == nullptr || !c->val.equal_value(0)) + return false; + } + return true; + } + return false; + } + + // Returns true if any stack was eliminated this pass. + static bool run_one_pass(Block *ib) { + // Collect every AdStackAllocaStmt anywhere within the IB. The IB is the root of a contiguous AD scope, so + // a single walk is sufficient. + std::vector stacks; + auto collected = irpass::analysis::gather_statements(ib, [&](Stmt *s) { return s->is(); }); + for (auto *s : collected) { + stacks.push_back(s->as()); + } + + bool modified = false; + std::unordered_map recomputable_cache; + + for (auto *stack : stacks) { + // Re-classify users of `stack` each iteration: a previous elimination on a downstream stack may have + // removed users that previously disqualified `stack`. + std::vector pushes; + std::vector load_tops; + bool disqualified = false; + + auto users = irpass::analysis::gather_statements(ib, [&](Stmt *s) { + if (auto *p = s->cast()) + return p->stack == stack; + if (auto *lt = s->cast()) + return lt->stack == stack; + if (auto *po = s->cast()) + return po->stack == stack; + if (auto *aa = s->cast()) + return aa->stack == stack; + if (auto *la = s->cast()) + return la->stack == stack; + return false; + }); + for (auto *user : users) { + if (auto *p = user->cast()) { + pushes.push_back(p); + } else if (auto *lt = user->cast()) { + if (lt->return_ptr) { + // return_ptr=true returns a pointer into the slot; replacing with an SSA value loses the + // pointer-identity contract. Keep the stack. + disqualified = true; + break; + } + load_tops.push_back(lt); + } else if (user->is() || user->is() || + user->is()) { + // Pre-MakeAdjoint, pop / adj-acc / load-top-adj should not appear yet for this stack. If they do, + // some upstream pass has already touched it - keep as-is to avoid double-rewrites. + disqualified = true; + break; + } + } + if (disqualified || pushes.empty()) { + continue; + } + + // ReplaceLocalVarWithStacks emits a const-zero "init" AdStackPushStmt immediately after the + // AdStackAllocaStmt as the stack's initial value (auto_diff.cpp:867-872 in pre-fix tree). Real + // user-level allocas in a loop body have at most ONE additional "body" push per iteration. Treat + // const-zero pushes whose value is a `ConstStmt` with all-zero `TypedConstant` as inits and require at + // most one non-init "body" push for elimination eligibility. Multi-body-push allocas hold loop-carried + // state (each iteration's push depends on the previous), so the reverse pass cannot reconstruct + // iteration k from iteration k-1 and the stack must stay. + AdStackPushStmt *body_push = nullptr; + std::vector init_pushes; + for (auto *p : pushes) { + if (is_zero_init_value(p->v)) { + init_pushes.push_back(p); + } else { + if (body_push != nullptr) { + disqualified = true; + break; + } + body_push = p; + } + } + if (disqualified || body_push == nullptr) { + continue; + } + + Stmt *pushed_val = body_push->v; + if (!RecomputableChainAnalyzer::is_recomputable(pushed_val, recomputable_cache)) { + continue; + } + // Loop-carried self-reference: when the pushed value's transitive operand DAG includes an + // AdStackLoadTopStmt of the SAME stack we are about to eliminate, the value is `read prev iter from + // stack -> compute -> push next iter`. Rewriting load_top($S) -> pushed_value_SSA leaves a self-cycle + // in the SSA graph (`acc_new = acc_new + sin(a)`) which both corrupts the IR and loses the iteration + // recurrence the stack was carrying. The recomputable-chain analyzer happily accepts such chains + // because AdStackLoadTopStmt is a recomputable leaf in general; the additional check below specifically + // disqualifies self-loaded stacks. + // + // The classic shape this protects: `acc = 0.0; for j in range(N): acc += sin(...)`. After + // PromoteSSA2LocalVar + ReplaceLocalVarWithStacks the IR has `$acc = stack_alloc; init push; + // for j: $tmp = stack_load_top $acc; $new = $tmp + sin(...); push $acc, val=$new`. The chain `$tmp + + // sin(...)` is recomputable per the leaf rules (LoadTop is a leaf), but rewriting load_top($acc) to + // ($tmp + sin(...)) makes the SSA graph self-reference $new and the reverse pass loses every + // iteration's accumulator state. + // Read-before-write protection: substituting `AdStackLoadTopStmt(S)` with the body push's value SSA + // chain only preserves semantics when the body push DOMINATES every load_top in the forward IR - + // that is, every load_top reads what the body push just wrote in the same iteration. The + // PromoteSSA2LocalVar + ReplaceLocalVarWithStacks pipeline emits exactly this shape for required-def + // spills: `alloca; push val=def; load_top` in document order, with the push immediately following + // the def and loads occurring later in the same iteration. + // + // Loop-carried recurrences violate the dominance rule. Take the canonical Fibonacci shape + // `p, q = q, p + q` lowered to IR: + // + // $tmp_p = load_top($p) # reads PREVIOUS iter's push into $p + // $tmp_q = load_top($p) + load_top($q) + // push $p, val=$tmp_p # writes NEXT iter's value into $p + // push $q, val=$tmp_q + // + // Here load_top($p) reads BEFORE the body push of $p within the same iter, so it observes iter + // (k-1)'s value, not iter k's. Substituting `load_top($p) -> $tmp_p`'s SSA chain (which reads + // load_top($q) at iter k) gives iter k's q-value instead of iter (k-1)'s p-value, off by one + // iteration and producing zero gradients on the Fibonacci-style regression test pinned by + // `test_ad_fibonacci_index`. + // + // The check below asserts that for every load_top in `load_tops`, the body push is its predecessor + // within the same containing block (or, equivalently, load_top comes AFTER the body push in + // document order at the body push's block level). Loads that live inside nested control flow + // (`IfStmt`, `RangeForStmt`) under the body push's block are also fine because the body push + // dominates them. We approximate dominance with the lexical "push's block contains load_top's + // ancestor block AND push position < load_top's ancestor's position in that block" check. + auto block_position = [](Stmt *s) -> int { + Block *b = s->parent; + if (b == nullptr) + return -1; + for (size_t i = 0; i < b->statements.size(); i++) { + if (b->statements[i].get() == s) + return static_cast(i); + } + return -1; + }; + Block *push_block = body_push->parent; + int push_pos = block_position(body_push); + bool dominates_all_loads = true; + for (auto *lt : load_tops) { + Stmt *cursor = lt; + Block *cursor_block = cursor->parent; + while (cursor_block != nullptr && cursor_block != push_block) { + cursor = cursor_block->parent_stmt(); + if (cursor == nullptr) + break; + cursor_block = cursor->parent; + } + if (cursor_block != push_block) { + // load_top is outside the push's block scope: cannot establish dominance. + dominates_all_loads = false; + break; + } + int cursor_pos = block_position(cursor); + if (cursor_pos <= push_pos) { + // load_top precedes the body push (or its enclosing container does): it reads the previous + // iteration's value, not iter k's. Substitution would shift iterations by one. + dominates_all_loads = false; + break; + } + } + if (!dominates_all_loads) { + continue; + } + + // Reverse-position correctness for chain leaves pushed in the same block as S. + // + // After elimination, every reverse stmt that consumed `load_top(S)` instead consumes the chain + // (cloned by `BackupSSA::generic_visit` at the consumer's reverse position). The cloned chain reads + // `load_top(T)` for each chain leaf T, where the read happens at the consumer's reverse cursor. + // + // `MakeAdjoint` visits forward stmts in REVERSE order, emitting at a cursor that advances per visit. + // For a forward stmt at position P, its reverse emissions land "later" in reverse block iff P is + // smaller. T's pop is emitted when `MakeAdjoint` visits T's body push at position P_T. The cloned + // chain at the consumer's reverse position reads `load_top(T)` POST-pop iff T's pop has fired by then, + // i.e. iff visit(P_T) precedes visit(C_pos), i.e. iff P_T > C_pos for every consumer C of S's + // load_tops. + // + // Forward chain evaluates at S's push position P_S, before any of T's same-block body pushes (since + // we already required P_T > P_S via the dominance check above for S's loads). So forward chain reads + // T's pre-iter-k-push value. Reverse clone post-pop also returns T's pre-iter-k-push value (= iter k + // INPUT for loop-carried T). Match. + // + // Without this check, the canonical Fibonacci-style shape (`p, q = q, p + q; b[q] += a[q]`) where S + // stores `p + q` and chain leaves p_stack / q_stack are pushed AFTER S but BEFORE the GlobalPtr + // index consumer would silently corrupt gradients: the reverse clone at the GlobalPtr's adjoint + // emission reads p_stack and q_stack PRE-pop (iter k POST-push values), summing to a different value + // than the forward chain at P_S used. + // + // Chain leaves T whose stacks have NO body push in S's containing block are stable within the iter + // (their tops do not change during the iter), so neither forward nor reverse evaluation order shifts + // their values. Excluded from this check. + auto find_consumers_of_load_tops = [&](std::vector &out_positions) { + std::function walk = [&](Block *b, int container_pos_in_push_block) { + for (size_t i = 0; i < b->statements.size(); i++) { + Stmt *s = b->statements[i].get(); + int effective_pos = (b == push_block) ? static_cast(i) : container_pos_in_push_block; + for (auto *op : s->get_operands()) { + if (op == nullptr) + continue; + for (auto *lt : load_tops) { + if (op == lt) { + out_positions.push_back(effective_pos); + break; + } + } + } + if (auto *if_s = s->cast()) { + if (if_s->true_statements) + walk(if_s->true_statements.get(), effective_pos); + if (if_s->false_statements) + walk(if_s->false_statements.get(), effective_pos); + } else if (auto *rf = s->cast()) { + if (rf->body) + walk(rf->body.get(), effective_pos); + } else if (auto *sf = s->cast()) { + if (sf->body) + walk(sf->body.get(), effective_pos); + } + } + }; + walk(push_block, -1); + }; + std::vector consumer_positions; + find_consumers_of_load_tops(consumer_positions); + int max_consumer_pos = -1; + for (int p : consumer_positions) { + if (p > max_consumer_pos) + max_consumer_pos = p; + } + + // Collect all distinct AdStackAllocaStmt referenced via AdStackLoadTopStmt leaves in pushed_val's chain. + std::unordered_set chain_leaf_stacks; + std::unordered_set walked_for_leaves; + std::function collect_leaves = [&](Stmt *s) { + if (walked_for_leaves.count(s)) + return; + walked_for_leaves.insert(s); + if (auto *lt = s->cast()) { + if (auto *as = lt->stack->cast()) + chain_leaf_stacks.insert(as); + return; + } + if (s->is() || s->is() || s->is()) + return; + for (auto *op : s->get_operands()) { + if (op != nullptr) + collect_leaves(op); + } + }; + collect_leaves(pushed_val); + + bool reverse_safe = true; + for (auto *T : chain_leaf_stacks) { + // Find T's body pushes (non-init) in push_block. Only same-block pushes can interfere: pushes in + // outer blocks happen once per outer iteration, so their tops are stable across the inner iter. + for (size_t i = 0; i < push_block->statements.size(); i++) { + Stmt *s = push_block->statements[i].get(); + if (auto *p = s->cast()) { + if (p->stack == T && !is_zero_init_value(p->v)) { + if (static_cast(i) <= max_consumer_pos) { + reverse_safe = false; + break; + } + } + } + } + if (!reverse_safe) + break; + } + if (!reverse_safe) { + continue; + } + + bool is_self_loaded = false; + std::unordered_set visited; + std::function walk = [&](Stmt *s) { + if (is_self_loaded || visited.count(s)) + return; + visited.insert(s); + if (auto *lt = s->cast()) { + if (lt->stack == stack) { + is_self_loaded = true; + return; + } + // load_top of a different stack: the operand chain stops here at this leaf for the self-check. + return; + } + if (s->is() || s->is() || s->is()) { + return; // leaf, not a self-load + } + for (auto *op : s->get_operands()) { + if (op != nullptr) + walk(op); + } + }; + walk(pushed_val); + if (is_self_loaded) { + continue; + } + + // Control-flow-cond consumers: by the time the IR reaches this pass, every IfStmt cond and RangeFor begin/end + // is either a bare `AdStackLoadTopStmt` (loop-carried local promoted via `PromoteSSA2LocalVar` -> + // `AdStackAllocaJudger::visit(IfStmt|RangeForStmt)` -> `ReplaceLocalVarWithStacks`) or a stmt outside any + // adstack-bearing alloca (no load_top in the chain at all). `MakeAdjoint::visit(IfStmt)` (auto_diff.cpp around + // line 2452) caps the bare-load_top case with a 1-push-per-execution snap-stack when the if body itself pushes + // to the cond's backing stack so the reverse cond reads the forward-time value, not a stack top mutated by the + // body. + // + // If we eliminate a stack whose load_top IS the bare cond of an IfStmt / inner RangeFor, the rewrite turns the + // cond / loop-bound into an inlined recomputed SSA chain. The snap-stack guard at the `AdStackLoadTopStmt`-cond + // check in `MakeAdjoint::visit(IfStmt)` stops firing (the cond is no longer bare), and `BackupSSA`'s clone path + // positions a fresh `AdStackLoadTopStmt` of an enclosing loop-carried stack at the consumer's IR location - + // which sits BEFORE the per-iter pops, returning the post-iteration value instead of the iteration-k value the + // cond was originally computed against. Net effect: silent gradient corruption in any if/loop nested inside a + // dynamic for-loop with a loop-carried alloca. + // + // Keep stacks whose load_tops have such consumers. The consumer check below trips only on a load_top of THIS + // stack feeding DIRECTLY into a control-flow-shaping operand of another stmt; that is also the exhaustive set + // of unsafe-elim shapes here, because compound-cond cases (arithmetic over a load_top feeding a cond) reach + // this pass with the cond rewritten into a separate adstack-promoted value via the alloca-promotion pipeline + // above and so look like the bare case from this guard's POV. + // `irpass::analysis::gather_statements` walks via `BasicStmtVisitor`, whose visit overrides for + // IfStmt / RangeForStmt / StructForStmt do not invoke the per-stmt test predicate on the container + // itself - only on stmts inside their bodies. Walk container stmts manually here. + bool has_control_flow_consumer = false; + std::function walk_block = [&](Block *b) { + if (has_control_flow_consumer) + return; + for (auto &owned : b->statements) { + Stmt *s = owned.get(); + if (auto *if_s = s->cast()) { + for (auto *lt : load_tops) { + if (if_s->cond == lt) { + has_control_flow_consumer = true; + return; + } + } + if (if_s->true_statements) + walk_block(if_s->true_statements.get()); + if (has_control_flow_consumer) + return; + if (if_s->false_statements) + walk_block(if_s->false_statements.get()); + } else if (auto *rf = s->cast()) { + for (auto *lt : load_tops) { + if (rf->begin == lt || rf->end == lt) { + has_control_flow_consumer = true; + return; + } + } + if (rf->body) + walk_block(rf->body.get()); + } else if (auto *sf = s->cast()) { + if (sf->body) + walk_block(sf->body.get()); + } else if (auto *off = s->cast()) { + if (off->body) + walk_block(off->body.get()); + if (off->tls_prologue) + walk_block(off->tls_prologue.get()); + if (off->bls_prologue) + walk_block(off->bls_prologue.get()); + if (off->bls_epilogue) + walk_block(off->bls_epilogue.get()); + if (off->tls_epilogue) + walk_block(off->tls_epilogue.get()); + } + } + }; + walk_block(ib); + if (has_control_flow_consumer) { + continue; + } + + // Eligible: rewrite each load_top to use the pushed value directly. The pushed value lives in the + // forward IR and dominates each load_top by SSA construction (the load_top reads what the push wrote; + // both are inside the same loop body in the dynamic-loop case, with the push preceding the load_top). + for (auto *lt : load_tops) { + irpass::replace_all_usages_with(ib, lt, pushed_val); + lt->parent->erase(lt); + } + // Erase init-zero pushes (they only matter if a load could observe them, but the rewriting above just + // routed every load to the body-pushed SSA value), the body push, and the alloca itself. + for (auto *p : init_pushes) { + p->parent->erase(p); + } + body_push->parent->erase(body_push); + stack->parent->erase(stack); + + // Invalidate the cache: the eliminated stack might have been referenced by other recomputable chains + // we evaluated earlier in this pass, but those evaluations only matter for stacks we eliminate THIS + // pass; re-running from scratch on the next iteration recomputes them. + recomputable_cache.clear(); + modified = true; + } + return modified; + } +}; + +} // namespace + +void eliminate_recomputable_ad_stack_pushes(Block *ib) { + EliminateRecomputableAdStackPushes::run(ib); +} + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/forward_state_spill.cpp b/quadrants/transforms/auto_diff/forward_state_spill.cpp new file mode 100644 index 0000000000..37ea0f5f40 --- /dev/null +++ b/quadrants/transforms/auto_diff/forward_state_spill.cpp @@ -0,0 +1,667 @@ +#include "quadrants/transforms/auto_diff/auto_diff_common.h" +#include "quadrants/transforms/auto_diff/forward_state_spill.h" + +namespace quadrants::lang { + +namespace { + +// ============================================================================ +// PromoteSSA2LocalVar: hoist forward-pass SSA defs that the reverse pass will +// re-read into AllocaStmt + LocalStore + LocalLoad. Demand-driven. +// +// Note that SSA does not mean the instruction will be executed at most once. +// For instructions that may be executed multiple times, we treat them as a +// mutable local variables. +// ============================================================================ +class PromoteSSA2LocalVar : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + + explicit PromoteSSA2LocalVar(Block *block) { + alloca_block_ = block; + invoke_default_visitor = true; + execute_once_ = true; + } + + // Demand-driven `required_defs_` set: the SSA defining stmts that downstream consumers actually require to be + // available at every iteration. A consumer requires its operand iff its adjoint formula reads it - the precise + // set is the operands of any non-linear unary / binary / ternary op (per `NonLinearOps::*_collections`), the + // indices of any `GlobalPtrStmt` / `ExternalPtrStmt` (the reverse pass replays the load), and the `cond` of any + // `IfStmt` / the `begin`/`end` of any `RangeForStmt` (the reverse pass clones the control flow). Stmts outside + // this set are left in pure SSA form: their values stay register-resident inside the forward pass, and + // `MakeAdjoint`'s reverse-pass formulas (which never read those values directly) generate correct adjoint + // accumulations against the adstack-backed defs that ARE in the set. Skipping promotion for the rest of the + // body collapses the alloca + LocalStore + LocalLoad triple-multiplier on unrolled IR that emitted a spill + + // reload pair per non-required arithmetic op with no reverse-pass consumer for the spilled value. + // + // Operands of `LocalStoreStmt` are not added because `MakeAdjoint::visit(LocalStoreStmt)` reads only + // `adjoint(stmt->dest)`, not the forward `stmt->val`. Linear ops (add / sub / mod / cmp / neg / floor / ceil / + // cast / logic_not / bit-ops) are likewise excluded because their adjoint formulas read only `adjoint(stmt)`. + static void compute_required_defs(Block *block, std::unordered_set &out) { + std::function walk = [&](Block *b) { + for (auto &owned : b->statements) { + Stmt *stmt = owned.get(); + if (auto *u = stmt->cast()) { + if (NonLinearOps::unary_collections.find(u->op_type) != NonLinearOps::unary_collections.end()) { + out.insert(u->operand); + } + } else if (auto *bin = stmt->cast()) { + if (NonLinearOps::binary_collections.find(bin->op_type) != NonLinearOps::binary_collections.end()) { + out.insert(bin->lhs); + out.insert(bin->rhs); + } + } else if (auto *tern = stmt->cast()) { + if (NonLinearOps::ternary_collections.find(tern->op_type) != NonLinearOps::ternary_collections.end()) { + out.insert(tern->op1); + out.insert(tern->op2); + out.insert(tern->op3); + } + } else if (auto *gp = stmt->cast()) { + for (auto *idx : gp->indices) { + out.insert(idx); + } + } else if (auto *ep = stmt->cast()) { + for (auto *idx : ep->indices) { + out.insert(idx); + } + } else if (auto *mp = stmt->cast()) { + // Reverse-mode `MakeAdjoint::visit(MatrixPtrStmt)` reads `stmt->offset` for dynamic-index adjoint + // routing, so a per-iteration-varying offset producer left in pure SSA would be backed by BackupSSA's + // single overwrite-each-iteration alloca and the reverse pass would read the last forward offset for + // every iteration. Promote it to alloca so `AdStackAllocaJudger::visit(MatrixPtrStmt)` can + // adstack-promote loop-varying offsets. Speculative fix: not exhibited by any failing test in the AD + // suite today, but the analysis is sound and the cost of leaving it unfixed is silent gradient + // corruption on `tensor[i + j]`-style local-tensor indexing inside a serial range-for. + out.insert(mp->offset); + } else if (auto *if_s = stmt->cast()) { + out.insert(if_s->cond); + if (if_s->true_statements) { + walk(if_s->true_statements.get()); + } + if (if_s->false_statements) { + walk(if_s->false_statements.get()); + } + } else if (auto *rf = stmt->cast()) { + out.insert(rf->begin); + out.insert(rf->end); + walk(rf->body.get()); + } else if (auto *sf = stmt->cast()) { + if (sf->body) { + walk(sf->body.get()); + } + } else if (auto *while_s = stmt->cast()) { + if (while_s->body) { + walk(while_s->body.get()); + } + } + } + }; + walk(block); + } + + void visit(Stmt *stmt) override { + if (execute_once_) + return; + if (!(stmt->is() || stmt->is() || stmt->is() || + stmt->is() || stmt->is() || stmt->is())) { + // TODO: this list may be incomplete + return; + } + + // `AllocaStmt`s always need to be hoisted to the top of the IB regardless of consumer analysis: a user-level + // `var = ...` construct inside a loop body must own a fixed slot at the IB's entry so every iteration shares it + // (cross-iteration accumulators are exactly the shape that drives the hoist). The demand-driven gate only + // applies to value-producing stmts (UnaryOp / BinaryOp / TernaryOp / GlobalLoad / LoopIndex) where the + // alloca + LocalStore + LocalLoad triple is purely a reverse-pass-readable spill - those skip when no consumer + // requires the value. The `LocalStoreStmt`s emitted here are placeholders that `ReplaceLocalVarWithStacks` + // rewrites into `AdStackPushStmt`s downstream. + if (stmt->is()) { + auto dtype = stmt->ret_type.ptr_removed(); + auto alloc = Stmt::make(dtype); + auto alloc_ptr = alloc.get(); + QD_ASSERT(alloca_block_); + alloca_block_->insert(std::move(alloc), 0); + immediate_modifier_->replace_usages_with(stmt, alloc_ptr); + + auto zero = insert_const(dtype, stmt, 0); + zero->insert_after_me(Stmt::make(alloc_ptr, zero)); + stmt->parent->erase(stmt); + return; + } + + if (required_defs_.find(stmt) == required_defs_.end()) { + return; + } + + auto alloc = Stmt::make(stmt->ret_type.ptr_removed()); + auto alloc_ptr = alloc.get(); + QD_ASSERT(alloca_block_); + alloca_block_->insert(std::move(alloc), 0); + auto load = stmt->insert_after_me(Stmt::make(alloc_ptr)); + immediate_modifier_->replace_usages_with(stmt, load); + // Create the load first so that the operand of the store does not get rewritten to point at the load (the + // SSA value `stmt` is still the right thing to spill; only the downstream consumers see the load). + stmt->insert_after_me(Stmt::make(alloc_ptr, stmt)); + } + + void visit(RangeForStmt *stmt) override { + auto old_execute_once = execute_once_; + execute_once_ = false; // loop body may be executed many times + stmt->body->accept(this); + execute_once_ = old_execute_once; + } + + static void run(Block *block) { + PromoteSSA2LocalVar pass(block); + compute_required_defs(block, pass.required_defs_); + pass.immediate_modifier_ = std::make_unique(block); + block->accept(&pass); + } + + private: + Block *alloca_block_{nullptr}; + bool execute_once_; + std::unordered_set required_defs_; + // ImmediateIRModifier collapses each `replace_usages_with` from a whole-tree walk (O(N)) to a constant-time + // operand-pointer rewrite: the modifier gathers every (consumer, operand_index) pair feeding any existing stmt + // once at construction (one O(N) pass), and per-replacement just looks up the table for `old_stmt` and rewrites + // each consumer's operand pointer. Without it, `PromoteSSA2LocalVar` ran K (= number of promoted defs) full IR + // walks - O(K*N), the dominant Quadrants-IR-side compile cost on unrolled reverse-mode bodies. + std::unique_ptr immediate_modifier_; +}; + +// ============================================================================ +// AdStackAllocaJudger: per-AllocaStmt analysis. Decides whether an alloca's +// runtime sequence of values must be preserved on an AdStack so the reverse +// pass can re-read every forward-time value, or whether a single overwrite- +// each-iteration backing slot suffices. +// ============================================================================ +class AdStackAllocaJudger : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + // Find the usage of the stmt recursively along the LocalLoadStmt + void visit(LocalLoadStmt *stmt) override { + if (stmt->src == target_alloca_) { + local_loaded_ = true; + target_alloca_ = stmt; + } + } + + // Track whether the alloca has any store at all so a load-only alloca (no adstack-relevant data flow either + // direction) can short-circuit `run()` regardless of what the per-op visitors below find. The decision of + // whether the alloca needs adstack promotion is made entirely by the precise visitors: non-linear unary / + // binary / ternary, GlobalPtr / ExternalPtr index, and IfStmt / RangeForStmt bound. Plus an alloca that is + // both loaded and stored anywhere in the IB is treated as loop-carried, which is needed for kernels like + // `for j: p, q = q, p + q` where the reverse pass routes the gradient through the cross-iteration + // recurrence and BackupSSA's single overwrite-each-iteration alloca cannot back the read-after-write across + // iterations. The visit-order-dependent load+store evidence here is conservative: any alloca with both a + // load and a store inside the IB triggers it, including pure accumulators whose adjoint formulas don't + // actually need per-iteration values - the slight over-promotion cost is the price of correctness on + // Fibonacci-style recurrences (silent gradient corruption otherwise). + void visit(LocalStoreStmt *stmt) override { + if (stmt->dest == target_alloca_backup_) { + load_only_ = false; + // Gate the load+store-implies-stack-needed rule on actually being inside a dynamic RangeForStmt at the + // point this evidence accumulates. The rule's purpose is to preserve cross-iteration RAW dependencies + // (`for j: p, q = q, p + q` Fibonacci-style) that BackupSSA's single overwrite-each-iteration alloca + // cannot back. With no enclosing dynamic for-loop the IB body executes once: there is no cross- + // iteration RAW to preserve, and the "load+store" pattern is just an in-block accumulator that the + // reverse pass handles via plain SSA cloning. Promoting such allocas under a static-unrolled loop body + // wastes one AdStack per accumulator (one push per unrolled-iter store + one load_top per unrolled- + // iter load) without any reverse-pass consumer needing per-iter replay - that is the unrolled-overhead + // bug Plan B targets. + // + // `dynamic_for_depth_` is incremented in `visit(RangeForStmt)` and decremented on exit. The judger + // walks the IB tree from the alloca's enclosing block, so depth here reflects exactly the nesting of + // *dynamic* for-loops between the alloca and the current load/store. StructFor / WhileStmt do not + // increment because their bodies still execute per-iter and need the same RAW protection (StructFor + // is the kernel-level offload-loop in some cases, but its body is a per-thread independent block; + // load+store there is the same shape as for a top-level alloca). + if (local_loaded_ && dynamic_for_depth_ > 0) { + is_stack_needed_ = true; + } + } + } + + // Check if the alloca is load only + void visit(AtomicOpStmt *stmt) override { + if (stmt->dest == target_alloca_backup_) + load_only_ = false; + } + + // The stack is needed if the alloca serves as the index of any global variables. Same cursor-vs-backup + // pattern as visit(IfStmt)/visit(RangeForStmt) below: `index` is always a value-producing stmt (typically a + // `LocalLoadStmt` reading the alloca, or a `ConstStmt`), never the alloca itself. The raw `index == + // target_alloca_` comparison only matches the first load's instance the `visit(LocalLoadStmt)` cursor + // advanced to - any subsequent load of the same alloca used as a different GlobalPtr index slips through. + // Resolve the LocalLoad chain and compare `ll->src` against `target_alloca_backup_` to catch every load. + void visit(GlobalPtrStmt *stmt) override { + if (is_stack_needed_) + return; + for (const auto &index : stmt->indices) { + auto *index_ll = index->cast(); + if (index_ll && index_ll->src == target_alloca_backup_) + is_stack_needed_ = true; + } + } + + void visit(ExternalPtrStmt *stmt) override { + if (is_stack_needed_) + return; + for (const auto &index : stmt->indices) { + auto *index_ll = index->cast(); + if (index_ll && index_ll->src == target_alloca_backup_) + is_stack_needed_ = true; + } + } + + // Reverse-mode `MakeAdjoint::visit(MatrixPtrStmt)` reads `stmt->offset` for dynamic-index adjoint routing, so + // a runtime-varying offset whose value comes from this alloca needs adstack promotion - otherwise BackupSSA + // backs the offset with a single overwrite-each-iteration slot and the reverse pass routes every iteration's + // adjoint into the last forward offset's slot. Same cursor-vs-backup pattern as the index visitors above. + void visit(MatrixPtrStmt *stmt) override { + if (is_stack_needed_) + return; + auto *offset_ll = stmt->offset->cast(); + if (offset_ll && offset_ll->src == target_alloca_backup_) + is_stack_needed_ = true; + } + + // Check whether the target alloca is fed into a non-linear unary op. Same cursor-vs-backup pattern as + // visit(GlobalPtrStmt) above: `stmt->operand` is a value-producing stmt (typically LocalLoad), never the + // alloca itself, so resolve the LocalLoad chain and compare against the backup. + void visit(UnaryOpStmt *stmt) override { + if (is_stack_needed_) + return; + if (NonLinearOps::unary_collections.find(stmt->op_type) != NonLinearOps::unary_collections.end()) { + auto *operand_ll = stmt->operand->cast(); + if (operand_ll && operand_ll->src == target_alloca_backup_) + is_stack_needed_ = true; + } + } + + // Check whether the target alloca is fed into a non-linear binary op. Same cursor-vs-backup pattern. + void visit(BinaryOpStmt *stmt) override { + if (is_stack_needed_) + return; + if (NonLinearOps::binary_collections.find(stmt->op_type) != NonLinearOps::binary_collections.end()) { + auto *lhs_ll = stmt->lhs->cast(); + auto *rhs_ll = stmt->rhs->cast(); + if ((lhs_ll && lhs_ll->src == target_alloca_backup_) || (rhs_ll && rhs_ll->src == target_alloca_backup_)) + is_stack_needed_ = true; + } + } + + // Check whether the target alloca is fed into a non-linear ternary op. Same cursor-vs-backup pattern. + void visit(TernaryOpStmt *stmt) override { + if (is_stack_needed_) + return; + if (NonLinearOps::ternary_collections.find(stmt->op_type) != NonLinearOps::ternary_collections.end()) { + auto *op1_ll = stmt->op1->cast(); + auto *op2_ll = stmt->op2->cast(); + auto *op3_ll = stmt->op3->cast(); + if ((op1_ll && op1_ll->src == target_alloca_backup_) || (op2_ll && op2_ll->src == target_alloca_backup_) || + (op3_ll && op3_ll->src == target_alloca_backup_)) + is_stack_needed_ = true; + } + } + + // Check whether the target alloca feeds the condition of an if stmt. `stmt->cond` is always a + // value-producing stmt - typically a direct `LocalLoadStmt` reading the alloca, but also commonly a + // `BinaryOpStmt` wrapping such a load (e.g. `j < i+1`). Walk the expression chain to catch every load of + // the target alloca: the raw `stmt->cond == target_alloca_` comparison the old code used only matched the + // first-visited load's instance, and a direct `cast` still misses the BinaryOp case that + // `visit(BinaryOpStmt)` cannot catch (comparison ops are linear and so not in `NonLinearOps`). Covers the + // shape defensively: IR simplification currently collapses most BinaryOp-wrapped conds before the judger + // sees them, so no failing regression test pins it today, but the fix is structurally correct for future + // IR changes that preserve the BinaryOp wrapping. + void visit(IfStmt *stmt) override { + if (is_stack_needed_) + return; + + if (feeds_target_alloca(stmt->cond)) { + is_stack_needed_ = true; + return; + } + + if (stmt->true_statements) + stmt->true_statements->accept(this); + if (stmt->false_statements) + stmt->false_statements->accept(this); + } + + // Check whether the target alloca feeds the begin or end of a range-for bound. Under reverse-mode AD, if an + // inner for-loop's bound is an enclosing loop-carried counter (the canonical triangular-nested + // `for k in range(j)` shape, or the `range(j+1)` / `range(n-i)` shapes where the bound is a linear arithmetic + // expression of a loop-carried alloca), its reverse clone must read the bound from the per-iteration forward + // value; without an adstack the reverse pass sees only the last forward value and the inner loop over- or + // under-runs, silently corrupting gradients for the earliest inner indices (those visited most often across + // outer iterations). This check is the only thing that promotes such a loop-counter alloca - + // `visit(LocalStoreStmt)`'s `local_loaded_` short-circuit does not fire because the counter is only LOAD-ed + // inside the inner-loop bound, not LOAD-then-STORE-ed. Walk the expression chain through + // `feeds_target_alloca` so both direct LocalLoads (`range(j)`) and LocalLoads nested under linear ops + // (`range(j+1)`, `range(n-i)`, ...) trigger promotion. The BinaryOp-wrapped case is defensively covered: IR + // simplification currently collapses most such bounds before the judger sees them, so no failing regression + // test pins it today, but the walker is structurally correct for future IR changes that preserve the + // wrapping. The raw-cast direct `LocalLoadStmt` variant pinned by `test_adstack_inner_for_bound_is_enclosing + // _loop_index` remains covered - that shape takes the first branch of the walker trivially. + void visit(RangeForStmt *stmt) override { + if (is_stack_needed_) + return; + + if (feeds_target_alloca(stmt->begin) || feeds_target_alloca(stmt->end)) { + is_stack_needed_ = true; + return; + } + + dynamic_for_depth_++; + stmt->body->accept(this); + dynamic_for_depth_--; + } + + static bool run(AllocaStmt *target_alloca) { + AdStackAllocaJudger judger; + judger.target_alloca_ = target_alloca; + judger.target_alloca_backup_ = target_alloca; + judger.dynamic_for_depth_ = 0; + target_alloca->parent->accept(&judger); + return (!judger.load_only_) && judger.is_stack_needed_; + } + + private: + // Recursively walk a value expression to decide whether it transitively reads `target_alloca_backup_` via a + // `LocalLoadStmt`. Used by `visit(IfStmt)` and `visit(RangeForStmt)` to detect the target alloca feeding a + // bound or condition even when wrapped by linear ops (e.g. `range(j+1)`, `j < i+1`). Linear binary/unary + // ops are traversed because `visit(BinaryOpStmt)`/`visit(UnaryOpStmt)` only flag *non-linear* ops - their + // linear-op path does not otherwise promote the alloca. `ConstStmt`s and unrelated values return false and + // terminate the recursion; the walker is always finite because SSA IR guarantees acyclic operand graphs. + bool feeds_target_alloca(Stmt *expr) const { + if (auto *ll = expr->cast()) { + return ll->src == target_alloca_backup_; + } + if (auto *bop = expr->cast()) { + return feeds_target_alloca(bop->lhs) || feeds_target_alloca(bop->rhs); + } + if (auto *uop = expr->cast()) { + return feeds_target_alloca(uop->operand); + } + if (auto *top = expr->cast()) { + return feeds_target_alloca(top->op1) || feeds_target_alloca(top->op2) || feeds_target_alloca(top->op3); + } + return false; + } + + Stmt *target_alloca_; + Stmt *target_alloca_backup_; + bool is_stack_needed_ = false; + bool local_loaded_ = false; + bool load_only_ = true; + // Nesting depth of dynamic `RangeForStmt` containers between the alloca's enclosing block and the current + // visit cursor. Static-unrolled `qd.static(range(...))` loops are removed by the AST transformer before the + // judger sees the IR, so they do not contribute to depth. The load+store-implies-stack-needed rule fires + // only when this depth is positive; see the rationale in `visit(LocalStoreStmt)`. + int dynamic_for_depth_ = 0; +}; + +// ============================================================================ +// ReplaceLocalVarWithStacks: rewrite each AllocaStmt that AdStackAllocaJudger +// approved into AdStackAllocaStmt + AdStackPush + AdStackLoadTop. Tensor-typed +// allocas need extra care because per-slot stack writes are not a reaching +// def for the store-to-load forwarding walker. +// ============================================================================ +class ReplaceLocalVarWithStacks : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + int ad_stack_size; + DelayedIRModifier delayed_modifier_; + + explicit ReplaceLocalVarWithStacks(int ad_stack_size) : ad_stack_size(ad_stack_size) { + } + + void visit(AllocaStmt *alloc) override { + bool is_stack_needed = AdStackAllocaJudger::run(alloc); + if (is_stack_needed) { + auto dtype = alloc->ret_type.ptr_removed(); + auto stack_alloca = Stmt::make(dtype, ad_stack_size); + auto stack_alloca_ptr = stack_alloca.get(); + + alloc->replace_with(VecStatement(std::move(stack_alloca))); + + // Note that unlike AllocaStmt, AdStackAllocaStmt does NOT have an 0 as + // initial value. Therefore here we push an initial 0 value. + auto zero = insert_const(dtype, stack_alloca_ptr, 0); + zero->insert_after_me(Stmt::make(stack_alloca_ptr, zero)); + } + } + + void visit(LocalLoadStmt *stmt) override { + if (stmt->src->is()) { + auto stack_load = Stmt::make(stmt->src); + stack_load->ret_type = stmt->ret_type; + + stmt->replace_with(std::move(stack_load)); + return; + } + + // Slot load from a stack-backed tensor. After `visit(MatrixPtrStmt)`, `stmt->src` is of the form + // `MatrixPtrStmt(AdStackLoadTopStmt(stack, return_ptr=true), offset)`. A direct load through that pointer + // leaves the store-to-load forwarding walker in `ir/control_flow_graph.cpp` with no reaching definition, + // because the only producer for the stack's top slots is an `AdStackPushStmt` (tagged `ir_traits::Load`, + // invisible to `get_store_destination`). Replace the load with a full-tensor `AdStackLoadTopStmt` + // materialized into a fresh regular `AllocaStmt`, then re-subscript it - a plain alloca + LocalStore + // sequence is a shape the reach-in walker can trace end-to-end. + if (stmt->src->is()) { + auto matrix_ptr = stmt->src->as(); + if (matrix_ptr->origin->is() && matrix_ptr->origin->as()->return_ptr) { + auto stack = matrix_ptr->origin->as()->stack; + QD_ASSERT(stack->is()); + auto tensor_type = stack->ret_type.ptr_removed(); + + auto full_load = Stmt::make(stack); + full_load->ret_type = tensor_type; + auto full_load_ptr = full_load.get(); + + auto fresh_alloca = Stmt::make(tensor_type); + auto fresh_alloca_ptr = fresh_alloca.get(); + fresh_alloca->ret_type = tensor_type; + fresh_alloca->ret_type.set_is_pointer(true); + + auto fresh_store = Stmt::make(fresh_alloca_ptr, full_load_ptr); + + auto new_matrix_ptr = Stmt::make(fresh_alloca_ptr, matrix_ptr->offset); + new_matrix_ptr->ret_type = stmt->ret_type; + + auto new_load = Stmt::make(new_matrix_ptr.get()); + new_load->ret_type = stmt->ret_type; + + stmt->insert_before_me(std::move(full_load)); + stmt->insert_before_me(std::move(fresh_alloca)); + stmt->insert_before_me(std::move(fresh_store)); + stmt->insert_before_me(std::move(new_matrix_ptr)); + stmt->replace_with(std::move(new_load)); + } + } + } + + void visit(LocalStoreStmt *stmt) override { + if (stmt->dest->is()) { + auto matrix_ptr_stmt = stmt->dest->as(); + if (matrix_ptr_stmt->origin->is()) { + auto stack_top_stmt = matrix_ptr_stmt->origin->as(); + QD_ASSERT(stack_top_stmt->return_ptr == true); + + if (!stack_top_stmt->ret_type.ptr_removed()->is()) { + return; + } + + auto tensor_type = stack_top_stmt->ret_type.ptr_removed()->as(); + auto num_elements = tensor_type->get_num_elements(); + + if (matrix_ptr_stmt->offset->is()) { + /* + [Static index] + Load the full current top as a tensor via `AdStackLoadTopStmt` and merge the new value at `offset` + using a boolean mask + `select`. Mirrors the dynamic-index lowering below so that every slot of the + new pushed tensor derives from either `stmt->val` or the loaded top tensor and the IR contains no + per-slot `LocalLoadStmt` on a stack-backed `MatrixPtrStmt`. + + Why that invariant matters: the store-to-load forwarding walker in `ir/control_flow_graph.cpp` does + not treat `AdStackPushStmt` as a reaching definition (it is tagged `ir_traits::Load`, so + `get_store_destination` returns nothing for it), so a `LocalLoadStmt(MatrixPtrStmt(stack_top_ptr, + i))` inserted here has no reaching def and ends up reading an uninitialized adjoint slot in the + reverse kernel. Keep the `AdStackLoadTopStmt(stack)` + mask-select shape when touching this path. + + Fwd: + $1 = alloca <4 x i32> + $2 = matrix ptr $1, 2 // offset = 2 + $3 : local store $2, $val + + Replaced: + $1 = alloca <4 x i32> + + $2 = matrix init [$val, $val, $val, $val] + $3 = matrix init [false, false, true, false] // mask with `offset == i` + + $4 = ad stack load top (full tensor) $1 + $5 = select $3, $2, $4 + + $6 : stack push $1, $5 + */ + int offset = matrix_ptr_stmt->offset->as()->val.val_int32(); + + QD_ASSERT(offset < num_elements); + + auto tensor_shape = tensor_type->get_shape(); + auto cmp_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::u1); + + std::vector val_values(num_elements, stmt->val); + std::vector mask_values(num_elements); + for (int i = 0; i < num_elements; i++) { + mask_values[i] = insert_const(PrimitiveType::u1, stmt, i == offset ? 1 : 0, true); + } + + auto matrix_val = Stmt::make(val_values); + matrix_val->ret_type = tensor_type; + + auto matrix_mask = Stmt::make(mask_values); + matrix_mask->ret_type = cmp_tensor_type; + + auto matrix_alloca_value = Stmt::make(stack_top_stmt->stack); + matrix_alloca_value->ret_type = tensor_type; + + auto matrix_select = Stmt::make(TernaryOpType::select, matrix_mask.get(), matrix_val.get(), + matrix_alloca_value.get()); + matrix_select->ret_type = tensor_type; + + auto stack_push = Stmt::make(stack_top_stmt->stack, matrix_select.get()); + + stmt->insert_before_me(std::move(matrix_val)); + stmt->insert_before_me(std::move(matrix_mask)); + stmt->insert_before_me(std::move(matrix_alloca_value)); + stmt->insert_before_me(std::move(matrix_select)); + stmt->replace_with(std::move(stack_push)); + + return; + + } else { + /* + [Dynamic index] + Fwd: + $1 = alloca <4 x i32> + $2 = matrix ptr $1, $offset // offset = 2 + $3 : local store $2, $val + + Replaced: + $1 = alloca <4 x i32> + + $2 = matrix init [$val, $val, $val, $val] + + $3 = matrix init [$offset, $offset, $offset, $offset] + $4 = matrix init [0, 1, 2, 3] + + $5 = bin_eq $3, $4 + $6 = select $5, $2, $1 + + $7 : store $1, $6 + */ + auto tensor_type = stack_top_stmt->ret_type.ptr_removed()->as(); + auto num_elements = tensor_type->get_num_elements(); + + auto tensor_shape = tensor_type->get_shape(); + auto index_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::i32); + + std::vector val_values(num_elements, stmt->val); + std::vector offset_values(num_elements, matrix_ptr_stmt->offset); + std::vector index_values(num_elements); + for (int i = 0; i < num_elements; i++) { + index_values[i] = insert_const(PrimitiveType::i32, stmt, i, true); + } + + auto matrix_val = Stmt::make(val_values); + matrix_val->ret_type = tensor_type; + + auto matrix_offset = Stmt::make(offset_values); + matrix_offset->ret_type = index_tensor_type; + + auto matrix_index = Stmt::make(index_values); + matrix_index->ret_type = index_tensor_type; + + auto cmp_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::u1); + auto matrix_eq = Stmt::make(BinaryOpType::cmp_eq, matrix_offset.get(), matrix_index.get()); + matrix_eq->ret_type = cmp_tensor_type; + + auto matrix_alloca_value = Stmt::make(stack_top_stmt->stack); + matrix_alloca_value->ret_type = tensor_type; + + auto matrix_select = Stmt::make(TernaryOpType::select, matrix_eq.get(), matrix_val.get(), + matrix_alloca_value.get()); + matrix_select->ret_type = tensor_type; + + auto stack_push = Stmt::make(stack_top_stmt->stack, matrix_select.get()); + + stmt->insert_before_me(std::move(matrix_val)); + stmt->insert_before_me(std::move(matrix_offset)); + stmt->insert_before_me(std::move(matrix_index)); + stmt->insert_before_me(std::move(matrix_eq)); + stmt->insert_before_me(std::move(matrix_alloca_value)); + stmt->insert_before_me(std::move(matrix_select)); + stmt->replace_with(std::move(stack_push)); + + return; + } + } + } + + // Non Tensor-type + if (stmt->dest->is()) + stmt->replace_with(Stmt::make(stmt->dest, stmt->val)); + } + + void visit(MatrixPtrStmt *stmt) override { + if (stmt->origin->is()) { + auto stack_top = Stmt::make(stmt->origin, true /*is_ptr*/); + stack_top->ret_type = stmt->origin->ret_type; + stack_top->ret_type.set_is_pointer(true); + + Stmt *stack_top_stmt = stack_top.get(); + stmt->insert_before_me(std::move(stack_top)); + + auto new_matrix_ptr_stmt = Stmt::make(stack_top_stmt, stmt->offset); + new_matrix_ptr_stmt->ret_type = stmt->ret_type; + stmt->replace_with(std::move(new_matrix_ptr_stmt)); + } + } +}; + +} // namespace + +void promote_ssa_to_local_var(Block *block) { + PromoteSSA2LocalVar::run(block); +} + +void replace_local_var_with_stacks(Block *block, int ad_stack_size) { + ReplaceLocalVarWithStacks pass(ad_stack_size); + block->accept(&pass); +} + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/forward_state_spill.h b/quadrants/transforms/auto_diff/forward_state_spill.h new file mode 100644 index 0000000000..8fd696b8d0 --- /dev/null +++ b/quadrants/transforms/auto_diff/forward_state_spill.h @@ -0,0 +1,30 @@ +#pragma once + +namespace quadrants::lang { + +class Block; + +// Stage: Forward-state spill. Pre-MakeAdjoint passes that prepare the forward +// IR so the reverse pass can re-read the per-iteration values it needs. + +// Hoist forward-pass SSA defs that the reverse pass will re-read (operands of +// non-linear ops, indices of GlobalPtr / ExternalPtr, range-for bounds, if +// conds) into AllocaStmt + LocalStore + LocalLoad. Demand-driven: defs no +// reverse formula reads stay in pure SSA form. +void promote_ssa_to_local_var(Block *block); + +// Replace each AllocaStmt the reverse pass needs to read across iterations +// with an AdStackAllocaStmt + AdStackPushStmt / AdStackLoadTopStmt. The +// AllocaJudger (file-private) decides which allocas need the promotion. +void replace_local_var_with_stacks(Block *block, int ad_stack_size); + +// Drop AdStackAllocaStmts whose pushed value is recomputable from already- +// stack-backed allocas + kernel args + constants + loop indices. Trades +// forward-pass adstack memory traffic for cloned arithmetic in the reverse +// scope, which `BackupSSA::generic_visit` synthesizes on demand via the +// shared RecomputableChainCloner. Must run after `replace_local_var_with_stacks` +// (so the analyzer sees the AdStack shape) and before `make_adjoint` (so the +// reverse pass is generated against the cleaned forward IR). +void eliminate_recomputable_ad_stack_pushes(Block *ib); + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/ir_shaping.cpp b/quadrants/transforms/auto_diff/ir_shaping.cpp new file mode 100644 index 0000000000..1b0a746be7 --- /dev/null +++ b/quadrants/transforms/auto_diff/ir_shaping.cpp @@ -0,0 +1,646 @@ +#include "quadrants/transforms/auto_diff/auto_diff_common.h" +#include "quadrants/transforms/auto_diff/ir_shaping.h" + +namespace quadrants::lang { + +namespace { + +// ============================================================================ +// RegulateTensorTypedStatements: rewrite tensor-typed local/global stores that +// touch a sub-tensor through MatrixPtr into an explicit gather + matrix-init +// + scalar store, so downstream passes never see a partial-tensor store. +// ============================================================================ + +class RegulateTensorTypedStatements : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + DelayedIRModifier delayed_modifier_; + + explicit RegulateTensorTypedStatements() { + } + + template + void process_store_stmt(Store *stmt) { + QD_ASSERT(stmt->template is() || stmt->template is()); + + if (stmt->dest->template is()) { + auto matrix_ptr_stmt = stmt->dest->template as(); + auto orig_stmt = matrix_ptr_stmt->origin; + + if (!orig_stmt->ret_type.ptr_removed()->template is()) { + return; + } + + auto tensor_type = orig_stmt->ret_type.ptr_removed()->template as(); + auto num_elements = tensor_type->get_num_elements(); + + if (matrix_ptr_stmt->offset->template is()) { + /* + [Static index] + Fwd: + $0 = alloca <4 x i32> + $1 = load $0 + $2 = matrix ptr $1, 2 // offset = 2 + $3 : local store $2, $val + + Replaced: + $0 = alloca <4 x i32> + $1 = load $0 + $2 = matrix ptr $1, 2 // --> erase + + $3 = matrix ptr $1, 0 + $4 = load $3 + + $5 = matrix ptr $1, 1 + $6 = load $5 + + $7 = matrix ptr $1, 3 + $8 = load $7 + + $9 = matrix init [$4, $6, $val, $8] + + $10 : store $0, $9 + */ + int offset = matrix_ptr_stmt->offset->template as()->val.val_int32(); + + QD_ASSERT(offset < num_elements); + + std::vector values; + for (int i = 0; i < num_elements; i++) { + if (i == offset) { + values.push_back(stmt->val); + continue; + } + + auto const_i = insert_const(PrimitiveType::i32, stmt, i, true); + auto matrix_ptr_stmt_i = Stmt::make(orig_stmt, const_i); + matrix_ptr_stmt_i->ret_type = tensor_type->get_element_type(); + + auto local_load_stmt_i = Stmt::make(matrix_ptr_stmt_i.get()); + local_load_stmt_i->ret_type = tensor_type->get_element_type(); + + values.push_back(local_load_stmt_i.get()); + + stmt->insert_before_me(std::move(matrix_ptr_stmt_i)); + stmt->insert_before_me(std::move(local_load_stmt_i)); + } + + auto matrix_init_stmt = Stmt::make(values); + matrix_init_stmt->ret_type = tensor_type; + + auto store_stmt = Stmt::make(orig_stmt, matrix_init_stmt.get()); + stmt->insert_before_me(std::move(matrix_init_stmt)); + stmt->replace_with(std::move(store_stmt)); + + return; + + } else { + /* + [Dynamic index] + Fwd: + $0 = alloca <4 x i32> + $1 = load $0 + $2 = matrix ptr $1, $offset // offset = 2 + $3 : local store $2, $val + + Replaced: + $0 = alloca <4 x i32> + + $1 = load $0 + $2 = matrix init [$val, $val, $val, $val] + + $3 = matrix init [$offset, $offset, $offset, $offset] + $4 = matrix init [0, 1, 2, 3] + + $5 = bin_eq $3, $4 + $6 = select $5, $2, $1 + + $7 : store $0, $6 + */ + auto tensor_type = orig_stmt->ret_type.ptr_removed()->template as(); + auto num_elements = tensor_type->get_num_elements(); + + auto tensor_shape = tensor_type->get_shape(); + auto index_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::i32); + + std::vector val_values(num_elements, stmt->val); + std::vector offset_values(num_elements, matrix_ptr_stmt->offset); + std::vector index_values(num_elements); + for (int i = 0; i < num_elements; i++) { + index_values[i] = insert_const(PrimitiveType::i32, stmt, i, true); + } + + auto matrix_val = Stmt::make(val_values); + matrix_val->ret_type = tensor_type; + + auto matrix_offset = Stmt::make(offset_values); + matrix_offset->ret_type = index_tensor_type; + + auto matrix_index = Stmt::make(index_values); + matrix_index->ret_type = index_tensor_type; + auto cmp_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::u1); + auto matrix_eq = Stmt::make(BinaryOpType::cmp_eq, matrix_offset.get(), matrix_index.get()); + matrix_eq->ret_type = cmp_tensor_type; + + auto orig_value = Stmt::make(orig_stmt); + orig_value->ret_type = tensor_type; + + auto matrix_select = + Stmt::make(TernaryOpType::select, matrix_eq.get(), matrix_val.get(), orig_value.get()); + matrix_select->ret_type = tensor_type; + + auto store_stmt = Stmt::make(orig_stmt, matrix_select.get()); + + stmt->insert_before_me(std::move(matrix_val)); + stmt->insert_before_me(std::move(matrix_offset)); + stmt->insert_before_me(std::move(matrix_index)); + stmt->insert_before_me(std::move(matrix_eq)); + stmt->insert_before_me(std::move(orig_value)); + stmt->insert_before_me(std::move(matrix_select)); + stmt->replace_with(std::move(store_stmt)); + return; + } + } + } + + void visit(LocalStoreStmt *stmt) override { + process_store_stmt(stmt); + } + + void visit(GlobalStoreStmt *stmt) override { + process_store_stmt(stmt); + } + + static void run(IRNode *root) { + RegulateTensorTypedStatements pass; + root->accept(&pass); + } +}; + +// ============================================================================ +// Independent-Blocks discovery: IBJudger / DupCleaner / IdentifyIBs. +// +// Independent Block (IB): a block (i.e. loop body) whose iterations are +// independent of previous iterations and outer scopes. IBs are where +// MakeAdjoint emits the reverse pass; outside an IB only iteration order +// matters and ReverseOuterLoops handles that. +// ============================================================================ + +class IndependentBlocksJudger : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + + void visit(LocalLoadStmt *stmt) override { + QD_ASSERT(stmt->src->is() || stmt->src->is() || stmt->src->is()); + touched_allocas_.insert(stmt->src); + } + + void visit(LocalStoreStmt *stmt) override { + QD_ASSERT(stmt->dest->is() || stmt->dest->is() || + stmt->dest->is()); + touched_allocas_.insert(stmt->dest); + } + + void visit(AtomicOpStmt *stmt) override { + // We don't need to check the global atomics inside the range for-loops + // because + // 1. If the range for-loop is innermost, they will be captured by + // MakeAdjoint anyway + // 2. If the range for-loop is not innermost, they will be processed by + // another IndependentBlocksJudger + if (is_inside_loop_) + return; + + Stmt *dest = stmt->dest; + if (dest->is()) { + dest = dest->as()->origin; + } + + if (dest->is()) { + if (dest->as() + ->base_ptr->as() + ->ret_type.ptr_removed() + ->as() + ->elements() + .size() > TypeFactory::GRAD_PTR_POS_IN_NDARRAY) { + qualified_glb_operations_ = true; + } + } else { + QD_ASSERT(dest->is()); + if (dest->as()->snode->has_adjoint()) { + qualified_glb_operations_ = true; + } + } + } + + void visit(GlobalLoadStmt *stmt) override { + // We don't need to check the global load inside the range for-loops + // because + // 1. If the range for-loop is innermost, they will be captured by + // MakeAdjoint anyway + // 2. If the range for-loop is not innermost, they will be processed by + // another IndependentBlocksJudger + if (is_inside_loop_) + return; + + Stmt *src = stmt->src; + if (src->is()) { + src = src->as()->origin; + } + + if ((src->is() && src->as() + ->base_ptr->as() + ->ret_type.ptr_removed() + ->as() + ->elements() + .size() > TypeFactory::GRAD_PTR_POS_IN_NDARRAY) || + (src->is() && src->as()->snode->has_adjoint())) { + qualified_glb_operations_ = true; + } + } + + void visit(RangeForStmt *stmt) override { + inner_most_loop_ = false; + is_inside_loop_ = true; + stmt->body->accept(this); + is_inside_loop_ = false; + } + + static void run(IRNode *root, IndependentBlockMetaData &ib_meta_data) { + IndependentBlocksJudger Judger; + Block *block = root->as(); + root->accept(&Judger); + std::set outside_blocks; + // Collect all parent blocks (i.e. outside blocks) of the current block for + // local load/store stmt checks + for (auto b = block->parent_block(); b; b = b->parent_block()) { + if (b) + outside_blocks.insert(b); + } + for (const auto &alloca : Judger.touched_allocas_) { + // Test if the alloca belongs to the current block + if (outside_blocks.find(alloca->parent) != outside_blocks.end()) { + // This block is not an IB since it loads/modifies outside variables + ib_meta_data.is_ib = false; + } + } + + // To judge whether a block is an IB + // - No local load/store to allocas *outside* itself has been strictly + // enforced + + // To judge whether a block is a smallest IB + // - If the #1 is satisfied, either an inner most loop or a block without + // global atomics / global load is an IB + ib_meta_data.is_smallest_ib = ib_meta_data.is_ib && (Judger.qualified_glb_operations_ || Judger.inner_most_loop_); + } + + private: + std::set touched_allocas_; + bool qualified_glb_operations_ = false; + bool inner_most_loop_ = true; + bool is_inside_loop_ = false; +}; + +// Remove the duplicated IBs, remove blocks who are others' children because +// each block should only be processed once +class DuplicateIndependentBlocksCleaner : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + + void check_children_ib(Block *target_block) { + // Remove the block if it is the child of the block being visiting + if (independent_blocks_cleaned_.find(target_block) != independent_blocks_cleaned_.end()) { + independent_blocks_cleaned_.erase(target_block); + } + } + + void visit(StructForStmt *stmt) override { + check_children_ib(stmt->body.get()); + stmt->body->accept(this); + } + void visit(RangeForStmt *stmt) override { + check_children_ib(stmt->body.get()); + stmt->body->accept(this); + } + + static std::set run(const std::vector> &raw_IBs) { + DuplicateIndependentBlocksCleaner cleaner; + // Remove duplicate IBs + for (auto const &item : raw_IBs) { + cleaner.independent_blocks_cleaned_.insert(item.second); + } + // No clean is needed if only one IB exists + if (cleaner.independent_blocks_cleaned_.size() > 1) { + // Check from the block with smallest depth, ensure no duplicate visit + // happens + for (const auto &block : cleaner.independent_blocks_cleaned_) { + block->accept(&cleaner); + } + } + return cleaner.independent_blocks_cleaned_; + } + + private: + std::set independent_blocks_cleaned_; +}; + +class IdentifyIndependentBlocks : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + + void visit(WhileStmt *stmt) override { + QD_ERROR("WhileStmt is not supported in AutoDiff."); + } + + void visit(ContinueStmt *stmt) override { + QD_ERROR("ContinueStmt is not supported in AutoDiff."); + } + + void visit(WhileControlStmt *stmt) override { + QD_ERROR("WhileControlStmt (break) is not supported in AutoDiff."); + } + + void visit_loop_body(Block *block) { + auto ib_meta_data = IndependentBlockMetaData(); + // An IB has no local load/store to allocas *outside* itself + // Note: + // - Local atomics should have been demoted before this pass. + // - It is OK for an IB to have more than two for loops. + // - No global load/atomics operations to the global variables which + // require gradient + if (block->statements.empty()) { + // A empty block shoud be a smallest IB + ib_meta_data.is_ib = true; + ib_meta_data.is_smallest_ib = true; + } else { + IndependentBlocksJudger::run(block, ib_meta_data); + } + + if (ib_meta_data.is_smallest_ib) { + independent_blocks_.push_back({depth_, block}); + } else if (ib_meta_data.is_ib) { + current_ib_ = block; + block->accept(this); + } else { + if (depth_ <= 1) { + QD_ASSERT(depth_ == 1); + // The top level block is already not an IB, store it + independent_blocks_.push_back({depth_ - 1, block}); + } else { + independent_blocks_.push_back({depth_ - 1, block->parent_block()}); + } + } + } + + void visit(StructForStmt *stmt) override { + QD_ASSERT(depth_ == 0); + depth_++; + current_ib_ = stmt->body.get(); + visit_loop_body(stmt->body.get()); + depth_--; + } + + void visit(RangeForStmt *stmt) override { + if (depth_ == 0) { + current_ib_ = stmt->body.get(); + } + depth_++; + visit_loop_body(stmt->body.get()); + depth_--; + } + + static std::set run(IRNode *root) { + IdentifyIndependentBlocks pass; + Block *block = root->as(); + bool has_for = false; + for (auto &s : block->statements) { + if (s->is() || s->is()) { + has_for = true; + } + } + if (!has_for) { + // The whole block is an IB + pass.independent_blocks_.push_back({0, block}); + } else { + root->accept(&pass); + } + // Sort the IBs by their depth from shallow to deep + std::sort( + pass.independent_blocks_.begin(), pass.independent_blocks_.end(), + [](const std::pair &a, const std::pair &b) -> bool { return a.first < b.first; }); + + QD_ASSERT(!pass.independent_blocks_.empty()); + return DuplicateIndependentBlocksCleaner::run(pass.independent_blocks_); + } + + private: + std::vector> independent_blocks_; + int depth_{0}; + Block *current_ib_{nullptr}; +}; + +// ============================================================================ +// ReverseOuterLoops: flip iteration direction on outer (non-IB) for-loops and +// reorder sibling for-loops in non-IB container blocks so the reverse pass +// walks the iteration trace backward. +// ============================================================================ + +class ReverseOuterLoops : public BasicStmtVisitor { + using BasicStmtVisitor::visit; + + private: + explicit ReverseOuterLoops(const std::set &IB) : loop_depth_(0), ib_(IB) { + } + + bool is_ib(Block *block) const { + return std::find(ib_.begin(), ib_.end(), block) != ib_.end(); + } + + // Sibling for-loops inside a non-IB container block execute their reverse-mode companions + // in the container's forward order by default, because MakeAdjoint only touches IB-level bodies + // and nothing else permutes the enclosing order. Reverse-mode AD requires the opposite: if the + // forward body runs `for_A; for_B` and for_B's reverse depends on reads produced by for_A's + // forward run, the reverse pass must execute `rev-for_B; rev-for_A` so for_A's reverse sees the + // adjoints for_B has populated (e.g. `cdof[i]=x[i]; cdofvel[i]=cdof[i]*vel[i]` silently returns + // x.grad=0 otherwise: rev-for_A clears cdof.grad before rev-for_B has populated it). + // + // Naive pairwise swap of for-loop positions is unsafe whenever a non-loop stmt between two + // for-loops feeds the later sibling's SSA operand chain (e.g. a GlobalLoad that supplies a + // dynamic trip count): after the swap, the consumer for-loop ends up before its producer and + // the IR verifier rejects the block. Before swapping, hoist any such producer (and its + // transitive in-block dependencies) to the slot just before the first sibling for-loop. Non-loop + // stmts unrelated to for-loop operands stay at their original indices; memory ordering between + // non-loop stmts is preserved because the hoist keeps them in their original relative order and + // only moves them upward over for-loops (which produce no SSA value and cannot be the source of + // a missed memory read for a non-loop that gets hoisted above them). + // + // The top-level kernel block is handled by `reverse_segments` before this pass, so we only + // reorder inside nested non-IB blocks here. + static void reverse_for_loop_order_in_place(Block *block) { + const int n = (int)block->statements.size(); + std::vector for_indices; + for (int i = 0; i < n; ++i) { + Stmt *s = block->statements[i].get(); + if (s->is() || s->is()) { + for_indices.push_back(i); + } + } + if (for_indices.size() < 2) { + return; + } + const int first_for = for_indices.front(); + + std::unordered_map pos_of; + pos_of.reserve(n); + for (int i = 0; i < n; ++i) { + pos_of[block->statements[i].get()] = i; + } + + // Walk the SSA operand graph of every for-loop (restricted to this block). Any in-block stmt + // that (a) the operand closure reaches and (b) sits at or after `first_for` gets flagged for + // hoisting: after swap, that stmt must precede every for-loop, not just the ones it feeds. + std::unordered_set must_hoist; + std::vector stack; + auto push_if_internal = [&](Stmt *s) { + if (s == nullptr) { + return; + } + auto it = pos_of.find(s); + if (it == pos_of.end() || it->second < first_for) { + return; + } + if (must_hoist.insert(s).second) { + stack.push_back(s); + } + }; + // Seed the hoist frontier from both the for-loop's direct SSA operands (`begin`, `end`) and + // from every stmt nested inside the for-loop's body that references an outer-block stmt as a + // free variable. The body-use gather is what catches the case where the later sibling + // for-loop consumes a non-loop outer-block stmt `S` inside its body (e.g. `for_B: body reads + // S`) rather than through `for_B`'s range bound: `RangeForStmt::get_operands()` returns only + // `{begin, end}`, so without walking the body `S` would miss `must_hoist`, the pairwise swap + // would place `for_B` ahead of `S`, and the IR verifier would reject the SSA violation. + for (int fi : for_indices) { + for (Stmt *op : block->statements[fi]->get_operands()) { + push_if_internal(op); + } + Stmt *for_stmt = block->statements[fi].get(); + irpass::analysis::gather_statements(for_stmt, [&](Stmt *body_stmt) { + for (Stmt *op : body_stmt->get_operands()) { + push_if_internal(op); + } + return false; + }); + } + while (!stack.empty()) { + Stmt *s = stack.back(); + stack.pop_back(); + for (Stmt *op : s->get_operands()) { + push_if_internal(op); + } + } + // For-loops themselves end up in `must_hoist` only because their own operand-closure reached + // them; they do not get hoisted as non-loop producers - strip them here to keep `must_hoist` + // to "non-loop stmts that need to move above all for-loops". + for (int fi : for_indices) { + must_hoist.erase(block->statements[fi].get()); + } + + std::vector> new_stmts; + new_stmts.reserve(n); + // Stmts strictly before `first_for` keep their original slot. + for (int i = 0; i < first_for; ++i) { + new_stmts.push_back(std::move(block->statements[i])); + } + // Hoisted non-loop stmts slot in here, in their original relative order. + for (int i = first_for; i < n; ++i) { + if (must_hoist.count(block->statements[i].get()) != 0) { + new_stmts.push_back(std::move(block->statements[i])); + } + } + // Remainder (for-loops and non-hoisted non-loops) in original order, with for-loops swapped + // pairwise inside this suffix. + std::vector> suffix; + std::vector suffix_for_positions; + for (int i = first_for; i < n; ++i) { + auto &sp = block->statements[i]; + if (!sp) { + continue; + } + bool is_for = sp->is() || sp->is(); + if (is_for) { + suffix_for_positions.push_back((int)suffix.size()); + } + suffix.push_back(std::move(sp)); + } + for (int lo = 0, hi = (int)suffix_for_positions.size() - 1; lo < hi; ++lo, --hi) { + std::swap(suffix[suffix_for_positions[lo]], suffix[suffix_for_positions[hi]]); + } + for (auto &s : suffix) { + new_stmts.push_back(std::move(s)); + } + + QD_ASSERT((int)new_stmts.size() == n); + block->statements.clear(); + for (auto &s : new_stmts) { + block->statements.push_back(std::move(s)); + } + } + + void visit(StructForStmt *stmt) override { + loop_depth_ += 1; + if (!is_ib(stmt->body.get())) { + stmt->body->accept(this); + reverse_for_loop_order_in_place(stmt->body.get()); + } + loop_depth_ -= 1; + } + + void visit(RangeForStmt *stmt) override { + if (loop_depth_ >= 1) { + stmt->reversed = !stmt->reversed; + } + loop_depth_ += 1; + if (!is_ib(stmt->body.get())) { + stmt->body->accept(this); + reverse_for_loop_order_in_place(stmt->body.get()); + } + loop_depth_ -= 1; + } + + // Deliberately no `visit(IfStmt *)` override, although sibling for-loops can live directly inside an if-branch + // block (`true_statements` / `false_statements`) the same way they live inside a for-body. The default + // `BasicStmtVisitor::visit(IfStmt *)` recurses into both branches so inner `RangeForStmt::body`s still get the + // sibling-reorder treatment via the range-for visitor above, but `reverse_for_loop_order_in_place` is never + // invoked on the branch block itself. That is intentional: `MakeAdjoint::visit(IfStmt *)` below emits the adjoint + // if-stmt by iterating each branch's statements in reverse order (`for (int i = size - 1; i >= 0; --i)` in its + // `true_statements` / `false_statements` loops), which achieves the same sibling-for reordering effect that the + // missing override here would provide. Overriding `visit(IfStmt)` in this class is therefore a no-op on the + // generated adjoint code. Keep the comment rather than the override so the visitor-coverage gap is documented. + + int loop_depth_; + std::set ib_; + + public: + static void run(IRNode *root, const std::set &IB) { + ReverseOuterLoops pass(IB); + root->accept(&pass); + } +}; + +} // namespace + +void regulate_tensor_typed_statements(IRNode *root) { + RegulateTensorTypedStatements::run(root); +} + +std::set identify_independent_blocks(IRNode *root) { + return IdentifyIndependentBlocks::run(root); +} + +void reverse_outer_loops(IRNode *root, const std::set &IB) { + ReverseOuterLoops::run(root, IB); +} + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/ir_shaping.h b/quadrants/transforms/auto_diff/ir_shaping.h new file mode 100644 index 0000000000..9a54aac2a6 --- /dev/null +++ b/quadrants/transforms/auto_diff/ir_shaping.h @@ -0,0 +1,28 @@ +#pragma once + +#include + +namespace quadrants::lang { + +class IRNode; +class Block; + +// Stage: IR shaping. Pre-MakeAdjoint passes that mutate or analyze the +// forward IR so the downstream reverse-mode codegen sees the expected shape. + +// Rewrite tensor-typed LocalStore / GlobalStore through MatrixPtr into an +// explicit gather + matrix-init + scalar store - downstream passes assume +// stores never touch a sub-tensor in place. +void regulate_tensor_typed_statements(IRNode *root); + +// Discover every Independent Block (loop body whose iterations carry no +// dependency on previous iterations or outer scopes) under `root`. Returns +// the set of blocks where MakeAdjoint will emit the reverse pass. +std::set identify_independent_blocks(IRNode *root); + +// Flip iteration direction of every outer (non-IB) RangeForStmt so the reverse +// pass walks the iteration trace backward; reorders sibling for-loops inside +// non-IB container blocks accordingly. +void reverse_outer_loops(IRNode *root, const std::set &IB); + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/make_adjoint.cpp b/quadrants/transforms/auto_diff/make_adjoint.cpp new file mode 100644 index 0000000000..0b041fd876 --- /dev/null +++ b/quadrants/transforms/auto_diff/make_adjoint.cpp @@ -0,0 +1,857 @@ +#include "quadrants/transforms/auto_diff/auto_diff_common.h" +#include "quadrants/transforms/auto_diff/make_adjoint.h" + +namespace quadrants::lang { + +namespace { + +// Generate the adjoint version of an independent block +class MakeAdjoint : public ADTransform { + public: + using ADTransform::visit; + Block *current_block; + Block *alloca_block; + // Backup the forward pass (the forward pass might be modified during the + // MakeAdjoint) for search whether a GlobalLoadStmt is inside a for-loop when + // allocating adjoint (see the function `adjoint`) Should be stored + // 1. Before entering a for-loop body + // 2. Before entering a if-stmt + // Should be restored after processing every statement in the two cases above + Block *forward_backup; + // IB root: stays constant across visitor recursion. Used when we need to allocate + // persistent storage that must survive enclosing for-loop iterations (e.g. the + // dedicated ad-stacks that snapshot IfStmt conds in visit(IfStmt)). + Block *ib_root; + std::map adjoint_stmt; + + explicit MakeAdjoint(Block *block) { + current_block = nullptr; + alloca_block = block; + forward_backup = block; + ib_root = block; + } + + static void run(Block *block) { + auto p = MakeAdjoint(block); + block->accept(&p); + } + + // Does `if_stmt`'s true/false body contain any AdStackPushStmt targeting `stack`? Recursive to + // catch pushes nested inside further control flow (if-in-if, if-in-for). Used by visit(IfStmt) + // to gate cond-snapshotting. Must be narrow: snapshotting every if-stmt would add an + // AdStackAllocaStmt per if, and determine_ad_stack_size cannot size stacks whose push/pop pair + // is only reachable through branches its Bellman-Ford walk considers "unreached" -- codegen then + // aborts with "Adaptive autodiff stack's size should have been determined" and the extras also + // spam "Unused autodiff stack should have been eliminated" for every untouched snap stack. Only + // when the body actually pushes onto the cond's backing stack does BackupSSA's reverse-time + // clone of load_top read a post-body value rather than the forward cond (the real bug); in every + // other case the clone is already correct and a snapshot would be dead weight. + static bool block_pushes_to_stack(Block *block, Stmt *stack) { + if (!block) + return false; + for (auto &stmt : block->statements) { + if (auto *push = stmt->cast()) { + if (push->stack == stack) + return true; + } + if (auto *inner_if = stmt->cast()) { + if (block_pushes_to_stack(inner_if->true_statements.get(), stack)) + return true; + if (block_pushes_to_stack(inner_if->false_statements.get(), stack)) + return true; + } + if (auto *inner_for = stmt->cast()) { + if (block_pushes_to_stack(inner_for->body.get(), stack)) + return true; + } + if (auto *inner_for = stmt->cast()) { + if (block_pushes_to_stack(inner_for->body.get(), stack)) + return true; + } + } + return false; + } + + static bool body_pushes_to_stack(IfStmt *if_stmt, Stmt *stack) { + return block_pushes_to_stack(if_stmt->true_statements.get(), stack) || + block_pushes_to_stack(if_stmt->false_statements.get(), stack); + } + + // TODO: current block might not be the right block to insert adjoint + // instructions! + void visit(Block *block) override { + std::vector statements; + // always make a copy since the list can be modified. + for (auto &stmt : block->statements) { + statements.push_back(stmt.get()); + } + std::reverse(statements.begin(), statements.end()); // reverse-mode AD... + for (auto stmt : statements) { + current_block = block; + stmt->accept(this); + } + } + + Stmt *insert_grad_stmt(std::unique_ptr &&stmt) override { + auto ptr = stmt.get(); + current_block->insert(std::move(stmt), -1); + return ptr; + } + + // Accumulate [value] to the adjoint of [primal] + void accumulate(Stmt *primal, Stmt *value) { + auto alloca_ = adjoint(primal); + if (!alloca_ || alloca_->is()) { + return; // primal may be int variable + } + if (alloca_->is()) { + auto alloca = alloca_->cast(); + if (is_real(alloca->ret_type.get_element_type())) { + insert(alloca, load(value)); + } + } else { + QD_ASSERT(alloca_->is()); + auto alloca = alloca_->as(); + auto local_load = insert(alloca); + local_load->ret_type = alloca->ret_type.ptr_removed(); + insert(alloca, add(local_load, value)); + } + } + + Stmt *adjoint(Stmt *stmt) { + DataType adjoint_dtype = stmt->ret_type.ptr_removed(); + if (adjoint_dtype->is()) { + DataType prim_dtype = PrimitiveType::f32; + if (is_real(adjoint_dtype.get_element_type())) { + prim_dtype = adjoint_dtype.get_element_type(); + } + adjoint_dtype = + TypeFactory::get_instance().get_tensor_type(adjoint_dtype->as()->get_shape(), prim_dtype); + } else if (stmt->is()) { + // pass + } else if (!is_real(adjoint_dtype) || stmt->is()) { + return constant(0); + } + + if (adjoint_stmt.find(stmt) == adjoint_stmt.end()) { + // normal SSA cases + + // create the alloca + // auto alloca = + // Stmt::make(get_current_program().config.gradient_dt); + // maybe it's better to use the statement data type than the default type + auto alloca = Stmt::make(adjoint_dtype); + adjoint_stmt[stmt] = alloca.get(); + + // We need to insert the alloca in the block of GlobalLoadStmt when the + // GlobalLoadStmt is not inside the currently-processed range-for. + // Code sample: + // a and b require grad + // Case 1 (GlobalLoadStmt is outside the for-loop, compute 5 times and + // accumulate once, alloca history value is needed): + // for i in range(5): + // p = a[i] + // q = b[i] + // for _ in range(5) + // q += p + + // Case 2 (GlobalLoadStmt is inside the for-loop, compute once and + // accumulate immediately, alloca history value can be discarded): + // for i in range(5): + // q = b[i] + // for _ in range(5) + // q += a[i] + if (stmt->is() && forward_backup->locate(stmt->as()) == -1) { + // Case 1: the GlobalLoadStmt lives in a block outside the currently-processed range-for iteration. Its + // adjoint must persist across all iterations of the inner reversed loop, so the alloca cannot live in the + // current alloca_block (which would be the inner reversed loop body). Walk up from the primal's enclosing + // block until we hit one whose owning statement unconditionally dominates both the forward and the reverse + // code (a loop / offloaded / kernel body, not an if / while body): visit(IfStmt) emits the reverse code + // into a brand new sibling IfStmt, not back into the forward if-body, so an alloca placed inside the + // forward branch is SSA-invalid from the reverse branch's point of view and gets DCE'd into silently-zero + // gradients. + Block *target = stmt->as()->parent; + while (target != nullptr) { + Stmt *parent_stmt = target->parent_stmt(); + if (parent_stmt == nullptr || parent_stmt->is() || parent_stmt->is() || + parent_stmt->is() || parent_stmt->is()) { + break; + } + target = parent_stmt->parent; + } + // Reaching a null target means the primal's enclosing-block chain is broken (an unparented block). Falling + // back to alloca_block here would silently restore the pre-fix bug (adjoint eliminated as DCE on the + // reverse branch); hard-assert instead so malformed IR surfaces loudly. + QD_ASSERT(target != nullptr); + target->insert(std::move(alloca), 0); + } else { + alloca_block->insert(std::move(alloca), 0); + } + } + return adjoint_stmt[stmt]; + } + + // For ops in `NonLinearOps::unary_collections` the reverse formula must NOT read the forward `stmt` + // directly - only `stmt->operand` (adstack-backed), `adjoint(stmt)`, and constants. BackupSSA spills the + // forward stmt to a single plain alloca overwritten each iteration, so reading `stmt` from a reversed + // dynamic loop would use the last-iteration value regardless of which reverse iteration is running. This + // helper walks the value-tree at IR-transform time and asserts the invariant; paired with the Python + // meta-test `test_unary_collections_audit` (which catches "forgot to add to unary_collections"), it + // covers both halves of the class of bugs the per-op audit used to miss. + void accumulate_unary_operand_checked(UnaryOpStmt *stmt, Stmt *value) { + if (NonLinearOps::unary_collections.find(stmt->op_type) != NonLinearOps::unary_collections.end()) { + std::unordered_set visited; + std::function walk = [&](const Stmt *s) { + if (!s || visited.count(s)) + return; + visited.insert(s); + QD_ASSERT_INFO(s != stmt, + "MakeAdjoint adjoint formula for UnaryOpType::{} reads the forward stmt directly. It " + "must read `stmt->operand` (adstack-backed) instead - see the tan/tanh/exp branches " + "for the recompute pattern.", + unary_op_type_name(stmt->op_type)); + for (auto *op : s->get_operands()) + walk(op); + }; + walk(value); + } + accumulate(stmt->operand, value); + } + + void visit(UnaryOpStmt *stmt) override { + auto acc = [&](Stmt *value) { accumulate_unary_operand_checked(stmt, value); }; + if (stmt->op_type == UnaryOpType::floor || stmt->op_type == UnaryOpType::ceil) { + // do nothing + } else if (stmt->op_type == UnaryOpType::neg) { + accumulate(stmt->operand, negate(adjoint(stmt))); + } else if (stmt->op_type == UnaryOpType::abs) { + acc(mul(adjoint(stmt), sgn(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::sin) { + acc(mul(adjoint(stmt), cos(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::cos) { + acc(negate(mul(adjoint(stmt), sin(stmt->operand)))); + } else if (stmt->op_type == UnaryOpType::tan) { + // d/dx tan(x) = 1 + tan(x)^2. Recompute tan(operand) rather than reusing the forward value: the primal is + // per-iteration inside dynamic loops but BackupSSA only spills forward values to a single plain alloca, so + // reading the forward tan would use the last-iteration value in the reversed loop. The operand, in contrast, + // rides the adstack through its LocalLoad, so a fresh tan on it is per-iteration correct. + acc(mul(adjoint(stmt), add(constant(1, stmt->ret_type), sqr(tan(stmt->operand))))); + } else if (stmt->op_type == UnaryOpType::tanh) { + // Recompute tanh(operand) in the reverse pass instead of reusing the forward stmt value. In dynamic loops + // BackupSSA spills the forward stmt to a single plain alloca overwritten each iteration, so the reversed + // loop would read the last-iteration tanh for every backward step. The operand rides the adstack through + // LocalLoad, so a fresh tanh on it is per-iteration correct. Trade-off: tanh is evaluated twice per + // iteration (once forward, once backward); caching the forward value on the adstack is a future optimization. + acc(mul(adjoint(stmt), sub(constant(1, stmt->ret_type), sqr(tanh(stmt->operand))))); + } else if (stmt->op_type == UnaryOpType::asin) { + acc(mul(adjoint(stmt), + div(constant(1, stmt->ret_type), sqrt(sub(constant(1, stmt->ret_type), sqr(stmt->operand)))))); + } else if (stmt->op_type == UnaryOpType::acos) { + acc(mul(adjoint(stmt), + negate(div(constant(1, stmt->ret_type), sqrt(sub(constant(1, stmt->ret_type), sqr(stmt->operand))))))); + } else if (stmt->op_type == UnaryOpType::exp) { + // See the tanh case above: recompute exp on the adstack-backed operand so the reversed loop sees the + // per-iteration value rather than the last-forward value spilled by BackupSSA. Same trade-off as tanh: exp + // is evaluated twice per iteration (once forward, once backward); caching the forward value on the adstack + // is a future optimization. + acc(mul(adjoint(stmt), exp(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::log) { + // No recompute workaround needed: the reverse formula `1 / operand` reads `stmt->operand` directly (which + // is adstack-backed via LocalLoad inside dynamic loops), not the forward `log(operand)` stmt value. + acc(div(adjoint(stmt), stmt->operand)); + } else if (stmt->op_type == UnaryOpType::sqrt) { + // No recompute workaround needed: the reverse formula reads `stmt->operand` (adstack-backed via LocalLoad + // inside dynamic loops, gated on `unary_collections` membership) and recomputes `sqrt(operand)` from it, + // not the forward `sqrt(operand)` stmt value. Unlike tanh/exp this case was already operand-based before + // the recompute fix landed; the structure mirrors log above. + acc(mul(adjoint(stmt), div(constant(0.5f, stmt->ret_type), sqrt(stmt->operand)))); + } else if (stmt->op_type == UnaryOpType::rsqrt) { + acc(mul(adjoint(stmt), + mul(constant(-0.5f, stmt->ret_type), pow(rsqrt(stmt->operand), constant(3, stmt->ret_type))))); + } else if (stmt->op_type == UnaryOpType::cast_value) { + if (is_real(stmt->cast_type.get_element_type()) && is_real(stmt->operand->ret_type.get_element_type())) { + accumulate(stmt->operand, adjoint(stmt)); + } + } else if (stmt->op_type == UnaryOpType::logic_not) { + // do nothing + } else { + QD_P(unary_op_type_name(stmt->op_type)); + QD_NOT_IMPLEMENTED; + } + } + + void visit(BinaryOpStmt *bin) override { + if (bin->op_type == BinaryOpType::add) { + accumulate(bin->lhs, adjoint(bin)); + accumulate(bin->rhs, adjoint(bin)); + } else if (bin->op_type == BinaryOpType::sub) { + accumulate(bin->lhs, adjoint(bin)); + accumulate(bin->rhs, negate(adjoint(bin))); + } else if (bin->op_type == BinaryOpType::mul) { + // d (x * y) = y * dx + x * dy + accumulate(bin->lhs, mul(adjoint(bin), bin->rhs)); + accumulate(bin->rhs, mul(adjoint(bin), bin->lhs)); + } else if (bin->op_type == BinaryOpType::mod) { + // Do nothing + } else if (bin->op_type == BinaryOpType::div) { + accumulate(bin->lhs, div(adjoint(bin), bin->rhs)); + accumulate(bin->rhs, negate(div(mul(adjoint(bin), bin->lhs), mul(bin->rhs, bin->rhs)))); + } else if (bin->op_type == BinaryOpType::atan2) { + auto numerator = add(sqr(bin->lhs), sqr(bin->rhs)); + accumulate(bin->lhs, div(mul(adjoint(bin), bin->rhs), numerator)); + accumulate(bin->rhs, negate(div(mul(adjoint(bin), bin->lhs), numerator))); + } else if (bin->op_type == BinaryOpType::pow) { + // d (x ^ y) = x ^ (y-1) * (y * dx + log(x) * x * dy) + auto common_coeff = pow(bin->lhs, sub(bin->rhs, constant(1, bin->ret_type))); // x ^ (y-1) + accumulate(bin->lhs, mul(adjoint(bin), mul(bin->rhs, common_coeff))); + accumulate(bin->rhs, mul(adjoint(bin), mul(log(bin->lhs), mul(bin->lhs, common_coeff)))); + } else if (bin->op_type == BinaryOpType::min || bin->op_type == BinaryOpType::max) { + auto cmp = bin->op_type == BinaryOpType::min ? cmp_lt(bin->lhs, bin->rhs) : cmp_lt(bin->rhs, bin->lhs); + auto zero = insert_const_for_grad(bin->ret_type, bin, 0); + accumulate(bin->lhs, sel(cmp, adjoint(bin), zero)); + accumulate(bin->rhs, sel(cmp, zero, adjoint(bin))); + } else if (bin->op_type == BinaryOpType::floordiv) { + // do nothing + } else if (is_comparison(bin->op_type) || is_bit_op(bin->op_type) || binary_is_logical(bin->op_type)) { + // do nothing + + } else { + QD_WARN("gradient of binary op {}\n{}", binary_op_type_name(bin->op_type), bin->get_tb()); + QD_NOT_IMPLEMENTED; + } + } + + void visit(TernaryOpStmt *stmt) override { + QD_ASSERT(stmt->op_type == TernaryOpType::select); + auto zero = insert_const_for_grad(stmt->ret_type, stmt, 0); + accumulate(stmt->op2, insert(TernaryOpType::select, stmt->op1, load(adjoint(stmt)), zero)); + accumulate(stmt->op3, insert(TernaryOpType::select, stmt->op1, zero, load(adjoint(stmt)))); + } + + void visit(IfStmt *if_stmt) override { + // Snapshot a stack-backed forward cond into a dedicated 1-push-per-if-execution ad-stack, + // but only when the cond's backing stack is also pushed inside the if body (e.g. short-circuit + // lowering pushes the rhs of `&&` onto the same stack that holds the cond). Without this, + // BackupSSA's clone of `if_stmt->cond` in the reverse block reads the cond stack AFTER the + // body-pushes rather than the forward-time cond value - the reverse IfStmt flips, pop counts + // drift, and gradients come out silently zero. A dedicated stack has exactly one push per + // forward if-execution, so the reverse load_top matches the forward cond. + // + // Guarded by the body-push check because snapshotting indiscriminately adds AdStackAllocaStmts + // that go through `determine_ad_stack_size` unused on every other if-stmt in the kernel - the + // adaptive-size pass emits "Unused autodiff stack should have been eliminated" warnings and + // the codegen step then fails with "Adaptive autodiff stack's size should have been determined". + Stmt *reverse_cond = if_stmt->cond; + AdStackAllocaStmt *snap_stack_ptr = nullptr; + // Narrow guard: only the bare `AdStackLoadTopStmt` shape needs the explicit snapshot below. A compound cond (e.g. + // `cmp_lt(load_top(x_stack) + 0.1, threshold)` from `if x + 0.1 < threshold` when `x` has been promoted to an + // adstack by `ReplaceLocalVarWithStacks`) reaches this visitor as a BARE `AdStackLoadTopStmt` cond anyway - the + // cmp / arithmetic value goes through `PromoteSSA2LocalVar`'s required-defs set (the IfStmt cond path adds the + // cond's value-producing op), then `AdStackAllocaJudger::visit(IfStmt)` (around line 763) marks its alloca + // stack-needed because it feeds the cond, then `ReplaceLocalVarWithStacks` promotes the alloca to an adstack. By + // the time control reaches here, `if_stmt->cond` is `AdStackLoadTopStmt` of that snap-promoted adstack and the + // body of the if does NOT push to that stack (the cmp value is pushed once just before the IfStmt, never inside), + // so the `body_pushes_to_stack` guard below is false and we correctly skip the additional snap-stack. Per-iter + // cond values are preserved by the alloca-promotion pipeline. + // + // The bare-`AdStackLoadTopStmt` case the snap-stack below handles is the OTHER shape: a load_top whose backing + // stack IS pushed to inside the if body (e.g. short-circuit lowering of `&&` pushes the rhs onto the same stack + // that holds the cond). Without the snap-stack, `BackupSSA::generic_visit`'s clone-branch for `AdStackLoadTopStmt` + // emits a fresh `AdStackLoadTopStmt` at the reverse cursor, which then reads the post-body-push top and sees the + // wrong cond value. The snap-stack here decouples the cond value from the body's pushes by capturing it once just + // before the IfStmt and reading it back at the matching reverse cursor. + // + // Earlier comments here claimed that compound conds rely on `BackupSSA::generic_visit`'s `load(op)` else-branch + // ("single alloca, last-write-wins" spill) for correctness. That description was incomplete: the alloca-promotion- + // to-adstack pipeline catches compound conds before they reach BackupSSA, so the spill is a FALLBACK for shapes + // the alloca-promotion missed, not the load-bearing path for compound conds in general. Verified empirically on + // 2026-05-04 with a loop-carried local v + compound cond `v + 0.1 > threshold` + nonlinear use of v in the if + // body: gradients match analytic on origin/main and on C3. + if (if_stmt->cond->is()) { + auto *cond_stack = if_stmt->cond->as()->stack->as(); + if (body_pushes_to_stack(if_stmt, cond_stack)) { + auto cond_type = if_stmt->cond->ret_type.ptr_removed(); + // Size the snap stack the same way as the cond stack it mirrors: one forward push per + // if-execution matched by one reverse pop. Reusing cond_stack->max_size keeps the snap + // stack exempt from `determine_ad_stack_size` when the cond stack itself was built with a + // fixed size, which is always true when ReplaceLocalVarWithStacks ran with a non-zero + // `ad_stack_size` (the only configuration we currently support for stack-based reverse AD). + auto snap_stack = Stmt::make(cond_type, cond_stack->max_size); + snap_stack_ptr = snap_stack->as(); + // Allocate at the IB root so the stack persists across enclosing for-loop iterations. + ib_root->insert(std::move(snap_stack), 0); + // Per-execution forward push of the cond value, just before the forward if-stmt. No + // initial zero push: the reverse load_top always runs after a matching forward push, so + // leaving the stack empty at entry is both correct and avoids a dead store that DSE would + // otherwise drop (and that `determine_ad_stack_size` would then miscount). + if_stmt->insert_before_me(Stmt::make(snap_stack_ptr, if_stmt->cond)); + // Per-execution reverse load of the snapshotted cond, emitted in the current reverse block. + reverse_cond = insert(snap_stack_ptr); + reverse_cond->ret_type = cond_type; + } + } + + auto new_if = Stmt::make_typed(reverse_cond); + if (if_stmt->true_statements) { + new_if->set_true_statements(std::make_unique()); + auto old_current_block = current_block; + // Backup forward pass + forward_backup = if_stmt->true_statements.get(); + + current_block = new_if->true_statements.get(); + for (int i = if_stmt->true_statements->statements.size() - 1; i >= 0; i--) { + if_stmt->true_statements->statements[i]->accept(this); + // Restore forward pass + forward_backup = if_stmt->true_statements.get(); + } + + current_block = old_current_block; + } + if (if_stmt->false_statements) { + new_if->set_false_statements(std::make_unique()); + auto old_current_block = current_block; + + // Backup forward pass + forward_backup = if_stmt->false_statements.get(); + + current_block = new_if->false_statements.get(); + for (int i = if_stmt->false_statements->statements.size() - 1; i >= 0; i--) { + if_stmt->false_statements->statements[i]->accept(this); + // Restore forward pass + forward_backup = if_stmt->false_statements.get(); + } + current_block = old_current_block; + } + insert_grad_stmt(std::move(new_if)); + if (snap_stack_ptr) { + // One pop per reverse if-execution, paired with the forward push above. + insert(snap_stack_ptr); + } + } + + void visit(RangeForStmt *for_stmt) override { + auto new_for = for_stmt->clone(); + auto new_for_ptr = new_for->as(); + new_for_ptr->reversed = !new_for_ptr->reversed; + insert_grad_stmt(std::move(new_for)); + const int len = new_for_ptr->body->size(); + + for (int i = 0; i < len; i++) { + new_for_ptr->body->erase(0); + } + + std::vector statements; + // always make a copy since the list can be modified. + for (auto &stmt : for_stmt->body->statements) { + statements.push_back(stmt.get()); + } + std::reverse(statements.begin(), statements.end()); // reverse-mode AD... + auto old_alloca_block = alloca_block; + auto old_current_block = current_block; + auto old_forward_backup = forward_backup; // store the block which is not inside the current IB, + // such as outer most loop + // Backup the forward pass + forward_backup = for_stmt->body.get(); + for (auto stmt : statements) { + alloca_block = new_for_ptr->body.get(); + current_block = new_for_ptr->body.get(); + stmt->accept(this); + // Restore the forward pass + forward_backup = for_stmt->body.get(); + } + // Restore current_block. Missing here before: if this RangeForStmt is visited from within another compound + // stmt (notably visit(IfStmt)), that outer visitor will continue iterating its own body in reverse after we + // return and emit further reverse stmts. Without this restore those emissions land in the reversed-for's + // body instead of the outer block, producing silently-wrong gradients whenever a runtime-guarded if wraps a + // for-loop with loop-carried variables (the reverse loop body ends up over-popping the adstack and emitting + // the x.grad accumulation on every iteration instead of once). + current_block = old_current_block; + forward_backup = old_forward_backup; + alloca_block = old_alloca_block; + } + + void visit(StructForStmt *for_stmt) override { + // Save/restore mirrors visit(RangeForStmt) above. Rationale: visit(Block) inside `body->accept(this)` + // sets current_block = for_stmt->body at the start of every iteration, so on return current_block + // points at the struct-for's body. An enclosing compound visitor (e.g. visit(IfStmt)) that resumes + // iterating its children in reverse after this StructForStmt needs current_block and alloca_block to + // still be its own, not this for's; otherwise subsequent reverse emissions land inside the struct-for + // body and any adjoint alloca lives in a block the enclosing if-branch cannot reach. forward_backup + // must be saved too because `visit(IfStmt)` mutates it without restoring, so a nested if inside the + // struct-for body leaves it pointing at the if-branch block, which then survives past this visitor and + // mis-routes `adjoint()` on GlobalLoadStmts for later siblings at the enclosing scope. + auto old_alloca_block = alloca_block; + auto old_current_block = current_block; + auto old_forward_backup = forward_backup; + alloca_block = for_stmt->body.get(); + for_stmt->body->accept(this); + current_block = old_current_block; + alloca_block = old_alloca_block; + forward_backup = old_forward_backup; + } + + // Equivalent to AdStackLoadTopStmt when no stack is needed + void visit(LocalLoadStmt *stmt) override { + // QD_ASSERT(!needs_grad(stmt->ret_type)); + if (is_real(stmt->ret_type.get_element_type())) + accumulate(stmt->src, load(adjoint(stmt))); + } + + // Equivalent to AdStackPushStmt when no stack is needed + void visit(LocalStoreStmt *stmt) override { + accumulate(stmt->val, load(adjoint(stmt->dest))); + + // Clear the adjoint of the dest after local store, + // Because LocalStoreStmt overwrites the dest, + // 1. If the alloca is inside a loop, the adjoint of this alloca of this + // iteration should be cleared after this iteration has been done + // 2. If the alloca serves as the dest of multiple LocalStoreStmt, only the + // last LocalStoreStmt should be taken account of + auto dest_type = stmt->dest->ret_type.ptr_removed(); + if (is_real(dest_type.get_element_type())) { + auto dtype = dest_type; + auto zero = insert_const_for_grad(dtype, stmt, 0); + insert(adjoint(stmt->dest), zero); + } + } + + void visit(AdStackLoadTopStmt *stmt) override { + if (is_real(stmt->ret_type.get_element_type())) + insert(stmt->stack, load(adjoint(stmt))); + } + + void visit(AdStackPushStmt *stmt) override { + accumulate(stmt->v, insert(stmt->stack)); + insert(stmt->stack); + } + + void visit(GlobalLoadStmt *stmt) override { + // issue global store to adjoint + + if (stmt->src->is() || + (stmt->src->is() && stmt->src->as()->origin->is())) { + ExternalPtrStmt *src = nullptr; + bool is_ptr_offset = false; + if (stmt->src->is()) { + src = stmt->src->as()->origin->as(); + is_ptr_offset = true; + } else { + src = stmt->src->as(); + } + auto arg = src->base_ptr->as(); + if (arg->ret_type.ptr_removed()->as()->elements().size() > TypeFactory::GRAD_PTR_POS_IN_NDARRAY) { + QD_ASSERT_INFO(!src->is_grad, + "Cannot automatically differentiate through a grad " + "tensor, if you really want to do that, pass the grad " + "tensor into the kernel directly"); + auto adj_ptr = + insert(src->base_ptr, src->indices, src->ndim, src->element_shape, /*is_grad=*/true); + adj_ptr->ret_type = src->ret_type; + + if (is_ptr_offset) { + adj_ptr = insert(adj_ptr, stmt->src->as()->offset); + adj_ptr->ret_type = stmt->src->ret_type; + adj_ptr->ret_type.set_is_pointer(true); + } + insert(AtomicOpType::add, adj_ptr, load(adjoint(stmt))); + } + return; + } + + if (stmt->src->is() || + (stmt->src->is() && stmt->src->as()->origin->is())) { + GlobalPtrStmt *src = nullptr; + bool is_ptr_offset = false; + if (stmt->src->is()) { + is_ptr_offset = true; + src = stmt->src->as()->origin->as(); + } else { + src = stmt->src->as(); + } + + auto snode = src->snode; + if (!snode->has_adjoint()) { + // No adjoint SNode. Do nothing + return; + } + if (gradients_stopped(stmt, snode)) { + // gradients stopped, do nothing. + return; + } + QD_ASSERT(snode->get_adjoint() != nullptr); + snode = snode->get_adjoint(); + auto adj_ptr = insert(snode, src->indices); + adj_ptr->ret_type = src->ret_type; + if (is_ptr_offset) { + adj_ptr = insert(adj_ptr, stmt->src->as()->offset); + } + insert(AtomicOpType::add, adj_ptr, load(adjoint(stmt))); + return; + } + } + + void visit(GlobalStoreStmt *stmt) override { + // erase and replace with global load adjoint + + Stmt *adjoint_ptr{nullptr}; + if (stmt->dest->is() || + (stmt->dest->is() && stmt->dest->as()->origin->is())) { + ExternalPtrStmt *dest = nullptr; + bool is_ptr_offset = false; + if (stmt->dest->is()) { + is_ptr_offset = true; + dest = stmt->dest->as()->origin->as(); + } else { + dest = stmt->dest->as(); + } + + auto arg = dest->base_ptr->as(); + if (arg->ret_type.ptr_removed()->as()->elements().size() <= TypeFactory::GRAD_PTR_POS_IN_NDARRAY) { + return; + } + QD_ASSERT_INFO(!dest->is_grad, + "Cannot automatically differentiate through a grad " + "tensor, if you really want to do that, pass the grad " + "tensor into the kernel directly"); + adjoint_ptr = insert(dest->base_ptr, dest->indices, dest->ndim, dest->element_shape, + /*is_grad=*/true); + adjoint_ptr->ret_type = dest->ret_type; + + if (is_ptr_offset) { + adjoint_ptr = insert(adjoint_ptr, stmt->dest->as()->offset); + adjoint_ptr->ret_type = stmt->dest->ret_type; + adjoint_ptr->ret_type.set_is_pointer(true); + } + + accumulate(stmt->val, insert(adjoint_ptr)); + } + + if (stmt->dest->is() || + (stmt->dest->is() && stmt->dest->as()->origin->is())) { + GlobalPtrStmt *dest = nullptr; + bool is_ptr_offset = false; + if (stmt->dest->is()) { + is_ptr_offset = true; + dest = stmt->dest->as()->origin->as(); + } else { + dest = stmt->dest->as(); + } + + auto snode = dest->snode; + if (!snode->has_adjoint()) { + // no gradient (likely integer types) + return; + } + QD_ASSERT(snode->get_adjoint() != nullptr); + snode = snode->get_adjoint(); + adjoint_ptr = insert(snode, dest->indices); + adjoint_ptr->ret_type = dest->ret_type; + if (is_ptr_offset) { + adjoint_ptr = insert(adjoint_ptr, stmt->dest->as()->offset); + } + accumulate(stmt->val, insert(adjoint_ptr)); + } + + // Clear the gradient after accumulation finished. + auto zero = insert_const_for_grad(adjoint_ptr->ret_type.ptr_removed(), stmt, 0); + insert(adjoint_ptr, zero); + + stmt->parent->erase(stmt); + } + + void visit(AtomicOpStmt *stmt) override { + if (stmt->dest->is() || + (stmt->dest->is() && stmt->dest->as()->origin->is())) { + ExternalPtrStmt *dest = nullptr; + bool is_ptr_offset = false; + if (stmt->dest->is()) { + is_ptr_offset = true; + dest = stmt->dest->as()->origin->as(); + } else { + dest = stmt->dest->as(); + } + + auto arg = dest->base_ptr->as(); + if (arg->ret_type.ptr_removed()->as()->elements().size() > TypeFactory::GRAD_PTR_POS_IN_NDARRAY) { + QD_ASSERT_INFO(!dest->is_grad, + "Cannot automatically differentiate through a grad " + "tensor, if you really want to do that, pass the grad " + "tensor into the kernel directly"); + auto adjoint_ptr = + insert(dest->base_ptr, dest->indices, dest->ndim, dest->element_shape, /*is_grad=*/true); + adjoint_ptr->ret_type = dest->ret_type; + + if (is_ptr_offset) { + adjoint_ptr = insert(adjoint_ptr, stmt->dest->as()->offset); + + adjoint_ptr->ret_type = stmt->dest->ret_type; + adjoint_ptr->ret_type.set_is_pointer(true); + } + adjoint_ptr->ret_type = dest->ret_type; + accumulate(stmt->val, insert(adjoint_ptr)); + stmt->parent->erase(stmt); + } + return; + } + + if (stmt->dest->is() || + (stmt->dest->is() && stmt->dest->as()->origin->is())) { + GlobalPtrStmt *dest = nullptr; + bool is_ptr_offset = false; + if (stmt->dest->is()) { + is_ptr_offset = true; + dest = stmt->dest->as()->origin->as(); + } else { + dest = stmt->dest->as(); + } + + auto snode = dest->snode; + if (!snode->has_adjoint()) { + // no gradient (likely integer types) + return; + } + + QD_ASSERT(snode->get_adjoint() != nullptr); + snode = snode->get_adjoint(); + auto adjoint_ptr = insert(snode, dest->indices); + adjoint_ptr->ret_type = dest->ret_type; + if (is_ptr_offset) { + adjoint_ptr = insert(adjoint_ptr, stmt->dest->as()->offset); + } + accumulate(stmt->val, insert(adjoint_ptr)); + stmt->parent->erase(stmt); + return; + } + } + + void visit(MatrixPtrStmt *stmt) override { + if (stmt->origin->is() || stmt->origin->is()) { + /* + The case of MatrixPtrStmt(GlobalPtrStmt, ...) is already handled in + GlobalPtrStmt, GlobalStoreStmt and AtomicStmt + + TODO(zhanlue): Try to separate out the chain rule for MatrixPtrStmt from + GlobalPtrStmt, GlobalStoreStmt and AtomicStmt and migrate the logics + here. + */ + return; + } + + DataType prim_dtype = PrimitiveType::f32; + if (is_real(stmt->ret_type.ptr_removed().get_element_type())) { + prim_dtype = stmt->ret_type.ptr_removed().get_element_type(); + } + + Stmt *adjoint_value = nullptr; + if (stmt->offset->is()) { + /* + [Static index] + Fwd: + $0 = alloca <4 x i32> + $1 = matrix ptr $0, 2 // offset = 2 + + Adjoint: + $3 = matrix init [0, 0, $1_adj, 0] // adjoint_value + + accumulate($0_adj, $3) + */ + int offset = stmt->offset->as()->val.val_int32(); + + auto tensor_type = stmt->origin->ret_type.ptr_removed()->as(); + int num_elements = tensor_type->get_num_elements(); + + auto zero = insert_const_for_grad(prim_dtype, stmt, 0); + std::vector values; + for (int i = 0; i < num_elements; i++) { + if (i == offset) { + values.push_back(load(adjoint(stmt))); + } else { + values.push_back(zero); + } + } + auto matrix_init_stmt = insert(values); + matrix_init_stmt->ret_type = tensor_type; + + adjoint_value = matrix_init_stmt; + + } else { + /* + [Dynamic index] + Fwd: + $0 = alloca <4 x i32> + $1 = matrix ptr $0, $offset + + Adjoint: + $3 = matrix init [0.0, 0.0, 0.0, 0.0] + $4 = matrix init [$1_adj, $1_adj, $1_adj, $1_adj] + + $5 = matrix init [0, 1, 2, 3] + $6 = matrix init [offset, offset, offset, offset] + $7 = bin_eq $6, $5 + $8 = select $7, $4, $3 // adjoint_value + + accumulate($0_adj, $7) + */ + auto tensor_type = stmt->origin->ret_type.ptr_removed()->as(); + auto tensor_shape = tensor_type->get_shape(); + int num_elements = tensor_type->get_num_elements(); + + auto zero = insert_const_for_grad(prim_dtype, stmt, 0); + auto stmt_adj = load(adjoint(stmt)); + + std::vector zero_values(num_elements, zero); + std::vector stmt_adj_values(num_elements, stmt_adj); + std::vector offset_values(num_elements, stmt->offset); + std::vector indices_values(num_elements); + for (size_t i = 0; i < num_elements; i++) { + indices_values[i] = insert(TypedConstant((int32)i)); + } + + auto zero_matrix_init_stmt = insert(zero_values); + zero_matrix_init_stmt->ret_type = tensor_type; + auto stmt_adj_matrix_init_stmt = insert(stmt_adj_values); + stmt_adj_matrix_init_stmt->ret_type = tensor_type; + + auto index_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::i32); + auto indices_matrix_init_stmt = insert(indices_values); + indices_matrix_init_stmt->ret_type = index_tensor_type; + + auto offset_matrix_init_stmt = insert(offset_values); + offset_matrix_init_stmt->ret_type = index_tensor_type; + auto cmp_tensor_type = TypeFactory::get_instance().get_tensor_type(tensor_shape, PrimitiveType::u1); + auto bin_eq_stmt = insert(BinaryOpType::cmp_eq, offset_matrix_init_stmt, indices_matrix_init_stmt); + bin_eq_stmt->ret_type = cmp_tensor_type; + + auto select_stmt = + insert(TernaryOpType::select, bin_eq_stmt, stmt_adj_matrix_init_stmt, zero_matrix_init_stmt); + adjoint_value = select_stmt; + } + + accumulate(stmt->origin, adjoint_value); + } + + void visit(MatrixInitStmt *stmt) override { + auto adjoint_ptr = adjoint(stmt); + + auto tensor_type = stmt->ret_type->as(); + int num_elements = tensor_type->get_num_elements(); + + for (size_t i = 0; i < num_elements; i++) { + auto const_i = insert_const_for_grad(PrimitiveType::i32, stmt, i); + + auto matrix_ptr_stmt_i = insert(adjoint_ptr, const_i); + matrix_ptr_stmt_i->ret_type = tensor_type->get_element_type(); + matrix_ptr_stmt_i->ret_type.set_is_pointer(true); + + accumulate(stmt->values[i], load(matrix_ptr_stmt_i)); + } + } +}; + +} // namespace + +void make_adjoint(Block *ib) { + MakeAdjoint::run(ib); +} + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/make_adjoint.h b/quadrants/transforms/auto_diff/make_adjoint.h new file mode 100644 index 0000000000..7694ee5ee5 --- /dev/null +++ b/quadrants/transforms/auto_diff/make_adjoint.h @@ -0,0 +1,12 @@ +#pragma once + +namespace quadrants::lang { + +class Block; + +// Stage: Reverse-mode codegen. Emit the reverse pass IR in `ib` (an +// Independent Block produced by identify_independent_blocks). Forward IR must +// already have been shaped by ir_shaping + forward_state_spill. +void make_adjoint(Block *ib); + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/make_dual.cpp b/quadrants/transforms/auto_diff/make_dual.cpp new file mode 100644 index 0000000000..f60318b132 --- /dev/null +++ b/quadrants/transforms/auto_diff/make_dual.cpp @@ -0,0 +1,345 @@ +#include "quadrants/transforms/auto_diff/auto_diff_common.h" +#include "quadrants/transforms/auto_diff/make_dual.h" + +namespace quadrants::lang { + +namespace { + +// Forward mode autodiff +class MakeDual : public ADTransform { + public: + using ADTransform::visit; + Stmt *current_stmt; + Block *current_block; + Block *alloca_block; + std::map dual_stmt; + + explicit MakeDual(Block *block) { + current_stmt = nullptr; + alloca_block = block; + current_block = block; + } + + static void run(Block *block) { + auto p = MakeDual(block); + block->accept(&p); + } + + Stmt *insert_grad_stmt(std::unique_ptr &&stmt) override { + auto ptr = stmt.get(); + current_stmt = current_stmt->insert_after_me(std::move(stmt)); + return ptr; + } + + void visit(Block *block) override { + std::vector statements; + // always make a copy since the list can be modified. + for (auto &stmt : block->statements) { + statements.push_back(stmt.get()); + } + for (auto stmt : statements) { + current_stmt = stmt; + stmt->accept(this); + } + } + + // Accumulate [value] to the dual of [primal] + void accumulate(Stmt *primal, Stmt *value) { + auto alloca_ = dual(primal); + if (!alloca_ || alloca_->is()) + return; // primal may be int variable + + QD_ASSERT(alloca_->is()); + auto alloca = alloca_->as(); + auto local_load = insert(alloca); + insert(alloca, add(local_load, value)); + } + + Stmt *dual(Stmt *stmt) { + auto dual_type = stmt->ret_type.ptr_removed(); + if (!is_real(dual_type.get_element_type()) || stmt->is()) { + return constant(0); + } + if (dual_stmt.find(stmt) == dual_stmt.end()) { + // normal SSA cases + + // create the alloca + // auto alloca = + // Stmt::make(get_current_program().config.gradient_dt); + // maybe it's better to use the statement data type than the default type + auto alloca = Stmt::make(dual_type); + dual_stmt[stmt] = alloca.get(); + + // TODO: check whether there are any edge cases for the alloca_block + alloca_block->insert(std::move(alloca), 0); + } + return dual_stmt[stmt]; + } + + void visit(UnaryOpStmt *stmt) override { + if (stmt->op_type == UnaryOpType::neg) { + accumulate(stmt, negate(dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::abs) { + accumulate(stmt, mul(sgn(stmt->operand), dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::sin) { + accumulate(stmt, mul(cos(stmt->operand), dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::cos) { + accumulate(stmt, negate(mul(sin(stmt->operand), dual(stmt->operand)))); + } else if (stmt->op_type == UnaryOpType::tan) { + // d/dx tan(x) = 1 + tan(x)^2. Forward mode executes in primal order, so `stmt` is the + // current-iteration tan value - no BackupSSA stale-value concern; reusing it is per-iteration + // correct. + accumulate(stmt, mul(add(constant(1), sqr(stmt)), dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::tanh) { + accumulate(stmt, mul(sub(constant(1), sqr(stmt)), dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::asin) { + accumulate(stmt, mul(div(constant(1), sqrt(sub(constant(1), sqr(stmt->operand)))), dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::acos) { + accumulate(stmt, mul(negate(div(constant(1), sqrt(sub(constant(1), sqr(stmt->operand))))), dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::exp) { + accumulate(stmt, mul(stmt, dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::log) { + accumulate(stmt, div(dual(stmt->operand), stmt->operand)); + } else if (stmt->op_type == UnaryOpType::sqrt) { + accumulate(stmt, mul(div(constant(0.5f), sqrt(stmt->operand)), dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::rsqrt) { + accumulate(stmt, mul(mul(constant(-0.5f), pow(rsqrt(stmt->operand), constant(3))), dual(stmt->operand))); + } else if (stmt->op_type == UnaryOpType::cast_value) { + if (is_real(stmt->cast_type.get_element_type()) && is_real(stmt->operand->ret_type.get_element_type())) { + accumulate(stmt, dual(stmt->operand)); + } + } else if (stmt->op_type == UnaryOpType::logic_not) { + // do nothing + } else { + QD_P(unary_op_type_name(stmt->op_type)); + QD_NOT_IMPLEMENTED + } + } + + void visit(BinaryOpStmt *bin) override { + if (bin->op_type == BinaryOpType::add) { + accumulate(bin, dual(bin->lhs)); + accumulate(bin, dual(bin->rhs)); + } else if (bin->op_type == BinaryOpType::sub) { + accumulate(bin, dual(bin->lhs)); + accumulate(bin, negate(dual(bin->rhs))); + } else if (bin->op_type == BinaryOpType::mul) { + // d (x * y) = y * dx + x * dy + accumulate(bin, mul(bin->lhs, dual(bin->rhs))); + accumulate(bin, mul(bin->rhs, dual(bin->lhs))); + } else if (bin->op_type == BinaryOpType::mod) { + // Do nothing + } else if (bin->op_type == BinaryOpType::div) { + accumulate(bin, div(dual(bin->lhs), bin->rhs)); + accumulate(bin, negate(div(mul(dual(bin->rhs), bin->lhs), mul(bin->rhs, bin->rhs)))); + } else if (bin->op_type == BinaryOpType::atan2) { + auto numerator = add(sqr(bin->lhs), sqr(bin->rhs)); + accumulate(bin, div(mul(bin->rhs, dual(bin->lhs)), numerator)); + accumulate(bin, negate(div(mul(bin->lhs, dual(bin->rhs)), numerator))); + } else if (bin->op_type == BinaryOpType::pow) { + // d (x ^ y) = x ^ (y-1) * (y * dx + log(x) * x * dy) + auto common_coeff = pow(bin->lhs, sub(bin->rhs, constant(1))); // x ^ (y-1) + accumulate(bin, mul(dual(bin->lhs), mul(bin->rhs, common_coeff))); + accumulate(bin, mul(dual(bin->rhs), mul(log(bin->lhs), mul(bin->lhs, common_coeff)))); + } else if (bin->op_type == BinaryOpType::min || bin->op_type == BinaryOpType::max) { + auto cmp = bin->op_type == BinaryOpType::min ? cmp_lt(bin->lhs, bin->rhs) : cmp_lt(bin->rhs, bin->lhs); + auto zero = insert_const_for_grad(bin->ret_type, bin, 0); + accumulate(bin, sel(cmp, dual(bin->lhs), zero)); + accumulate(bin, sel(cmp, zero, dual(bin->rhs))); + } else if (bin->op_type == BinaryOpType::floordiv) { + // do nothing + } else if (is_comparison(bin->op_type) || is_bit_op(bin->op_type)) { + // do nothing + } else { + QD_WARN("gradient of binary op {}\n{}", binary_op_type_name(bin->op_type), bin->get_tb()); + QD_NOT_IMPLEMENTED + } + } + + void visit(TernaryOpStmt *stmt) override { + QD_ASSERT(stmt->op_type == TernaryOpType::select); + auto zero = insert_const_for_grad(stmt->ret_type, stmt, 0); + accumulate(stmt, insert(TernaryOpType::select, stmt->op1, load(dual(stmt->op2)), zero)); + accumulate(stmt, insert(TernaryOpType::select, stmt->op1, zero, load(dual(stmt->op3)))); + } + + void visit(IfStmt *if_stmt) override { + if (if_stmt->true_statements) { + std::vector true_statements; + for (auto &stmt : if_stmt->true_statements->statements) { + true_statements.push_back(stmt.get()); + } + + for (auto stmt : true_statements) { + current_stmt = stmt; + stmt->accept(this); + } + } + if (if_stmt->false_statements) { + std::vector false_statements; + for (auto &stmt : if_stmt->false_statements->statements) { + false_statements.push_back(stmt.get()); + } + + for (auto stmt : false_statements) { + current_stmt = stmt; + stmt->accept(this); + } + } + } + + void visit(RangeForStmt *for_stmt) override { + std::vector statements; + // always make a copy since the list can be modified. + for (auto &stmt : for_stmt->body->statements) { + statements.push_back(stmt.get()); + } + auto previous_alloca_block = alloca_block; + alloca_block = for_stmt->body.get(); + for (auto stmt : statements) { + current_stmt = stmt; + stmt->accept(this); + } + alloca_block = previous_alloca_block; + } + + void visit(StructForStmt *for_stmt) override { + // Save/restore mirrors visit(RangeForStmt) above and MakeAdjoint::visit(StructForStmt). An enclosing + // compound visitor that resumes iterating its body after this StructForStmt needs alloca_block to still + // point at its own block, not the sparse-for body, so new dual allocas land where the enclosing reverse + // code can reach them. + auto previous_alloca_block = alloca_block; + alloca_block = for_stmt->body.get(); + for_stmt->body->accept(this); + alloca_block = previous_alloca_block; + } + + void visit(LocalLoadStmt *stmt) override { + // QD_ASSERT(!needs_grad(stmt->ret_type)); + accumulate(stmt, dual(stmt->src)); + } + + void visit(LocalStoreStmt *stmt) override { + // Clear the dual of the dest before local store, + // Because LocalStoreStmt overwrites the dest, + // If the alloca serves as the dest of multiple LocalStoreStmt, only the + // last LocalStoreStmt should be taken account of, i.e, its history should + // be cleared + auto dtype = stmt->dest->ret_type.ptr_removed(); + if (is_real(dtype.get_element_type())) { + auto zero = insert_const_for_grad(dtype, stmt, 0); + insert(dual(stmt->dest), zero); + } + + accumulate(stmt->dest, dual(stmt->val)); + } + + void visit(GlobalLoadStmt *stmt) override { + // issue global store to dual + GlobalPtrStmt *src = nullptr; + bool is_ptr_offset = false; + if (stmt->src->is()) { + is_ptr_offset = true; + src = stmt->src->as()->origin->as(); + } else { + src = stmt->src->as(); + } + auto snode = src->snode; + if (!snode->has_dual()) { + // No dual SNode. Do nothing + return; + } + if (gradients_stopped(stmt, snode)) { + // gradients stopped, do nothing. + return; + } + QD_ASSERT(snode->get_dual() != nullptr); + snode = snode->get_dual(); + auto dual_ptr = insert(snode, src->indices); + dual_ptr->ret_type = src->ret_type; + if (is_ptr_offset) { + dual_ptr = insert(dual_ptr, stmt->src->as()->offset); + } + accumulate(stmt, insert(dual_ptr)); + } + + void visit(GlobalStoreStmt *stmt) override { + GlobalPtrStmt *dest = nullptr; + bool is_ptr_offset = false; + if (stmt->dest->is()) { + is_ptr_offset = true; + dest = stmt->dest->as()->origin->as(); + } else { + dest = stmt->dest->as(); + } + auto snode = dest->snode; + if (!snode->has_dual()) { + // no gradient (likely integer types) + return; + } + QD_ASSERT(snode->get_dual() != nullptr); + snode = snode->get_dual(); + auto dual_ptr = insert(snode, dest->indices); + dual_ptr->ret_type = dest->ret_type; + if (is_ptr_offset) { + dual_ptr = insert(dual_ptr, stmt->dest->as()->offset); + } + insert(AtomicOpType::add, dual_ptr, load(dual(stmt->val))); + } + + void visit(AtomicOpStmt *stmt) override { + GlobalPtrStmt *dest = nullptr; + bool is_ptr_offset = false; + if (stmt->dest->is()) { + is_ptr_offset = true; + dest = stmt->dest->as()->origin->as(); + } else { + dest = stmt->dest->as(); + } + auto snode = dest->snode; + if (!snode->has_dual()) { + // no gradient (likely integer types) + return; + } + QD_ASSERT(snode->get_dual() != nullptr); + snode = snode->get_dual(); + auto dual_ptr = insert(snode, dest->indices); + dual_ptr->ret_type = dest->ret_type; + if (is_ptr_offset) { + dual_ptr = insert(dual_ptr, stmt->dest->as()->offset); + } + insert(AtomicOpType::add, dual_ptr, load(dual(stmt->val))); + } + + void visit(MatrixInitStmt *stmt) override { + std::vector duals; + for (auto &s : stmt->values) { + duals.push_back(dual(s)); + } + auto dual_stmt = insert(duals); + dual_stmt->ret_type = stmt->ret_type; + + accumulate(stmt, dual_stmt); + } + + void visit(MatrixPtrStmt *stmt) override { + if (stmt->origin->is()) { + // Handled in GlobalLoadStmt and GlobalStoreStmt + return; + } + + auto origin_dual = dual(stmt->origin); + auto origin_dual_ptr = insert(origin_dual, stmt->offset); + origin_dual_ptr->ret_type = stmt->ret_type; + + accumulate(stmt, origin_dual_ptr); + } +}; + +} // namespace + +void make_dual(Block *block) { + MakeDual::run(block); +} + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/make_dual.h b/quadrants/transforms/auto_diff/make_dual.h new file mode 100644 index 0000000000..5daeb40c00 --- /dev/null +++ b/quadrants/transforms/auto_diff/make_dual.h @@ -0,0 +1,11 @@ +#pragma once + +namespace quadrants::lang { + +class Block; + +// Stage: Forward-mode codegen. Emit a tangent-propagating dual pass for +// `block` (the kernel root for forward-mode autodiff). +void make_dual(Block *block); + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp b/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp new file mode 100644 index 0000000000..9521ab70dd --- /dev/null +++ b/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp @@ -0,0 +1,324 @@ +#include "quadrants/transforms/auto_diff/auto_diff_common.h" +#include "quadrants/transforms/auto_diff/post_adjoint_cleanup.h" + +namespace quadrants::lang { + +namespace { + +// ============================================================================ +// BackupSSA: spill cross-block SSA operands the reverse-mode clones reference. +// +// MakeAdjoint clones forward stmts into the reverse scope but shares operand +// pointers with the forward graph; when those operands live inside the +// forward body (e.g. an inner-for's `begin`/`end`), the reverse clone's +// operand no longer dominates its use. This pass rewrites such operands: +// AdStackLoadTopStmt / ArgLoadStmt get cloned in place, AdStackAllocaStmt +// gets re-rooted at the IB, generic SSA values get a per-IB backup +// AllocaStmt + LocalStore + LocalLoad chain. +// ============================================================================ +class BackupSSA : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + + Block *independent_block; + std::map backup_alloca; + + explicit BackupSSA(Block *independent_block) : independent_block(independent_block) { + allow_undefined_visitor = true; + invoke_default_visitor = true; + } + + Stmt *load(Stmt *stmt) { + if (backup_alloca.find(stmt) == backup_alloca.end()) { + auto alloca = Stmt::make(stmt->ret_type.ptr_removed()); + auto alloca_ptr = alloca.get(); + independent_block->insert(std::move(alloca), 0); + auto local_store = Stmt::make(alloca_ptr, stmt); + stmt->insert_after_me(std::move(local_store)); + backup_alloca[stmt] = alloca_ptr; + } + return backup_alloca[stmt]; + } + + void generic_visit(Stmt *stmt) { + std::vector leaf_to_root; + auto t = stmt->parent; + while (t != nullptr) { + leaf_to_root.push_back(t); + t = t->parent_block(); + } + int num_operands = stmt->get_operands().size(); + for (int i = 0; i < num_operands; i++) { + auto op = stmt->operand(i); + if (op == nullptr) { + continue; + } + if (std::find(leaf_to_root.begin(), leaf_to_root.end(), op->parent) == leaf_to_root.end() && + !op->is()) { + if (op->is()) { + // Just create another AdStackLoadTopStmt + stmt->set_operand(i, stmt->insert_before_me(op->clone())); + } else if (op->is()) { + // Backup AdStackAllocaStmt because it should not be local stored and + // local loaded + auto stack_alloca = op->as(); + if (backup_alloca.find(op) == backup_alloca.end()) { + auto backup_stack_alloca = Stmt::make(stack_alloca->dt, stack_alloca->max_size); + auto backup_stack_alloca_ptr = backup_stack_alloca.get(); + independent_block->insert(std::move(backup_stack_alloca), 0); + backup_alloca[op] = backup_stack_alloca_ptr; + // Replace usages of all blocks i.e., the entry point for the + // replace is the top level block + irpass::replace_all_usages_with(leaf_to_root.back(), op, backup_stack_alloca_ptr); + // Erase the outdated AdStackAllocaStmt + op->parent->erase(op); + } + } else if (op->is()) { + stmt->set_operand(i, stmt->insert_before_me(op->clone())); + } else { + // Recomputable-chain fallback before the last-iter `load(op)` spill: when the cross-block SSA op + // is rooted in a DAG of side-effect-free arithmetic over already-stack-backed allocas, kernel-args, + // constants, and loop indices, clone the chain into the reverse scope at the consumer site. The + // cloned chain reads stack tops live (matching `MakeAdjoint`'s pop ordering, which fires pops + // AFTER all uses of a stack within one reverse iter) and reads kernel-args / constants / loop + // indices via direct clones - exactly the path that was already correct for AdStackLoadTopStmt and + // ArgLoadStmt operands. + // + // The pre-existing `load(op)` fallback below remains correct for genuinely non-recomputable + // cross-block ops (a forward `GlobalLoadStmt` of a needs_grad SNode whose value the reverse must + // read at last-write rather than recompute, for instance) and for shapes outside one of our + // independent blocks. Adding the recomputable path above the fallback strictly subsets the + // previous behaviour: a chain that fails the predicate falls through to the old `load(op)` line. + if (RecomputableChainAnalyzer::is_recomputable(op, recomputable_cache_)) { + std::unordered_map clone_cache; + Stmt *cloned = RecomputableChainCloner::clone_at(op, stmt, clone_cache); + stmt->set_operand(i, cloned); + } else { + auto alloca = load(op); + stmt->set_operand(i, stmt->insert_before_me(Stmt::make(alloca))); + } + } + } + } + } + + // Memoization cache for `RecomputableChainAnalyzer::is_recomputable` queries within one BackupSSA run. + // Re-used across all generic_visit calls; invariant during the visit because forward IR is read-only here. + std::unordered_map recomputable_cache_; + + void visit(Stmt *stmt) override { + generic_visit(stmt); + } + + void visit(IfStmt *stmt) override { + generic_visit(stmt); + BasicStmtVisitor::visit(stmt); + } + + // generic_visit spills cross-block operands (the for-loop's `begin` and `end`) the same way it does for an + // IfStmt's cond. MakeAdjoint clones a forward for-loop into the reverse scope and shares the clone's + // `begin`/`end` pointers with the forward stmt; when those operands live inside the forward for's body (e.g. + // inner `for k in range(j)` where `j` is an enclosing loop's index promoted to a per-iter adstack), the reverse + // clone's operand no longer dominates its use. generic_visit's AdStackLoadTopStmt branch handles this by + // inserting a fresh AdStackLoadTop in the reverse scope, which reads the correct per-iteration value. + void visit(RangeForStmt *stmt) override { + generic_visit(stmt); + stmt->body->accept(this); + } + + void visit(StructForStmt *stmt) override { + generic_visit(stmt); + stmt->body->accept(this); + } + + void visit(WhileStmt *stmt) override { + QD_ERROR("WhileStmt not supported in AutoDiff for now."); + } + + void visit(Block *block) override { + std::vector statements; + // always make a copy since the list can be modified. + for (auto &stmt : block->statements) { + statements.push_back(stmt.get()); + } + for (auto stmt : statements) { + QD_ASSERT(!stmt->erased); + stmt->accept(this); + } + } + + static void run(Block *block) { + BackupSSA pass(block); + block->accept(&pass); + } +}; + +// ============================================================================ +// CoalesceAdStackLoads: dedup redundant AdStackLoadTop / AdStackLoadTopAdj +// reads emitted by the reverse pass. +// +// Within a single straight-line block, multiple `AdStackLoadTopStmt` reads of the same stack with no +// intervening `AdStackPushStmt` / `AdStackPopStmt` for that stack return the same value, and multiple +// `AdStackLoadTopAdjStmt` reads are equivalent under the same conditions plus no intervening +// `AdStackAccAdjointStmt`. After Python-side static unrolling collapses an inner loop into straight-line IR, +// every iteration's read of an outer-loop-invariant adstack value emits a separate LoadTop in the same block; +// each individual load is cheap (read u64 count + GEP) but unrolled-loop counts of hundreds to thousands +// inflate PTX size and ptxas register-allocator cost. This pass walks each block, caches the most recent +// LoadTop / LoadTopAdj per stack, replaces subsequent same-stack reads with the cached SSA value until a +// Push / Pop / AccAdjoint invalidates the cache for that stack, and conservatively clears the cache when +// crossing into nested control flow (IfStmt / RangeForStmt / StructForStmt / WhileStmt) where unseen +// push/pop/AccAdjoint may appear. +// ============================================================================ + +// clang-format off +/* +Support for TensorType: How to handle MatrixPtrStmt & MatrixInitStmt + +[Original Quadrants Code] + +@ti.kernel +def test(...): + b = ti.Vector([0, 1, 2, 3]) + b[2] = 100 + y = b[3] * b[2] * x + + +[Forward] [Forward-Replaced] [Backward] +$b = alloca Tensor<4 x i32> --> $b = adstack alloca <4 x i32> +$1 = matrix init [0, 1, 2, 3] --> $1 = matrix init [0, 1, 2, 3] + adstack pop +$2: local store $b, $1 --> adstack push $1 --> acc($1_adj, adstack top adj()) + +$3 = matrix ptr $b, 2 --> $2 = adstack top(is_ptr=True) --> adstack acc adj($2_adj) + + acc($2_adj, $14) + $3 = matrix ptr $2, 0 --> $14 = matrix_init({$3_adj, 0, 0, 0}) + + acc($2_adj, $13) + $4 = matrix ptr $2, 2 --> $13 = matrix_init({0, 0, $4_adj, 0}) + + acc($2_adj, $12) + $5 = matrix ptr $2, 3 --> $12 = matrix_init({0, 0, 0, $5_adj}) + + $6 = load($3) --> acc($3_adj, $6_adj) + $7 = load($4) --> acc($4_adj, $7_adj) + $8 = load($5) --> acc($5_adj, $8_adj) + + acc($8_adj, matrix ptr($9_adj, 3)) + acc($7_adj, matrix ptr($9_adj, 2)) + acc(100_adj, matrix ptr($9_adj, 1)) + $9 = matrix_init($6,100,$7,$8) --> acc($6_adj, matrix ptr($9_adj, 0)) + + adstack pop +$4 = local store $3, 100 --> adstack push $9 --> acc($9_adj, adstack top adj()) + + $10 = adstack top(is_ptr=True) --> adstack acc adj($10_adj) + + acc($10_adj, $18) +$5 = matrix ptr $b, 3 --> $11 = matrix ptr $10, 3 --> $18 = matrix_init({0, 0, 0, $11_adj}) +$b3 = local load $5 --> $b3 = local load $11 --> acc($11_adj, $b3_adj) + + $12 = adstack top(is_ptr=True) --> adstack acc adj($12_adj) + + acc($12_adj, $17) +$6 = matrix ptr $b, 2 --> $13 = matrix ptr $12, 2 --> $17 = matrix_init({0, 0, $13_adj, 0}) +$b2 = local load $6 --> $b2 = local load $13 --> acc($13_adj, $b2_adj) + + acc($b3_adj, $15) + acc($b2_adj, $16) + $16 = mul($tmp_adj, $b3) +$tmp = mul b3, b2 --> $tmp = mul $b3, $b2 --> $15 = mul($tmp_adj, $b2) + + acc($tmp_adj, $14) +$y = mul $tmp, $x --> $y = mul $tmp, $x --> $14 = mul($y_adj, $x) +*/ +// clang-format on +class CoalesceAdStackLoads : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + + CoalesceAdStackLoads() { + invoke_default_visitor = true; + allow_undefined_visitor = true; + } + + void visit(Block *block) override { + std::unordered_map primal_cache; + std::unordered_map adjoint_cache; + // Iterate over a snapshot so erase()s during the walk do not invalidate the iteration. The + // DelayedIRModifier still applies erases at the end via `modify_ir()` so SSA users get rewritten in one + // pass; the snapshot only protects the visitor's own walk. + std::vector stmts; + stmts.reserve(block->statements.size()); + for (auto &s : block->statements) { + stmts.push_back(s.get()); + } + for (Stmt *s : stmts) { + if (auto *lt = s->cast()) { + // `return_ptr=true` returns a pointer into the stack slot rather than loading a value; coalescing + // those is unsound because subsequent stores via the pointer would alias with each other through + // the cached SSA value. The non-pointer variant is the common case driven by reverse-pass formulas. + if (lt->return_ptr) { + continue; + } + auto it = primal_cache.find(lt->stack); + if (it != primal_cache.end()) { + irpass::replace_all_usages_with(lt->parent, lt, it->second); + modifier_.erase(lt); + } else { + primal_cache[lt->stack] = lt; + } + } else if (auto *la = s->cast()) { + auto it = adjoint_cache.find(la->stack); + if (it != adjoint_cache.end()) { + irpass::replace_all_usages_with(la->parent, la, it->second); + modifier_.erase(la); + } else { + adjoint_cache[la->stack] = la; + } + } else if (auto *p = s->cast()) { + primal_cache.erase(p->stack); + adjoint_cache.erase(p->stack); + } else if (auto *po = s->cast()) { + primal_cache.erase(po->stack); + adjoint_cache.erase(po->stack); + } else if (auto *aa = s->cast()) { + // Accumulating into the top adjoint slot does not alter the primal half, so the primal cache for + // the same stack stays valid; only the adjoint cache must be invalidated. + adjoint_cache.erase(aa->stack); + } else if (s->is() || s->is() || s->is() || s->is()) { + // The pass is intentionally intra-block: any push / pop hidden inside a nested block would + // invalidate caches asymmetrically across branches and make a sound merge complex. Conservatively + // drop both caches before descending so reads after the nested block see no stale entries. + primal_cache.clear(); + adjoint_cache.clear(); + s->accept(this); + } else { + // Any other statement is assumed inert with respect to the AdStack count headers; recurse into + // potential nested blocks (defensively) but leave the caches intact. + s->accept(this); + } + } + } + + static bool run(IRNode *root) { + CoalesceAdStackLoads pass; + root->accept(&pass); + return pass.modifier_.modify_ir(); + } + + private: + DelayedIRModifier modifier_; +}; + +} // namespace + +void backup_ssa(Block *block) { + BackupSSA::run(block); +} + +bool coalesce_ad_stack_loads(IRNode *root) { + return CoalesceAdStackLoads::run(root); +} + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/post_adjoint_cleanup.h b/quadrants/transforms/auto_diff/post_adjoint_cleanup.h new file mode 100644 index 0000000000..c80453bc62 --- /dev/null +++ b/quadrants/transforms/auto_diff/post_adjoint_cleanup.h @@ -0,0 +1,21 @@ +#pragma once + +namespace quadrants::lang { + +class IRNode; +class Block; + +// Stage: Post-MakeAdjoint cleanup. + +// Spill cross-block SSA operands the reverse-mode clones reference, so the +// reverse IR is verifier-clean: AdStackLoadTop / ArgLoad get cloned in place, +// AdStackAlloca gets re-rooted, generic SSA values get a per-IB backup +// AllocaStmt + LocalStore + LocalLoad chain. +void backup_ssa(Block *block); + +// Within a single straight-line block, dedup repeated AdStackLoadTop / +// AdStackLoadTopAdj reads of the same stack with no intervening +// Push / Pop / AccAdjoint. Returns whether any IR mutation took place. +bool coalesce_ad_stack_loads(IRNode *root); + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/validation.cpp b/quadrants/transforms/auto_diff/validation.cpp new file mode 100644 index 0000000000..2676e7a40a --- /dev/null +++ b/quadrants/transforms/auto_diff/validation.cpp @@ -0,0 +1,264 @@ +#include "quadrants/transforms/auto_diff/auto_diff_common.h" +#include "quadrants/transforms/auto_diff/validation.h" + +namespace quadrants::lang { + +namespace { + +// Detect cross-iteration read-after-write through a `needs_grad` global field at the body of an offload-level +// range-for. `MakeAdjoint` treats each iteration of an offload-level loop as independent and only chases +// loop-carried dependencies through `AdStackAllocaJudger`, which inspects local `AllocaStmt`s - cross-iteration +// dependencies through `GlobalStoreStmt` / `GlobalLoadStmt` on a `has_adjoint()` SNode are silently dropped. +// A kernel like `out[i] = out[i-1] + x[None]` therefore returns wrong gradients with no diagnostic, even with +// `ad_stack_experimental_enabled`. Pair every same-SNode (store, load) and raise when the index Stmt*s differ. +// Same-Stmt indices (`out[i] = out[i] + x[None]`) are in-place accumulation and the per-iteration backward +// pass handles those correctly; pointer equality on the SSA index Stmts after the `pre_autodiff` simplify +// (`compile_to_offloads.cpp:121`) is enough to distinguish the two cases. Walks the body of every direct-child +// `RangeForStmt` of `root`; nested for-loops' bodies are skipped via `inside_nested_` because their cross-iter +// behavior is governed by the existing IB / adstack machinery, not by this offload-level guard. +class OffloadLevelGlobalCrossIterRAWChecker : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + + void visit(GlobalStoreStmt *stmt) override { + if (inside_nested_) + return; + if (auto *gp = extract_global_ptr(stmt->dest)) { + if (gp->snode->has_adjoint()) + stores_[gp->snode].push_back(gp); + } + } + + void visit(GlobalLoadStmt *stmt) override { + if (inside_nested_) + return; + if (auto *gp = extract_global_ptr(stmt->src)) { + if (gp->snode->has_adjoint()) + loads_[gp->snode].push_back(gp); + } + } + + void visit(RangeForStmt *stmt) override { + walk_nested_loop_body(stmt->body.get()); + } + + void visit(StructForStmt *stmt) override { + walk_nested_loop_body(stmt->body.get()); + } + + void visit(MeshForStmt *stmt) override { + walk_nested_loop_body(stmt->body.get()); + } + + void visit(WhileStmt *stmt) override { + walk_nested_loop_body(stmt->body.get()); + } + + static void run(Block *offload_body) { + OffloadLevelGlobalCrossIterRAWChecker checker; + offload_body->accept(&checker); + for (const auto &kv : checker.stores_) { + auto *snode = kv.first; + auto load_it = checker.loads_.find(snode); + if (load_it == checker.loads_.end()) + continue; + for (auto *store : kv.second) { + // The store's index must reference the offload's `LoopIndexStmt` on at least one axis - otherwise it is + // iteration-independent (writes the same slot every iteration) and there is no cross-iter chain to drop. + if (!any_axis_references_loop_index(store)) + continue; + for (auto *load : load_it->second) { + // Skip loads whose index is iteration-independent on every axis (constants, captures from outer + // scopes, ...). Such loads read from a slice of the SNode that the loop never writes to in this + // iteration, so they are not a cross-iter RAW hazard. + if (!any_axis_references_loop_index(load)) + continue; + if (!indices_have_no_cross_iter_dependency(store, load)) { + QD_ERROR( + "Cross-iteration read-after-write on the `needs_grad` global field `{}` at the offload-level " + "loop body is not supported in reverse-mode autodiff: the cross-iteration adjoint chain is " + "not generated and the backward pass returns wrong gradients with no diagnostic. Wrap the " + "loop in an outer for-loop in the kernel so the inner range becomes adstack-eligible, or " + "restructure the kernel so the offload-level loop has no global field cross-iteration " + "read-after-write. Sibling-component access on a disjoint axis like `out[i, 0] = out[i, 1]` " + "is allowed; in-place accumulation `out[i] = out[i] + x` (same index Stmt on every axis) is " + "also fine.", + snode->get_node_type_name_hinted()); + } + } + } + } + } + + private: + void walk_nested_loop_body(Block *body) { + bool prev = inside_nested_; + inside_nested_ = true; + body->accept(this); + inside_nested_ = prev; + } + + static GlobalPtrStmt *extract_global_ptr(Stmt *s) { + if (auto *gp = s->cast()) + return gp; + if (auto *mp = s->cast()) { + if (auto *gp = mp->origin->cast()) + return gp; + } + return nullptr; + } + + // Decide whether two `GlobalPtrStmt`s on the same SNode index iter-independent slices of that SNode (no + // cross-iter dependency). The rules are per-axis: + // - axis where neither index references `LoopIndexStmt`: constants or outer-scope captures, may differ + // (sibling-component access on a disjoint axis like `state[i, 0]` vs `state[i, 1]`); never a cross-iter + // dependency. + // - axis where both indices reference `LoopIndexStmt`: must be the same SSA `Stmt*` (e.g. both bare + // `LoopIndexStmt(i)`). If one is `LoopIndex` and the other is `LoopIndex - 1` or similar, the load reads + // a slot the store wrote in a different iteration: this is cross-iter. + // - axis where exactly one references `LoopIndexStmt`: load and store have iter-dependent vs iter- + // independent indexing on the same axis; conservatively flag as cross-iter, since whether the constant + // index aliases an iteration of the loop is not provable here. + static bool indices_have_no_cross_iter_dependency(GlobalPtrStmt *a, GlobalPtrStmt *b) { + if (a->indices.size() != b->indices.size()) + return false; + for (std::size_t i = 0; i < a->indices.size(); ++i) { + bool a_refs = references_loop_index(a->indices[i]); + bool b_refs = references_loop_index(b->indices[i]); + if (a_refs != b_refs) + return false; + if (a_refs && b_refs && a->indices[i] != b->indices[i]) + return false; + } + return true; + } + + // Walk an index `Stmt*` and decide whether it (transitively) references a `LoopIndexStmt` of any enclosing + // loop. Recurses through linear arithmetic ops because an index like `i - 1` lowers to + // `BinaryOpStmt(LoopIndexStmt(i), Sub, ConstStmt(1))` and we want to recognise the `LoopIndexStmt` underneath. + // Conservatively returns false for anything else (constants, ndarray loads from outer scope, ...) - those + // express iteration-independent indexing that the cross-iter guard intentionally exempts. + static bool references_loop_index(Stmt *s) { + if (s == nullptr) + return false; + if (s->is()) + return true; + if (auto *bo = s->cast()) + return references_loop_index(bo->lhs) || references_loop_index(bo->rhs); + if (auto *uo = s->cast()) + return references_loop_index(uo->operand); + if (auto *to = s->cast()) + return references_loop_index(to->op1) || references_loop_index(to->op2) || references_loop_index(to->op3); + return false; + } + + static bool any_axis_references_loop_index(GlobalPtrStmt *gp) { + for (auto *idx : gp->indices) { + if (references_loop_index(idx)) + return true; + } + return false; + } + + std::unordered_map> stores_; + std::unordered_map> loads_; + bool inside_nested_ = false; +}; + +class GloablDataAccessRuleChecker : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + + void visit(GlobalLoadStmt *stmt) override { + GlobalPtrStmt *src = nullptr; + if (stmt->src->is()) { + src = stmt->src->as(); + } else { + QD_ASSERT(stmt->src->is()); + src = stmt->src->as()->origin->as(); + } + auto snode = src->snode; + if (!snode->has_adjoint_checkbit()) { + return; + } + QD_ASSERT(snode->get_adjoint_checkbit() != nullptr); + snode = snode->get_adjoint_checkbit(); + auto global_ptr = stmt->insert_after_me(Stmt::make(snode, src->indices)); + auto dtype = global_ptr->ret_type.ptr_removed(); + + auto one = insert_const(dtype, global_ptr, 1, false /*insert_before_me*/); + one->insert_after_me(Stmt::make(global_ptr, one)); + } + + void visit_gloabl_store_stmt_and_atomic_add(Stmt *stmt, GlobalPtrStmt *dest) { + auto snode = dest->snode; + if (!snode->has_adjoint_checkbit()) { + return; + } + QD_ASSERT(snode->get_adjoint_checkbit() != nullptr); + snode = snode->get_adjoint_checkbit(); + auto global_ptr = stmt->insert_before_me(Stmt::make(snode, dest->indices)); + auto global_load = stmt->insert_before_me(Stmt::make(global_ptr)); + auto dtype = global_ptr->ret_type.ptr_removed(); + auto zero = insert_const(dtype, stmt, 0, /*insert_before_me=*/true); + auto check_equal = stmt->insert_before_me(Stmt::make(BinaryOpType::cmp_eq, global_load, zero)); + std::string msg = fmt::format( + "(kernel={}) Breaks the global data access rule. Snode {} is " + "overwritten unexpectedly.", + kernel_name_, dest->snode->get_node_type_name()); + msg += "\n" + stmt->get_tb(); + + stmt->insert_before_me(Stmt::make(check_equal, msg, std::vector())); + } + + void visit(GlobalStoreStmt *stmt) override { + GlobalPtrStmt *dest = nullptr; + if (stmt->dest->is()) { + dest = stmt->dest->as(); + } else { + QD_ASSERT(stmt->dest->is()); + dest = stmt->dest->as()->origin->as(); + } + visit_gloabl_store_stmt_and_atomic_add(stmt, dest); + } + + void visit(AtomicOpStmt *stmt) override { + GlobalPtrStmt *dest = nullptr; + if (stmt->dest->is()) { + dest = stmt->dest->as(); + } else { + QD_ASSERT(stmt->dest->is()); + dest = stmt->dest->as()->origin->as(); + } + visit_gloabl_store_stmt_and_atomic_add(stmt, dest); + } + + static void run(IRNode *root, const std::string &kernel_name) { + GloablDataAccessRuleChecker checker; + checker.kernel_name_ = kernel_name; + root->accept(&checker); + } + + private: + std::string kernel_name_; +}; + +} // namespace + +void offload_level_global_cross_iter_raw_check(Block *offload_body) { + OffloadLevelGlobalCrossIterRAWChecker::run(offload_body); +} + +void global_data_access_rule_check(IRNode *root, const std::string &kernel_name) { + GloablDataAccessRuleChecker::run(root, kernel_name); +} + +namespace irpass { + +void differentiation_validation_check(IRNode *root, const CompileConfig &config, const std::string &kernel_name) { + global_data_access_rule_check(root, kernel_name); +} + +} // namespace irpass + +} // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/validation.h b/quadrants/transforms/auto_diff/validation.h new file mode 100644 index 0000000000..59b9822381 --- /dev/null +++ b/quadrants/transforms/auto_diff/validation.h @@ -0,0 +1,23 @@ +#pragma once + +#include + +namespace quadrants::lang { + +class IRNode; +class Block; + +// Stage: Validation. Diagnostic-only passes that refuse autodiff inputs which +// would silently produce wrong gradients, and inserted runtime assertions for +// global-data access rules. + +// Refuses cross-iteration read-after-write through a needs_grad global field +// at an offload-level loop body. Run pre-transform from irpass::auto_diff. +void offload_level_global_cross_iter_raw_check(Block *offload_body); + +// Inserts runtime check-bit accumulators around GlobalStore / AtomicOp on +// snodes whose adjoint check-bit is enabled. Run from +// irpass::differentiation_validation_check. +void global_data_access_rule_check(IRNode *root, const std::string &kernel_name); + +} // namespace quadrants::lang From 836847a271545193884e017f8ec74618ed48f2fe Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 4 May 2026 22:47:05 +0200 Subject: [PATCH 4/5] [Fixup] Point test_unary_collections_audit at the new auto_diff/{auto_diff_common.h,make_adjoint.cpp} paths --- tests/python/test_adstack.py | 39 ++++++++++++++++++------------------ 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index 9c476ee979..4d61aa537b 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -184,33 +184,34 @@ def kern(): def test_unary_collections_audit(): - # Prevents drift between the Python unary op registry and the C++ `unary_collections` set in - # `quadrants/transforms/auto_diff.cpp`. Every unary op whose `MakeAdjoint` branch accumulates onto - # `stmt->operand` must be either in `unary_collections` (nonlinear: needs per-iteration operand spilling on - # the adstack inside dynamic loops) or in the local `_KNOWN_LINEAR_UNARY_OPS` allow-list (reverse formula - # uses a compile-time constant coefficient, so the single-slot spill path is correct for it). Forgetting to - # classify a new diffable unary op falls back to the single-slot spill and produces silently wrong gradients - # in dynamic loops. + # Prevents drift between the Python unary op registry and the C++ `unary_collections` set declared in + # `quadrants/transforms/auto_diff/auto_diff_common.h`. Every unary op whose `MakeAdjoint` branch (defined in + # `quadrants/transforms/auto_diff/make_adjoint.cpp`) accumulates onto `stmt->operand` must be either in + # `unary_collections` (nonlinear: needs per-iteration operand spilling on the adstack inside dynamic loops) or + # in the local `_KNOWN_LINEAR_UNARY_OPS` allow-list (reverse formula uses a compile-time constant coefficient, + # so the single-slot spill path is correct for it). Forgetting to classify a new diffable unary op falls back + # to the single-slot spill and produces silently wrong gradients in dynamic loops. # # Four invariants are checked, all of them symmetric: # (1) Every diffable-math op detected in MakeAdjoint is in `cpp_nonlinear` OR in `_KNOWN_LINEAR_UNARY_OPS`. # (2) Every op in `cpp_nonlinear` has a matching diffable-math branch in MakeAdjoint. # (3) Every op in `_KNOWN_LINEAR_UNARY_OPS` has a matching diffable-math branch in MakeAdjoint. # (4) `cpp_nonlinear` and `_KNOWN_LINEAR_UNARY_OPS` are disjoint (an op cannot be both nonlinear and linear). - src_path = pathlib.Path(__file__).resolve().parents[2] / "quadrants" / "transforms" / "auto_diff.cpp" - src = src_path.read_text() + auto_diff_dir = pathlib.Path(__file__).resolve().parents[2] / "quadrants" / "transforms" / "auto_diff" + common_src = (auto_diff_dir / "auto_diff_common.h").read_text() + make_adjoint_src = (auto_diff_dir / "make_adjoint.cpp").read_text() - cc_match = re.search(r"unary_collections\s*\{([^}]+)\}", src) - assert cc_match is not None, "unary_collections not located in auto_diff.cpp" + cc_match = re.search(r"unary_collections\s*\{([^}]+)\}", common_src) + assert cc_match is not None, "unary_collections not located in auto_diff_common.h" cpp_nonlinear = set(re.findall(r"UnaryOpType::(\w+)", cc_match.group(1))) - make_adjoint_start = src.find("class MakeAdjoint") - assert make_adjoint_start != -1, "class MakeAdjoint not located in auto_diff.cpp" - adj_start = src.find("void visit(UnaryOpStmt *stmt) override", make_adjoint_start) - assert adj_start != -1, "MakeAdjoint::visit(UnaryOpStmt*) not located in auto_diff.cpp" - adj_end = src.find("void visit(", adj_start + 10) + make_adjoint_start = make_adjoint_src.find("class MakeAdjoint") + assert make_adjoint_start != -1, "class MakeAdjoint not located in make_adjoint.cpp" + adj_start = make_adjoint_src.find("void visit(UnaryOpStmt *stmt) override", make_adjoint_start) + assert adj_start != -1, "MakeAdjoint::visit(UnaryOpStmt*) not located in make_adjoint.cpp" + adj_end = make_adjoint_src.find("void visit(", adj_start + 10) assert adj_end != -1, "next visitor method after MakeAdjoint::visit(UnaryOpStmt*) not found" - adj_block = src[adj_start:adj_end] + adj_block = make_adjoint_src[adj_start:adj_end] # Split the visitor's if/else-if chain into per-op segments, then classify each segment as "diffable math" # iff its body accumulates onto `stmt->operand`. Two accumulate entry points are recognised: the raw # `accumulate(stmt->operand, ...)` call (used by `neg` and `cast_value`) and the `acc(...)` lambda (used by @@ -242,8 +243,8 @@ def test_unary_collections_audit(): missing = diffable_math - cpp_nonlinear - _KNOWN_LINEAR_UNARY_OPS assert not missing, ( f"Diffable unary ops not classified as nonlinear or linear: {sorted(missing)}. " - f"Add each one to `unary_collections` in quadrants/transforms/auto_diff.cpp (if nonlinear) or to " - f"`_KNOWN_LINEAR_UNARY_OPS` in this file (if linear)." + f"Add each one to `unary_collections` in quadrants/transforms/auto_diff/auto_diff_common.h (if " + f"nonlinear) or to `_KNOWN_LINEAR_UNARY_OPS` in this file (if linear)." ) stray_nonlinear = cpp_nonlinear - diffable_math assert not stray_nonlinear, ( From bb5a62d4024c740c763c150f00ddb06ad0f0813c Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Tue, 5 May 2026 00:24:02 +0200 Subject: [PATCH 5/5] [Fixup] Elim-pushes: detect prologue init push by position, not by zero value, to preserve real x = 0.0 body stores --- quadrants/transforms/auto_diff/auto_diff.cpp | 33 +-- .../transforms/auto_diff/auto_diff_common.h | 115 ++++---- .../eliminate_recomputable_pushes.cpp | 278 +++++++++--------- .../auto_diff/forward_state_spill.cpp | 232 +++++++-------- .../auto_diff/forward_state_spill.h | 28 +- quadrants/transforms/auto_diff/ir_shaping.cpp | 148 ++++------ quadrants/transforms/auto_diff/ir_shaping.h | 19 +- .../transforms/auto_diff/make_adjoint.cpp | 228 +++++++------- quadrants/transforms/auto_diff/make_adjoint.h | 5 +- quadrants/transforms/auto_diff/make_dual.cpp | 26 +- quadrants/transforms/auto_diff/make_dual.h | 4 +- .../auto_diff/post_adjoint_cleanup.cpp | 106 +++---- .../auto_diff/post_adjoint_cleanup.h | 12 +- quadrants/transforms/auto_diff/validation.cpp | 36 +-- quadrants/transforms/auto_diff/validation.h | 14 +- tests/python/test_adstack.py | 61 ++++ 16 files changed, 660 insertions(+), 685 deletions(-) diff --git a/quadrants/transforms/auto_diff/auto_diff.cpp b/quadrants/transforms/auto_diff/auto_diff.cpp index 177a42a345..77556ba947 100644 --- a/quadrants/transforms/auto_diff/auto_diff.cpp +++ b/quadrants/transforms/auto_diff/auto_diff.cpp @@ -17,12 +17,11 @@ void auto_diff(IRNode *root, const CompileConfig &config, AutodiffMode autodiff_ QD_AUTO_PROF; if (autodiff_mode == AutodiffMode::kReverse) { if (auto *root_block = root->cast()) { - // Top-level range-fors at the kernel root become the offload-level loop. Walk each direct child once - // and refuse cross-iteration global RAW on a needs_grad SNode before the AD transformation runs - any - // diagnostic from inside `MakeAdjoint` would already point at the partially-rewritten reverse IR and - // be much harder to read. StructFor offloads are out of scope: their loop variable is an SNode index - // produced by the runtime, not a user-controlled integer arithmeticed against itself, so the - // `out[i-1]` shape this guard targets does not arise. + // Top-level range-fors at the kernel root become the offload-level loop. Walk each direct child once and refuse + // cross-iteration global RAW on a needs_grad SNode before the AD transformation runs - any diagnostic from inside + // `MakeAdjoint` would already point at the partially-rewritten reverse IR and be much harder to read. StructFor + // offloads are out of scope: their loop variable is an SNode index produced by the runtime, not a user-controlled + // integer arithmeticed against itself, so the `out[i-1]` shape this guard targets does not arise. for (auto &s : root_block->statements) { if (auto *rf = s->cast()) offload_level_global_cross_iter_raw_check(rf->body.get()); @@ -38,23 +37,23 @@ void auto_diff(IRNode *root, const CompileConfig &config, AutodiffMode autodiff_ replace_local_var_with_stacks(ib, config.ad_stack_size); type_check(root, config); - // Drop AdStackAllocas whose pushed value is recomputable from already-stack-backed allocas + args + - // const + loop-index. Trades forward-pass adstack memory traffic (one push + one load per iter per - // intermediate spill) for cloned arithmetic in the reverse scope, which `BackupSSA::generic_visit` - // generates on demand via the same RecomputableChainCloner path. Must run after - // `replace_local_var_with_stacks` (so the analyzer sees the AdStackAlloca shape, not the alloca-with-store - // shape promote_ssa_to_local_var emits) and before `make_adjoint` (so the reverse pass is generated - // against the cleaned forward IR with no spurious push/load scaffolding). + // Drop AdStackAllocas whose pushed value is recomputable from already-stack-backed allocas + args + const + + // loop-index. Trades forward-pass adstack memory traffic (one push + one load per iter per intermediate spill) + // for cloned arithmetic in the reverse scope, which `BackupSSA::generic_visit` generates on demand via the same + // RecomputableChainCloner path. Must run after `replace_local_var_with_stacks` (so the analyzer sees the + // AdStackAlloca shape, not the alloca-with-store shape promote_ssa_to_local_var emits) and before + // `make_adjoint` (so the reverse pass is generated against the cleaned forward IR with no spurious push/load + // scaffolding). eliminate_recomputable_ad_stack_pushes(ib); type_check(root, config); make_adjoint(ib); type_check(root, config); backup_ssa(ib); - // After MakeAdjoint emits the reverse-pass body, an outer-loop-invariant value pulled into the - // reverse direction by `accumulate_unary_operand_checked` becomes a fresh `AdStackLoadTopStmt` per - // user-side use; in straight-line unrolled IR those reads coalesce to one load per stack per block, - // which the dedicated pass handles before the IR reaches `irpass::analysis::verify_if_debug`. + // After MakeAdjoint emits the reverse-pass body, an outer-loop-invariant value pulled into the reverse + // direction by `accumulate_unary_operand_checked` becomes a fresh `AdStackLoadTopStmt` per user-side use; in + // straight-line unrolled IR those reads coalesce to one load per stack per block, which the dedicated pass + // handles before the IR reaches `irpass::analysis::verify_if_debug`. coalesce_ad_stack_loads(ib); irpass::analysis::verify_if_debug(root, config); } diff --git a/quadrants/transforms/auto_diff/auto_diff_common.h b/quadrants/transforms/auto_diff/auto_diff_common.h index ca067045b7..87a2023e0d 100644 --- a/quadrants/transforms/auto_diff/auto_diff_common.h +++ b/quadrants/transforms/auto_diff/auto_diff_common.h @@ -61,31 +61,28 @@ class NonLinearOps { }; // ---------------------------------------------------------------------------- -// Recomputable chain analyzer + cloner: decide whether a forward SSA value can -// be reconstructed in the reverse-pass scope from already-stack-backed allocas, -// kernel args, constants, and loop indices, and clone the DAG at a target -// insertion point. Cross-stage shared infrastructure: used by -// `EliminateRecomputableAdStackPushes` (forward_state_spill stage) to drop -// pushes whose values are recomputable, and by `BackupSSA::generic_visit` -// (post_adjoint_cleanup stage) to clone such chains in place of cross-block -// SSA reads. +// Recomputable chain analyzer + cloner: decide whether a forward SSA value can be reconstructed in the reverse-pass +// scope from already-stack-backed allocas, kernel args, constants, and loop indices, and clone the DAG at a target +// insertion point. Cross-stage shared infrastructure: used by `EliminateRecomputableAdStackPushes` +// (forward_state_spill stage) to drop pushes whose values are recomputable, and by `BackupSSA::generic_visit` +// (post_adjoint_cleanup stage) to clone such chains in place of cross-block SSA reads. // ---------------------------------------------------------------------------- -// Returns true iff `stmt`'s transitive operand DAG terminates at recomputable leaves via side-effect-free -// interior ops only. Used by `EliminateRecomputableAdStackPushes` and `BackupSSA::generic_visit` to decide -// whether a forward SSA value can be reconstructed in the reverse-pass scope from already-stack-backed allocas -// + kernel-args + constants + loop indices, instead of being spilled to a per-iteration adstack or to -// `BackupSSA::load`'s last-iteration plain alloca. +// Returns true iff `stmt`'s transitive operand DAG terminates at recomputable leaves via side-effect-free interior ops +// only. Used by `EliminateRecomputableAdStackPushes` and `BackupSSA::generic_visit` to decide whether a forward SSA +// value can be reconstructed in the reverse-pass scope from already-stack-backed allocas + kernel-args + constants + +// loop indices, instead of being spilled to a per-iteration adstack or to `BackupSSA::load`'s last-iteration plain +// alloca. // -// Recomputable leaves: AdStackLoadTopStmt (re-readable via cloned load), AdStackAllocaStmt (the stack itself, -// shared not cloned), ArgLoadStmt (kernel-arg, immutable within the launch), ConstStmt, LoopIndexStmt (clonable -// to read the reverse-direction loop's index, which matches the forward iteration the reverse is currently -// processing). Side-effect-free interior ops: UnaryOp, BinaryOp, TernaryOp, MatrixPtr, GlobalPtr, ExternalPtr. +// Recomputable leaves: AdStackLoadTopStmt (re-readable via cloned load), AdStackAllocaStmt (the stack itself, shared +// not cloned), ArgLoadStmt (kernel-arg, immutable within the launch), ConstStmt, LoopIndexStmt (clonable to read the +// reverse-direction loop's index, which matches the forward iteration the reverse is currently processing). +// Side-effect-free interior ops: UnaryOp, BinaryOp, TernaryOp, MatrixPtr, GlobalPtr, ExternalPtr. // -// `GlobalLoadStmt` and `ExternalPtrAccessStmt`-style reads are intentionally not recomputable: the underlying -// global memory may be mutated mid-kernel by a sibling task, so reading at a different IR position can yield a -// different value. `LocalLoadStmt` similarly aliases mutable allocas; the reverse pass must read its forward -// value through the dedicated spill machinery. +// `GlobalLoadStmt` and `ExternalPtrAccessStmt`-style reads are intentionally not recomputable: the underlying global +// memory may be mutated mid-kernel by a sibling task, so reading at a different IR position can yield a different +// value. `LocalLoadStmt` similarly aliases mutable allocas; the reverse pass must read its forward value through the +// dedicated spill machinery. // // Caller passes a `cache` to share memoization across multiple queries on the same DAG (diamond shapes). class RecomputableChainAnalyzer { @@ -94,8 +91,8 @@ class RecomputableChainAnalyzer { auto it = cache.find(stmt); if (it != cache.end()) return it->second; - // Tentatively false to break cycles in pathological IR (real SSA DAGs are acyclic, but the cache also - // serves as a visited set during recursion). + // Tentatively false to break cycles in pathological IR (real SSA DAGs are acyclic, but the cache also serves as a + // visited set during recursion). cache[stmt] = false; bool result = check(stmt, cache); cache[stmt] = result; @@ -105,14 +102,12 @@ class RecomputableChainAnalyzer { private: static bool check(Stmt *stmt, std::unordered_map &cache) { // Recomputable leaves: ConstStmt, ArgLoadStmt, AdStackLoadTopStmt, AdStackAllocaStmt. LoopIndexStmt is - // intentionally excluded - cloning a LoopIndexStmt copies the reference to the forward RangeForStmt, - // but the cloned consumer lives inside the reverse RangeForStmt (a separate stmt with its own loop - // index), so the cloned read points at undefined state and silently double-accumulates gradients - // (`test_adstack_sum_linear`). A future MakeAdjoint coordination pass that tracks - // forward-RangeFor-to-reverse-RangeFor mapping could lift this restriction. + // intentionally excluded - cloning a LoopIndexStmt copies the reference to the forward RangeForStmt, but the cloned + // consumer lives inside the reverse RangeForStmt (a separate stmt with its own loop index), so the cloned read + // points at undefined state and silently double-accumulates gradients. // - // AdStackLoadTopStmt as a leaf is correct under the dominance + control-flow-consumer + self-load - // guards in `EliminateRecomputableAdStackPushes::run_one_pass`. The reasoning: + // AdStackLoadTopStmt as a leaf is correct under the dominance + control-flow-consumer + self-load guards in + // `EliminateRecomputableAdStackPushes::run_one_pass`. The reasoning: // // - In FORWARD: the eliminated stack S's body push dominates each `AdStackLoadTopStmt(S)` (the // dominance guard), and the chain leaves' pushes dominate the chain's evaluation point. Each @@ -130,27 +125,24 @@ class RecomputableChainAnalyzer { // been popped yet, so load_top(T) returns T's iter-k-push value - matching what the original // load_top(S) returned in forward iter k. // - // The control-flow-consumer guard in `run_one_pass` covers a separate issue: stacks whose load_tops - // are direct operands of IfStmt cond / RangeFor begin/end. `MakeAdjoint::visit(IfStmt)` runs a - // dedicated snap-stack fixup that ONLY triggers when the cond is a bare `AdStackLoadTopStmt` (line - // 2168-2202). Eliminating the cond stack converts the cond into a compound stmt, the snap-stack - // does not trigger, and the reverse cond falls back to BackupSSA's load(op) which is single-slot - // last-iter only - silent gradient corruption on multi-iter loops. + // The control-flow-consumer guard in `run_one_pass` covers a separate issue: stacks whose load_tops are direct + // operands of IfStmt cond / RangeFor begin/end. `MakeAdjoint::visit(IfStmt)` runs a dedicated snap-stack fixup that + // ONLY triggers when the cond is a bare `AdStackLoadTopStmt` (line 2168-2202). Eliminating the cond stack converts + // the cond into a compound stmt, the snap-stack does not trigger, and the reverse cond falls back to BackupSSA's + // load(op) which is single-slot last-iter only - silent gradient corruption on multi-iter loops. if (stmt->is() || stmt->is() || stmt->is() || stmt->is()) { return true; } - // GlobalLoadStmt as recomputable: the load reads a SNode value via GlobalPtrStmt. The cloned chain in - // the reverse pass re-issues the same load - safe iff the global is not mutated between the forward - // chain evaluation and the cloned re-read. Within a single kernel execution, forward writes complete - // before reverse runs, so the reverse re-read sees the kernel's final post-write state. If a global is - // mutated by the forward and then re-read by the reverse clone, the values can differ. + // GlobalLoadStmt as recomputable: the load reads a SNode value via GlobalPtrStmt. The cloned chain in the reverse + // pass re-issues the same load - safe iff the global is not mutated between the forward chain evaluation and the + // cloned re-read. Within a single kernel execution, forward writes complete before reverse runs, so the reverse + // re-read sees the kernel's final post-write state. If a global is mutated by the forward and then re-read by the + // reverse clone, the values can differ. // - // Most rigid-step kernel chain leaves are reads of input parameters (mass, inertia, joint params, - // morphology) that the kernel does NOT write. For those, the re-read returns the same value. - // Conservative analysis: a future safety check could prove the SNode is read-only in the kernel; until - // then this path relies on the test suite's gradient-correctness asserts to surface any mutated-global - // miscompilation. + // Chain leaves that pass this branch are typically reads of kernel-input parameters not written by the kernel; for + // those, the re-read returns the same value as the forward read. The pass does not statically verify SNode + // read-only-ness, and mutated-global cases would silently produce wrong gradients. bool is_interior = stmt->is() || stmt->is() || stmt->is() || stmt->is() || stmt->is() || stmt->is() || stmt->is(); @@ -168,18 +160,18 @@ class RecomputableChainAnalyzer { } }; -// Clones the SSA chain rooted at `src` into the IR, inserting cloned stmts before `insert_point`. Returns the -// cloned root. Per-stmt cache shared across one resolution materializes each SSA value at most once: diamond -// DAGs see two consumers but get one shared clone. `AdStackAllocaStmt` is treated as a leaf and shared (not -// cloned) - the stack itself is a unique storage handle that must not be duplicated. +// Clones the SSA chain rooted at `src` into the IR, inserting cloned stmts before `insert_point`. Returns the cloned +// root. Per-stmt cache shared across one resolution materializes each SSA value at most once: diamond DAGs see two +// consumers but get one shared clone. `AdStackAllocaStmt` is treated as a leaf and shared (not cloned) - the stack +// itself is a unique storage handle that must not be duplicated. // -// Pop-ordering safety: cloned `AdStackLoadTopStmt`s read the live top at the cloned position. `MakeAdjoint` -// emits `AdStackPopStmt` for each surviving `AdStackPushStmt`, and the existing reverse-pass scheme places -// the pop AFTER all uses of that stack within the reverse iteration (uses include both the original -// `AdStackLoadTopStmt`s emitted by `ReplaceLocalVarWithStacks` and the consumers' clones). For loop-carried -// allocas the pop fires early to expose the iteration's INPUT primal as the new top, which is exactly the -// value the recomputed chain needs - the existing per-consumer clone path at `BackupSSA::generic_visit` line -// ~2697 relies on this same property and has been correct in production. +// Pop-ordering safety: cloned `AdStackLoadTopStmt`s read the live top at the cloned position. `MakeAdjoint` emits +// `AdStackPopStmt` for each surviving `AdStackPushStmt`, and the existing reverse-pass scheme places the pop AFTER all +// uses of that stack within the reverse iteration (uses include both the original `AdStackLoadTopStmt`s emitted by +// `ReplaceLocalVarWithStacks` and the consumers' clones). For loop-carried allocas the pop fires early to expose the +// iteration's INPUT primal as the new top, which is exactly the value the recomputed chain needs - the existing +// per-consumer clone path at `BackupSSA::generic_visit` line ~2697 relies on this same property and has been correct in +// production. class RecomputableChainCloner { public: static Stmt *clone_at(Stmt *src, Stmt *insert_point, std::unordered_map &cache) { @@ -193,8 +185,8 @@ class RecomputableChainCloner { } else if (src->is() || src->is() || src->is()) { auto cloned_unique = src->clone(); cloned = insert_point->insert_before_me(std::move(cloned_unique)); - // For AdStackLoadTopStmt clones, the cloned stmt's `stack` operand still points at the original - // AdStackAllocaStmt - that's the desired sharing. + // For AdStackLoadTopStmt clones, the cloned stmt's `stack` operand still points at the original AdStackAllocaStmt + // - that's the desired sharing. } else { // Compound op: clone first, then walk operands and rewire each to a recursive clone. auto cloned_unique = src->clone(); @@ -214,9 +206,8 @@ class RecomputableChainCloner { }; // ---------------------------------------------------------------------------- -// ADTransform: shared base for reverse-mode (MakeAdjoint) and forward-mode -// (MakeDual) IR builders. All methods are inline so derived classes in -// separate translation units can use them without ODR concerns. +// ADTransform: shared base for reverse-mode (MakeAdjoint) and forward-mode (MakeDual) IR builders. Methods are inline +// so derived classes in separate translation units can use them without ODR concerns. // ---------------------------------------------------------------------------- class ADTransform : public IRVisitor { protected: diff --git a/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp b/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp index 6230174365..e826db5b92 100644 --- a/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp +++ b/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp @@ -5,10 +5,10 @@ namespace quadrants::lang { namespace { -// Eliminate AdStackAllocaStmts whose pushed value is recomputable from already-stack-backed allocas, kernel -// args, constants, and loop indices. Runs between `ReplaceLocalVarWithStacks` and `MakeAdjoint`, so the reverse -// pass is generated against the cleaned IR (no spurious AdStackPushStmt / AdStackLoadTopStmt scaffolding for -// values the reverse can reconstruct on the fly via cloned forward DAGs). +// Eliminate AdStackAllocaStmts whose pushed value is recomputable from already-stack-backed allocas, kernel args, +// constants, and loop indices. Runs between `ReplaceLocalVarWithStacks` and `MakeAdjoint`, so the reverse pass is +// generated against the cleaned IR (no spurious AdStackPushStmt / AdStackLoadTopStmt scaffolding for values the reverse +// can reconstruct on the fly via cloned forward DAGs). // // Eligibility per AdStackAllocaStmt S in an independent block: // 1. S is written by exactly one AdStackPushStmt (single-push pattern). Multi-push allocas hold loop-carried @@ -17,31 +17,30 @@ namespace { // 2. The pushed value's transitive operand DAG is recomputable per RecomputableChainAnalyzer (leaves at // AdStackLoadTop / AdStackAlloca / ArgLoad / Const / LoopIndex; interior side-effect-free ops only). // 3. S has no AdStackLoadTopStmt with `return_ptr=true` consumers (those return a pointer aliasing the slot -// and a follow-up store would not be modeled by an SSA replacement). +// and a subsequent store would not be modeled by an SSA replacement). // // Action on eligibility: replace every `AdStackLoadTopStmt(S)` with the original pushed SSA value (the -// `AdStackPushStmt::v`), then erase the push and the alloca. The pushed SSA value's chain stays in the forward -// IR; downstream consumers in the forward pass now reference it directly as SSA, and `MakeAdjoint` plus -// `BackupSSA` reconstruct the chain on demand in the reverse pass via the same DAG-clone path that -// `BackupSSA::generic_visit` exercises for cross-block SSA references. +// `AdStackPushStmt::v`), then erase the push and the alloca. The pushed SSA value's chain stays in the forward IR; +// downstream consumers in the forward pass now reference it directly as SSA, and `MakeAdjoint` plus `BackupSSA` +// reconstruct the chain on demand in the reverse pass via the same DAG-clone path that `BackupSSA::generic_visit` +// exercises for cross-block SSA references. // -// Iterates to fixed point: eliminating one stack can newly expose another stack's chain as recomputable -// (S2's pushed value chained through `AdStackLoadTopStmt(S1)`; once S1 is gone, the chain may collapse into -// pure SSA chained through S1's pushed-value chain instead). Each pass eliminates at least one stack or -// terminates. +// Iterates to fixed point: eliminating one stack can newly expose another stack's chain as recomputable (S2's pushed +// value chained through `AdStackLoadTopStmt(S1)`; once S1 is gone, the chain may collapse into pure SSA chained through +// S1's pushed-value chain instead). Each pass eliminates at least one stack or terminates. // -// Cost model: this pass trades forward-pass adstack pushes (memory write + top-pointer bump per push) for -// extra arithmetic in the reverse pass (the cloned DAG re-executes the forward chain). For pure-arithmetic -// chains rooted at a single loop-carried alloca - the dominant shape on Genesis-style rigid-step kernels - -// the win is typically 5-10x: each push is a memory op crossing the L1 boundary on every iteration, while -// the recomputed arithmetic stays in registers and reuses warmed-up sin/cos/exp pipeline state. +// Cost model: this pass trades forward-pass adstack pushes (memory write + top-pointer bump per push) for extra +// arithmetic in the reverse pass (the cloned DAG re-executes the forward chain). For pure-arithmetic chains rooted at a +// single loop-carried alloca - the dominant shape in nonlinear-recurrence kernels - the win is typically 5-10x: each +// push is a memory op crossing the L1 boundary on every iteration, while the recomputed arithmetic stays in registers +// and reuses warmed-up sin/cos/exp pipeline state. class EliminateRecomputableAdStackPushes { public: static void run(Block *ib) { - // Iterate to fixed point. Each pass either eliminates at least one stack or terminates. The hard cap - // protects against analysis bugs (e.g. an eligibility check that returns true on the same stmt twice - // in a row); under a correct implementation each pass strictly reduces the AdStackAllocaStmt count, so - // the actual iteration bound is the number of stacks at entry. 1024 is well above any realistic kernel. + // Iterate to fixed point. Each pass either eliminates at least one stack or terminates. The hard cap protects + // against analysis bugs (e.g. an eligibility check that returns true on the same stmt twice in a row); under a + // correct implementation each pass strictly reduces the AdStackAllocaStmt count, so the actual iteration bound is + // the number of stacks at entry. 1024 is well above any realistic kernel. for (int i = 0; i < 1024; i++) { if (!run_one_pass(ib)) return; @@ -49,32 +48,39 @@ class EliminateRecomputableAdStackPushes { } private: - // True iff `s` is a literal-zero ConstStmt or a MatrixInitStmt whose every element is a literal-zero - // ConstStmt. Matches the init-zero push pattern that ReplaceLocalVarWithStacks emits right after each - // AdStackAllocaStmt: scalar allocas get a ConstStmt(0), tensor-typed allocas get a MatrixInitStmt of - // ConstStmt(0)s. Both are erased along with the alloca during elimination - they only matter if a load - // could observe them before the body push fires, which the standard PromoteSSA2LocalVar + - // ReplaceLocalVarWithStacks codegen never produces (loads always follow the body push within the loop - // iter that emits them). - static bool is_zero_init_value(Stmt *s) { - if (auto *c = s->cast()) { - return c->val.equal_value(0); - } - if (auto *mi = s->cast()) { - for (auto *elem : mi->values) { - auto *c = elem->cast(); - if (c == nullptr || !c->val.equal_value(0)) - return false; + // Identify the prologue init push for `stack`: the AdStackPushStmt that ReplaceLocalVarWithStacks emits in the + // alloca's parent block, immediately after the alloca and the ConstStmt (plus MatrixInitStmt for tensor allocas) that + // produces its zero initial value. By construction it is the first AdStackPushStmt for `stack` in document order + // within `stack->parent`, so a forward scan from the alloca's position locates it deterministically. Detecting the + // init push by VALUE (literal zero) instead would also match real user-level `x = 0.0` body stores lowered to + // zero-valued AdStackPushStmts: routing such a store into `init_pushes` lets the pass accept the stack as "single + // body push", erase the user's zero store, and rewire every load_top to the non-zero body push's SSA chain - silently + // corrupting forward values and gradients in iterations the user explicitly zeroed. + static AdStackPushStmt *find_init_push(AdStackAllocaStmt *stack) { + Block *stack_block = stack->parent; + if (stack_block == nullptr) + return nullptr; + bool past_alloca = false; + for (auto &owned : stack_block->statements) { + Stmt *s = owned.get(); + if (s == stack) { + past_alloca = true; + continue; + } + if (!past_alloca) + continue; + if (auto *p = s->cast()) { + if (p->stack == stack) + return p; } - return true; } - return false; + return nullptr; } // Returns true if any stack was eliminated this pass. static bool run_one_pass(Block *ib) { - // Collect every AdStackAllocaStmt anywhere within the IB. The IB is the root of a contiguous AD scope, so - // a single walk is sufficient. + // Collect every AdStackAllocaStmt anywhere within the IB. The IB is the root of a contiguous AD scope, so a single + // walk is sufficient. std::vector stacks; auto collected = irpass::analysis::gather_statements(ib, [&](Stmt *s) { return s->is(); }); for (auto *s : collected) { @@ -85,8 +91,8 @@ class EliminateRecomputableAdStackPushes { std::unordered_map recomputable_cache; for (auto *stack : stacks) { - // Re-classify users of `stack` each iteration: a previous elimination on a downstream stack may have - // removed users that previously disqualified `stack`. + // Re-classify users of `stack` on each pass iteration: an elimination on a downstream stack within the same outer + // fixed-point round may have removed users that would otherwise disqualify `stack`. std::vector pushes; std::vector load_tops; bool disqualified = false; @@ -109,16 +115,16 @@ class EliminateRecomputableAdStackPushes { pushes.push_back(p); } else if (auto *lt = user->cast()) { if (lt->return_ptr) { - // return_ptr=true returns a pointer into the slot; replacing with an SSA value loses the - // pointer-identity contract. Keep the stack. + // return_ptr=true returns a pointer into the slot; replacing with an SSA value loses the pointer-identity + // contract. Keep the stack. disqualified = true; break; } load_tops.push_back(lt); } else if (user->is() || user->is() || user->is()) { - // Pre-MakeAdjoint, pop / adj-acc / load-top-adj should not appear yet for this stack. If they do, - // some upstream pass has already touched it - keep as-is to avoid double-rewrites. + // Pre-MakeAdjoint, pop / adj-acc / load-top-adj should not appear yet for this stack. If they do, some + // upstream pass has already touched it - keep as-is to avoid double-rewrites. disqualified = true; break; } @@ -127,17 +133,17 @@ class EliminateRecomputableAdStackPushes { continue; } - // ReplaceLocalVarWithStacks emits a const-zero "init" AdStackPushStmt immediately after the - // AdStackAllocaStmt as the stack's initial value (auto_diff.cpp:867-872 in pre-fix tree). Real - // user-level allocas in a loop body have at most ONE additional "body" push per iteration. Treat - // const-zero pushes whose value is a `ConstStmt` with all-zero `TypedConstant` as inits and require at - // most one non-init "body" push for elimination eligibility. Multi-body-push allocas hold loop-carried - // state (each iteration's push depends on the previous), so the reverse pass cannot reconstruct - // iteration k from iteration k-1 and the stack must stay. + // ReplaceLocalVarWithStacks emits exactly one prologue "init" AdStackPushStmt next to the AdStackAllocaStmt to + // seed the stack with a zero initial value. Real user-level allocas in a loop body have at most ONE additional + // "body" push per iteration. Find the prologue push by position (`find_init_push`) and require at most one + // non-init "body" push for elimination eligibility. Multi-body-push allocas hold loop-carried state (each + // iteration's push depends on the previous), so the reverse pass cannot reconstruct iteration k from iteration + // k-1 and the stack must stay. + AdStackPushStmt *init_push = find_init_push(stack); AdStackPushStmt *body_push = nullptr; std::vector init_pushes; for (auto *p : pushes) { - if (is_zero_init_value(p->v)) { + if (p == init_push) { init_pushes.push_back(p); } else { if (body_push != nullptr) { @@ -155,47 +161,44 @@ class EliminateRecomputableAdStackPushes { if (!RecomputableChainAnalyzer::is_recomputable(pushed_val, recomputable_cache)) { continue; } - // Loop-carried self-reference: when the pushed value's transitive operand DAG includes an - // AdStackLoadTopStmt of the SAME stack we are about to eliminate, the value is `read prev iter from - // stack -> compute -> push next iter`. Rewriting load_top($S) -> pushed_value_SSA leaves a self-cycle - // in the SSA graph (`acc_new = acc_new + sin(a)`) which both corrupts the IR and loses the iteration - // recurrence the stack was carrying. The recomputable-chain analyzer happily accepts such chains - // because AdStackLoadTopStmt is a recomputable leaf in general; the additional check below specifically - // disqualifies self-loaded stacks. + // Loop-carried self-reference: when the pushed value's transitive operand DAG includes an AdStackLoadTopStmt of + // the SAME stack we are about to eliminate, the value is `read prev iter from stack -> compute -> push next + // iter`. Rewriting load_top($S) -> pushed_value_SSA leaves a self-cycle in the SSA graph (`acc_new = acc_new + + // sin(a)`) which both corrupts the IR and loses the iteration recurrence the stack was carrying. The + // recomputable-chain analyzer happily accepts such chains because AdStackLoadTopStmt is a recomputable leaf in + // general; the additional check below specifically disqualifies self-loaded stacks. + // + // The classic shape this protects: `acc = 0.0; for j in range(N): acc += sin(...)`. After PromoteSSA2LocalVar + // + ReplaceLocalVarWithStacks the IR has `$acc = stack_alloc; init push; for j: $tmp = stack_load_top $acc; + // $new = $tmp + sin(...); push $acc, val=$new`. The chain `$tmp + sin(...)` is recomputable per the leaf rules + // (LoadTop is a leaf), but rewriting load_top($acc) to ($tmp + sin(...)) makes the SSA graph self-reference + // $new and the reverse pass loses every iteration's accumulator state. // - // The classic shape this protects: `acc = 0.0; for j in range(N): acc += sin(...)`. After - // PromoteSSA2LocalVar + ReplaceLocalVarWithStacks the IR has `$acc = stack_alloc; init push; - // for j: $tmp = stack_load_top $acc; $new = $tmp + sin(...); push $acc, val=$new`. The chain `$tmp + - // sin(...)` is recomputable per the leaf rules (LoadTop is a leaf), but rewriting load_top($acc) to - // ($tmp + sin(...)) makes the SSA graph self-reference $new and the reverse pass loses every - // iteration's accumulator state. - // Read-before-write protection: substituting `AdStackLoadTopStmt(S)` with the body push's value SSA - // chain only preserves semantics when the body push DOMINATES every load_top in the forward IR - - // that is, every load_top reads what the body push just wrote in the same iteration. The - // PromoteSSA2LocalVar + ReplaceLocalVarWithStacks pipeline emits exactly this shape for required-def - // spills: `alloca; push val=def; load_top` in document order, with the push immediately following - // the def and loads occurring later in the same iteration. + // Read-before-write protection: substituting `AdStackLoadTopStmt(S)` with the body push's value SSA chain only + // preserves semantics when the body push DOMINATES every load_top in the forward IR - that is, every load_top + // reads what the body push just wrote in the same iteration. The PromoteSSA2LocalVar + ReplaceLocalVarWithStacks + // pipeline emits exactly this shape for required-def spills: `alloca; push val=def; load_top` in document order, + // with the push immediately following the def and loads occurring later in the same iteration. // - // Loop-carried recurrences violate the dominance rule. Take the canonical Fibonacci shape - // `p, q = q, p + q` lowered to IR: + // Loop-carried recurrences violate the dominance rule. Take the canonical Fibonacci shape `p, q = q, p + q` + // lowered to IR: // // $tmp_p = load_top($p) # reads PREVIOUS iter's push into $p // $tmp_q = load_top($p) + load_top($q) // push $p, val=$tmp_p # writes NEXT iter's value into $p // push $q, val=$tmp_q // - // Here load_top($p) reads BEFORE the body push of $p within the same iter, so it observes iter - // (k-1)'s value, not iter k's. Substituting `load_top($p) -> $tmp_p`'s SSA chain (which reads - // load_top($q) at iter k) gives iter k's q-value instead of iter (k-1)'s p-value, off by one - // iteration and producing zero gradients on the Fibonacci-style regression test pinned by - // `test_ad_fibonacci_index`. + // Here load_top($p) reads BEFORE the body push of $p within the same iter, so it observes iter (k-1)'s value, not + // iter k's. Substituting `load_top($p) -> $tmp_p`'s SSA chain (which reads load_top($q) at iter k) gives iter k's + // q-value instead of iter (k-1)'s p-value, off by one iteration and producing zero gradients on Fibonacci-style + // recurrences. // - // The check below asserts that for every load_top in `load_tops`, the body push is its predecessor - // within the same containing block (or, equivalently, load_top comes AFTER the body push in - // document order at the body push's block level). Loads that live inside nested control flow - // (`IfStmt`, `RangeForStmt`) under the body push's block are also fine because the body push - // dominates them. We approximate dominance with the lexical "push's block contains load_top's - // ancestor block AND push position < load_top's ancestor's position in that block" check. + // The check below asserts that for every load_top in `load_tops`, the body push is its predecessor within the + // same containing block (or, equivalently, load_top comes AFTER the body push in document order at the body + // push's block level). Loads that live inside nested control flow (`IfStmt`, `RangeForStmt`) under the body + // push's block are also fine because the body push dominates them. We approximate dominance with the lexical + // "push's block contains load_top's ancestor block AND push position < load_top's ancestor's position in that + // block" check. auto block_position = [](Stmt *s) -> int { Block *b = s->parent; if (b == nullptr) @@ -225,8 +228,8 @@ class EliminateRecomputableAdStackPushes { } int cursor_pos = block_position(cursor); if (cursor_pos <= push_pos) { - // load_top precedes the body push (or its enclosing container does): it reads the previous - // iteration's value, not iter k's. Substitution would shift iterations by one. + // load_top precedes the body push (or its enclosing container does): it reads the previous iteration's value, + // not iter k's. Substitution would shift iterations by one. dominates_all_loads = false; break; } @@ -237,31 +240,29 @@ class EliminateRecomputableAdStackPushes { // Reverse-position correctness for chain leaves pushed in the same block as S. // - // After elimination, every reverse stmt that consumed `load_top(S)` instead consumes the chain - // (cloned by `BackupSSA::generic_visit` at the consumer's reverse position). The cloned chain reads - // `load_top(T)` for each chain leaf T, where the read happens at the consumer's reverse cursor. + // After elimination, every reverse stmt that consumed `load_top(S)` instead consumes the chain (cloned by + // `BackupSSA::generic_visit` at the consumer's reverse position). The cloned chain reads `load_top(T)` for each + // chain leaf T, where the read happens at the consumer's reverse cursor. // - // `MakeAdjoint` visits forward stmts in REVERSE order, emitting at a cursor that advances per visit. - // For a forward stmt at position P, its reverse emissions land "later" in reverse block iff P is - // smaller. T's pop is emitted when `MakeAdjoint` visits T's body push at position P_T. The cloned - // chain at the consumer's reverse position reads `load_top(T)` POST-pop iff T's pop has fired by then, - // i.e. iff visit(P_T) precedes visit(C_pos), i.e. iff P_T > C_pos for every consumer C of S's - // load_tops. + // `MakeAdjoint` visits forward stmts in REVERSE order, emitting at a cursor that advances per visit. For a + // forward stmt at position P, its reverse emissions land "later" in reverse block iff P is smaller. T's pop is + // emitted when `MakeAdjoint` visits T's body push at position P_T. The cloned chain at the consumer's reverse + // position reads `load_top(T)` POST-pop iff T's pop has fired by then, i.e. iff visit(P_T) precedes visit(C_pos), + // i.e. iff P_T > C_pos for every consumer C of S's load_tops. // - // Forward chain evaluates at S's push position P_S, before any of T's same-block body pushes (since - // we already required P_T > P_S via the dominance check above for S's loads). So forward chain reads - // T's pre-iter-k-push value. Reverse clone post-pop also returns T's pre-iter-k-push value (= iter k - // INPUT for loop-carried T). Match. + // Forward chain evaluates at S's push position P_S, before any of T's same-block body pushes (since we already + // required P_T > P_S via the dominance check above for S's loads). So forward chain reads T's pre-iter-k-push + // value. Reverse clone post-pop also returns T's pre-iter-k-push value (= iter k INPUT for loop-carried T). + // Match. // - // Without this check, the canonical Fibonacci-style shape (`p, q = q, p + q; b[q] += a[q]`) where S - // stores `p + q` and chain leaves p_stack / q_stack are pushed AFTER S but BEFORE the GlobalPtr - // index consumer would silently corrupt gradients: the reverse clone at the GlobalPtr's adjoint - // emission reads p_stack and q_stack PRE-pop (iter k POST-push values), summing to a different value - // than the forward chain at P_S used. + // Without this check, the canonical Fibonacci-style shape (`p, q = q, p + q; b[q] += a[q]`) where S stores `p + + // q` and chain leaves p_stack / q_stack are pushed AFTER S but BEFORE the GlobalPtr index consumer would silently + // corrupt gradients: the reverse clone at the GlobalPtr's adjoint emission reads p_stack and q_stack PRE-pop + // (iter k POST-push values), summing to a different value than the forward chain at P_S used. // - // Chain leaves T whose stacks have NO body push in S's containing block are stable within the iter - // (their tops do not change during the iter), so neither forward nor reverse evaluation order shifts - // their values. Excluded from this check. + // Chain leaves T whose stacks have NO body push in S's containing block are stable within the iter (their tops do + // not change during the iter), so neither forward nor reverse evaluation order shifts their values. Excluded from + // this check. auto find_consumers_of_load_tops = [&](std::vector &out_positions) { std::function walk = [&](Block *b, int container_pos_in_push_block) { for (size_t i = 0; i < b->statements.size(); i++) { @@ -324,12 +325,13 @@ class EliminateRecomputableAdStackPushes { bool reverse_safe = true; for (auto *T : chain_leaf_stacks) { - // Find T's body pushes (non-init) in push_block. Only same-block pushes can interfere: pushes in - // outer blocks happen once per outer iteration, so their tops are stable across the inner iter. + // Find T's body pushes (non-init) in push_block. Only same-block pushes can interfere: pushes in outer blocks + // happen once per outer iteration, so their tops are stable across the inner iter. + AdStackPushStmt *T_init = find_init_push(T); for (size_t i = 0; i < push_block->statements.size(); i++) { Stmt *s = push_block->statements[i].get(); if (auto *p = s->cast()) { - if (p->stack == T && !is_zero_init_value(p->v)) { + if (p->stack == T && p != T_init) { if (static_cast(i) <= max_consumer_pos) { reverse_safe = false; break; @@ -371,30 +373,28 @@ class EliminateRecomputableAdStackPushes { continue; } - // Control-flow-cond consumers: by the time the IR reaches this pass, every IfStmt cond and RangeFor begin/end - // is either a bare `AdStackLoadTopStmt` (loop-carried local promoted via `PromoteSSA2LocalVar` -> + // Control-flow-cond consumers: by the time the IR reaches this pass, every IfStmt cond and RangeFor begin/end is + // either a bare `AdStackLoadTopStmt` (loop-carried local promoted via `PromoteSSA2LocalVar` -> // `AdStackAllocaJudger::visit(IfStmt|RangeForStmt)` -> `ReplaceLocalVarWithStacks`) or a stmt outside any - // adstack-bearing alloca (no load_top in the chain at all). `MakeAdjoint::visit(IfStmt)` (auto_diff.cpp around - // line 2452) caps the bare-load_top case with a 1-push-per-execution snap-stack when the if body itself pushes - // to the cond's backing stack so the reverse cond reads the forward-time value, not a stack top mutated by the - // body. + // adstack-bearing alloca (no load_top in the chain at all). `MakeAdjoint::visit(IfStmt)` caps the bare-load_top + // case with a 1-push-per-execution snap-stack when the if body itself pushes to the cond's backing stack so the + // reverse cond reads the forward-time value, not a stack top mutated by the body. // // If we eliminate a stack whose load_top IS the bare cond of an IfStmt / inner RangeFor, the rewrite turns the // cond / loop-bound into an inlined recomputed SSA chain. The snap-stack guard at the `AdStackLoadTopStmt`-cond // check in `MakeAdjoint::visit(IfStmt)` stops firing (the cond is no longer bare), and `BackupSSA`'s clone path - // positions a fresh `AdStackLoadTopStmt` of an enclosing loop-carried stack at the consumer's IR location - - // which sits BEFORE the per-iter pops, returning the post-iteration value instead of the iteration-k value the - // cond was originally computed against. Net effect: silent gradient corruption in any if/loop nested inside a - // dynamic for-loop with a loop-carried alloca. + // positions a fresh `AdStackLoadTopStmt` of an enclosing loop-carried stack at the consumer's IR location - which + // sits BEFORE the per-iter pops, returning the post-iteration value instead of the iteration-k value the cond was + // originally computed against. Net effect: silent gradient corruption in any if/loop nested inside a dynamic + // for-loop with a loop-carried alloca. // // Keep stacks whose load_tops have such consumers. The consumer check below trips only on a load_top of THIS - // stack feeding DIRECTLY into a control-flow-shaping operand of another stmt; that is also the exhaustive set - // of unsafe-elim shapes here, because compound-cond cases (arithmetic over a load_top feeding a cond) reach - // this pass with the cond rewritten into a separate adstack-promoted value via the alloca-promotion pipeline - // above and so look like the bare case from this guard's POV. - // `irpass::analysis::gather_statements` walks via `BasicStmtVisitor`, whose visit overrides for - // IfStmt / RangeForStmt / StructForStmt do not invoke the per-stmt test predicate on the container - // itself - only on stmts inside their bodies. Walk container stmts manually here. + // stack feeding DIRECTLY into a control-flow-shaping operand of another stmt; that is also the exhaustive set of + // unsafe-elim shapes here, because compound-cond cases (arithmetic over a load_top feeding a cond) reach this + // pass with the cond rewritten into a separate adstack-promoted value via the alloca-promotion pipeline above and + // so look like the bare case from this guard's POV. `irpass::analysis::gather_statements` walks via + // `BasicStmtVisitor`, whose visit overrides for IfStmt / RangeForStmt / StructForStmt do not invoke the per-stmt + // test predicate on the container itself - only on stmts inside their bodies. Walk container stmts manually here. bool has_control_flow_consumer = false; std::function walk_block = [&](Block *b) { if (has_control_flow_consumer) @@ -445,24 +445,24 @@ class EliminateRecomputableAdStackPushes { continue; } - // Eligible: rewrite each load_top to use the pushed value directly. The pushed value lives in the - // forward IR and dominates each load_top by SSA construction (the load_top reads what the push wrote; - // both are inside the same loop body in the dynamic-loop case, with the push preceding the load_top). + // Eligible: rewrite each load_top to use the pushed value directly. The pushed value lives in the forward IR and + // dominates each load_top by SSA construction (the load_top reads what the push wrote; both are inside the same + // loop body in the dynamic-loop case, with the push preceding the load_top). for (auto *lt : load_tops) { irpass::replace_all_usages_with(ib, lt, pushed_val); lt->parent->erase(lt); } - // Erase init-zero pushes (they only matter if a load could observe them, but the rewriting above just - // routed every load to the body-pushed SSA value), the body push, and the alloca itself. + // Erase init-zero pushes (they only matter if a load could observe them, but the rewriting above just routed + // every load to the body-pushed SSA value), the body push, and the alloca itself. for (auto *p : init_pushes) { p->parent->erase(p); } body_push->parent->erase(body_push); stack->parent->erase(stack); - // Invalidate the cache: the eliminated stack might have been referenced by other recomputable chains - // we evaluated earlier in this pass, but those evaluations only matter for stacks we eliminate THIS - // pass; re-running from scratch on the next iteration recomputes them. + // Invalidate the cache: the eliminated stack might have been referenced by other recomputable chains we evaluated + // earlier in this pass, but those evaluations only matter for stacks we eliminate THIS pass; re-running from + // scratch on the next iteration recomputes them. recomputable_cache.clear(); modified = true; } diff --git a/quadrants/transforms/auto_diff/forward_state_spill.cpp b/quadrants/transforms/auto_diff/forward_state_spill.cpp index 37ea0f5f40..963a5099db 100644 --- a/quadrants/transforms/auto_diff/forward_state_spill.cpp +++ b/quadrants/transforms/auto_diff/forward_state_spill.cpp @@ -6,12 +6,11 @@ namespace quadrants::lang { namespace { // ============================================================================ -// PromoteSSA2LocalVar: hoist forward-pass SSA defs that the reverse pass will -// re-read into AllocaStmt + LocalStore + LocalLoad. Demand-driven. +// PromoteSSA2LocalVar: hoist forward-pass SSA defs that the reverse pass will re-read into AllocaStmt + LocalStore + +// LocalLoad. Demand-driven. // -// Note that SSA does not mean the instruction will be executed at most once. -// For instructions that may be executed multiple times, we treat them as a -// mutable local variables. +// Note that SSA does not mean the instruction will be executed at most once. For instructions that may be executed +// multiple times, we treat them as a mutable local variables. // ============================================================================ class PromoteSSA2LocalVar : public BasicStmtVisitor { public: @@ -24,19 +23,19 @@ class PromoteSSA2LocalVar : public BasicStmtVisitor { } // Demand-driven `required_defs_` set: the SSA defining stmts that downstream consumers actually require to be - // available at every iteration. A consumer requires its operand iff its adjoint formula reads it - the precise - // set is the operands of any non-linear unary / binary / ternary op (per `NonLinearOps::*_collections`), the - // indices of any `GlobalPtrStmt` / `ExternalPtrStmt` (the reverse pass replays the load), and the `cond` of any - // `IfStmt` / the `begin`/`end` of any `RangeForStmt` (the reverse pass clones the control flow). Stmts outside - // this set are left in pure SSA form: their values stay register-resident inside the forward pass, and - // `MakeAdjoint`'s reverse-pass formulas (which never read those values directly) generate correct adjoint - // accumulations against the adstack-backed defs that ARE in the set. Skipping promotion for the rest of the - // body collapses the alloca + LocalStore + LocalLoad triple-multiplier on unrolled IR that emitted a spill + - // reload pair per non-required arithmetic op with no reverse-pass consumer for the spilled value. + // available at every iteration. A consumer requires its operand iff its adjoint formula reads it - the precise set is + // the operands of any non-linear unary / binary / ternary op (per `NonLinearOps::*_collections`), the indices of any + // `GlobalPtrStmt` / `ExternalPtrStmt` (the reverse pass replays the load), and the `cond` of any `IfStmt` / the + // `begin`/`end` of any `RangeForStmt` (the reverse pass clones the control flow). Stmts outside this set are left in + // pure SSA form: their values stay register-resident inside the forward pass, and `MakeAdjoint`'s reverse-pass + // formulas (which never read those values directly) generate correct adjoint accumulations against the adstack-backed + // defs that ARE in the set. Skipping promotion for the rest of the body collapses the alloca + LocalStore + LocalLoad + // triple-multiplier on unrolled IR that emitted a spill + reload pair per non-required arithmetic op with no + // reverse-pass consumer for the spilled value. // // Operands of `LocalStoreStmt` are not added because `MakeAdjoint::visit(LocalStoreStmt)` reads only - // `adjoint(stmt->dest)`, not the forward `stmt->val`. Linear ops (add / sub / mod / cmp / neg / floor / ceil / - // cast / logic_not / bit-ops) are likewise excluded because their adjoint formulas read only `adjoint(stmt)`. + // `adjoint(stmt->dest)`, not the forward `stmt->val`. Linear ops (add / sub / mod / cmp / neg / floor / ceil / cast / + // logic_not / bit-ops) are likewise excluded because their adjoint formulas read only `adjoint(stmt)`. static void compute_required_defs(Block *block, std::unordered_set &out) { std::function walk = [&](Block *b) { for (auto &owned : b->statements) { @@ -65,13 +64,12 @@ class PromoteSSA2LocalVar : public BasicStmtVisitor { out.insert(idx); } } else if (auto *mp = stmt->cast()) { - // Reverse-mode `MakeAdjoint::visit(MatrixPtrStmt)` reads `stmt->offset` for dynamic-index adjoint - // routing, so a per-iteration-varying offset producer left in pure SSA would be backed by BackupSSA's - // single overwrite-each-iteration alloca and the reverse pass would read the last forward offset for - // every iteration. Promote it to alloca so `AdStackAllocaJudger::visit(MatrixPtrStmt)` can - // adstack-promote loop-varying offsets. Speculative fix: not exhibited by any failing test in the AD - // suite today, but the analysis is sound and the cost of leaving it unfixed is silent gradient - // corruption on `tensor[i + j]`-style local-tensor indexing inside a serial range-for. + // Reverse-mode `MakeAdjoint::visit(MatrixPtrStmt)` reads `stmt->offset` for dynamic-index adjoint routing, so + // a per-iteration-varying offset producer left in pure SSA would be backed by BackupSSA's single + // overwrite-each-iteration alloca and the reverse pass would read the last forward offset for every + // iteration. Promote it to alloca so `AdStackAllocaJudger::visit(MatrixPtrStmt)` can adstack-promote + // loop-varying offsets. The cost of skipping this promotion is silent gradient corruption on `tensor[i + + // j]`-style local-tensor indexing inside a serial range-for. out.insert(mp->offset); } else if (auto *if_s = stmt->cast()) { out.insert(if_s->cond); @@ -108,13 +106,13 @@ class PromoteSSA2LocalVar : public BasicStmtVisitor { return; } - // `AllocaStmt`s always need to be hoisted to the top of the IB regardless of consumer analysis: a user-level - // `var = ...` construct inside a loop body must own a fixed slot at the IB's entry so every iteration shares it - // (cross-iteration accumulators are exactly the shape that drives the hoist). The demand-driven gate only - // applies to value-producing stmts (UnaryOp / BinaryOp / TernaryOp / GlobalLoad / LoopIndex) where the - // alloca + LocalStore + LocalLoad triple is purely a reverse-pass-readable spill - those skip when no consumer - // requires the value. The `LocalStoreStmt`s emitted here are placeholders that `ReplaceLocalVarWithStacks` - // rewrites into `AdStackPushStmt`s downstream. + // `AllocaStmt`s always need to be hoisted to the top of the IB regardless of consumer analysis: a user-level `var = + // ...` construct inside a loop body must own a fixed slot at the IB's entry so every iteration shares it + // (cross-iteration accumulators are exactly the shape that drives the hoist). The demand-driven gate only applies + // to value-producing stmts (UnaryOp / BinaryOp / TernaryOp / GlobalLoad / LoopIndex) where the alloca + LocalStore + // + LocalLoad triple is purely a reverse-pass-readable spill - those skip when no consumer requires the value. The + // `LocalStoreStmt`s emitted here are placeholders that `ReplaceLocalVarWithStacks` rewrites into `AdStackPushStmt`s + // downstream. if (stmt->is()) { auto dtype = stmt->ret_type.ptr_removed(); auto alloc = Stmt::make(dtype); @@ -139,8 +137,8 @@ class PromoteSSA2LocalVar : public BasicStmtVisitor { alloca_block_->insert(std::move(alloc), 0); auto load = stmt->insert_after_me(Stmt::make(alloc_ptr)); immediate_modifier_->replace_usages_with(stmt, load); - // Create the load first so that the operand of the store does not get rewritten to point at the load (the - // SSA value `stmt` is still the right thing to spill; only the downstream consumers see the load). + // Create the load first so that the operand of the store does not get rewritten to point at the load (the SSA value + // `stmt` is still the right thing to spill; only the downstream consumers see the load). stmt->insert_after_me(Stmt::make(alloc_ptr, stmt)); } @@ -163,18 +161,17 @@ class PromoteSSA2LocalVar : public BasicStmtVisitor { bool execute_once_; std::unordered_set required_defs_; // ImmediateIRModifier collapses each `replace_usages_with` from a whole-tree walk (O(N)) to a constant-time - // operand-pointer rewrite: the modifier gathers every (consumer, operand_index) pair feeding any existing stmt - // once at construction (one O(N) pass), and per-replacement just looks up the table for `old_stmt` and rewrites - // each consumer's operand pointer. Without it, `PromoteSSA2LocalVar` ran K (= number of promoted defs) full IR - // walks - O(K*N), the dominant Quadrants-IR-side compile cost on unrolled reverse-mode bodies. + // operand-pointer rewrite: the modifier gathers every (consumer, operand_index) pair feeding any existing stmt once + // at construction (one O(N) pass), and per-replacement just looks up the table for `old_stmt` and rewrites each + // consumer's operand pointer. Without it, `PromoteSSA2LocalVar` would run K (= number of promoted defs) full IR walks + // - O(K*N), which dominates the Quadrants-IR-side compile cost on unrolled reverse-mode bodies. std::unique_ptr immediate_modifier_; }; // ============================================================================ -// AdStackAllocaJudger: per-AllocaStmt analysis. Decides whether an alloca's -// runtime sequence of values must be preserved on an AdStack so the reverse -// pass can re-read every forward-time value, or whether a single overwrite- -// each-iteration backing slot suffices. +// AdStackAllocaJudger: per-AllocaStmt analysis. Decides whether an alloca's runtime sequence of values must be +// preserved on an AdStack so the reverse pass can re-read every forward-time value, or whether a single +// overwrite-each-iteration backing slot suffices. // ============================================================================ class AdStackAllocaJudger : public BasicStmtVisitor { public: @@ -188,35 +185,33 @@ class AdStackAllocaJudger : public BasicStmtVisitor { } // Track whether the alloca has any store at all so a load-only alloca (no adstack-relevant data flow either - // direction) can short-circuit `run()` regardless of what the per-op visitors below find. The decision of - // whether the alloca needs adstack promotion is made entirely by the precise visitors: non-linear unary / - // binary / ternary, GlobalPtr / ExternalPtr index, and IfStmt / RangeForStmt bound. Plus an alloca that is - // both loaded and stored anywhere in the IB is treated as loop-carried, which is needed for kernels like - // `for j: p, q = q, p + q` where the reverse pass routes the gradient through the cross-iteration - // recurrence and BackupSSA's single overwrite-each-iteration alloca cannot back the read-after-write across - // iterations. The visit-order-dependent load+store evidence here is conservative: any alloca with both a - // load and a store inside the IB triggers it, including pure accumulators whose adjoint formulas don't - // actually need per-iteration values - the slight over-promotion cost is the price of correctness on - // Fibonacci-style recurrences (silent gradient corruption otherwise). + // direction) can short-circuit `run()` regardless of what the per-op visitors below find. The decision of whether the + // alloca needs adstack promotion is made entirely by the precise visitors: non-linear unary / binary / ternary, + // GlobalPtr / ExternalPtr index, and IfStmt / RangeForStmt bound. Plus an alloca that is both loaded and stored + // anywhere in the IB is treated as loop-carried, which is needed for kernels like `for j: p, q = q, p + q` where the + // reverse pass routes the gradient through the cross-iteration recurrence and BackupSSA's single + // overwrite-each-iteration alloca cannot back the read-after-write across iterations. The visit-order-dependent + // load+store evidence here is conservative: any alloca with both a load and a store inside the IB triggers it, + // including pure accumulators whose adjoint formulas don't actually need per-iteration values - the slight + // over-promotion cost is the price of correctness on Fibonacci-style recurrences (silent gradient corruption + // otherwise). void visit(LocalStoreStmt *stmt) override { if (stmt->dest == target_alloca_backup_) { load_only_ = false; - // Gate the load+store-implies-stack-needed rule on actually being inside a dynamic RangeForStmt at the - // point this evidence accumulates. The rule's purpose is to preserve cross-iteration RAW dependencies - // (`for j: p, q = q, p + q` Fibonacci-style) that BackupSSA's single overwrite-each-iteration alloca - // cannot back. With no enclosing dynamic for-loop the IB body executes once: there is no cross- - // iteration RAW to preserve, and the "load+store" pattern is just an in-block accumulator that the - // reverse pass handles via plain SSA cloning. Promoting such allocas under a static-unrolled loop body - // wastes one AdStack per accumulator (one push per unrolled-iter store + one load_top per unrolled- - // iter load) without any reverse-pass consumer needing per-iter replay - that is the unrolled-overhead - // bug Plan B targets. + // Gate the load+store-implies-stack-needed rule on actually being inside a dynamic RangeForStmt at the point this + // evidence accumulates. The rule's purpose is to preserve cross-iteration RAW dependencies (`for j: p, q = q, p + + // q` Fibonacci-style) that BackupSSA's single overwrite-each-iteration alloca cannot back. With no enclosing + // dynamic for-loop the IB body executes once: there is no cross-iteration RAW to preserve, and the "load+store" + // pattern is just an in-block accumulator that the reverse pass handles via plain SSA cloning. Promoting such + // allocas under a static-unrolled loop body wastes one AdStack per accumulator (one push per unrolled-iter store + // + one load_top per unrolled-iter load) without any reverse-pass consumer needing per-iter replay. // - // `dynamic_for_depth_` is incremented in `visit(RangeForStmt)` and decremented on exit. The judger - // walks the IB tree from the alloca's enclosing block, so depth here reflects exactly the nesting of - // *dynamic* for-loops between the alloca and the current load/store. StructFor / WhileStmt do not - // increment because their bodies still execute per-iter and need the same RAW protection (StructFor - // is the kernel-level offload-loop in some cases, but its body is a per-thread independent block; - // load+store there is the same shape as for a top-level alloca). + // `dynamic_for_depth_` is incremented in `visit(RangeForStmt)` and decremented on exit. The judger walks the IB + // tree from the alloca's enclosing block, so depth here reflects exactly the nesting of *dynamic* for-loops + // between the alloca and the current load/store. StructFor / WhileStmt do not increment because their bodies + // still execute per-iter and need the same RAW protection (StructFor is the kernel-level offload-loop in some + // cases, but its body is a per-thread independent block; load+store there is the same shape as for a top-level + // alloca). if (local_loaded_ && dynamic_for_depth_ > 0) { is_stack_needed_ = true; } @@ -229,12 +224,12 @@ class AdStackAllocaJudger : public BasicStmtVisitor { load_only_ = false; } - // The stack is needed if the alloca serves as the index of any global variables. Same cursor-vs-backup - // pattern as visit(IfStmt)/visit(RangeForStmt) below: `index` is always a value-producing stmt (typically a - // `LocalLoadStmt` reading the alloca, or a `ConstStmt`), never the alloca itself. The raw `index == - // target_alloca_` comparison only matches the first load's instance the `visit(LocalLoadStmt)` cursor - // advanced to - any subsequent load of the same alloca used as a different GlobalPtr index slips through. - // Resolve the LocalLoad chain and compare `ll->src` against `target_alloca_backup_` to catch every load. + // The stack is needed if the alloca serves as the index of any global variables. Same cursor-vs-backup pattern as + // visit(IfStmt) / visit(RangeForStmt) below: `index` is always a value-producing stmt (typically a `LocalLoadStmt` + // reading the alloca, or a `ConstStmt`), never the alloca itself. The raw `index == target_alloca_` comparison only + // matches the first load's instance the `visit(LocalLoadStmt)` cursor advanced to - any subsequent load of the same + // alloca used as a different GlobalPtr index slips through. Resolve the LocalLoad chain and compare `ll->src` against + // `target_alloca_backup_` to catch every load. void visit(GlobalPtrStmt *stmt) override { if (is_stack_needed_) return; @@ -255,10 +250,10 @@ class AdStackAllocaJudger : public BasicStmtVisitor { } } - // Reverse-mode `MakeAdjoint::visit(MatrixPtrStmt)` reads `stmt->offset` for dynamic-index adjoint routing, so - // a runtime-varying offset whose value comes from this alloca needs adstack promotion - otherwise BackupSSA - // backs the offset with a single overwrite-each-iteration slot and the reverse pass routes every iteration's - // adjoint into the last forward offset's slot. Same cursor-vs-backup pattern as the index visitors above. + // Reverse-mode `MakeAdjoint::visit(MatrixPtrStmt)` reads `stmt->offset` for dynamic-index adjoint routing, so a + // runtime-varying offset whose value comes from this alloca needs adstack promotion - otherwise BackupSSA backs the + // offset with a single overwrite-each-iteration slot and the reverse pass routes every iteration's adjoint into the + // last forward offset's slot. Same cursor-vs-backup pattern as the index visitors above. void visit(MatrixPtrStmt *stmt) override { if (is_stack_needed_) return; @@ -268,8 +263,8 @@ class AdStackAllocaJudger : public BasicStmtVisitor { } // Check whether the target alloca is fed into a non-linear unary op. Same cursor-vs-backup pattern as - // visit(GlobalPtrStmt) above: `stmt->operand` is a value-producing stmt (typically LocalLoad), never the - // alloca itself, so resolve the LocalLoad chain and compare against the backup. + // visit(GlobalPtrStmt) above: `stmt->operand` is a value-producing stmt (typically LocalLoad), never the alloca + // itself, so resolve the LocalLoad chain and compare against the backup. void visit(UnaryOpStmt *stmt) override { if (is_stack_needed_) return; @@ -306,15 +301,13 @@ class AdStackAllocaJudger : public BasicStmtVisitor { } } - // Check whether the target alloca feeds the condition of an if stmt. `stmt->cond` is always a - // value-producing stmt - typically a direct `LocalLoadStmt` reading the alloca, but also commonly a - // `BinaryOpStmt` wrapping such a load (e.g. `j < i+1`). Walk the expression chain to catch every load of - // the target alloca: the raw `stmt->cond == target_alloca_` comparison the old code used only matched the - // first-visited load's instance, and a direct `cast` still misses the BinaryOp case that - // `visit(BinaryOpStmt)` cannot catch (comparison ops are linear and so not in `NonLinearOps`). Covers the - // shape defensively: IR simplification currently collapses most BinaryOp-wrapped conds before the judger - // sees them, so no failing regression test pins it today, but the fix is structurally correct for future - // IR changes that preserve the BinaryOp wrapping. + // Check whether the target alloca feeds the condition of an if stmt. `stmt->cond` is always a value-producing stmt - + // typically a direct `LocalLoadStmt` reading the alloca, but also commonly a `BinaryOpStmt` wrapping such a load + // (e.g. `j < i+1`). Walk the expression chain via `feeds_target_alloca` to catch every load of the target alloca, + // including loads nested under comparison or other linear ops (which `visit(BinaryOpStmt)` does not flag because + // comparison ops are not in `NonLinearOps`). The walk is defensive: IR simplification currently collapses most + // BinaryOp-wrapped conds before the judger sees them, so most input IR uses the bare-LocalLoad shape; the walker + // stays sound when the wrapping is preserved. void visit(IfStmt *stmt) override { if (is_stack_needed_) return; @@ -330,21 +323,19 @@ class AdStackAllocaJudger : public BasicStmtVisitor { stmt->false_statements->accept(this); } - // Check whether the target alloca feeds the begin or end of a range-for bound. Under reverse-mode AD, if an - // inner for-loop's bound is an enclosing loop-carried counter (the canonical triangular-nested - // `for k in range(j)` shape, or the `range(j+1)` / `range(n-i)` shapes where the bound is a linear arithmetic - // expression of a loop-carried alloca), its reverse clone must read the bound from the per-iteration forward - // value; without an adstack the reverse pass sees only the last forward value and the inner loop over- or - // under-runs, silently corrupting gradients for the earliest inner indices (those visited most often across - // outer iterations). This check is the only thing that promotes such a loop-counter alloca - - // `visit(LocalStoreStmt)`'s `local_loaded_` short-circuit does not fire because the counter is only LOAD-ed - // inside the inner-loop bound, not LOAD-then-STORE-ed. Walk the expression chain through - // `feeds_target_alloca` so both direct LocalLoads (`range(j)`) and LocalLoads nested under linear ops - // (`range(j+1)`, `range(n-i)`, ...) trigger promotion. The BinaryOp-wrapped case is defensively covered: IR - // simplification currently collapses most such bounds before the judger sees them, so no failing regression - // test pins it today, but the walker is structurally correct for future IR changes that preserve the - // wrapping. The raw-cast direct `LocalLoadStmt` variant pinned by `test_adstack_inner_for_bound_is_enclosing - // _loop_index` remains covered - that shape takes the first branch of the walker trivially. + // Check whether the target alloca feeds the begin or end of a range-for bound. Under reverse-mode AD, if an inner + // for-loop's bound is an enclosing loop-carried counter (the canonical triangular-nested `for k in range(j)` shape, + // or the `range(j+1)` / `range(n-i)` shapes where the bound is a linear arithmetic expression of a loop-carried + // alloca), its reverse clone must read the bound from the per-iteration forward value; without an adstack the reverse + // pass sees only the last forward value and the inner loop over- or under-runs, silently corrupting gradients for the + // earliest inner indices (those visited most often across outer iterations). This check is the only thing that + // promotes such a loop-counter alloca - `visit(LocalStoreStmt)`'s `local_loaded_` short-circuit does not fire because + // the counter is only LOAD-ed inside the inner-loop bound, not LOAD-then-STORE-ed. Walk the expression chain through + // `feeds_target_alloca` so both direct LocalLoads (`range(j)`) and LocalLoads nested under linear ops (`range(j+1)`, + // `range(n-i)`, ...) trigger promotion. The BinaryOp-wrapped case is covered defensively: IR simplification currently + // collapses most such bounds before the judger sees them, so most input IR uses the bare-LocalLoad shape; the walker + // stays sound when the wrapping is preserved. The bare-LocalLoad shape takes the first branch of the walker + // trivially. void visit(RangeForStmt *stmt) override { if (is_stack_needed_) return; @@ -370,11 +361,11 @@ class AdStackAllocaJudger : public BasicStmtVisitor { private: // Recursively walk a value expression to decide whether it transitively reads `target_alloca_backup_` via a - // `LocalLoadStmt`. Used by `visit(IfStmt)` and `visit(RangeForStmt)` to detect the target alloca feeding a - // bound or condition even when wrapped by linear ops (e.g. `range(j+1)`, `j < i+1`). Linear binary/unary - // ops are traversed because `visit(BinaryOpStmt)`/`visit(UnaryOpStmt)` only flag *non-linear* ops - their - // linear-op path does not otherwise promote the alloca. `ConstStmt`s and unrelated values return false and - // terminate the recursion; the walker is always finite because SSA IR guarantees acyclic operand graphs. + // `LocalLoadStmt`. Used by `visit(IfStmt)` and `visit(RangeForStmt)` to detect the target alloca feeding a bound or + // condition even when wrapped by linear ops (e.g. `range(j+1)`, `j < i+1`). Linear binary/unary ops are traversed + // because `visit(BinaryOpStmt)`/`visit(UnaryOpStmt)` only flag *non-linear* ops - their linear-op path does not + // otherwise promote the alloca. `ConstStmt`s and unrelated values return false and terminate the recursion; the + // walker is always finite because SSA IR guarantees acyclic operand graphs. bool feeds_target_alloca(Stmt *expr) const { if (auto *ll = expr->cast()) { return ll->src == target_alloca_backup_; @@ -396,17 +387,16 @@ class AdStackAllocaJudger : public BasicStmtVisitor { bool is_stack_needed_ = false; bool local_loaded_ = false; bool load_only_ = true; - // Nesting depth of dynamic `RangeForStmt` containers between the alloca's enclosing block and the current - // visit cursor. Static-unrolled `qd.static(range(...))` loops are removed by the AST transformer before the - // judger sees the IR, so they do not contribute to depth. The load+store-implies-stack-needed rule fires - // only when this depth is positive; see the rationale in `visit(LocalStoreStmt)`. + // Nesting depth of dynamic `RangeForStmt` containers between the alloca's enclosing block and the current visit + // cursor. Static-unrolled `qd.static(range(...))` loops are removed by the AST transformer before the judger sees the + // IR, so they do not contribute to depth. The load+store-implies-stack-needed rule fires only when this depth is + // positive; see the rationale in `visit(LocalStoreStmt)`. int dynamic_for_depth_ = 0; }; // ============================================================================ -// ReplaceLocalVarWithStacks: rewrite each AllocaStmt that AdStackAllocaJudger -// approved into AdStackAllocaStmt + AdStackPush + AdStackLoadTop. Tensor-typed -// allocas need extra care because per-slot stack writes are not a reaching +// ReplaceLocalVarWithStacks: rewrite each AllocaStmt that AdStackAllocaJudger approved into AdStackAllocaStmt + +// AdStackPush + AdStackLoadTop. Tensor-typed allocas need extra care because per-slot stack writes are not a reaching // def for the store-to-load forwarding walker. // ============================================================================ class ReplaceLocalVarWithStacks : public BasicStmtVisitor { @@ -427,8 +417,8 @@ class ReplaceLocalVarWithStacks : public BasicStmtVisitor { alloc->replace_with(VecStatement(std::move(stack_alloca))); - // Note that unlike AllocaStmt, AdStackAllocaStmt does NOT have an 0 as - // initial value. Therefore here we push an initial 0 value. + // Note that unlike AllocaStmt, AdStackAllocaStmt does NOT have a 0 as initial value, so we push an initial 0 + // here. auto zero = insert_const(dtype, stack_alloca_ptr, 0); zero->insert_after_me(Stmt::make(stack_alloca_ptr, zero)); } @@ -444,12 +434,12 @@ class ReplaceLocalVarWithStacks : public BasicStmtVisitor { } // Slot load from a stack-backed tensor. After `visit(MatrixPtrStmt)`, `stmt->src` is of the form - // `MatrixPtrStmt(AdStackLoadTopStmt(stack, return_ptr=true), offset)`. A direct load through that pointer - // leaves the store-to-load forwarding walker in `ir/control_flow_graph.cpp` with no reaching definition, - // because the only producer for the stack's top slots is an `AdStackPushStmt` (tagged `ir_traits::Load`, - // invisible to `get_store_destination`). Replace the load with a full-tensor `AdStackLoadTopStmt` - // materialized into a fresh regular `AllocaStmt`, then re-subscript it - a plain alloca + LocalStore - // sequence is a shape the reach-in walker can trace end-to-end. + // `MatrixPtrStmt(AdStackLoadTopStmt(stack, return_ptr=true), offset)`. A direct load through that pointer leaves + // the store-to-load forwarding walker in `ir/control_flow_graph.cpp` with no reaching definition, because the only + // producer for the stack's top slots is an `AdStackPushStmt` (tagged `ir_traits::Load`, invisible to + // `get_store_destination`). Replace the load with a full-tensor `AdStackLoadTopStmt` materialized into a fresh + // regular `AllocaStmt`, then re-subscript it - a plain alloca + LocalStore sequence is a shape the reach-in walker + // can trace end-to-end. if (stmt->src->is()) { auto matrix_ptr = stmt->src->as(); if (matrix_ptr->origin->is() && matrix_ptr->origin->as()->return_ptr) { diff --git a/quadrants/transforms/auto_diff/forward_state_spill.h b/quadrants/transforms/auto_diff/forward_state_spill.h index 8fd696b8d0..edb1387f81 100644 --- a/quadrants/transforms/auto_diff/forward_state_spill.h +++ b/quadrants/transforms/auto_diff/forward_state_spill.h @@ -4,27 +4,23 @@ namespace quadrants::lang { class Block; -// Stage: Forward-state spill. Pre-MakeAdjoint passes that prepare the forward -// IR so the reverse pass can re-read the per-iteration values it needs. +// Stage: Forward-state spill. Pre-MakeAdjoint passes that prepare the forward IR so the reverse pass can re-read the +// per-iteration values it needs. -// Hoist forward-pass SSA defs that the reverse pass will re-read (operands of -// non-linear ops, indices of GlobalPtr / ExternalPtr, range-for bounds, if -// conds) into AllocaStmt + LocalStore + LocalLoad. Demand-driven: defs no -// reverse formula reads stay in pure SSA form. +// Hoist forward-pass SSA defs that the reverse pass will re-read (operands of non-linear ops, indices of GlobalPtr / +// ExternalPtr, range-for bounds, if conds) into AllocaStmt + LocalStore + LocalLoad. Demand-driven: defs no reverse +// formula reads stay in pure SSA form. void promote_ssa_to_local_var(Block *block); -// Replace each AllocaStmt the reverse pass needs to read across iterations -// with an AdStackAllocaStmt + AdStackPushStmt / AdStackLoadTopStmt. The -// AllocaJudger (file-private) decides which allocas need the promotion. +// Replace each AllocaStmt the reverse pass needs to read across iterations with an AdStackAllocaStmt + AdStackPushStmt +// / AdStackLoadTopStmt. The AllocaJudger (file-private) decides which allocas need the promotion. void replace_local_var_with_stacks(Block *block, int ad_stack_size); -// Drop AdStackAllocaStmts whose pushed value is recomputable from already- -// stack-backed allocas + kernel args + constants + loop indices. Trades -// forward-pass adstack memory traffic for cloned arithmetic in the reverse -// scope, which `BackupSSA::generic_visit` synthesizes on demand via the -// shared RecomputableChainCloner. Must run after `replace_local_var_with_stacks` -// (so the analyzer sees the AdStack shape) and before `make_adjoint` (so the -// reverse pass is generated against the cleaned forward IR). +// Drop AdStackAllocaStmts whose pushed value is recomputable from already-stack-backed allocas + kernel args + +// constants + loop indices. Trades forward-pass adstack memory traffic for cloned arithmetic in the reverse scope, +// which `BackupSSA::generic_visit` synthesizes on demand via the shared RecomputableChainCloner. Must run after +// `replace_local_var_with_stacks` (so the analyzer sees the AdStack shape) and before `make_adjoint` (so the reverse +// pass is generated against the cleaned forward IR). void eliminate_recomputable_ad_stack_pushes(Block *ib); } // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/ir_shaping.cpp b/quadrants/transforms/auto_diff/ir_shaping.cpp index 1b0a746be7..7b58eb1569 100644 --- a/quadrants/transforms/auto_diff/ir_shaping.cpp +++ b/quadrants/transforms/auto_diff/ir_shaping.cpp @@ -6,9 +6,8 @@ namespace quadrants::lang { namespace { // ============================================================================ -// RegulateTensorTypedStatements: rewrite tensor-typed local/global stores that -// touch a sub-tensor through MatrixPtr into an explicit gather + matrix-init -// + scalar store, so downstream passes never see a partial-tensor store. +// RegulateTensorTypedStatements: rewrite tensor-typed local/global stores that touch a sub-tensor through MatrixPtr +// into an explicit gather + matrix-init + scalar store, so downstream passes never see a partial-tensor store. // ============================================================================ class RegulateTensorTypedStatements : public BasicStmtVisitor { @@ -180,10 +179,9 @@ class RegulateTensorTypedStatements : public BasicStmtVisitor { // ============================================================================ // Independent-Blocks discovery: IBJudger / DupCleaner / IdentifyIBs. // -// Independent Block (IB): a block (i.e. loop body) whose iterations are -// independent of previous iterations and outer scopes. IBs are where -// MakeAdjoint emits the reverse pass; outside an IB only iteration order -// matters and ReverseOuterLoops handles that. +// Independent Block (IB): a block (i.e. loop body) whose iterations are independent of previous iterations and outer +// scopes. IBs are where MakeAdjoint emits the reverse pass; outside an IB only iteration order matters and +// ReverseOuterLoops handles that. // ============================================================================ class IndependentBlocksJudger : public BasicStmtVisitor { @@ -202,12 +200,9 @@ class IndependentBlocksJudger : public BasicStmtVisitor { } void visit(AtomicOpStmt *stmt) override { - // We don't need to check the global atomics inside the range for-loops - // because - // 1. If the range for-loop is innermost, they will be captured by - // MakeAdjoint anyway - // 2. If the range for-loop is not innermost, they will be processed by - // another IndependentBlocksJudger + // We don't need to check the global atomics inside the range for-loops because: + // 1. If the range for-loop is innermost, they are captured by MakeAdjoint anyway. + // 2. If the range for-loop is not innermost, they are processed by another IndependentBlocksJudger. if (is_inside_loop_) return; @@ -234,12 +229,9 @@ class IndependentBlocksJudger : public BasicStmtVisitor { } void visit(GlobalLoadStmt *stmt) override { - // We don't need to check the global load inside the range for-loops - // because - // 1. If the range for-loop is innermost, they will be captured by - // MakeAdjoint anyway - // 2. If the range for-loop is not innermost, they will be processed by - // another IndependentBlocksJudger + // We don't need to check the global load inside the range for-loops because: + // 1. If the range for-loop is innermost, they are captured by MakeAdjoint anyway. + // 2. If the range for-loop is not innermost, they are processed by another IndependentBlocksJudger. if (is_inside_loop_) return; @@ -271,27 +263,23 @@ class IndependentBlocksJudger : public BasicStmtVisitor { Block *block = root->as(); root->accept(&Judger); std::set outside_blocks; - // Collect all parent blocks (i.e. outside blocks) of the current block for - // local load/store stmt checks + // Collect all parent blocks (i.e. outside blocks) of the current block for local load/store stmt checks. for (auto b = block->parent_block(); b; b = b->parent_block()) { if (b) outside_blocks.insert(b); } for (const auto &alloca : Judger.touched_allocas_) { - // Test if the alloca belongs to the current block + // Test if the alloca belongs to the current block. if (outside_blocks.find(alloca->parent) != outside_blocks.end()) { - // This block is not an IB since it loads/modifies outside variables + // This block is not an IB since it loads/modifies outside variables. ib_meta_data.is_ib = false; } } - // To judge whether a block is an IB - // - No local load/store to allocas *outside* itself has been strictly - // enforced - - // To judge whether a block is a smallest IB - // - If the #1 is satisfied, either an inner most loop or a block without - // global atomics / global load is an IB + // IB classification rules: + // - To be an IB, the block must have no local load/store to allocas outside itself (enforced above). + // - To be a smallest IB on top of that, it must be an inner-most loop or a block without qualified global + // atomics / global loads. ib_meta_data.is_smallest_ib = ib_meta_data.is_ib && (Judger.qualified_glb_operations_ || Judger.inner_most_loop_); } @@ -302,8 +290,7 @@ class IndependentBlocksJudger : public BasicStmtVisitor { bool is_inside_loop_ = false; }; -// Remove the duplicated IBs, remove blocks who are others' children because -// each block should only be processed once +// Remove duplicated IBs and remove blocks that are others' children, so each block is processed at most once. class DuplicateIndependentBlocksCleaner : public BasicStmtVisitor { public: using BasicStmtVisitor::visit; @@ -332,8 +319,7 @@ class DuplicateIndependentBlocksCleaner : public BasicStmtVisitor { } // No clean is needed if only one IB exists if (cleaner.independent_blocks_cleaned_.size() > 1) { - // Check from the block with smallest depth, ensure no duplicate visit - // happens + // Check from the block with smallest depth, ensure no duplicate visit happens. for (const auto &block : cleaner.independent_blocks_cleaned_) { block->accept(&cleaner); } @@ -363,12 +349,10 @@ class IdentifyIndependentBlocks : public BasicStmtVisitor { void visit_loop_body(Block *block) { auto ib_meta_data = IndependentBlockMetaData(); - // An IB has no local load/store to allocas *outside* itself - // Note: - // - Local atomics should have been demoted before this pass. - // - It is OK for an IB to have more than two for loops. - // - No global load/atomics operations to the global variables which - // require gradient + // An IB has no local load/store to allocas *outside* itself. Note: + // - Local atomics must have been demoted before this pass. + // - It is OK for an IB to have more than two for loops. + // - No global load/atomics operations to global variables that require gradient. if (block->statements.empty()) { // A empty block shoud be a smallest IB ib_meta_data.is_ib = true; @@ -441,9 +425,8 @@ class IdentifyIndependentBlocks : public BasicStmtVisitor { }; // ============================================================================ -// ReverseOuterLoops: flip iteration direction on outer (non-IB) for-loops and -// reorder sibling for-loops in non-IB container blocks so the reverse pass -// walks the iteration trace backward. +// ReverseOuterLoops: flip iteration direction on outer (non-IB) for-loops and reorder sibling for-loops in non-IB +// container blocks so the reverse pass walks the iteration trace backward. // ============================================================================ class ReverseOuterLoops : public BasicStmtVisitor { @@ -457,26 +440,24 @@ class ReverseOuterLoops : public BasicStmtVisitor { return std::find(ib_.begin(), ib_.end(), block) != ib_.end(); } - // Sibling for-loops inside a non-IB container block execute their reverse-mode companions - // in the container's forward order by default, because MakeAdjoint only touches IB-level bodies - // and nothing else permutes the enclosing order. Reverse-mode AD requires the opposite: if the - // forward body runs `for_A; for_B` and for_B's reverse depends on reads produced by for_A's - // forward run, the reverse pass must execute `rev-for_B; rev-for_A` so for_A's reverse sees the - // adjoints for_B has populated (e.g. `cdof[i]=x[i]; cdofvel[i]=cdof[i]*vel[i]` silently returns - // x.grad=0 otherwise: rev-for_A clears cdof.grad before rev-for_B has populated it). + // Sibling for-loops inside a non-IB container block execute their reverse-mode companions in the container's forward + // order by default, because MakeAdjoint only touches IB-level bodies and nothing else permutes the enclosing order. + // Reverse-mode AD requires the opposite: if the forward body runs `for_A; for_B` and for_B's reverse depends on reads + // produced by for_A's forward run, the reverse pass must execute `rev-for_B; rev-for_A` so for_A's reverse sees the + // adjoints for_B has populated (e.g. `a[i]=x[i]; b[i]=a[i]*y[i]` silently returns x.grad=0 otherwise: rev-for_A + // clears a.grad before rev-for_B has populated it). // - // Naive pairwise swap of for-loop positions is unsafe whenever a non-loop stmt between two - // for-loops feeds the later sibling's SSA operand chain (e.g. a GlobalLoad that supplies a - // dynamic trip count): after the swap, the consumer for-loop ends up before its producer and - // the IR verifier rejects the block. Before swapping, hoist any such producer (and its - // transitive in-block dependencies) to the slot just before the first sibling for-loop. Non-loop - // stmts unrelated to for-loop operands stay at their original indices; memory ordering between - // non-loop stmts is preserved because the hoist keeps them in their original relative order and - // only moves them upward over for-loops (which produce no SSA value and cannot be the source of - // a missed memory read for a non-loop that gets hoisted above them). + // Naive pairwise swap of for-loop positions is unsafe whenever a non-loop stmt between two for-loops feeds the later + // sibling's SSA operand chain (e.g. a GlobalLoad that supplies a dynamic trip count): after the swap, the consumer + // for-loop ends up before its producer and the IR verifier rejects the block. Before swapping, hoist any such + // producer (and its transitive in-block dependencies) to the slot just before the first sibling for-loop. Non-loop + // stmts unrelated to for-loop operands stay at their original indices; memory ordering between non-loop stmts is + // preserved because the hoist keeps them in their original relative order and only moves them upward over for-loops + // (which produce no SSA value and cannot be the source of a missed memory read for a non-loop that gets hoisted above + // them). // - // The top-level kernel block is handled by `reverse_segments` before this pass, so we only - // reorder inside nested non-IB blocks here. + // The top-level kernel block is handled by `reverse_segments` before this pass, so we only reorder inside nested + // non-IB blocks here. static void reverse_for_loop_order_in_place(Block *block) { const int n = (int)block->statements.size(); std::vector for_indices; @@ -497,9 +478,9 @@ class ReverseOuterLoops : public BasicStmtVisitor { pos_of[block->statements[i].get()] = i; } - // Walk the SSA operand graph of every for-loop (restricted to this block). Any in-block stmt - // that (a) the operand closure reaches and (b) sits at or after `first_for` gets flagged for - // hoisting: after swap, that stmt must precede every for-loop, not just the ones it feeds. + // Walk the SSA operand graph of every for-loop (restricted to this block). Any in-block stmt that (a) the operand + // closure reaches and (b) sits at or after `first_for` gets flagged for hoisting: after swap, that stmt must + // precede every for-loop, not just the ones it feeds. std::unordered_set must_hoist; std::vector stack; auto push_if_internal = [&](Stmt *s) { @@ -514,13 +495,12 @@ class ReverseOuterLoops : public BasicStmtVisitor { stack.push_back(s); } }; - // Seed the hoist frontier from both the for-loop's direct SSA operands (`begin`, `end`) and - // from every stmt nested inside the for-loop's body that references an outer-block stmt as a - // free variable. The body-use gather is what catches the case where the later sibling - // for-loop consumes a non-loop outer-block stmt `S` inside its body (e.g. `for_B: body reads - // S`) rather than through `for_B`'s range bound: `RangeForStmt::get_operands()` returns only - // `{begin, end}`, so without walking the body `S` would miss `must_hoist`, the pairwise swap - // would place `for_B` ahead of `S`, and the IR verifier would reject the SSA violation. + // Seed the hoist frontier from both the for-loop's direct SSA operands (`begin`, `end`) and from every stmt nested + // inside the for-loop's body that references an outer-block stmt as a free variable. The body-use gather catches + // the case where the later sibling for-loop consumes a non-loop outer-block stmt `S` inside its body (e.g. `for_B: + // body reads S`) rather than through `for_B`'s range bound: `RangeForStmt::get_operands()` returns only `{begin, + // end}`, so without walking the body `S` would miss `must_hoist`, the pairwise swap would place `for_B` ahead of + // `S`, and the IR verifier would reject the SSA violation. for (int fi : for_indices) { for (Stmt *op : block->statements[fi]->get_operands()) { push_if_internal(op); @@ -540,9 +520,9 @@ class ReverseOuterLoops : public BasicStmtVisitor { push_if_internal(op); } } - // For-loops themselves end up in `must_hoist` only because their own operand-closure reached - // them; they do not get hoisted as non-loop producers - strip them here to keep `must_hoist` - // to "non-loop stmts that need to move above all for-loops". + // For-loops themselves end up in `must_hoist` only because their own operand-closure reached them; they do not get + // hoisted as non-loop producers - strip them here to keep `must_hoist` to "non-loop stmts that need to move above + // all for-loops". for (int fi : for_indices) { must_hoist.erase(block->statements[fi].get()); } @@ -559,8 +539,8 @@ class ReverseOuterLoops : public BasicStmtVisitor { new_stmts.push_back(std::move(block->statements[i])); } } - // Remainder (for-loops and non-hoisted non-loops) in original order, with for-loops swapped - // pairwise inside this suffix. + // Remainder (for-loops and non-hoisted non-loops) in original order, with for-loops swapped pairwise inside this + // suffix. std::vector> suffix; std::vector suffix_for_positions; for (int i = first_for; i < n; ++i) { @@ -609,15 +589,15 @@ class ReverseOuterLoops : public BasicStmtVisitor { loop_depth_ -= 1; } - // Deliberately no `visit(IfStmt *)` override, although sibling for-loops can live directly inside an if-branch - // block (`true_statements` / `false_statements`) the same way they live inside a for-body. The default + // Deliberately no `visit(IfStmt *)` override, although sibling for-loops can live directly inside an if-branch block + // (`true_statements` / `false_statements`) the same way they live inside a for-body. The default // `BasicStmtVisitor::visit(IfStmt *)` recurses into both branches so inner `RangeForStmt::body`s still get the - // sibling-reorder treatment via the range-for visitor above, but `reverse_for_loop_order_in_place` is never - // invoked on the branch block itself. That is intentional: `MakeAdjoint::visit(IfStmt *)` below emits the adjoint - // if-stmt by iterating each branch's statements in reverse order (`for (int i = size - 1; i >= 0; --i)` in its - // `true_statements` / `false_statements` loops), which achieves the same sibling-for reordering effect that the - // missing override here would provide. Overriding `visit(IfStmt)` in this class is therefore a no-op on the - // generated adjoint code. Keep the comment rather than the override so the visitor-coverage gap is documented. + // sibling-reorder treatment via the range-for visitor above, but `reverse_for_loop_order_in_place` is never invoked + // on the branch block itself. That is intentional: `MakeAdjoint::visit(IfStmt *)` below emits the adjoint if-stmt by + // iterating each branch's statements in reverse order (`for (int i = size - 1; i >= 0; --i)` in its `true_statements` + // / `false_statements` loops), which achieves the same sibling-for reordering effect that the missing override here + // would provide. Overriding `visit(IfStmt)` in this class is therefore a no-op on the generated adjoint code. Keep + // the comment rather than the override so the visitor-coverage gap is documented. int loop_depth_; std::set ib_; diff --git a/quadrants/transforms/auto_diff/ir_shaping.h b/quadrants/transforms/auto_diff/ir_shaping.h index 9a54aac2a6..855354fb81 100644 --- a/quadrants/transforms/auto_diff/ir_shaping.h +++ b/quadrants/transforms/auto_diff/ir_shaping.h @@ -7,22 +7,19 @@ namespace quadrants::lang { class IRNode; class Block; -// Stage: IR shaping. Pre-MakeAdjoint passes that mutate or analyze the -// forward IR so the downstream reverse-mode codegen sees the expected shape. +// Stage: IR shaping. Pre-MakeAdjoint passes that mutate or analyze the forward IR so the downstream reverse-mode +// codegen sees the expected shape. -// Rewrite tensor-typed LocalStore / GlobalStore through MatrixPtr into an -// explicit gather + matrix-init + scalar store - downstream passes assume -// stores never touch a sub-tensor in place. +// Rewrite tensor-typed LocalStore / GlobalStore through MatrixPtr into an explicit gather + matrix-init + scalar store +// - downstream passes assume stores never touch a sub-tensor in place. void regulate_tensor_typed_statements(IRNode *root); -// Discover every Independent Block (loop body whose iterations carry no -// dependency on previous iterations or outer scopes) under `root`. Returns -// the set of blocks where MakeAdjoint will emit the reverse pass. +// Discover every Independent Block (loop body whose iterations carry no dependency on previous iterations or outer +// scopes) under `root`. Returns the set of blocks where MakeAdjoint will emit the reverse pass. std::set identify_independent_blocks(IRNode *root); -// Flip iteration direction of every outer (non-IB) RangeForStmt so the reverse -// pass walks the iteration trace backward; reorders sibling for-loops inside -// non-IB container blocks accordingly. +// Flip iteration direction of every outer (non-IB) RangeForStmt so the reverse pass walks the iteration trace backward; +// reorders sibling for-loops inside non-IB container blocks accordingly. void reverse_outer_loops(IRNode *root, const std::set &IB); } // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/make_adjoint.cpp b/quadrants/transforms/auto_diff/make_adjoint.cpp index 0b041fd876..aacdd299f7 100644 --- a/quadrants/transforms/auto_diff/make_adjoint.cpp +++ b/quadrants/transforms/auto_diff/make_adjoint.cpp @@ -11,16 +11,12 @@ class MakeAdjoint : public ADTransform { using ADTransform::visit; Block *current_block; Block *alloca_block; - // Backup the forward pass (the forward pass might be modified during the - // MakeAdjoint) for search whether a GlobalLoadStmt is inside a for-loop when - // allocating adjoint (see the function `adjoint`) Should be stored - // 1. Before entering a for-loop body - // 2. Before entering a if-stmt - // Should be restored after processing every statement in the two cases above + // Backup the forward pass (the forward pass might be modified during the MakeAdjoint) for search whether a + // GlobalLoadStmt is inside a for-loop when allocating adjoint (see the function `adjoint`). Should be stored before + // entering a for-loop body or an if-stmt, and restored after processing every statement in those two cases. Block *forward_backup; - // IB root: stays constant across visitor recursion. Used when we need to allocate - // persistent storage that must survive enclosing for-loop iterations (e.g. the - // dedicated ad-stacks that snapshot IfStmt conds in visit(IfStmt)). + // IB root: stays constant across visitor recursion. Used when we need to allocate persistent storage that must + // survive enclosing for-loop iterations (e.g. the dedicated ad-stacks that snapshot IfStmt conds in visit(IfStmt)). Block *ib_root; std::map adjoint_stmt; @@ -36,16 +32,15 @@ class MakeAdjoint : public ADTransform { block->accept(&p); } - // Does `if_stmt`'s true/false body contain any AdStackPushStmt targeting `stack`? Recursive to - // catch pushes nested inside further control flow (if-in-if, if-in-for). Used by visit(IfStmt) - // to gate cond-snapshotting. Must be narrow: snapshotting every if-stmt would add an - // AdStackAllocaStmt per if, and determine_ad_stack_size cannot size stacks whose push/pop pair - // is only reachable through branches its Bellman-Ford walk considers "unreached" -- codegen then - // aborts with "Adaptive autodiff stack's size should have been determined" and the extras also - // spam "Unused autodiff stack should have been eliminated" for every untouched snap stack. Only - // when the body actually pushes onto the cond's backing stack does BackupSSA's reverse-time - // clone of load_top read a post-body value rather than the forward cond (the real bug); in every - // other case the clone is already correct and a snapshot would be dead weight. + // Does `if_stmt`'s true/false body contain any AdStackPushStmt targeting `stack`? Recursive to catch pushes nested + // inside further control flow (if-in-if, if-in-for). Used by visit(IfStmt) to gate cond-snapshotting. Must be narrow: + // snapshotting every if-stmt would add an AdStackAllocaStmt per if, and determine_ad_stack_size cannot size stacks + // whose push/pop pair is only reachable through branches its Bellman-Ford walk considers "unreached" - codegen then + // aborts with "Adaptive autodiff stack's size should have been determined" and the extras also spam "Unused autodiff + // stack should have been eliminated" for every untouched snap stack. Only when the body actually pushes onto the + // cond's backing stack does BackupSSA's reverse-time clone of load_top read a post-body value rather than the forward + // cond (the case this snapshot guards against); in every other case the clone is correct and a snapshot would be dead + // weight. static bool block_pushes_to_stack(Block *block, Stmt *stack) { if (!block) return false; @@ -77,8 +72,7 @@ class MakeAdjoint : public ADTransform { block_pushes_to_stack(if_stmt->false_statements.get(), stack); } - // TODO: current block might not be the right block to insert adjoint - // instructions! + // TODO: current block might not be the right block to insert adjoint instructions! void visit(Block *block) override { std::vector statements; // always make a copy since the list can be modified. @@ -136,40 +130,36 @@ class MakeAdjoint : public ADTransform { if (adjoint_stmt.find(stmt) == adjoint_stmt.end()) { // normal SSA cases - // create the alloca - // auto alloca = - // Stmt::make(get_current_program().config.gradient_dt); - // maybe it's better to use the statement data type than the default type + // Create the alloca. Using the statement's own `ret_type` tends to fit better than the kernel-wide + // `get_current_program().config.gradient_dt` default. auto alloca = Stmt::make(adjoint_dtype); adjoint_stmt[stmt] = alloca.get(); - // We need to insert the alloca in the block of GlobalLoadStmt when the - // GlobalLoadStmt is not inside the currently-processed range-for. - // Code sample: - // a and b require grad - // Case 1 (GlobalLoadStmt is outside the for-loop, compute 5 times and - // accumulate once, alloca history value is needed): + // We need to insert the alloca in the block of GlobalLoadStmt when the GlobalLoadStmt is not inside the + // currently-processed range-for. Code sample (a and b require grad): + // + // Case 1 (GlobalLoadStmt is outside the for-loop, compute 5 times and accumulate once, alloca history value + // is needed): // for i in range(5): // p = a[i] // q = b[i] // for _ in range(5) // q += p - // Case 2 (GlobalLoadStmt is inside the for-loop, compute once and - // accumulate immediately, alloca history value can be discarded): + // Case 2 (GlobalLoadStmt is inside the for-loop, compute once and accumulate immediately, alloca history + // value can be discarded): // for i in range(5): // q = b[i] // for _ in range(5) // q += a[i] if (stmt->is() && forward_backup->locate(stmt->as()) == -1) { - // Case 1: the GlobalLoadStmt lives in a block outside the currently-processed range-for iteration. Its - // adjoint must persist across all iterations of the inner reversed loop, so the alloca cannot live in the - // current alloca_block (which would be the inner reversed loop body). Walk up from the primal's enclosing - // block until we hit one whose owning statement unconditionally dominates both the forward and the reverse - // code (a loop / offloaded / kernel body, not an if / while body): visit(IfStmt) emits the reverse code - // into a brand new sibling IfStmt, not back into the forward if-body, so an alloca placed inside the - // forward branch is SSA-invalid from the reverse branch's point of view and gets DCE'd into silently-zero - // gradients. + // Case 1: the GlobalLoadStmt lives in a block outside the currently-processed range-for iteration. Its adjoint + // must persist across all iterations of the inner reversed loop, so the alloca cannot live in the current + // alloca_block (which would be the inner reversed loop body). Walk up from the primal's enclosing block until + // we hit one whose owning statement unconditionally dominates both the forward and the reverse code (a loop / + // offloaded / kernel body, not an if / while body): visit(IfStmt) emits the reverse code into a brand new + // sibling IfStmt, not back into the forward if-body, so an alloca placed inside the forward branch is + // SSA-invalid from the reverse branch's point of view and gets DCE'd into silently-zero gradients. Block *target = stmt->as()->parent; while (target != nullptr) { Stmt *parent_stmt = target->parent_stmt(); @@ -179,9 +169,9 @@ class MakeAdjoint : public ADTransform { } target = parent_stmt->parent; } - // Reaching a null target means the primal's enclosing-block chain is broken (an unparented block). Falling - // back to alloca_block here would silently restore the pre-fix bug (adjoint eliminated as DCE on the - // reverse branch); hard-assert instead so malformed IR surfaces loudly. + // Reaching a null target means the primal's enclosing-block chain is broken (an unparented block). Falling back + // to alloca_block here would place the adjoint inside a branch that gets DCE'd on the reverse side (silently + // zeroed gradients); hard-assert instead so malformed IR surfaces loudly. QD_ASSERT(target != nullptr); target->insert(std::move(alloca), 0); } else { @@ -191,13 +181,13 @@ class MakeAdjoint : public ADTransform { return adjoint_stmt[stmt]; } - // For ops in `NonLinearOps::unary_collections` the reverse formula must NOT read the forward `stmt` - // directly - only `stmt->operand` (adstack-backed), `adjoint(stmt)`, and constants. BackupSSA spills the - // forward stmt to a single plain alloca overwritten each iteration, so reading `stmt` from a reversed - // dynamic loop would use the last-iteration value regardless of which reverse iteration is running. This - // helper walks the value-tree at IR-transform time and asserts the invariant; paired with the Python - // meta-test `test_unary_collections_audit` (which catches "forgot to add to unary_collections"), it - // covers both halves of the class of bugs the per-op audit used to miss. + // For ops in `NonLinearOps::unary_collections` the reverse formula must NOT read the forward `stmt` directly - only + // `stmt->operand` (adstack-backed), `adjoint(stmt)`, and constants. BackupSSA spills the forward stmt to a single + // plain alloca overwritten each iteration, so reading `stmt` from a reversed dynamic loop would use the + // last-iteration value regardless of which reverse iteration is running. This helper walks the value-tree at + // IR-transform time and asserts the invariant. It covers the formula-reads-forward-stmt half of the per-op + // classification check; the missing-from-unary_collections half is covered by a Python-side audit of the + // unary_collections set. void accumulate_unary_operand_checked(UnaryOpStmt *stmt, Stmt *value) { if (NonLinearOps::unary_collections.find(stmt->op_type) != NonLinearOps::unary_collections.end()) { std::unordered_set visited; @@ -238,10 +228,10 @@ class MakeAdjoint : public ADTransform { acc(mul(adjoint(stmt), add(constant(1, stmt->ret_type), sqr(tan(stmt->operand))))); } else if (stmt->op_type == UnaryOpType::tanh) { // Recompute tanh(operand) in the reverse pass instead of reusing the forward stmt value. In dynamic loops - // BackupSSA spills the forward stmt to a single plain alloca overwritten each iteration, so the reversed - // loop would read the last-iteration tanh for every backward step. The operand rides the adstack through - // LocalLoad, so a fresh tanh on it is per-iteration correct. Trade-off: tanh is evaluated twice per - // iteration (once forward, once backward); caching the forward value on the adstack is a future optimization. + // BackupSSA spills the forward stmt to a single plain alloca overwritten each iteration, so the reversed loop + // would read the last-iteration tanh for every backward step. The operand rides the adstack through LocalLoad, so + // a fresh tanh on it is per-iteration correct. Trade-off: tanh is evaluated twice per iteration (once forward, + // once backward). acc(mul(adjoint(stmt), sub(constant(1, stmt->ret_type), sqr(tanh(stmt->operand))))); } else if (stmt->op_type == UnaryOpType::asin) { acc(mul(adjoint(stmt), @@ -251,19 +241,17 @@ class MakeAdjoint : public ADTransform { negate(div(constant(1, stmt->ret_type), sqrt(sub(constant(1, stmt->ret_type), sqr(stmt->operand))))))); } else if (stmt->op_type == UnaryOpType::exp) { // See the tanh case above: recompute exp on the adstack-backed operand so the reversed loop sees the - // per-iteration value rather than the last-forward value spilled by BackupSSA. Same trade-off as tanh: exp - // is evaluated twice per iteration (once forward, once backward); caching the forward value on the adstack - // is a future optimization. + // per-iteration value rather than the last-forward value spilled by BackupSSA. Same trade-off as tanh: exp is + // evaluated twice per iteration (once forward, once backward). acc(mul(adjoint(stmt), exp(stmt->operand))); } else if (stmt->op_type == UnaryOpType::log) { - // No recompute workaround needed: the reverse formula `1 / operand` reads `stmt->operand` directly (which - // is adstack-backed via LocalLoad inside dynamic loops), not the forward `log(operand)` stmt value. + // No recompute workaround needed: the reverse formula `1 / operand` reads `stmt->operand` directly (which is + // adstack-backed via LocalLoad inside dynamic loops), not the forward `log(operand)` stmt value. acc(div(adjoint(stmt), stmt->operand)); } else if (stmt->op_type == UnaryOpType::sqrt) { - // No recompute workaround needed: the reverse formula reads `stmt->operand` (adstack-backed via LocalLoad - // inside dynamic loops, gated on `unary_collections` membership) and recomputes `sqrt(operand)` from it, - // not the forward `sqrt(operand)` stmt value. Unlike tanh/exp this case was already operand-based before - // the recompute fix landed; the structure mirrors log above. + // No recompute workaround needed: the reverse formula reads `stmt->operand` (adstack-backed via LocalLoad inside + // dynamic loops, gated on `unary_collections` membership) and recomputes `sqrt(operand)` from it, not the forward + // `sqrt(operand)` stmt value. Structure mirrors log above. acc(mul(adjoint(stmt), div(constant(0.5f, stmt->ret_type), sqrt(stmt->operand)))); } else if (stmt->op_type == UnaryOpType::rsqrt) { acc(mul(adjoint(stmt), @@ -329,30 +317,29 @@ class MakeAdjoint : public ADTransform { } void visit(IfStmt *if_stmt) override { - // Snapshot a stack-backed forward cond into a dedicated 1-push-per-if-execution ad-stack, - // but only when the cond's backing stack is also pushed inside the if body (e.g. short-circuit - // lowering pushes the rhs of `&&` onto the same stack that holds the cond). Without this, - // BackupSSA's clone of `if_stmt->cond` in the reverse block reads the cond stack AFTER the - // body-pushes rather than the forward-time cond value - the reverse IfStmt flips, pop counts - // drift, and gradients come out silently zero. A dedicated stack has exactly one push per - // forward if-execution, so the reverse load_top matches the forward cond. + // Snapshot a stack-backed forward cond into a dedicated 1-push-per-if-execution ad-stack, but only when the cond's + // backing stack is also pushed inside the if body (e.g. short-circuit lowering pushes the rhs of `&&` onto the same + // stack that holds the cond). Without this, BackupSSA's clone of `if_stmt->cond` in the reverse block reads the + // cond stack AFTER the body-pushes rather than the forward-time cond value - the reverse IfStmt flips, pop counts + // drift, and gradients come out silently zero. A dedicated stack has exactly one push per forward if-execution, so + // the reverse load_top matches the forward cond. // - // Guarded by the body-push check because snapshotting indiscriminately adds AdStackAllocaStmts - // that go through `determine_ad_stack_size` unused on every other if-stmt in the kernel - the - // adaptive-size pass emits "Unused autodiff stack should have been eliminated" warnings and - // the codegen step then fails with "Adaptive autodiff stack's size should have been determined". + // Guarded by the body-push check because snapshotting indiscriminately adds AdStackAllocaStmts that go through + // `determine_ad_stack_size` unused on every other if-stmt in the kernel - the adaptive-size pass emits "Unused + // autodiff stack should have been eliminated" warnings and the codegen step then fails with "Adaptive autodiff + // stack's size should have been determined". Stmt *reverse_cond = if_stmt->cond; AdStackAllocaStmt *snap_stack_ptr = nullptr; // Narrow guard: only the bare `AdStackLoadTopStmt` shape needs the explicit snapshot below. A compound cond (e.g. // `cmp_lt(load_top(x_stack) + 0.1, threshold)` from `if x + 0.1 < threshold` when `x` has been promoted to an - // adstack by `ReplaceLocalVarWithStacks`) reaches this visitor as a BARE `AdStackLoadTopStmt` cond anyway - the - // cmp / arithmetic value goes through `PromoteSSA2LocalVar`'s required-defs set (the IfStmt cond path adds the - // cond's value-producing op), then `AdStackAllocaJudger::visit(IfStmt)` (around line 763) marks its alloca - // stack-needed because it feeds the cond, then `ReplaceLocalVarWithStacks` promotes the alloca to an adstack. By - // the time control reaches here, `if_stmt->cond` is `AdStackLoadTopStmt` of that snap-promoted adstack and the - // body of the if does NOT push to that stack (the cmp value is pushed once just before the IfStmt, never inside), - // so the `body_pushes_to_stack` guard below is false and we correctly skip the additional snap-stack. Per-iter - // cond values are preserved by the alloca-promotion pipeline. + // adstack by `ReplaceLocalVarWithStacks`) reaches this visitor as a BARE `AdStackLoadTopStmt` cond anyway - the cmp + // / arithmetic value goes through `PromoteSSA2LocalVar`'s required-defs set (the IfStmt cond path adds the cond's + // value-producing op), then `AdStackAllocaJudger::visit(IfStmt)` marks its alloca stack-needed because it feeds the + // cond, then `ReplaceLocalVarWithStacks` promotes the alloca to an adstack. By the time control reaches here, + // `if_stmt->cond` is `AdStackLoadTopStmt` of that snap-promoted adstack and the body of the if does NOT push to + // that stack (the cmp value is pushed once just before the IfStmt, never inside), so the `body_pushes_to_stack` + // guard below is false and we correctly skip the additional snap-stack. Per-iter cond values are preserved by the + // alloca-promotion pipeline. // // The bare-`AdStackLoadTopStmt` case the snap-stack below handles is the OTHER shape: a load_top whose backing // stack IS pushed to inside the if body (e.g. short-circuit lowering of `&&` pushes the rhs onto the same stack @@ -361,29 +348,26 @@ class MakeAdjoint : public ADTransform { // wrong cond value. The snap-stack here decouples the cond value from the body's pushes by capturing it once just // before the IfStmt and reading it back at the matching reverse cursor. // - // Earlier comments here claimed that compound conds rely on `BackupSSA::generic_visit`'s `load(op)` else-branch - // ("single alloca, last-write-wins" spill) for correctness. That description was incomplete: the alloca-promotion- - // to-adstack pipeline catches compound conds before they reach BackupSSA, so the spill is a FALLBACK for shapes - // the alloca-promotion missed, not the load-bearing path for compound conds in general. Verified empirically on - // 2026-05-04 with a loop-carried local v + compound cond `v + 0.1 > threshold` + nonlinear use of v in the if - // body: gradients match analytic on origin/main and on C3. + // Compound conds whose value-producing op is not adstack-promoted by the alloca-promotion pipeline fall through to + // `BackupSSA::generic_visit`'s `load(op)` spill: a single alloca written immediately after the forward cond and + // reloaded at reverse-cond construction. That captures the forward-time cond value correctly within one iteration + // of an enclosing dynamic loop. if (if_stmt->cond->is()) { auto *cond_stack = if_stmt->cond->as()->stack->as(); if (body_pushes_to_stack(if_stmt, cond_stack)) { auto cond_type = if_stmt->cond->ret_type.ptr_removed(); - // Size the snap stack the same way as the cond stack it mirrors: one forward push per - // if-execution matched by one reverse pop. Reusing cond_stack->max_size keeps the snap - // stack exempt from `determine_ad_stack_size` when the cond stack itself was built with a - // fixed size, which is always true when ReplaceLocalVarWithStacks ran with a non-zero - // `ad_stack_size` (the only configuration we currently support for stack-based reverse AD). + // Size the snap stack the same way as the cond stack it mirrors: one forward push per if-execution matched by + // one reverse pop. Reusing cond_stack->max_size keeps the snap stack exempt from `determine_ad_stack_size` when + // the cond stack itself was built with a fixed size, which holds when ReplaceLocalVarWithStacks ran with a + // non-zero `ad_stack_size` (the only configuration supported for stack-based reverse AD). auto snap_stack = Stmt::make(cond_type, cond_stack->max_size); snap_stack_ptr = snap_stack->as(); // Allocate at the IB root so the stack persists across enclosing for-loop iterations. ib_root->insert(std::move(snap_stack), 0); - // Per-execution forward push of the cond value, just before the forward if-stmt. No - // initial zero push: the reverse load_top always runs after a matching forward push, so - // leaving the stack empty at entry is both correct and avoids a dead store that DSE would - // otherwise drop (and that `determine_ad_stack_size` would then miscount). + // Per-execution forward push of the cond value, just before the forward if-stmt. No initial zero push: the + // reverse load_top always runs after a matching forward push, so leaving the stack empty at entry is both + // correct and avoids a dead store that DSE would otherwise drop (and that `determine_ad_stack_size` would then + // miscount). if_stmt->insert_before_me(Stmt::make(snap_stack_ptr, if_stmt->cond)); // Per-execution reverse load of the snapshotted cond, emitted in the current reverse block. reverse_cond = insert(snap_stack_ptr); @@ -459,27 +443,26 @@ class MakeAdjoint : public ADTransform { // Restore the forward pass forward_backup = for_stmt->body.get(); } - // Restore current_block. Missing here before: if this RangeForStmt is visited from within another compound - // stmt (notably visit(IfStmt)), that outer visitor will continue iterating its own body in reverse after we - // return and emit further reverse stmts. Without this restore those emissions land in the reversed-for's - // body instead of the outer block, producing silently-wrong gradients whenever a runtime-guarded if wraps a - // for-loop with loop-carried variables (the reverse loop body ends up over-popping the adstack and emitting - // the x.grad accumulation on every iteration instead of once). + // Restore current_block: if this RangeForStmt is visited from within another compound stmt (notably visit(IfStmt)), + // that outer visitor continues iterating its own body in reverse after we return and emit further reverse stmts. + // Without the restore, those emissions would land in the reversed-for's body instead of the outer block, producing + // silently-wrong gradients whenever a runtime-guarded if wraps a for-loop with loop-carried variables (the reverse + // loop body would over-pop the adstack and emit the x.grad accumulation on every iteration instead of once). current_block = old_current_block; forward_backup = old_forward_backup; alloca_block = old_alloca_block; } void visit(StructForStmt *for_stmt) override { - // Save/restore mirrors visit(RangeForStmt) above. Rationale: visit(Block) inside `body->accept(this)` - // sets current_block = for_stmt->body at the start of every iteration, so on return current_block - // points at the struct-for's body. An enclosing compound visitor (e.g. visit(IfStmt)) that resumes - // iterating its children in reverse after this StructForStmt needs current_block and alloca_block to - // still be its own, not this for's; otherwise subsequent reverse emissions land inside the struct-for - // body and any adjoint alloca lives in a block the enclosing if-branch cannot reach. forward_backup - // must be saved too because `visit(IfStmt)` mutates it without restoring, so a nested if inside the - // struct-for body leaves it pointing at the if-branch block, which then survives past this visitor and - // mis-routes `adjoint()` on GlobalLoadStmts for later siblings at the enclosing scope. + // Save/restore mirrors visit(RangeForStmt) above. Rationale: visit(Block) inside `body->accept(this)` sets + // current_block = for_stmt->body at the start of every iteration, so on return current_block points at the + // struct-for's body. An enclosing compound visitor (e.g. visit(IfStmt)) that resumes iterating its children in + // reverse after this StructForStmt needs current_block and alloca_block to still be its own, not this for's; + // otherwise subsequent reverse emissions land inside the struct-for body and any adjoint alloca lives in a block + // the enclosing if-branch cannot reach. forward_backup must be saved too because `visit(IfStmt)` mutates it + // without restoring, so a nested if inside the struct-for body leaves it pointing at the if-branch block, which + // then survives past this visitor and mis-routes `adjoint()` on GlobalLoadStmts for later siblings at the + // enclosing scope. auto old_alloca_block = alloca_block; auto old_current_block = current_block; auto old_forward_backup = forward_backup; @@ -501,12 +484,10 @@ class MakeAdjoint : public ADTransform { void visit(LocalStoreStmt *stmt) override { accumulate(stmt->val, load(adjoint(stmt->dest))); - // Clear the adjoint of the dest after local store, - // Because LocalStoreStmt overwrites the dest, - // 1. If the alloca is inside a loop, the adjoint of this alloca of this - // iteration should be cleared after this iteration has been done - // 2. If the alloca serves as the dest of multiple LocalStoreStmt, only the - // last LocalStoreStmt should be taken account of + // Clear the adjoint of the dest after local store, because LocalStoreStmt overwrites the dest: + // 1. If the alloca is inside a loop, the adjoint of this alloca for this iteration must be cleared once the + // iteration completes. + // 2. If the alloca serves as the dest of multiple LocalStoreStmts, only the last LocalStoreStmt counts. auto dest_type = stmt->dest->ret_type.ptr_removed(); if (is_real(dest_type.get_element_type())) { auto dtype = dest_type; @@ -727,12 +708,11 @@ class MakeAdjoint : public ADTransform { void visit(MatrixPtrStmt *stmt) override { if (stmt->origin->is() || stmt->origin->is()) { /* - The case of MatrixPtrStmt(GlobalPtrStmt, ...) is already handled in - GlobalPtrStmt, GlobalStoreStmt and AtomicStmt + The case of MatrixPtrStmt(GlobalPtrStmt, ...) is already handled in GlobalPtrStmt, GlobalStoreStmt and + AtomicStmt. - TODO(zhanlue): Try to separate out the chain rule for MatrixPtrStmt from - GlobalPtrStmt, GlobalStoreStmt and AtomicStmt and migrate the logics - here. + TODO(zhanlue): Try to separate out the chain rule for MatrixPtrStmt from GlobalPtrStmt, GlobalStoreStmt and + AtomicStmt and migrate the logics here. */ return; } diff --git a/quadrants/transforms/auto_diff/make_adjoint.h b/quadrants/transforms/auto_diff/make_adjoint.h index 7694ee5ee5..e8828746d7 100644 --- a/quadrants/transforms/auto_diff/make_adjoint.h +++ b/quadrants/transforms/auto_diff/make_adjoint.h @@ -4,9 +4,8 @@ namespace quadrants::lang { class Block; -// Stage: Reverse-mode codegen. Emit the reverse pass IR in `ib` (an -// Independent Block produced by identify_independent_blocks). Forward IR must -// already have been shaped by ir_shaping + forward_state_spill. +// Stage: Reverse-mode codegen. Emit the reverse pass IR in `ib` (an Independent Block produced by +// identify_independent_blocks). Forward IR must already have been shaped by ir_shaping + forward_state_spill. void make_adjoint(Block *ib); } // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/make_dual.cpp b/quadrants/transforms/auto_diff/make_dual.cpp index f60318b132..1cbf9301b9 100644 --- a/quadrants/transforms/auto_diff/make_dual.cpp +++ b/quadrants/transforms/auto_diff/make_dual.cpp @@ -63,10 +63,8 @@ class MakeDual : public ADTransform { if (dual_stmt.find(stmt) == dual_stmt.end()) { // normal SSA cases - // create the alloca - // auto alloca = - // Stmt::make(get_current_program().config.gradient_dt); - // maybe it's better to use the statement data type than the default type + // Create the alloca. Using the statement's own `ret_type` tends to fit better than the kernel-wide + // `get_current_program().config.gradient_dt` default. auto alloca = Stmt::make(dual_type); dual_stmt[stmt] = alloca.get(); @@ -86,9 +84,8 @@ class MakeDual : public ADTransform { } else if (stmt->op_type == UnaryOpType::cos) { accumulate(stmt, negate(mul(sin(stmt->operand), dual(stmt->operand)))); } else if (stmt->op_type == UnaryOpType::tan) { - // d/dx tan(x) = 1 + tan(x)^2. Forward mode executes in primal order, so `stmt` is the - // current-iteration tan value - no BackupSSA stale-value concern; reusing it is per-iteration - // correct. + // d/dx tan(x) = 1 + tan(x)^2. Forward mode executes in primal order, so `stmt` is the current-iteration tan value + // - no BackupSSA stale-value concern; reusing it is per-iteration correct. accumulate(stmt, mul(add(constant(1), sqr(stmt)), dual(stmt->operand))); } else if (stmt->op_type == UnaryOpType::tanh) { accumulate(stmt, mul(sub(constant(1), sqr(stmt)), dual(stmt->operand))); @@ -204,10 +201,9 @@ class MakeDual : public ADTransform { } void visit(StructForStmt *for_stmt) override { - // Save/restore mirrors visit(RangeForStmt) above and MakeAdjoint::visit(StructForStmt). An enclosing - // compound visitor that resumes iterating its body after this StructForStmt needs alloca_block to still - // point at its own block, not the sparse-for body, so new dual allocas land where the enclosing reverse - // code can reach them. + // Save/restore mirrors visit(RangeForStmt) above and MakeAdjoint::visit(StructForStmt). An enclosing compound + // visitor that resumes iterating its body after this StructForStmt needs alloca_block to still point at its own + // block, not the sparse-for body, so new dual allocas land where the enclosing reverse code can reach them. auto previous_alloca_block = alloca_block; alloca_block = for_stmt->body.get(); for_stmt->body->accept(this); @@ -220,11 +216,9 @@ class MakeDual : public ADTransform { } void visit(LocalStoreStmt *stmt) override { - // Clear the dual of the dest before local store, - // Because LocalStoreStmt overwrites the dest, - // If the alloca serves as the dest of multiple LocalStoreStmt, only the - // last LocalStoreStmt should be taken account of, i.e, its history should - // be cleared + // Clear the dual of the dest before local store, because LocalStoreStmt overwrites the dest. If the alloca serves + // as the dest of multiple LocalStoreStmts, only the last one counts and the prior accumulated dual must be + // discarded. auto dtype = stmt->dest->ret_type.ptr_removed(); if (is_real(dtype.get_element_type())) { auto zero = insert_const_for_grad(dtype, stmt, 0); diff --git a/quadrants/transforms/auto_diff/make_dual.h b/quadrants/transforms/auto_diff/make_dual.h index 5daeb40c00..0205db7b29 100644 --- a/quadrants/transforms/auto_diff/make_dual.h +++ b/quadrants/transforms/auto_diff/make_dual.h @@ -4,8 +4,8 @@ namespace quadrants::lang { class Block; -// Stage: Forward-mode codegen. Emit a tangent-propagating dual pass for -// `block` (the kernel root for forward-mode autodiff). +// Stage: Forward-mode codegen. Emit a tangent-propagating dual pass for `block` (the kernel root for forward-mode +// autodiff). void make_dual(Block *block); } // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp b/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp index 9521ab70dd..6925b254a7 100644 --- a/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp +++ b/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp @@ -8,13 +8,11 @@ namespace { // ============================================================================ // BackupSSA: spill cross-block SSA operands the reverse-mode clones reference. // -// MakeAdjoint clones forward stmts into the reverse scope but shares operand -// pointers with the forward graph; when those operands live inside the -// forward body (e.g. an inner-for's `begin`/`end`), the reverse clone's -// operand no longer dominates its use. This pass rewrites such operands: -// AdStackLoadTopStmt / ArgLoadStmt get cloned in place, AdStackAllocaStmt -// gets re-rooted at the IB, generic SSA values get a per-IB backup -// AllocaStmt + LocalStore + LocalLoad chain. +// MakeAdjoint clones forward stmts into the reverse scope but shares operand pointers with the forward graph; when +// those operands live inside the forward body (e.g. an inner-for's `begin`/`end`), the reverse clone's operand no +// longer dominates its use. This pass rewrites such operands: AdStackLoadTopStmt / ArgLoadStmt get cloned in place, +// AdStackAllocaStmt gets re-rooted at the IB, generic SSA values get a per-IB backup AllocaStmt + LocalStore + +// LocalLoad chain. // ============================================================================ class BackupSSA : public BasicStmtVisitor { public: @@ -59,36 +57,32 @@ class BackupSSA : public BasicStmtVisitor { // Just create another AdStackLoadTopStmt stmt->set_operand(i, stmt->insert_before_me(op->clone())); } else if (op->is()) { - // Backup AdStackAllocaStmt because it should not be local stored and - // local loaded + // Backup AdStackAllocaStmt because it should not be local stored and local loaded. auto stack_alloca = op->as(); if (backup_alloca.find(op) == backup_alloca.end()) { auto backup_stack_alloca = Stmt::make(stack_alloca->dt, stack_alloca->max_size); auto backup_stack_alloca_ptr = backup_stack_alloca.get(); independent_block->insert(std::move(backup_stack_alloca), 0); backup_alloca[op] = backup_stack_alloca_ptr; - // Replace usages of all blocks i.e., the entry point for the - // replace is the top level block + // Replace usages of all blocks, i.e. the entry point for the replace is the top level block. irpass::replace_all_usages_with(leaf_to_root.back(), op, backup_stack_alloca_ptr); - // Erase the outdated AdStackAllocaStmt + // Erase the outdated AdStackAllocaStmt. op->parent->erase(op); } } else if (op->is()) { stmt->set_operand(i, stmt->insert_before_me(op->clone())); } else { - // Recomputable-chain fallback before the last-iter `load(op)` spill: when the cross-block SSA op - // is rooted in a DAG of side-effect-free arithmetic over already-stack-backed allocas, kernel-args, - // constants, and loop indices, clone the chain into the reverse scope at the consumer site. The - // cloned chain reads stack tops live (matching `MakeAdjoint`'s pop ordering, which fires pops - // AFTER all uses of a stack within one reverse iter) and reads kernel-args / constants / loop - // indices via direct clones - exactly the path that was already correct for AdStackLoadTopStmt and - // ArgLoadStmt operands. + // Recomputable-chain fallback before the last-iter `load(op)` spill: when the cross-block SSA op is rooted in + // a DAG of side-effect-free arithmetic over already-stack-backed allocas, kernel-args, constants, and loop + // indices, clone the chain into the reverse scope at the consumer site. The cloned chain reads stack tops + // live (matching `MakeAdjoint`'s pop ordering, which fires pops AFTER all uses of a stack within one reverse + // iter) and reads kernel-args / constants / loop indices via direct clones - matching the path used for + // AdStackLoadTopStmt and ArgLoadStmt operands. // - // The pre-existing `load(op)` fallback below remains correct for genuinely non-recomputable - // cross-block ops (a forward `GlobalLoadStmt` of a needs_grad SNode whose value the reverse must - // read at last-write rather than recompute, for instance) and for shapes outside one of our - // independent blocks. Adding the recomputable path above the fallback strictly subsets the - // previous behaviour: a chain that fails the predicate falls through to the old `load(op)` line. + // The `load(op)` fallback below covers genuinely non-recomputable cross-block ops (a forward `GlobalLoadStmt` + // of a needs_grad SNode whose value the reverse must read at last-write rather than recompute, for instance) + // and shapes outside one of our independent blocks. The recomputable path above takes precedence when its + // predicate matches; otherwise control falls through to the `load(op)` line. if (RecomputableChainAnalyzer::is_recomputable(op, recomputable_cache_)) { std::unordered_map clone_cache; Stmt *cloned = RecomputableChainCloner::clone_at(op, stmt, clone_cache); @@ -102,8 +96,8 @@ class BackupSSA : public BasicStmtVisitor { } } - // Memoization cache for `RecomputableChainAnalyzer::is_recomputable` queries within one BackupSSA run. - // Re-used across all generic_visit calls; invariant during the visit because forward IR is read-only here. + // Memoization cache for `RecomputableChainAnalyzer::is_recomputable` queries within one BackupSSA run. Re-used across + // all generic_visit calls; invariant during the visit because forward IR is read-only here. std::unordered_map recomputable_cache_; void visit(Stmt *stmt) override { @@ -115,12 +109,12 @@ class BackupSSA : public BasicStmtVisitor { BasicStmtVisitor::visit(stmt); } - // generic_visit spills cross-block operands (the for-loop's `begin` and `end`) the same way it does for an - // IfStmt's cond. MakeAdjoint clones a forward for-loop into the reverse scope and shares the clone's - // `begin`/`end` pointers with the forward stmt; when those operands live inside the forward for's body (e.g. - // inner `for k in range(j)` where `j` is an enclosing loop's index promoted to a per-iter adstack), the reverse - // clone's operand no longer dominates its use. generic_visit's AdStackLoadTopStmt branch handles this by - // inserting a fresh AdStackLoadTop in the reverse scope, which reads the correct per-iteration value. + // generic_visit spills cross-block operands (the for-loop's `begin` and `end`) the same way it does for an IfStmt's + // cond. MakeAdjoint clones a forward for-loop into the reverse scope and shares the clone's `begin`/`end` pointers + // with the forward stmt; when those operands live inside the forward for's body (e.g. inner `for k in range(j)` where + // `j` is an enclosing loop's index promoted to a per-iter adstack), the reverse clone's operand no longer dominates + // its use. generic_visit's AdStackLoadTopStmt branch handles this by inserting a fresh AdStackLoadTop in the reverse + // scope, which reads the correct per-iteration value. void visit(RangeForStmt *stmt) override { generic_visit(stmt); stmt->body->accept(this); @@ -154,19 +148,17 @@ class BackupSSA : public BasicStmtVisitor { }; // ============================================================================ -// CoalesceAdStackLoads: dedup redundant AdStackLoadTop / AdStackLoadTopAdj -// reads emitted by the reverse pass. +// CoalesceAdStackLoads: dedup redundant AdStackLoadTop / AdStackLoadTopAdj reads emitted by the reverse pass. // -// Within a single straight-line block, multiple `AdStackLoadTopStmt` reads of the same stack with no -// intervening `AdStackPushStmt` / `AdStackPopStmt` for that stack return the same value, and multiple -// `AdStackLoadTopAdjStmt` reads are equivalent under the same conditions plus no intervening -// `AdStackAccAdjointStmt`. After Python-side static unrolling collapses an inner loop into straight-line IR, -// every iteration's read of an outer-loop-invariant adstack value emits a separate LoadTop in the same block; -// each individual load is cheap (read u64 count + GEP) but unrolled-loop counts of hundreds to thousands -// inflate PTX size and ptxas register-allocator cost. This pass walks each block, caches the most recent -// LoadTop / LoadTopAdj per stack, replaces subsequent same-stack reads with the cached SSA value until a -// Push / Pop / AccAdjoint invalidates the cache for that stack, and conservatively clears the cache when -// crossing into nested control flow (IfStmt / RangeForStmt / StructForStmt / WhileStmt) where unseen +// Within a single straight-line block, multiple `AdStackLoadTopStmt` reads of the same stack with no intervening +// `AdStackPushStmt` / `AdStackPopStmt` for that stack return the same value, and multiple `AdStackLoadTopAdjStmt` +// reads are equivalent under the same conditions plus no intervening `AdStackAccAdjointStmt`. After Python-side +// static unrolling collapses an inner loop into straight-line IR, every iteration's read of an outer-loop-invariant +// adstack value emits a separate LoadTop in the same block; each individual load is cheap (read u64 count + GEP) but +// unrolled-loop counts of hundreds to thousands inflate PTX size and ptxas register-allocator cost. This pass walks +// each block, caches the most recent LoadTop / LoadTopAdj per stack, replaces subsequent same-stack reads with the +// cached SSA value until a Push / Pop / AccAdjoint invalidates the cache for that stack, and conservatively clears +// the cache when crossing into nested control flow (IfStmt / RangeForStmt / StructForStmt / WhileStmt) where unseen // push/pop/AccAdjoint may appear. // ============================================================================ @@ -245,9 +237,9 @@ class CoalesceAdStackLoads : public BasicStmtVisitor { void visit(Block *block) override { std::unordered_map primal_cache; std::unordered_map adjoint_cache; - // Iterate over a snapshot so erase()s during the walk do not invalidate the iteration. The - // DelayedIRModifier still applies erases at the end via `modify_ir()` so SSA users get rewritten in one - // pass; the snapshot only protects the visitor's own walk. + // Iterate over a snapshot so erase()s during the walk do not invalidate the iteration. The DelayedIRModifier still + // applies erases at the end via `modify_ir()` so SSA users get rewritten in one pass; the snapshot only protects + // the visitor's own walk. std::vector stmts; stmts.reserve(block->statements.size()); for (auto &s : block->statements) { @@ -255,9 +247,9 @@ class CoalesceAdStackLoads : public BasicStmtVisitor { } for (Stmt *s : stmts) { if (auto *lt = s->cast()) { - // `return_ptr=true` returns a pointer into the stack slot rather than loading a value; coalescing - // those is unsound because subsequent stores via the pointer would alias with each other through - // the cached SSA value. The non-pointer variant is the common case driven by reverse-pass formulas. + // `return_ptr=true` returns a pointer into the stack slot rather than loading a value; coalescing those is + // unsound because subsequent stores via the pointer would alias with each other through the cached SSA value. + // The non-pointer variant is the common case driven by reverse-pass formulas. if (lt->return_ptr) { continue; } @@ -283,19 +275,19 @@ class CoalesceAdStackLoads : public BasicStmtVisitor { primal_cache.erase(po->stack); adjoint_cache.erase(po->stack); } else if (auto *aa = s->cast()) { - // Accumulating into the top adjoint slot does not alter the primal half, so the primal cache for - // the same stack stays valid; only the adjoint cache must be invalidated. + // Accumulating into the top adjoint slot does not alter the primal half, so the primal cache for the same stack + // stays valid; only the adjoint cache must be invalidated. adjoint_cache.erase(aa->stack); } else if (s->is() || s->is() || s->is() || s->is()) { - // The pass is intentionally intra-block: any push / pop hidden inside a nested block would - // invalidate caches asymmetrically across branches and make a sound merge complex. Conservatively - // drop both caches before descending so reads after the nested block see no stale entries. + // The pass is intentionally intra-block: any push / pop hidden inside a nested block would invalidate caches + // asymmetrically across branches and make a sound merge complex. Conservatively drop both caches before + // descending so reads after the nested block see no stale entries. primal_cache.clear(); adjoint_cache.clear(); s->accept(this); } else { - // Any other statement is assumed inert with respect to the AdStack count headers; recurse into - // potential nested blocks (defensively) but leave the caches intact. + // Any other statement is assumed inert with respect to the AdStack count headers; recurse into potential nested + // blocks (defensively) but leave the caches intact. s->accept(this); } } diff --git a/quadrants/transforms/auto_diff/post_adjoint_cleanup.h b/quadrants/transforms/auto_diff/post_adjoint_cleanup.h index c80453bc62..636f50645e 100644 --- a/quadrants/transforms/auto_diff/post_adjoint_cleanup.h +++ b/quadrants/transforms/auto_diff/post_adjoint_cleanup.h @@ -7,15 +7,13 @@ class Block; // Stage: Post-MakeAdjoint cleanup. -// Spill cross-block SSA operands the reverse-mode clones reference, so the -// reverse IR is verifier-clean: AdStackLoadTop / ArgLoad get cloned in place, -// AdStackAlloca gets re-rooted, generic SSA values get a per-IB backup -// AllocaStmt + LocalStore + LocalLoad chain. +// Spill cross-block SSA operands the reverse-mode clones reference, so the reverse IR is verifier-clean: AdStackLoadTop +// / ArgLoad get cloned in place, AdStackAlloca gets re-rooted, generic SSA values get a per-IB backup AllocaStmt + +// LocalStore + LocalLoad chain. void backup_ssa(Block *block); -// Within a single straight-line block, dedup repeated AdStackLoadTop / -// AdStackLoadTopAdj reads of the same stack with no intervening -// Push / Pop / AccAdjoint. Returns whether any IR mutation took place. +// Within a single straight-line block, dedup repeated AdStackLoadTop / AdStackLoadTopAdj reads of the same stack with +// no intervening Push / Pop / AccAdjoint. Returns whether any IR mutation took place. bool coalesce_ad_stack_loads(IRNode *root); } // namespace quadrants::lang diff --git a/quadrants/transforms/auto_diff/validation.cpp b/quadrants/transforms/auto_diff/validation.cpp index 2676e7a40a..70d882ec15 100644 --- a/quadrants/transforms/auto_diff/validation.cpp +++ b/quadrants/transforms/auto_diff/validation.cpp @@ -6,16 +6,16 @@ namespace quadrants::lang { namespace { // Detect cross-iteration read-after-write through a `needs_grad` global field at the body of an offload-level -// range-for. `MakeAdjoint` treats each iteration of an offload-level loop as independent and only chases -// loop-carried dependencies through `AdStackAllocaJudger`, which inspects local `AllocaStmt`s - cross-iteration -// dependencies through `GlobalStoreStmt` / `GlobalLoadStmt` on a `has_adjoint()` SNode are silently dropped. -// A kernel like `out[i] = out[i-1] + x[None]` therefore returns wrong gradients with no diagnostic, even with -// `ad_stack_experimental_enabled`. Pair every same-SNode (store, load) and raise when the index Stmt*s differ. -// Same-Stmt indices (`out[i] = out[i] + x[None]`) are in-place accumulation and the per-iteration backward -// pass handles those correctly; pointer equality on the SSA index Stmts after the `pre_autodiff` simplify -// (`compile_to_offloads.cpp:121`) is enough to distinguish the two cases. Walks the body of every direct-child -// `RangeForStmt` of `root`; nested for-loops' bodies are skipped via `inside_nested_` because their cross-iter -// behavior is governed by the existing IB / adstack machinery, not by this offload-level guard. +// range-for. `MakeAdjoint` treats each iteration of an offload-level loop as independent and only chases loop-carried +// dependencies through `AdStackAllocaJudger`, which inspects local `AllocaStmt`s - cross-iteration dependencies through +// `GlobalStoreStmt` / `GlobalLoadStmt` on a `has_adjoint()` SNode are silently dropped. A kernel like `out[i] = +// out[i-1] + x[None]` therefore returns wrong gradients with no diagnostic, even with `ad_stack_experimental_enabled`. +// Pair every same-SNode (store, load) and raise when the index Stmt*s differ. Same-Stmt indices (`out[i] = out[i] + +// x[None]`) are in-place accumulation and the per-iteration backward pass handles those correctly; pointer equality on +// the SSA index Stmts after the `pre_autodiff` simplify (`compile_to_offloads.cpp:121`) is enough to distinguish the +// two cases. Walks the body of every direct-child `RangeForStmt` of `root`; nested for-loops' bodies are skipped via +// `inside_nested_` because their cross-iter behavior is governed by the existing IB / adstack machinery, not by this +// offload-level guard. class OffloadLevelGlobalCrossIterRAWChecker : public BasicStmtVisitor { public: using BasicStmtVisitor::visit; @@ -68,9 +68,9 @@ class OffloadLevelGlobalCrossIterRAWChecker : public BasicStmtVisitor { if (!any_axis_references_loop_index(store)) continue; for (auto *load : load_it->second) { - // Skip loads whose index is iteration-independent on every axis (constants, captures from outer - // scopes, ...). Such loads read from a slice of the SNode that the loop never writes to in this - // iteration, so they are not a cross-iter RAW hazard. + // Skip loads whose index is iteration-independent on every axis (constants, captures from outer scopes, ...). + // Such loads read from a slice of the SNode that the loop never writes to in this iteration, so they are not + // a cross-iter RAW hazard. if (!any_axis_references_loop_index(load)) continue; if (!indices_have_no_cross_iter_dependency(store, load)) { @@ -133,11 +133,11 @@ class OffloadLevelGlobalCrossIterRAWChecker : public BasicStmtVisitor { return true; } - // Walk an index `Stmt*` and decide whether it (transitively) references a `LoopIndexStmt` of any enclosing - // loop. Recurses through linear arithmetic ops because an index like `i - 1` lowers to - // `BinaryOpStmt(LoopIndexStmt(i), Sub, ConstStmt(1))` and we want to recognise the `LoopIndexStmt` underneath. - // Conservatively returns false for anything else (constants, ndarray loads from outer scope, ...) - those - // express iteration-independent indexing that the cross-iter guard intentionally exempts. + // Walk an index `Stmt*` and decide whether it (transitively) references a `LoopIndexStmt` of any enclosing loop. + // Recurses through linear arithmetic ops because an index like `i - 1` lowers to `BinaryOpStmt(LoopIndexStmt(i), Sub, + // ConstStmt(1))` and we want to recognise the `LoopIndexStmt` underneath. Conservatively returns false for anything + // else (constants, ndarray loads from outer scope, ...) - those express iteration-independent indexing that the + // cross-iter guard intentionally exempts. static bool references_loop_index(Stmt *s) { if (s == nullptr) return false; diff --git a/quadrants/transforms/auto_diff/validation.h b/quadrants/transforms/auto_diff/validation.h index 59b9822381..31712445f0 100644 --- a/quadrants/transforms/auto_diff/validation.h +++ b/quadrants/transforms/auto_diff/validation.h @@ -7,17 +7,15 @@ namespace quadrants::lang { class IRNode; class Block; -// Stage: Validation. Diagnostic-only passes that refuse autodiff inputs which -// would silently produce wrong gradients, and inserted runtime assertions for -// global-data access rules. +// Stage: Validation. Diagnostic-only passes that refuse autodiff inputs which would silently produce wrong gradients, +// and inserted runtime assertions for global-data access rules. -// Refuses cross-iteration read-after-write through a needs_grad global field -// at an offload-level loop body. Run pre-transform from irpass::auto_diff. +// Refuses cross-iteration read-after-write through a needs_grad global field at an offload-level loop body. Run +// pre-transform from irpass::auto_diff. void offload_level_global_cross_iter_raw_check(Block *offload_body); -// Inserts runtime check-bit accumulators around GlobalStore / AtomicOp on -// snodes whose adjoint check-bit is enabled. Run from -// irpass::differentiation_validation_check. +// Inserts runtime check-bit accumulators around GlobalStore / AtomicOp on snodes whose adjoint check-bit is enabled. +// Run from irpass::differentiation_validation_check. void global_data_access_rule_check(IRNode *root, const std::string &kernel_name); } // namespace quadrants::lang diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index 4d61aa537b..86d8ea33ec 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -414,6 +414,67 @@ def test_adstack_basic_gradient_f64(n_iter): _run_basic_gradient(qd.f64, n_iter=n_iter, rel_tol=1e-14, approx=pytest.approx, abs_tol=0) +@test_utils.test(require=qd.extension.adstack) +def test_eliminate_recomputable_pushes_preserves_zero_body_store(): + # Cross-checks `dloss/dc` against the analytic value `2 * cos(c) * cos(sin(c))` for an adstack-mode kernel where + # `tmp` receives one recomputable body push (`tmp = qd.sin(c[None])`) plus one conditional zero body store + # (`if reset[i] != 0: tmp = 0.0`). The reset mask `[0, 1, 0, 1]` zeros `tmp` for two of four iterations so only + # the other two contribute to the gradient. + # + # Internal details: the adstack promotion of `tmp` comes from `AdStackAllocaJudger::visit(UnaryOpStmt)` (`tmp` + # feeds the non-linear `qd.sin(tmp)` accumulator); the load+store rule does not fire because the offload-level + # for-loop sits outside the IB and `dynamic_for_depth_` stays 0 throughout the judger walk. The body push value + # `sin(GlobalLoad(c))` is a recomputable chain by `RecomputableChainAnalyzer`, which makes the stack a candidate + # for `EliminateRecomputableAdStackPushes`. The eligibility gate must count the conditional `tmp = 0.0` as a + # body push: two body pushes for one stack disqualify the stack and the original IR survives. A weaker gate that + # classifies the conditional zero as the prologue init - e.g. detecting init by literal-zero value rather than + # by position relative to the alloca - drops the user's zero store, rewires every `load_top` to `sin(c)`, and + # produces `c.grad` exactly 2x the analytic value because the gradient flows through every iteration regardless + # of the reset mask. + n = 4 + c = qd.field(qd.f32, shape=(), needs_grad=True) + reset = qd.field(qd.i32, shape=n) + loss = qd.field(qd.f32, shape=(), needs_grad=True) + + @qd.kernel + def compute(): + for i in range(n): + # Recomputable body push: sin of a globally-loaded scalar. The chain `sin(GlobalLoad(c))` has interior + # side-effect-free ops only and no LoopIndex / LocalLoad leaves, so RecomputableChainAnalyzer returns + # true. + tmp = qd.sin(c[None]) + if reset[i] != 0: + # Real `tmp = 0.0` body store. Lowers to AdStackPushStmt with value zero. A weaker eligibility gate + # misclassifies this push as an init prologue push and silently erases it. + tmp = 0.0 + # Use tmp as a non-linear unary op operand so AdStackAllocaJudger marks the alloca stack-needed via its + # visit(UnaryOpStmt) rule, which is what makes this an AdStack in the first place (the load+store rule + # does not fire here because the offload-level for-loop is outside the IB and `dynamic_for_depth_` stays + # 0). + loss[None] += qd.sin(tmp) + + c[None] = 0.5 + reset[0] = 0 + reset[1] = 1 + reset[2] = 0 + reset[3] = 1 + loss[None] = 0.0 + c.grad[None] = 0.0 + + compute() + loss.grad[None] = 1.0 + compute.grad() + + # Forward semantics: tmp_i = sin(c) when reset[i]==0, 0 when reset[i]!=0. Two of four iterations are not reset, + # so loss == 2 * sin(sin(c)). Gradient: dloss/dc contributes cos(c)*cos(sin(c)) per non-reset iteration, so + # dloss/dc == 2 * cos(c) * cos(sin(c)). + n_no_reset = 2 + expected_loss = n_no_reset * math.sin(math.sin(0.5)) + expected_grad = n_no_reset * math.cos(0.5) * math.cos(math.sin(0.5)) + assert loss[None] == pytest.approx(expected_loss, rel=1e-5) + assert c.grad[None] == pytest.approx(expected_grad, rel=1e-5) + + @pytest.mark.parametrize("n_iter", [1, 3, 10]) @pytest.mark.parametrize("wrap_inner_in_func", [False, True]) @test_utils.test(ad_stack_experimental_enabled=False)