[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,7 @@ 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.
+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 applies when each loop iteration touches a different field cell (`field[i]`, `field[I, J, K]`, `grid[i, 0]`, `grid[arg, i, j]` where `arg` is a kernel argument that slices alongside iterating axes); shapes where more than one iteration may hit the same cell (`field[i % K]` with K < n, `field[i / 2]` and similar affine folds, `field[42]`, `field[arg]` with no iterating axis, `field[other_field[i]]`) automatically fall back to the `num_threads * stack_size` heap and stay correct.
 
 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 +354,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,112 @@ 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). Runtime-bounded loops cannot be checked
+      // statically and fall through to the joint axis check below; the conservative reject for the unsafe
+      // single-axis-bare runtime-bounded shape happens after `n_iterating` is known.
+      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;
+        }
+      }
+      // Reject when the loop-to-cell mapping is not provably injective.
+      //   * `n_iterating == 0`: every axis is loop-invariant (`field[42]`, `field[arg]`, `field[other_field[i]]`).
+      //     The gate evaluates to a launch-constant cell, so the reducer reports a launch-constant count while the
+      //     n iterations claim n rows - rows past that count alias.
+      //   * `n_iterating == 1 && n_bare_iterating == 0`: the single iterating axis goes through arithmetic on the
+      //     loop index (`field[i / 2]`, `field[i % K]`, `field[i + 5]`). Affine folds collapse iterations onto
+      //     fewer cells; bijective shifts like `i + 5` are caught here too as conservative fallback - they would
+      //     need range analysis we do not have at this IR stage to admit safely.
+      // Multi-axis cases with `n_iterating >= 2` are accepted: the canonical sparse-grid shape decomposes a
+      // single linear loop index into multiple axes via floordiv / mod chains whose joint mapping is bijective
+      // when the divisors / moduli form a non-overlapping partition (the `qd.ndrange(*shape)` lowering); the cell
+      // visited per iteration is unique. The axis walk recurses through `BinaryOpStmt` / `UnaryOpStmt` and stops
+      // at load nodes so dynamic indices like `field[other_field[i]]` register as non-iterating.
+      auto *lookup = getch->input_ptr ? getch->input_ptr->cast<SNodeLookupStmt>() : nullptr;
+      if (lookup == nullptr || lookup->input_index == nullptr) {
+        return false;
+      }
+      // Collect the per-axis components of `lookup->input_index`. After `lower_access`, the input_index can be
+      // either a `LinearizeStmt` (whose `inputs` are the axes - the StructFor path preserves this) or an
+      // `add`/`mul` arithmetic tree of the form `axis_0 * stride_0 + axis_1 * stride_1 + ...` (the ndrange path
+      // expands the LinearizeStmt). Either shape encodes the same axis decomposition; the recursive collector
+      // walks `add` nodes and extracts the non-const operand of each `mul(axis, stride_const)` leaf, falling
+      // through to push any other shape as a single axis.
+      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);
+        }
+        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;
+      }
+      // Runtime-bounded loops with a single iterating axis are unsafe: even when the axis is a bare `LoopIndexStmt`
+      // (e.g. `for i in range(n_arg): if grid[i] > eps:`), the launch can dispatch `n_arg > snode_iter_count`
+      // threads and undershoot the heap. Multi-axis runtime-bounded ndrange shapes (Genesis-style `for ii, jj, kk,
+      // ib in qd.ndrange(*shape, batch)`) are accepted: the joint floordiv / mod decomposition from a single
+      // linear index is structurally bijective even when the per-axis sizes are runtime values, so excess
+      // iterations would manifest as out-of-bounds field access (a different error class) before fold aliasing.
+      if (!static_bound && n_iterating == 1) {
+        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;