diff --git a/quadrants/transforms/auto_diff/auto_diff.cpp b/quadrants/transforms/auto_diff/auto_diff.cpp index 157c275ef2..77556ba947 100644 --- a/quadrants/transforms/auto_diff/auto_diff.cpp +++ b/quadrants/transforms/auto_diff/auto_diff.cpp @@ -37,13 +37,14 @@ void auto_diff(IRNode *root, const CompileConfig &config, AutodiffMode autodiff_ replace_local_var_with_stacks(ib, config.ad_stack_size); type_check(root, config); - // Disabled: `RecomputableChainAnalyzer` admits `GlobalLoadStmt` as a recomputable interior node without - // verifying SNode read-only-ness. When the chain reaches a `GlobalLoadStmt` of an SNode written elsewhere in - // the same kernel execution, `BackupSSA`'s reverse-side re-clone reads post-write state instead of the iter-k - // value the original `AdStackLoadTopStmt` would have returned, silently corrupting gradients on MPM-style - // kernels that mix per-particle reads with grid writes. Re-enable once the chain analyzer rejects - // mutated-SNode loads (or the cost model gates on read-only-verified leaves). - // eliminate_recomputable_ad_stack_pushes(ib); + // 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); diff --git a/quadrants/transforms/auto_diff/auto_diff_common.h b/quadrants/transforms/auto_diff/auto_diff_common.h index 87a2023e0d..4b20df14c4 100644 --- a/quadrants/transforms/auto_diff/auto_diff_common.h +++ b/quadrants/transforms/auto_diff/auto_diff_common.h @@ -79,28 +79,69 @@ class NonLinearOps { // 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` is admitted as a recomputable interior node only when its source SNode is not mutated anywhere in +// the surrounding IB (per `mutated_snodes`); a re-clone at the reverse cursor of a load whose SNode is also written +// observes post-write state instead of the iter-k value the forward chain consumed. Loads whose source is not a +// `GlobalPtrStmt` (e.g. `ExternalPtrStmt`, `GlobalTemporaryStmt`) are rejected outright since their write set is not +// collected here. `LocalLoadStmt` aliases mutable allocas; the reverse pass reads its forward value through the +// dedicated spill path. // -// Caller passes a `cache` to share memoization across multiple queries on the same DAG (diamond shapes). +// Caller passes a `cache` to share memoization across multiple queries on the same DAG (diamond shapes) and the +// `mutated_snodes` set produced by `collect_mutated_snodes(ib)`. class RecomputableChainAnalyzer { public: - static bool is_recomputable(Stmt *stmt, std::unordered_map &cache) { + // Set of SNodes written somewhere in the IB. `GlobalStoreStmt`, `AtomicOpStmt`, and `SNodeOpStmt` destinations are + // tracked; `MatrixPtrStmt` is chased one level to its origin so post-`lower_matrix_ptr` shapes are captured. + static std::unordered_set collect_mutated_snodes(Block *ib) { + std::unordered_set result; + auto add_dest = [&](Stmt *dest) { + if (dest == nullptr) + return; + if (auto *mp = dest->cast()) + dest = mp->origin; + if (dest == nullptr) + return; + if (auto *gp = dest->cast()) { + if (gp->snode != nullptr) + result.insert(gp->snode); + } else if (auto *mgp = dest->cast()) { + for (auto *s : mgp->snodes) + if (s != nullptr) + result.insert(s); + } + }; + auto stmts = irpass::analysis::gather_statements(ib, [](Stmt *) { return true; }); + for (auto *s : stmts) { + if (auto *gs = s->cast()) { + add_dest(gs->dest); + } else if (auto *ao = s->cast()) { + add_dest(ao->dest); + } else if (auto *so = s->cast()) { + if (so->snode != nullptr) + result.insert(so->snode); + } + } + return result; + } + + static bool is_recomputable(Stmt *stmt, + std::unordered_map &cache, + const std::unordered_set &mutated_snodes) { 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); + bool result = check(stmt, cache, mutated_snodes); cache[stmt] = result; return result; } private: - static bool check(Stmt *stmt, std::unordered_map &cache) { + static bool check(Stmt *stmt, + std::unordered_map &cache, + const std::unordered_set &mutated_snodes) { // 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 @@ -134,15 +175,16 @@ class RecomputableChainAnalyzer { 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. - // - // 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. + // A `GlobalLoadStmt` chain leaf is recomputable only when its `GlobalPtrStmt` source resolves to an SNode the IB + // never writes; a reverse-side re-clone of a write-aliased SNode read returns post-write state. Non-`GlobalPtrStmt` + // sources (e.g. `ExternalPtrStmt`, `GlobalTemporaryStmt`) are rejected because their write set is not collected. + if (auto *gl = stmt->cast()) { + auto *src = gl->src; + auto *gp = src ? src->cast() : nullptr; + if (gp == nullptr || gp->snode == nullptr || mutated_snodes.count(gp->snode)) { + return false; + } + } bool is_interior = stmt->is() || stmt->is() || stmt->is() || stmt->is() || stmt->is() || stmt->is() || stmt->is(); @@ -153,7 +195,7 @@ class RecomputableChainAnalyzer { for (auto *op : operands) { if (op == nullptr) continue; - if (!is_recomputable(op, cache)) + if (!is_recomputable(op, cache, mutated_snodes)) return false; } return true; diff --git a/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp b/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp index e826db5b92..d90b6ce1c2 100644 --- a/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp +++ b/quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp @@ -87,6 +87,10 @@ class EliminateRecomputableAdStackPushes { stacks.push_back(s->as()); } + // SNodes written in the IB; gates `GlobalLoadStmt` chain leaves in `is_recomputable`. The pass body only erases + // adstack stmts, so the write-stmt set is invariant across one outer iteration and we can recompute it per pass. + auto mutated_snodes = RecomputableChainAnalyzer::collect_mutated_snodes(ib); + bool modified = false; std::unordered_map recomputable_cache; @@ -158,7 +162,7 @@ class EliminateRecomputableAdStackPushes { } Stmt *pushed_val = body_push->v; - if (!RecomputableChainAnalyzer::is_recomputable(pushed_val, recomputable_cache)) { + if (!RecomputableChainAnalyzer::is_recomputable(pushed_val, recomputable_cache, mutated_snodes)) { continue; } // Loop-carried self-reference: when the pushed value's transitive operand DAG includes an AdStackLoadTopStmt of @@ -238,80 +242,42 @@ class EliminateRecomputableAdStackPushes { 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. + // Reverse-position correctness for chain leaves. // - // 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. + // Each chain leaf `AdStackLoadTopStmt(T)` is re-cloned by `BackupSSA::generic_visit` at the consumer's reverse + // position; the cloned `load_top(T)` reads `T`'s top at that reverse cursor, which must equal what the original + // load_top read in forward iter k. `MakeAdjoint` emits `T`'s pop for each forward push at the reverse cursor + // mirroring that push's forward position, so by consumer-reverse only the pops for `T`-pushes with forward + // position > P_cons have fired. The reverse top therefore equals `T`'s last forward push at-or-before P_cons, + // while the forward chain at the load_top's own forward position P_lt reads `T`'s last forward push + // at-or-before P_lt. They match iff no `T`-push falls strictly between P_lt and P_cons in document order. // - // 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; + // A single preorder index over the IB lets the check compare positions across nested blocks in one integer + // compare; pushes inside `IfStmt`, `RangeForStmt`, and `StructForStmt` bodies are reached via the visitor walk + // and slot in between the container's preorder and the next sibling's preorder. + std::unordered_map preorder; + { + auto all = irpass::analysis::gather_statements(ib, [](Stmt *) { return true; }); + preorder.reserve(all.size()); + for (size_t i = 0; i < all.size(); i++) + preorder[all[i]] = static_cast(i); } + auto preorder_of = [&](Stmt *s) -> int { + auto it = preorder.find(s); + return it == preorder.end() ? -1 : it->second; + }; - // Collect all distinct AdStackAllocaStmt referenced via AdStackLoadTopStmt leaves in pushed_val's chain. - std::unordered_set chain_leaf_stacks; + // Collect every `AdStackLoadTopStmt` reachable from `pushed_val`'s recomputable chain. Each one has its own + // forward position P_lt - the position at which the cloned chain in reverse will read its target stack's top + // for the corresponding forward iter - so the check is per-load_top, not per-stack. + std::vector chain_load_tops; 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); + chain_load_tops.push_back(lt); return; } if (s->is() || s->is() || s->is()) @@ -323,20 +289,40 @@ class EliminateRecomputableAdStackPushes { }; collect_leaves(pushed_val); + int max_consumer_preorder = -1; + auto consumer_users = irpass::analysis::gather_statements(ib, [&](Stmt *s) { + for (auto *op : s->get_operands()) { + for (auto *lt : load_tops) { + if (op == lt) + return true; + } + } + return false; + }); + for (auto *u : consumer_users) { + int p = preorder_of(u); + if (p > max_consumer_preorder) + max_consumer_preorder = p; + } + 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 (auto *lt : chain_load_tops) { + auto *T = lt->stack ? lt->stack->cast() : nullptr; + if (T == nullptr) { + reverse_safe = false; + break; + } + int lt_pre = preorder_of(lt); 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 && p != T_init) { - if (static_cast(i) <= max_consumer_pos) { - reverse_safe = false; - break; - } - } + auto T_pushes = irpass::analysis::gather_statements(ib, [&](Stmt *s) { + auto *p = s->cast(); + return p != nullptr && p->stack == T && p != T_init; + }); + for (auto *p : T_pushes) { + int p_pre = preorder_of(p); + if (p_pre > lt_pre && p_pre <= max_consumer_preorder) { + reverse_safe = false; + break; } } if (!reverse_safe) diff --git a/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp b/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp index 6925b254a7..61ba25cb5b 100644 --- a/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp +++ b/quadrants/transforms/auto_diff/post_adjoint_cleanup.cpp @@ -20,10 +20,14 @@ class BackupSSA : public BasicStmtVisitor { Block *independent_block; std::map backup_alloca; + std::unordered_set mutated_snodes_; - explicit BackupSSA(Block *independent_block) : independent_block(independent_block) { + explicit BackupSSA(Block *independent_block) + : independent_block(independent_block), + mutated_snodes_(RecomputableChainAnalyzer::collect_mutated_snodes(independent_block)) { allow_undefined_visitor = true; invoke_default_visitor = true; + compute_forward_preorder(); } Stmt *load(Stmt *stmt) { @@ -54,8 +58,17 @@ class BackupSSA : public BasicStmtVisitor { 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())); + // Cloning an AdStackLoadTopStmt at the reverse cursor reads the stack's live top there. Correct iff every + // subsequent forward push to the same stack has had its matching reverse pop fire before the consumer. + // The unsafe case is a push whose parent block strictly contains the consumer's parent AND precedes the + // consumer's enclosing-scope stmt in document order: pop_P lands AFTER consumer's reverse code. On + // rejection, fall back to the per-IB alloca spill. + if (is_load_top_safe_for_consumer(op->as(), stmt)) { + 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))); + } } else if (op->is()) { // Backup AdStackAllocaStmt because it should not be local stored and local loaded. auto stack_alloca = op->as(); @@ -83,7 +96,11 @@ class BackupSSA : public BasicStmtVisitor { // 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_)) { + // `is_chain_clone_safe` adds a per-leaf check on top of `is_recomputable`: each `AdStackLoadTopStmt` chain + // leaf must satisfy the consumer-aware ancestor predicate (same rule as the bare-`AdStackLoadTopStmt` + // operand path). On rejection, fall back to the per-IB alloca spill. + if (RecomputableChainAnalyzer::is_recomputable(op, recomputable_cache_, mutated_snodes_) && + is_chain_clone_safe(op, stmt)) { std::unordered_map clone_cache; Stmt *cloned = RecomputableChainCloner::clone_at(op, stmt, clone_cache); stmt->set_operand(i, cloned); @@ -100,6 +117,149 @@ class BackupSSA : public BasicStmtVisitor { // all generic_visit calls; invariant during the visit because forward IR is read-only here. std::unordered_map recomputable_cache_; + // Document-order index of every stmt reachable from `independent_block` at construction time. Captures forward-IR + // positions before any reverse-side insertion occurs; reused by `is_load_top_safe_for_consumer` so the per-load_top + // ancestor check runs in O(1) per query after a single O(N) walk. + std::unordered_map forward_preorder_; + + void compute_forward_preorder() { + int idx = 0; + std::function walk = [&](Block *block) { + for (auto &owned : block->statements) { + Stmt *s = owned.get(); + forward_preorder_[s] = idx++; + if (auto *if_s = s->cast()) { + 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 *for_s = s->cast()) { + if (for_s->body) + walk(for_s->body.get()); + } else if (auto *for_s = s->cast()) { + if (for_s->body) + walk(for_s->body.get()); + } else if (auto *for_s = s->cast()) { + if (for_s->body) + walk(for_s->body.get()); + } + } + }; + walk(independent_block); + } + + static bool is_strict_ancestor(Block *ancestor, Block *descendant) { + if (ancestor == nullptr || descendant == nullptr || ancestor == descendant) { + return false; + } + Block *cur = descendant->parent_block(); + while (cur != nullptr) { + if (cur == ancestor) { + return true; + } + cur = cur->parent_block(); + } + return false; + } + + static Stmt *ancestor_in_block(Stmt *stmt, Block *target_block) { + Stmt *cur = stmt; + while (cur != nullptr) { + if (cur->parent == target_block) { + return cur; + } + Block *cur_parent = cur->parent; + if (cur_parent == nullptr) { + return nullptr; + } + cur = cur_parent->parent_stmt(); + } + return nullptr; + } + + // Returns true iff cloning `lt` at `consumer`'s position reads `lt`'s forward value. The unsafe case is a forward + // push P to the same stack such that (a) P's parent block is a strict ancestor of `consumer`'s parent and (b) P + // precedes consumer's enclosing-scope stmt in P's parent block. Both conditions together mean pop_P fires after + // `consumer` in execution order, so the cloned load_top reads the post-push top. + bool is_load_top_safe_for_consumer(AdStackLoadTopStmt *lt, Stmt *consumer) { + auto lt_it = forward_preorder_.find(lt); + if (lt_it == forward_preorder_.end()) { + return true; + } + auto *stack = lt->stack ? lt->stack->cast() : nullptr; + if (stack == nullptr) { + return true; + } + int lt_pos = lt_it->second; + Block *consumer_parent = consumer->parent; + auto pushes = irpass::analysis::gather_statements(independent_block, [&](Stmt *s) { + auto *p = s->cast(); + return p != nullptr && p->stack == stack; + }); + for (auto *p : pushes) { + auto pit = forward_preorder_.find(p); + if (pit == forward_preorder_.end() || pit->second <= lt_pos) { + continue; + } + if (!is_strict_ancestor(p->parent, consumer_parent)) { + continue; + } + Stmt *consumer_anc = ancestor_in_block(consumer, p->parent); + if (consumer_anc != nullptr) { + auto cit = forward_preorder_.find(consumer_anc); + if (cit != forward_preorder_.end() && pit->second > cit->second) { + continue; + } + } + return false; + } + return true; + } + + bool is_chain_clone_safe(Stmt *op, Stmt *consumer) { + Block *consumer_parent = consumer->parent; + auto cache_key = std::make_pair(op, consumer_parent); + auto it = chain_clone_safe_cache_.find(cache_key); + if (it != chain_clone_safe_cache_.end()) { + return it->second; + } + chain_clone_safe_cache_[cache_key] = false; + bool result = chain_clone_safe_check(op, consumer); + chain_clone_safe_cache_[cache_key] = result; + return result; + } + + bool chain_clone_safe_check(Stmt *op, Stmt *consumer) { + if (auto *lt = op->cast()) { + return is_load_top_safe_for_consumer(lt, consumer); + } + if (op->is() || op->is() || op->is()) { + return true; + } + bool is_interior = op->is() || op->is() || op->is() || + op->is() || op->is() || op->is() || + op->is(); + if (!is_interior) { + return false; + } + for (auto *child : op->get_operands()) { + if (child == nullptr) { + continue; + } + if (!is_chain_clone_safe(child, consumer)) { + return false; + } + } + return true; + } + + struct PairHash { + std::size_t operator()(const std::pair &p) const noexcept { + return std::hash{}(p.first) ^ (std::hash{}(p.second) << 1); + } + }; + std::unordered_map, bool, PairHash> chain_clone_safe_cache_; + void visit(Stmt *stmt) override { generic_visit(stmt); } diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index ba9fa5fa4a..35d9a9a33c 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -477,6 +477,123 @@ def compute(): assert c.grad[None] == pytest.approx(expected_grad, rel=1e-5) +@test_utils.test(require=qd.extension.adstack, cfg_optimization=False, ad_stack_experimental_enabled=True) +def test_backup_ssa_load_top_with_subsequent_pushes(): + # Pins `BackupSSA::is_load_top_stable` against the reverse-mode aliasing bug where cloning an + # `AdStackLoadTopStmt` at the reverse cursor reads the stack's live top - which differs from the original + # forward-time top whenever the same stack received another push between the original load_top and the end of + # the IB. The minimal IR shape that exhibits this: + # + # - A multi-axis `qd.ndrange(...)` whose flat loop index is decomposed into per-axis indices via repeated + # `floordiv` / `sub` over an adstack-backed value (each step pushes its result, then a subsequent + # `load_top` reads that pushed value to feed the next decomposition step). + # - A guarded division `out = in / mass` inside the kernel: `MakeAdjoint` emits two reverse formulas + # (`adj(in) += adj(out) / mass`, `adj(mass) += -adj(out) * in / mass^2`) that re-evaluate `mass` and `in` + # at the reverse cursor; that re-evaluation walks the per-axis index DAG, which BackupSSA was happy to + # reconstruct from `floordiv(load_top(stack), const)` clones - the load_top clones read the stack's + # final top, which is the LAST push (the inner axis) instead of the intermediate flat index. + # + # With the stable-load-top guard, BackupSSA detects the unsafe load_top, falls back to the per-IB alloca + # spill (`load(op)`), and the reverse re-evaluation reads the per-iteration forward-time index, producing + # the analytic gradients. Without the guard, the reverse pass reads `g_mass[j/2, j]` and `g_in[j/2, j]` + # instead of `g_mass[i, j]` / `g_in[i, j]`, so cells whose forward gate was true at `(i=1, j=1)` end up + # dividing `adj(out)` by `g_mass[0, 1]` (= 0) and the gradient comes out as `inf`/`nan`. + # + # Internal details: `cfg_optimization=False` keeps the compound `if true && (g_mass[i, j] > 0)` form alive + # (the `&&` lowering inserts the multi-push into the gate stack that the chain analyzer must clone through); + # `ad_stack_experimental_enabled=True` is the configuration where the cloned-load_top path actually fires. + g_mass = qd.field(qd.f32, shape=(2, 2), needs_grad=True) + g_in = qd.field(qd.f32, shape=(2, 2), needs_grad=True) + g_out = qd.field(qd.f32, shape=(2, 2), needs_grad=True) + + @qd.kernel + def gated_div(): + for i, j in qd.ndrange(2, 2): + if g_mass[i, j] > 0.0: + g_out[i, j] = g_in[i, j] / g_mass[i, j] + + g_mass[1, 1] = 2.0 + g_in[1, 1] = 4.0 + g_out.grad[1, 1] = 1.0 + gated_div.grad() + + # Forward: g_out[1, 1] = g_in[1, 1] / g_mass[1, 1] = 2.0; gate is false at every other cell so all other + # gradients are zero. With dy = 1.0 on the active cell: + # d(g_in[1, 1]) = 1 / g_mass[1, 1] = 0.5 + # d(g_mass[1, 1]) = -g_in[1, 1] / g_mass[1, 1]^2 = -1.0 + assert g_in.grad[1, 1] == pytest.approx(0.5, rel=1e-6) + assert g_mass.grad[1, 1] == pytest.approx(-1.0, rel=1e-6) + # Inactive cells must not have received any gradient. A regression where `BackupSSA` re-clones the + # load_top at the reverse cursor reads `g_mass[j/2, j] = g_mass[0, 1] = 0`, so 1/g_mass[0, 1] = inf + # would land on `g_in.grad[0, 1]` via the wrong index. Probing every other cell catches both the inf + # spill and any silent off-by-one accumulation. + for i in range(2): + for j in range(2): + if (i, j) == (1, 1): + continue + assert g_in.grad[i, j] == pytest.approx(0.0, abs=1e-6) + assert g_mass.grad[i, j] == pytest.approx(0.0, abs=1e-6) + + +@test_utils.test( + require=qd.extension.adstack, + cfg_optimization=False, + force_scalarize_matrix=True, + ad_stack_experimental_enabled=True, +) +def test_eliminate_recomputable_pushes_rejects_mutated_snode_chain_leaf(): + # Pins the `mutated_snodes` guard inside `RecomputableChainAnalyzer::is_recomputable` against the chain-clone + # post-write read miscompile: a forward `GlobalLoadStmt` chain leaf reading a SNode the same kernel mutates + # cannot be re-cloned at the reverse cursor - the cloned load re-issues the read after the forward writes have + # updated the SNode, producing wrong gradients (`nan` / `inf`). + # + # Internal details: + # - Kernel shape: three top-level for-loops over the same `field` SNode. + # 1. Atomic-add into `field[base, 0]` keyed by `base = floor(data[0]).cast(i32)`. + # 2. Gated divide-by-self over the whole `field`, gate predicate `field[I, 0][0] > 0.0`. + # 3. Consumer that reads `field[base, 0]` and accumulates into `data[1]` (the gradient endpoint). + # - `mutated_snodes(IB) = {field}`. + # - Without the guard: the analyzer admits the chain `GlobalLoad(GlobalPtr(field, base))` as recomputable; + # ERAP drops the adstack carrying that chain's value; `BackupSSA::generic_visit` re-clones the chain at + # the reverse cursor; the cloned `GlobalLoadStmt` reads `field` POST stage 2 (the divide-by-self has set + # every gated cell to the all-ones vector), so the adjoint flowing through `data[0]` blows up. + # - With the guard: the chain is rejected, ERAP keeps the original push/pop, reverse pops the iter-k value + # verbatim, and `data.grad[0]` matches the analytic all-ones vector. + # - `force_scalarize_matrix=True` is structural: the matrix-typed `field` value path is what ERAP latches + # onto; with scalarization off the chain's leaf shape changes and the bug no longer fires. + # - `cfg_optimization=False` keeps the gate's compound `&&` lowering alive so the reverse clone path + # actually triggers; with cfg-opt the cond folds and the clone never happens. + # - The trailing `1` axis on `field` matches `MakeAdjoint::visit(RangeForStmt)`'s reverse iteration; + # collapsing it folds the loop into a shape that misses ERAP's eligibility entirely. + vec3 = qd.types.vector(3, qd.f32) + data = qd.field(dtype=vec3, shape=(2,), needs_grad=True) + field = qd.field(dtype=vec3, shape=(2, 2, 2, 1), needs_grad=True) + + data[0] = qd.Vector([1.0, 1.0, 1.0]) + + @qd.kernel + def k(data: qd.template(), field: qd.template()): + for _ in qd.ndrange(1): + base = qd.floor(data[0]).cast(qd.i32) + field[base, 0] += data[0] + for ii, jj, kk, i_b in qd.ndrange(2, 2, 2, 1): + I = (ii, jj, kk) + if field[I, i_b][0] > 0.0: + field[I, i_b] = field[I, i_b] / field[I, i_b] + for _ in qd.ndrange(1): + base = qd.floor(data[0]).cast(qd.i32) + data[1] = data[0] + field[base, 0] + + field[1, 1, 1, 0] = qd.Vector([1.0, 1.0, 1.0]) + data.grad[1] = qd.Vector([1.0, 1.0, 1.0]) + k.grad(data, field) + + expected = [1.0, 1.0, 1.0] + for d in range(3): + assert math.isfinite(data.grad[0][d]), f"non-finite at axis {d}: {data.grad[0][d]}" + assert data.grad[0][d] == pytest.approx(expected[d], rel=1e-6, abs=1e-6) + + @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)