Skip to content

[AutoDiff] Adstack chain-clone safety: mutated-SNode leaf reject + load_top consumer-aware guard #634

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
duburcqa merged 2 commits into from
May 6, 2026
Merged

[AutoDiff] Adstack chain-clone safety: mutated-SNode leaf reject + load_top consumer-aware guard #634

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c51f218
[AutoDiff] BackupSSA: reject recomputable-chain clone when an AdStack…
duburcqa May 6, 2026
79091b7
[AutoDiff] Re-enable EliminateRecomputableAdStackPushes with mutated-…
duburcqa May 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 8 additions & 7 deletions quadrants/transforms/auto_diff/auto_diff.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
78 changes: 60 additions & 18 deletions quadrants/transforms/auto_diff/auto_diff_common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Stmt *, bool> &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<SNode *> collect_mutated_snodes(Block *ib) {
std::unordered_set<SNode *> result;
auto add_dest = [&](Stmt *dest) {
if (dest == nullptr)
return;
if (auto *mp = dest->cast<MatrixPtrStmt>())
dest = mp->origin;
if (dest == nullptr)
return;
if (auto *gp = dest->cast<GlobalPtrStmt>()) {
if (gp->snode != nullptr)
result.insert(gp->snode);
} else if (auto *mgp = dest->cast<MatrixOfGlobalPtrStmt>()) {
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<GlobalStoreStmt>()) {
add_dest(gs->dest);
} else if (auto *ao = s->cast<AtomicOpStmt>()) {
add_dest(ao->dest);
} else if (auto *so = s->cast<SNodeOpStmt>()) {
if (so->snode != nullptr)
result.insert(so->snode);
}
}
return result;
}

static bool is_recomputable(Stmt *stmt,
std::unordered_map<Stmt *, bool> &cache,
const std::unordered_set<SNode *> &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<Stmt *, bool> &cache) {
static bool check(Stmt *stmt,
std::unordered_map<Stmt *, bool> &cache,
const std::unordered_set<SNode *> &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
Expand Down Expand Up @@ -134,15 +175,16 @@ class RecomputableChainAnalyzer {
stmt->is<ConstStmt>()) {
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<GlobalLoadStmt>()) {
auto *src = gl->src;
auto *gp = src ? src->cast<GlobalPtrStmt>() : nullptr;
if (gp == nullptr || gp->snode == nullptr || mutated_snodes.count(gp->snode)) {
return false;
}
}
bool is_interior = stmt->is<UnaryOpStmt>() || stmt->is<BinaryOpStmt>() || stmt->is<TernaryOpStmt>() ||
stmt->is<MatrixPtrStmt>() || stmt->is<GlobalPtrStmt>() || stmt->is<ExternalPtrStmt>() ||
stmt->is<GlobalLoadStmt>();
Expand All @@ -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;
Expand Down
140 changes: 63 additions & 77 deletions quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,10 @@ class EliminateRecomputableAdStackPushes {
stacks.push_back(s->as<AdStackAllocaStmt>());
}

// 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<Stmt *, bool> recomputable_cache;

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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<int> &out_positions) {
std::function<void(Block *, int)> 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<int>(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<IfStmt>()) {
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<RangeForStmt>()) {
if (rf->body)
walk(rf->body.get(), effective_pos);
} else if (auto *sf = s->cast<StructForStmt>()) {
if (sf->body)
walk(sf->body.get(), effective_pos);
}
}
};
walk(push_block, -1);
};
std::vector<int> 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<Stmt *, int> 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<int>(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<AdStackAllocaStmt *> 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<AdStackLoadTopStmt *> chain_load_tops;
std::unordered_set<Stmt *> walked_for_leaves;
std::function<void(Stmt *)> collect_leaves = [&](Stmt *s) {
if (walked_for_leaves.count(s))
return;
walked_for_leaves.insert(s);
if (auto *lt = s->cast<AdStackLoadTopStmt>()) {
if (auto *as = lt->stack->cast<AdStackAllocaStmt>())
chain_leaf_stacks.insert(as);
chain_load_tops.push_back(lt);
return;
}
if (s->is<AdStackAllocaStmt>() || s->is<ArgLoadStmt>() || s->is<ConstStmt>())
Expand All @@ -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<AdStackAllocaStmt>() : 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<AdStackPushStmt>()) {
if (p->stack == T && p != T_init) {
if (static_cast<int>(i) <= max_consumer_pos) {
reverse_safe = false;
break;
}
}
auto T_pushes = irpass::analysis::gather_statements(ib, [&](Stmt *s) {
auto *p = s->cast<AdStackPushStmt>();
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)
Expand Down
Loading
Loading