[AutoDiff] Adstack heap: clip reducer count by static loop trip count to avoid oversize on small-loop large-snode shapes

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)). |
 
-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, so a workload whose gate matches 5% of iterations pays 5% of the float-adstack cost. See [Appendix B: gate-index shapes that capture vs fall back to the worst-case heap](#appendix-b-gate-index-shapes-that-capture-vs-fall-back-to-the-worst-case-heap) for the authoritative list of supported shapes.
+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. 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. So a workload whose gate matches 5% of iterations pays 5% of the float-adstack cost. See [Appendix B: gate-index shapes that capture vs fall back to the worst-case heap](#appendix-b-gate-index-shapes-that-capture-vs-fall-back-to-the-worst-case-heap) for the authoritative list of supported shapes.
 
 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.
 

quadrants/codegen/llvm/codegen_llvm.cpp

@@ -1850,8 +1850,16 @@ std::string TaskCodeGenLLVM::init_offloaded_task_function(OffloadedStmt *stmt, s
     desc.iter_count = static_cast<uint32_t>(iter_count);
     return desc;
   };
-  auto adstack_analysis =
-      analyze_adstack_static_bounds(stmt, snode_resolver, compile_config.ad_stack_sparse_threshold_bytes);
+  // CPU LLVM goes through `make_cpu_multithreaded_range_for` in `offload_to_executable`, which rewrites the user
+  // loop's `[begin_value, end_value)` into per-thread chunks before codegen runs. The atomic row counter the
+  // codegen emits is shared across every chunk of the same task, so the total claim count is the original
+  // pre-chunk loop trip count, not the per-chunk subrange. Signal that to the analyzer so it skips filling
+  // `bound_expr.loop_iter_static` on CPU and the runtime falls back to the unclipped reducer count there. CUDA
+  // and AMDGPU dispatch one thread per iteration without chunking, so their per-task `[begin_value, end_value)`
+  // matches the user loop and the analyzer can fill the field.
+  const bool task_range_is_original_loop = !arch_is_cpu(compile_config.arch);
+  auto adstack_analysis = analyze_adstack_static_bounds(
+      stmt, snode_resolver, compile_config.ad_stack_sparse_threshold_bytes, task_range_is_original_loop);
   ad_stack_bootstrap_pushes_ = std::move(adstack_analysis.bootstrap_pushes);
   ad_stack_lca_block_float_ir_ = adstack_analysis.lca_block_float;
   ad_stack_static_bound_expr_ = adstack_analysis.bound_expr;

quadrants/runtime/gfx/runtime.cpp

@@ -685,6 +685,22 @@ void GfxRuntime::launch_kernel(KernelHandle handle, LaunchContextBuilder &host_c
         auto bound_count_it = per_task_bound_count.find(i);
         if (bound_count_it != per_task_bound_count.end()) {
           effective_rows = bound_count_it->second;
+          // Clip by the captured compile-time loop trip count. Each loop iteration claims at most one row at the
+          // LCA-block (one `atomic_add` per gating iteration), so the heap needs at most `loop_iter_static` rows
+          // regardless of how many cells of the gating SNode the reducer counted; without this, an oversized SNode
+          // (1024-cell `selector` paired with a 64-iter loop) inflates the float heap to the SNode's cell count
+          // even though the loop can only claim 64 rows. Skipped when `loop_iter_static == 0` (runtime-bounded
+          // loops, where the analyzer cannot resolve the trip count from IR alone) - the unclipped reducer count
+          // is already a valid upper bound there. Skipped on `dynamic_gpu_range_for` shapes that exceed the GPU
+          // dispatch cap and serialize iterations across threads (each thread reaches the LCA-block multiple
+          // times, so `loop_iter_static` does not bound row claims) by capping to `dispatched_threads` whenever
+          // it is smaller, falling back cleanly to the unclipped reducer count.
+          if (attribs.ad_stack.bound_expr.has_value() && attribs.ad_stack.bound_expr->loop_iter_static > 0) {
+            const size_t loop_iter_static = static_cast<size_t>(attribs.ad_stack.bound_expr->loop_iter_static);
+            if (loop_iter_static <= dispatched_threads) {
+              effective_rows = std::min<size_t>(effective_rows, loop_iter_static);
+            }
+          }
         } else if (attribs.ad_stack.bound_expr.has_value()) {
           // Reaching here means the bound reducer skipped this `bound_expr`-captured task and `per_task_bound_count`
           // has no entry for slot `i`. The reducer's skip paths in `dispatch_adstack_bound_reducers` are: PSB

quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp

@@ -436,9 +436,19 @@ void LlvmRuntimeExecutor::ensure_per_task_float_heap_post_reducer(std::size_t ta
 
   // Floor at 1 row when the captured count is zero (no thread passed the gate this launch). The codegen-emitted bounds
   // clamp keeps `claimed_row` in [0, count-1] so threads that miss the gate never reach the LCA-block claim - the heap
-  // row stays unused. A 1-row allocation is cheap and keeps the heap pointer non-null.
-  const std::size_t effective_rows =
+  // row stays unused. A 1-row allocation is cheap and keeps the heap pointer non-null. Clip by the captured
+  // compile-time loop trip count when known: each iteration claims at most one row at the LCA-block (one `atomic_add`
+  // per gating iteration), so the heap needs at most `loop_iter_static` rows regardless of how many cells of an
+  // oversized gating SNode the reducer counted. The analyzer leaves `loop_iter_static == 0` for runtime-bounded loops
+  // and for CPU LLVM tasks whose `[begin_value, end_value)` is a post-chunking subrange (the unclipped reducer count is
+  // the right upper bound there).
+  std::size_t effective_rows =
       (count == std::numeric_limits<uint32_t>::max()) ? num_threads : std::max<std::size_t>(count, 1);
+  if (count != std::numeric_limits<uint32_t>::max() && ad_stack.bound_expr.has_value() &&
+      ad_stack.bound_expr->loop_iter_static > 0) {
+    effective_rows =
+        std::min<std::size_t>(effective_rows, static_cast<std::size_t>(ad_stack.bound_expr->loop_iter_static));
+  }
   // Read back the per-thread float stride (in bytes) that `publish_adstack_metadata` published into
   // `runtime->adstack_per_thread_stride_float`. `AdStackSizingInfo::per_thread_stride_float` from the analysis pre-pass
   // is in entry-count units (`2 * max_size`), not bytes, and would massively undersize the heap.

quadrants/transforms/static_adstack_analysis.cpp

@@ -140,7 +140,8 @@ void walk_ir(IRNode *node, Fn &&visit) {
 
 StaticAdStackAnalysisResult analyze_adstack_static_bounds(OffloadedStmt *task_ir,
                                                           const SNodeDescriptorResolver &snode_descriptor_resolver,
-                                                          std::size_t sparse_heap_threshold_bytes) {
+                                                          std::size_t sparse_heap_threshold_bytes,
+                                                          bool task_range_is_original_loop) {
   StaticAdStackAnalysisResult result;
   if (task_ir == nullptr || task_ir->body == nullptr) {
     return result;
@@ -228,6 +229,26 @@ StaticAdStackAnalysisResult analyze_adstack_static_bounds(OffloadedStmt *task_ir
     }
   });
 
+  // Compile-time loop trip count of the analyzed task, used by the runtime to clip the reducer's gate-passing-cell
+  // count down to the actual maximum row claim count: each iteration claims at most one row at the LCA-block, so
+  // `loop_iter` is a sound upper bound on heap rows regardless of how oversized the gating SNode / ndarray is.
+  // Zero when:
+  //   * the task's loop is runtime-bounded (`task_ir->end_stmt != nullptr` or non-const begin / end - the analyzer
+  //     cannot resolve the trip count from IR alone), or
+  //   * `task_range_is_original_loop` is false (the caller signals that the task's `[begin_value, end_value)` is a
+  //     post-chunking subrange, e.g. CPU LLVM after `make_cpu_multithreaded_range_for` rewrote a 256-iteration
+  //     range-for into 16 parallel chunks of 16; the chunked subrange would massively undersize the heap clip
+  //     because the atomic row counter is shared across all chunks of the same task).
+  // In both cases the runtime falls back to the unclipped reducer count.
+  const bool task_static_bound = task_ir->const_begin && task_ir->const_end && task_ir->end_stmt == nullptr;
+  const int64_t task_loop_iter =
+      task_static_bound ? (static_cast<int64_t>(task_ir->end_value) - static_cast<int64_t>(task_ir->begin_value)) : 0;
+  const uint32_t loop_iter_static_for_bound_expr =
+      (task_range_is_original_loop && task_loop_iter > 0 &&
+       static_cast<uint64_t>(task_loop_iter) <= std::numeric_limits<uint32_t>::max())
+          ? static_cast<uint32_t>(task_loop_iter)
+          : 0;
+
   // Resolve a `GlobalLoadStmt::src` chain to a captured field source. Returns true on a recognized shape (ndarray
   // ext-ptr or SNode root->dense->place(scalar)); on success populates the source-kind-specific fields of `out`.
   auto match_field_source = [&](Stmt *load_src, StaticAdStackBoundExpr &out) -> bool {
@@ -245,6 +266,7 @@ StaticAdStackAnalysisResult analyze_adstack_static_bounds(OffloadedStmt *task_ir
           }
         }
         out.field_source_kind = StaticAdStackBoundExpr::FieldSourceKind::NdArray;
+        out.loop_iter_static = loop_iter_static_for_bound_expr;
         out.ndarray_arg_id = base_arg->arg_id;
         // Capture the gating ndarray's ndim so the host launcher can walk shape[0..ndim) at dispatch time and product
         // them into the reducer's flat-element walk bound. Without this the launcher would have to fall back to
@@ -498,6 +520,7 @@ StaticAdStackAnalysisResult analyze_adstack_static_bounds(OffloadedStmt *task_ir
       out.snode_byte_base_offset = desc_opt->byte_base_offset;
       out.snode_byte_cell_stride = desc_opt->byte_cell_stride;
       out.snode_iter_count = desc_opt->iter_count;
+      out.loop_iter_static = loop_iter_static_for_bound_expr;
       return true;
     }
     return false;

quadrants/transforms/static_adstack_analysis.h

@@ -94,6 +94,14 @@ struct StaticAdStackBoundExpr {
   uint32_t snode_byte_cell_stride{0};
   uint32_t snode_iter_count{0};
 
+  // Compile-time loop trip count of the task whose pushes this gate dominates. Set when the task's
+  // `range_for` has a constant `[begin, end)` (the same `static_bound` predicate the iteration-count check
+  // uses). Zero means runtime-bound (analyzer cannot resolve the trip count from IR alone). The runtime
+  // uses this as an upper bound when sizing the float adstack heap from the reducer's gate-passing-cell
+  // count: each loop iteration claims at most one row at the LCA-block, so the heap needs at most
+  // `loop_iter_static` rows regardless of how many cells of an oversized SNode the reducer counted.
+  uint32_t loop_iter_static{0};
+
   QD_IO_DEF(cmp_op,
             field_dtype_is_float,
             field_dtype_is_double,
@@ -107,7 +115,8 @@ struct StaticAdStackBoundExpr {
             snode_root_id,
             snode_byte_base_offset,
             snode_byte_cell_stride,
-            snode_iter_count);
+            snode_iter_count,
+            loop_iter_static);
 };
 
 // SNode descriptor info the analysis needs to capture an SNode-backed gate. The resolver returns `std::nullopt` when
@@ -154,8 +163,14 @@ struct StaticAdStackAnalysisResult {
 // addressing and the launchers skip the per-launch reducer dispatch + DtoH; see `CompileConfig::
 // ad_stack_sparse_threshold_bytes` for the user-facing knob (default 100 MiB; 0 forces capture, useful for tests
 // that pin the reducer-backed sizing path).
+// `task_range_is_original_loop`: true when `task_ir->{begin_value, end_value}` is the original user-loop trip
+// count, false when an upstream pass (e.g. `make_cpu_multithreaded_range_for` on CPU LLVM) has rewritten it
+// into a per-chunk subrange. The analyzer only fills `bound_expr.loop_iter_static` when this is true; on
+// rewritten loops the original trip count is no longer recoverable from `task_ir` and the runtime falls back to
+// the unclipped reducer count for that task.
 StaticAdStackAnalysisResult analyze_adstack_static_bounds(OffloadedStmt *task_ir,
                                                           const SNodeDescriptorResolver &snode_descriptor_resolver,
-                                                          std::size_t sparse_heap_threshold_bytes);
+                                                          std::size_t sparse_heap_threshold_bytes,
+                                                          bool task_range_is_original_loop = true);
 
 }  // namespace quadrants::lang