[AutoDiff] Adstack heap: clip reducer count by per-task loop trip count (compile-time and SizeExpr-evaluated)

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/ir/statements.h

@@ -1353,6 +1353,15 @@ class OffloadedStmt : public Stmt {
   std::size_t bls_size{0};
   MemoryAccessOptions mem_access_opt;
 
+  // Pre-chunking loop trip-count `SizeExpr` captured by `determine_ad_stack_size`. Set on adstack-bearing
+  // range-for tasks before `make_cpu_multithreaded_range_for` rewrites the loop into per-thread chunks, so the
+  // SizeExpr still describes the original user-loop bound (handles both compile-time constants and
+  // runtime-bounded shapes like `for j in range(field[i])` via the same `FieldLoad` / `ExternalTensorRead` /
+  // `MaxOverRange` grammar `compute_bounded_adstack_size` already uses for per-thread stack sizing). Read by
+  // `analyze_adstack_static_bounds` at codegen time, serialised into `StaticAdStackBoundExpr::loop_iter_size_expr`,
+  // and evaluated at launch time as the per-task row-claim upper bound for the float-heap clip.
+  std::shared_ptr<SizeExpr> pre_chunk_loop_trip_count_expr;
+
   OffloadedStmt(TaskType task_type, Arch arch, Kernel *kernel);
 
   std::string task_name() const;

quadrants/runtime/amdgpu/kernel_launcher.cpp

@@ -103,7 +103,7 @@ void KernelLauncher::launch_offloaded_tasks(LaunchContextBuilder &ctx,
                                                       context_pointer);
         // Size the float heap from the published gate-passing count (DtoH'd per task). Mirrors the CUDA / CPU
         // launcher post-reducer sizing.
-        executor->ensure_per_task_float_heap_post_reducer(task_index, task.ad_stack, n_threads_amdgpu);
+        executor->ensure_per_task_float_heap_post_reducer(task_index, task.ad_stack, n_threads_amdgpu, &ctx);
       }
     }
     ++task_index;

quadrants/runtime/cpu/kernel_launcher.cpp

@@ -71,7 +71,7 @@ void KernelLauncher::launch_offloaded_tasks(LaunchContextBuilder &ctx,
         // allocas (in tasks with a captured `bound_expr`) address through `heap_float + row_id_var * stride_float +
         // float_offset`; sizing the heap at `count * stride_float` instead of the dispatched-threads worst case is
         // where the actual memory savings on sparse-grid workloads come from.
-        executor->ensure_per_task_float_heap_post_reducer(i, ad_stacks[i], num_threads_per_task[i]);
+        executor->ensure_per_task_float_heap_post_reducer(i, ad_stacks[i], num_threads_per_task[i], &ctx);
       }
     }
     task_funcs[i](&ctx.get_context());

quadrants/runtime/cuda/kernel_launcher.cpp

@@ -129,7 +129,7 @@ void KernelLauncher::launch_offloaded_tasks(LaunchContextBuilder &ctx,
         // Size the float heap from the published gate-passing count (DtoH'd per task). Mirrors the CPU launcher's
         // post-reducer sizing call - this is what shrinks the float slab to `count * stride_float` instead of the
         // dispatched-threads worst case on sparse-grid workloads.
-        executor->ensure_per_task_float_heap_post_reducer(task_index, task.ad_stack, n);
+        executor->ensure_per_task_float_heap_post_reducer(task_index, task.ad_stack, n, &ctx);
       }
     }
     ++task_index;

quadrants/runtime/gfx/runtime.cpp

