[Lang] SNode-arm bound-expr capture rejects fold-attack gate indices via iter-count + axis-classification checks

Author: duburcqa

docs/source/user_guide/autodiff.md

@@ -317,7 +317,19 @@ where each quantity means:
 | `bytes_per_slot` | Depends on `T` and on the backend (see table below). |
 | `num_buffers` | Number of adstacks the kernel allocates - one per loop-carried variable plus one per dependent branch flag (see [One adstack per variable](#one-adstack-per-variable)). |
 
-Kernels of the shape `for i in range(...): if field[i] cmp literal: <adstack work>` (a runtime gate directly above the adstack-using body, comparing one field entry to a constant) shrink further: the compiler counts gate-passing iterations at launch time and sizes the float adstack to that count instead of `num_threads * stack_size`. A workload whose gate matches 5% of iterations pays 5% of the float-adstack cost; the float heap grows on demand if a later launch matches more. Integer / boolean adstacks stay at `num_threads * stack_size` - their pushes fire unconditionally for control-flow replay. The shrinking is exact only when the gate's per-axis index is a bare loop variable (`field[i]`, `field[I, J, K]`); see [What can go wrong](#what-can-go-wrong) for a known limitation on `qd.field`-backed gates indexed by compound expressions.
+The float heap is by far the main reverse-mode memory bottleneck because a typical kernel allocates many float-typed adstacks - one per floating-point loop-carried scalar, each storing both primal and adjoint - and the total scales as `num_threads * stack_size * num_float_buffers * 8` bytes, dominating the integer / boolean heap. Advanced static IR analysis is used to further shrink the float adstack in some common gated-kernel shapes. When a runtime gate sits directly above the adstack-using body and compares a single field entry to a constant, the compiler counts the gate-passing iterations at launch time and sizes the float adstack to that count instead of `num_threads * stack_size`, so a workload whose gate matches 5% of iterations pays 5% of the float-adstack cost; the float heap grows on demand if a later launch matches more, and integer / boolean adstacks stay at `num_threads * stack_size` since their pushes fire unconditionally for control-flow replay. The shrinking kicks in only when the compiler can statically prove that each loop iteration touches a different field cell; patterns where uniqueness cannot be proven statically fall back to worst-case heap `num_threads * stack_size`.
+
+Patterns that capture:
+  - **Linear range loop**: `for i in range(n): if field[i] > eps: ...`
+  - **Multi-axis StructFor**: `for I, J, K in field3d: if field3d[I, J, K] > eps: ...`
+  - **Multi-axis ndrange**: `for ii, jj, kk in qd.ndrange(*shape): if grid[ii, jj, kk] > eps: ...`
+  - **Iterating axes plus a slice**: a kernel argument or constant on an extra axis, e.g. `grid[i, 0]`, `grid[arg, ii, jj, kk]`.
+
+Patterns that fall back to the worst-case heap:
+  - **Single-axis arithmetic on the loop variable**: `field[i % K]`, `field[i / 2]`, `field[i + 5]`, `field[2 * i]`, and similar.
+  - **Constant-index gate**: `field[42]`, or any axis that is a literal constant.
+  - **Kernel-argument index, no iterating axis**: `field[arg]` where every axis is launch-constant.
+  - **Indirect index via runtime load**: `field[other_field[i]]`; the compiler cannot prove `other_field` is injective.
 
 Every adstack slot always stores a *primal* value - the forward-pass value the reverse pass pops to recover the chain-rule step. Floating-point adstacks additionally store an *adjoint* slot where the reverse pass accumulates chain-rule contributions. Integer / boolean adstacks do not need an adjoint slot.
 
@@ -354,9 +366,6 @@ A large `ndrange` combined with several loop-carried variables multiplies quickl
   - pass `ad_stack_size=N` to `qd.init()` with `N` large enough to cover the real push count (bypasses the sizer).
 - **Out-of-memory before the kernel even runs.** A reverse pass through many loop-carried variables at a large ndrange can ask the runtime for more adstack memory than the device can physically back, even when the sizer's number is correct. Surfaces as an allocator OOM at launch time. Remedies are the ones listed under *Avoiding OOM on GPU* above: fewer loop-carried variables, a smaller ndrange, manual checkpointing, or more device-memory headroom.
 - **Loop bounds backed by a mutated ndarray.** A reverse-mode kernel with `for i in range(n[j])` requires `n[j]` to hold the same value at the forward call and at `.grad()`. If anything writes to `n[j]` between those two points - the differentiable kernel itself, or any other kernel call - the computed gradient may come out wrong, sometimes as an `Adstack overflow` exception at `qd.sync()`, sometimes silently. The safe rule: populate loop-bound ndarrays before the forward call and leave them untouched until `.grad()` returns. The reason for that is Quadrants' adstack sizer design: it reads the loop bound separately at each dispatch, which includes forward and backward calls. Tape-based eager AD like [PyTorch's autograd](https://pytorch.org/docs/stable/notes/autograd.html) is not affected, since the trip count is recorded as the forward runs and reused at backward time.
-- :warning: **Gate on a `qd.field` indexed by an expression that is not a plain loop variable.** A reverse-mode kernel of the shape `for i in range(n): if field[i % K] > eps: <adstack work>` (or any gate whose index is not a plain loop variable - `field[2 * i]`, `field[42]`, `field[other_field[i]]`) may produce silently wrong gradients. Workarounds:
-  - raise `ad_stack_sparse_threshold_bytes` in `qd.init()` past the kernel's conservative-heap byte size;
-  - use a `qd.ndarray` for the gating field instead of a `qd.field`.
 
 ## Performance characteristics
 

quadrants/transforms/static_adstack_analysis.cpp

@@ -274,20 +274,102 @@ StaticAdStackAnalysisResult analyze_adstack_static_bounds(OffloadedStmt *task_ir
       if (!desc_opt.has_value()) {
         return false;
       }
-      // KNOWN LIMITATION: the SNode arm trusts whatever index expression the codegen passed to `SNodeLookupStmt` and
-      // does not verify that it is a bijection with the kernel's loop iteration space. A pathological pattern like
-      // `for i in range(n): if field[i % K] > eps: <push f32>` (with `K < n`, `field` an SNode-backed `qd.field`)
-      // captures `iter_count = K` while the main pass walks `[0, n)` and claims a heap row for every iteration whose
-      // `i % K` cell passes - aliasing the n - K excess gated iterations onto the K-row heap and corrupting
-      // gradients (silent on LLVM, hard "reducer count diverged" overflow on SPIR-V). The ndarray arm above DOES
-      // validate that each axis is a `LoopIndexStmt` because at analysis time the ndarray's per-axis indices are
-      // still individual statements; the SNode case has no such per-axis information once `LinearizeStmt` is
-      // lowered into raw `add` / `mul` arithmetic, and any narrower walker we tried would also reject legitimate
-      // multi-axis kernels (e.g. `for I, J, K in grid: if grid[I, J, K].mass > eps: ...`, the canonical MPM-grid
-      // shape), where the lowered offset is `add(mul(I, sx), add(mul(J, sy), mul(K, sz)))` - the same affine shape
-      // a malicious manual linearisation can fake. Until the analysis runs at an earlier IR stage where
-      // `LinearizeStmt` is preserved (or a different bijection-witness is identified), this gap is documented
-      // rather than gated. Working assumption: production kernels don't use `field[i % K]` as a gate.
+      // Iteration-count check (statically-bounded loops only): reject when the task's loop iterates more times
+      // than the SNode has cells. The reducer walks `desc_opt->iter_count` cells and counts gate-passing ones,
+      // the main pass claims a heap row per gated iteration, so a mismatch under-allocates the heap and aliases
+      // excess iterations onto a shared row. Structural signature of `for i in range(n): if field[i % K] > eps:
+      // <push f32>` (K < n, field on a K-cell SNode): `loop_iter = n > K = snode_iter_count`. Legitimate
+      // multi-axis kernels of the shape `for ii, jj, kk, ib in qd.ndrange(...): if grid[f, ii, jj, kk, ib] >
+      // eps:` keep `loop_iter <= snode_iter_count` (slice axes like the kernel-arg `f` make the reducer
+      // over-count by the slice factor, which is benign over-allocation).
+      const bool static_bound = task_ir->const_begin && task_ir->const_end && task_ir->end_stmt == nullptr;
+      if (static_bound) {
+        const int64_t loop_iter = static_cast<int64_t>(task_ir->end_value) - static_cast<int64_t>(task_ir->begin_value);
+        if (loop_iter <= 0 || static_cast<uint64_t>(loop_iter) > static_cast<uint64_t>(desc_opt->iter_count)) {
+          return false;
+        }
+      }
+      // Per-axis classification on the gate's `LinearizeStmt::inputs`. After `lower_access` the input_index is
+      // either a `LinearizeStmt` (StructFor path preserves it) or an `add`/`mul` arithmetic tree (ndrange path
+      // expanded form); a recursive collector recovers per-axis components from either shape. For each axis,
+      // walk through `BinaryOpStmt` / `UnaryOpStmt` looking for a `LoopIndexStmt` to classify as iterating; pure
+      // `ConstStmt` / `ArgLoadStmt` axes are slices. Reject when:
+      //   * `n_iterating == 0`: every axis is loop-invariant (`field[42]`, `field[arg]`, `field[other_field[i]]`).
+      //   * `n_iterating == 1 && n_bare_iterating == 0`: single non-bare iterating axis (`field[i / 2]`,
+      //     `field[i % K]`, `field[i + 5]`). Conservatively rejects bijective shifts too; the worst-case heap
+      //     stays correct, just larger.
+      // Multi-axis cases with `n_iterating >= 2` are accepted: the canonical `qd.ndrange(*shape)` decomposes a
+      // single linear loop index into multiple axes via floordiv / mod chains whose joint mapping is bijective.
+      auto *lookup = getch->input_ptr ? getch->input_ptr->cast<SNodeLookupStmt>() : nullptr;
+      if (lookup == nullptr || lookup->input_index == nullptr) {
+        return false;
+      }
+      std::vector<Stmt *> axes;
+      if (auto *lin = lookup->input_index->cast<LinearizeStmt>()) {
+        axes.assign(lin->inputs.begin(), lin->inputs.end());
+      } else {
+        std::function<void(Stmt *)> collect_axes = [&](Stmt *s) {
+          if (auto *bin = s->cast<BinaryOpStmt>()) {
+            if (bin->op_type == BinaryOpType::add) {
+              collect_axes(bin->lhs);
+              collect_axes(bin->rhs);
+              return;
+            }
+            if (bin->op_type == BinaryOpType::mul) {
+              if (bin->rhs && bin->rhs->is<ConstStmt>()) {
+                axes.push_back(bin->lhs);
+                return;
+              }
+              if (bin->lhs && bin->lhs->is<ConstStmt>()) {
+                axes.push_back(bin->rhs);
+                return;
+              }
+            }
+          }
+          axes.push_back(s);
+        };
+        collect_axes(lookup->input_index);
+      }
+      std::function<bool(Stmt *, int)> contains_loop_index = [&](Stmt *s, int depth) -> bool {
+        if (s == nullptr || depth > 8) {
+          return false;
+        }
+        if (s->is<LoopIndexStmt>()) {
+          return true;
+        }
+        if (auto *bin = s->cast<BinaryOpStmt>()) {
+          return contains_loop_index(bin->lhs, depth + 1) || contains_loop_index(bin->rhs, depth + 1);
+        }
+        if (auto *un = s->cast<UnaryOpStmt>()) {
+          return contains_loop_index(un->operand, depth + 1);
+        }
+        // The reverse task replays multi-axis loop indices by spilling them onto per-axis adstacks during the forward
+        // pass and loading via `AdStackLoadTopStmt` in the reverse pass. The push (in the forward `OffloadedStmt`) and
+        // the load (in the reverse `OffloadedStmt`) sit in different tasks, so the per-task `per_stack_pushed_values`
+        // map cannot trace from this load back to the `LoopIndexStmt` source. Treat the load as iterating directly:
+        // autodiff only spills runtime-varying values onto adstacks, and per-axis adstacks - the only ones reaching a
+        // `LinearizeStmt::input` - carry replayed loop indices by construction.
+        if (s->is<AdStackLoadTopStmt>()) {
+          return true;
+        }
+        return false;
+      };
+      int n_iterating = 0;
+      int n_bare_iterating = 0;
+      for (Stmt *axis : axes) {
+        if (contains_loop_index(axis, 0)) {
+          n_iterating++;
+          if (axis->is<LoopIndexStmt>()) {
+            n_bare_iterating++;
+          }
+        }
+      }
+      if (n_iterating == 0) {
+        return false;
+      }
+      if (n_iterating == 1 && n_bare_iterating == 0) {
+        return false;
+      }
       out.field_source_kind = StaticAdStackBoundExpr::FieldSourceKind::SNode;
       out.snode_root_id = desc_opt->root_id;
       out.snode_byte_base_offset = desc_opt->byte_base_offset;