@@ -685,6 +685,45 @@ 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 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 `trip_count` 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. Two trip-count sources, both gated by the
+          // `dispatched_threads` ceiling so a `dynamic_gpu_range_for` that exceeds the SPIR-V dispatch cap and
+          // serialises iterations across threads (each thread reaches the LCA-block multiple times) does not
+          // accidentally undersize the heap:
+          //   * `loop_iter_static`: compile-time-known trip count (`for i in range(N)` with N a constant).
+          //     Cheap integer compare.
+          //   * `loop_iter_size_expr`: runtime-evaluated `SizeExpr` for shapes the static field cannot cover
+          //     (`for j in range(field[i])`, `for k in range(arr.shape[axis])`, two-arg ranges over either,
+          //     plus the multi-axis ndrange and bound-cast variants the SizeExpr grammar already supports).
+          //     Evaluator cost is the same per-launch host walk the per-thread `ad_stack_size` already runs;
+          //     returns -1 on shapes that aren't host-resolvable, in which case we skip the clip from this
+          //     source.
+          if (attribs.ad_stack.bound_expr.has_value()) {
+            const auto &bound_expr = *attribs.ad_stack.bound_expr;
+            if (bound_expr.loop_iter_static > 0) {
+              // Compile-time trip count: integer compare, no per-launch eval cost. Constant `SizeExpr` shapes
+              // are already collapsed into this field by the analyzer so they short-circuit the runtime eval
+              // below.
+              const size_t loop_iter_static = static_cast<size_t>(bound_expr.loop_iter_static);
+              if (loop_iter_static <= dispatched_threads) {
+                effective_rows = std::min<size_t>(effective_rows, loop_iter_static);
+              }
+            } else if (!bound_expr.loop_iter_size_expr.nodes.empty() && program_impl_ != nullptr &&
+                       program_impl_->program != nullptr) {
+              // Runtime-bounded clip: only evaluated when the static field is unset (the analyzer leaves it
+              // 0 for shapes the compile-time path cannot cover, e.g. `for j in range(field[i])` /
+              // `for k in range(arr.shape[axis])`). Cost = one tree walk per launch, same path the
+              // per-thread `ad_stack_size` host-eval already runs.
+              const int64_t evaluated =
+                  evaluate_adstack_size_expr(bound_expr.loop_iter_size_expr, program_impl_->program, &host_ctx);
+              if (evaluated > 0 && static_cast<size_t>(evaluated) <= dispatched_threads) {
+                effective_rows = std::min<size_t>(effective_rows, static_cast<size_t>(evaluated));
+              }
+            }
+          }
         } 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

@@ -404,7 +404,8 @@ void LlvmRuntimeExecutor::ensure_adstack_heap_int(std::size_t needed_bytes) {
 
 void LlvmRuntimeExecutor::ensure_per_task_float_heap_post_reducer(std::size_t task_index,
                                                                   const AdStackSizingInfo &ad_stack,
-                                                                  std::size_t num_threads) {
+                                                                  std::size_t num_threads,
+                                                                  LaunchContextBuilder *ctx) {
   // Skip when the task has no float heap need (no f32 allocas, or analysis didn't capture a gate so we wouldn't have
   // routed it through the lazy float path on the codegen side).
   if (!ad_stack.bound_expr.has_value() || ad_stack.per_thread_stride_float == 0) {
@@ -436,9 +437,37 @@ 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()) {
+    if (ad_stack.bound_expr->loop_iter_static > 0) {
+      // Compile-time trip count: integer compare, no per-launch eval cost. Constant `SizeExpr` shapes are
+      // already collapsed into this field by the analyzer so they short-circuit the runtime eval below.
+      effective_rows =
+          std::min<std::size_t>(effective_rows, static_cast<std::size_t>(ad_stack.bound_expr->loop_iter_static));
+    } else if (!ad_stack.bound_expr->loop_iter_size_expr.nodes.empty()) {
+      // Runtime-bounded clip: evaluate the captured trip-count `SizeExpr` only when the static field is unset
+      // (the analyzer leaves `loop_iter_static == 0` for shapes the compile-time path cannot cover, e.g.
+      // `for j in range(field[i])` / `for k in range(arr.shape[axis])`). Cost = one tree walk per launch,
+      // dominated by host scalar reads through `SNodeRwAccessorsBank` on `FieldLoad` / `ExternalTensorRead`
+      // nodes (CPU: a memory load; CUDA / AMDGPU: a 4-8 byte DtoH). The evaluator returns -1 when the tree
+      // references state that is not host-resolvable from `ctx`; in that case we leave `effective_rows`
+      // unclipped from this source.
+      Program *prog = (program_impl_ != nullptr) ? program_impl_->program : nullptr;
+      if (ctx != nullptr && prog != nullptr) {
+        const int64_t evaluated = evaluate_adstack_size_expr(ad_stack.bound_expr->loop_iter_size_expr, prog, ctx);
+        if (evaluated > 0) {
+          effective_rows = std::min<std::size_t>(effective_rows, static_cast<std::size_t>(evaluated));
+        }
+      }
+    }
+  }
   // 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/runtime/llvm/llvm_runtime_executor.h

@@ -170,7 +170,8 @@ class LlvmRuntimeExecutor {
   // is sized exactly to the reducer's count instead of the dispatched-threads worst case.
   void ensure_per_task_float_heap_post_reducer(std::size_t task_index,
                                                const AdStackSizingInfo &ad_stack,
-                                               std::size_t num_threads);
+                                               std::size_t num_threads,
+                                               LaunchContextBuilder *ctx);
 
   // Return (and lazily cache) the device pointer to `runtime->temporaries`, the global temporary buffer backing
   // `GlobalTemporaryStmt` loads and stores. GPU kernel launchers use this to read back dynamic range_for bounds (begin

quadrants/transforms/determine_ad_stack_size.cpp

@@ -1393,6 +1393,47 @@ bool determine_ad_stack_size(IRNode *root, const CompileConfig &config) {
         " or a shape the pre-pass already recognises, or extend the `SizeExpr` grammar to cover this shape.",
         alloca->get_last_tb(), dump_hint);
   }
+  // Build a per-task `SizeExpr` for each adstack-bearing parallel range-for offload's loop trip count, captured
+  // here BEFORE `make_cpu_multithreaded_range_for` rewrites the user loop into chunks. The same `SizeExpr`
+  // grammar `compute_bounded_adstack_size` uses for per-thread stack sizing applies: integer constants, scalar
+  // i32 / i64 field reads, ndarray shape axes, ndarray element reads at constant or loop indices, plus the
+  // arithmetic / max-over-range combinators. The runtime float-heap clip evaluates this expression at launch
+  // and uses it as an upper bound on row claims (each loop iteration claims at most one row at the LCA-block,
+  // so the heap needs at most `evaluate(loop_trip_count_expr)` rows regardless of how many cells of an
+  // oversized gating SNode the reducer counted). Skipped for non-range-for offloads (struct-for / mesh-for /
+  // serial / listgen / gc) because their iteration model does not fit the simple `[begin, end)` walk the
+  // expression captures; those tasks fall back to the unclipped reducer count.
+  auto offloaded_tasks = irpass::analysis::gather_statements(root, [&](Stmt *s) { return s->is<OffloadedStmt>(); });
+  for (Stmt *s : offloaded_tasks) {
+    auto *task = s->as<OffloadedStmt>();
+    if (task->task_type != OffloadedTaskType::range_for) {
+      continue;
+    }
+    bool has_adstack = false;
+    irpass::analysis::gather_statements(task, [&](Stmt *inner) {
+      if (inner->is<AdStackAllocaStmt>()) {
+        has_adstack = true;
+      }
+      return false;
+    });
+    if (!has_adstack) {
+      continue;
+    }
+    int32_t var_id_counter = 0;
+    auto end_e = resolve_loop_end(task, root, &var_id_counter);
+    if (!end_e) {
+      continue;
+    }
+    auto begin_e = resolve_loop_begin_lower_bound(task);
+    auto trip = expr_sub(std::move(end_e), std::move(begin_e));
+    if (!trip) {
+      continue;
+    }
+    if (trip->kind == SizeExpr::Kind::Const && trip->const_value < 0) {
+      trip = SizeExpr::make_const(0);
+    }
+    task->pre_chunk_loop_trip_count_expr = std::shared_ptr<SizeExpr>(std::move(trip));
+  }
   return true;
 }