[Perf] Adstack max-reducer: skip recognizer on CPU; lift host-eval cap to UINT32_MAX since CPU max-reducer is serial; restrict GPU-only tests; layer doc

Author: duburcqa

docs/source/user_guide/autodiff.md

@@ -465,27 +465,23 @@ Quadrants computes that worst case at launch time - in this example, the max of
 
 ### Evaluation paths
 
-The compiler picks one of two evaluation paths to compute the maximum based on the bound expression's structure:
+On GPU backends the compiler picks one of two evaluation paths to compute the maximum based on the bound expression's structure; on CPU the sequential path is always taken since the runtime's CPU max-reducer is single-threaded and the parallel dispatch's per-launch setup would be pure overhead:
 
-- **Parallel:** the maximum is computed with a tiny parallel reduction kernel for efficiency. The reducer accepts a common subset of bound expressions:
+- **Parallel:** the maximum is computed with a tiny parallel reduction kernel on the GPU for efficiency. The reducer accepts a common subset of bound expressions:
   - **Integer ndarray or field read** up to 32 bits wide, indexed by literal constants or outer-loop variables: `arr[i, j]`, `field[i]`.
   - **Shape term**: `arr.shape[k]`.
   - **Literal integer constant**: `42`.
   - **Arithmetic combinator**: any `+`, `-`, `*`, `max` of the above.
-- **Sequential:** the fallback path, used whenever the parallel path doesn't support the bound expression. Quadrants walks the bound expression one outer-loop iteration at a time on a single thread; the adstack is sized identically, only the upfront cost differs. This path accepts everything the parallel path does, plus:
+- **Sequential:** the fallback path, used whenever the parallel path doesn't support the bound expression. Quadrants walks the bound expression one outer-loop iteration at a time on a single thread (host-side on CPU, single-thread on-device kernel on GPU); the adstack is sized identically, only the upfront cost differs. This path accepts everything the parallel path does, plus:
   - **Arithmetic-indexed read**: `arr[i // 2]`, `arr[i % 4]`.
   - **Indirect / nested read**: `arr1[arr2[i]]`, `my_field[arr[i]]`.
 
 ### Nested loops
 
-Quadrants supports arbitrarily nested loops. When the bound expression itself contains another enclosed loop whose own bound expression must be reduced first, the enclosing bound expression takes the parallel path only if every nested bound expression also fits the parallel-path grammar; otherwise it falls back to the sequential walk. This keeps the runtime from mixing parallel and sequential evaluators inside a single bound expression, which would otherwise force per-iteration kernel launches.
+Quadrants supports arbitrarily nested loops. When the bound expression itself contains another enclosed loop whose own bound expression must be reduced first, the enclosing bound expression takes the parallel path only if every nested bound expression is also supported by the parallel path; otherwise it falls back to the sequential walk. This keeps the runtime from mixing parallel and sequential evaluators inside a single bound expression, which would otherwise force per-iteration kernel launches.
 
 ### Sequential walk cap
 
-The sequential walk's outer loop is artificially capped at 2^24 = 16 777 216 iterations to keep both the walk time and the read-tracking memory bounded; past that the kernel raises `RuntimeError: ... iteration count ... exceeds the 16777216 guard`.
+The sequential walk's outer loop is artificially capped at 2^24 = 16 777 216 iterations on GPU backends to keep both the walk time and the read-tracking memory bounded; past that the kernel raises `RuntimeError: ... iteration count ... exceeds the 16777216 guard`. In the example above, the iteration count of the enclosed loop takes the sequential path because of the `i // 2` index, so it would raise at launch on GPU backends if `arr.shape[0] > (1 << 24)`.
 
-In the example above, the iteration count of the enclosed loop takes the sequential path because of the `i // 2` index. As such, it would raise at launch if `arr.shape[0] > (1 << 24)`.
-
-### Workaround
-
-Rewrite the bound expression to unlock the parallel path (e.g. precompute `bounds[i] = arr[i // 2]` into a persistent separate buffer, pass `bounds` in as an input, and use `for j in range(bounds[i]):`), or keep the outer loop count below 2^24.
+To circumvent this limitation, rewrite the bound expression to unlock the parallel path (e.g. precompute `bounds[i] = arr[i // 2]` into a persistent separate buffer, pass `bounds` in as an input, and use `for j in range(bounds[i]):`), or keep the outer loop count below 2^24.

quadrants/codegen/llvm/codegen_llvm.cpp

@@ -1994,10 +1994,18 @@ void TaskCodeGenLLVM::finalize_offloaded_task_function() {
     current_task->ad_stack.allocas = ad_stack_allocas_info_;
     current_task->ad_stack.size_exprs = ad_stack_size_exprs_;
     current_task->ad_stack.bound_expr = ad_stack_static_bound_expr_;
-    // recognize `MaxOverRange` nodes that the runtime can reduce in parallel via the dedicated max-reducer dispatch
-    // instead of letting the per-thread sizer enumerate. Indexing matches `ad_stack_size_exprs_` (same iteration order
-    // as the pre-scan above).
-    current_task->ad_stack.max_reducer_specs = recognize_adstack_max_reducer_specs(ad_stack_size_exprs_);
+    // Recognize `MaxOverRange` nodes the runtime can reduce in parallel via the dedicated max-reducer dispatch instead
+    // of letting the per-thread sizer enumerate. Indexing matches `ad_stack_size_exprs_` (same iteration order as the
+    // pre-scan above). Skip on CPU: `runtime_eval_adstack_max_reduce_serial` walks single-threaded just like the host
+    // evaluator's `MaxOverRange` loop in `program/adstack/eval.cpp`, so the dispatch's per-launch setup overhead
+    // (params blob encode, body bytecode encode, observation bookkeeping, JIT call) is pure cost without compute
+    // parallelism to offset it - measured ~28 % wallclock regression on the rigid-step CPU bench. The host evaluator
+    // handles every iteration count up to its own cap (raised to UINT32_MAX on CPU in `eval.cpp`) so above-cap shapes
+    // still resolve correctly. On CUDA / AMDGPU the parallel reducer is the whole point of the dispatch and the
+    // recognizer stays active.
+    if (!arch_is_cpu(compile_config.arch)) {
+      current_task->ad_stack.max_reducer_specs = recognize_adstack_max_reducer_specs(ad_stack_size_exprs_);
+    }
     // Snodes the task body mutates. Persisted on `OffloadedTask::snode_writes` so the LLVM
     // launcher can invalidate the per-task adstack metadata cache when a kernel that runs in
     // between mutated a SNode an enclosing `size_expr::FieldLoad` reads. Mirrors the SPIR-V

quadrants/program/adstack/eval.cpp

@@ -1,6 +1,8 @@
 #include "quadrants/program/adstack/eval.h"
 
 #include <algorithm>
+#include <climits>
+#include <cstdint>
 #include <functional>
 #include <limits>
 #include <unordered_map>
@@ -10,6 +12,8 @@
 #include "quadrants/common/logging.h"
 #include "quadrants/ir/snode.h"
 #include "quadrants/ir/type.h"
+#include "quadrants/program/program.h"
+#include "quadrants/rhi/arch.h"
 #include "quadrants/ir/type_factory.h"
 #include "quadrants/ir/type_utils.h"
 #include "quadrants/program/launch_context_builder.h"
@@ -248,11 +252,18 @@ int64_t evaluate_node(const SerializedSizeExpr &expr,
     case SizeExpr::Kind::MaxOverRange: {
       int64_t begin = evaluate_node(expr, node.operand_a, bound_vars, prog, ctx, reads);
       int64_t end = evaluate_node(expr, node.operand_b, bound_vars, prog, ctx, reads);
-      // Guard against pathological trip counts. The evaluator walks `[begin, end)` linearly and re-evaluates the
-      // body at every i; a range of several million would stall the launch hot path for seconds. Real reverse-mode
-      // trip counts sit well below this cap (a few hundred to a few thousand in practice); anything above is
-      // almost certainly a pre-pass grammar bug the user should file, and a clear QD_ERROR beats a silent hang.
-      constexpr int64_t kMaxOverRangeIterations = int64_t{1} << 24;
+      // Guard against pathological trip counts. The evaluator walks `[begin, end)` linearly and re-evaluates the body
+      // at every i; on Metal / Vulkan / CUDA / AMDGPU the recognizer-captured `MaxOverRange` shapes are dispatched in
+      // parallel by the max-reducer and substituted to a `Const` before the host evaluator walks the tree, so any
+      // shape that lands here above the `1<<24` cap is out-of-grammar and a clear QD_ERROR beats a silent hang. On
+      // CPU the recognizer is intentionally skipped (see the matching `arch_is_cpu` gate in
+      // `codegen/llvm/codegen_llvm.cpp::finalize_offloaded_task_function`: the CPU max-reducer walks single-threaded
+      // just like this loop, so the dispatch's per-launch setup overhead is pure cost), and the CPU max-reducer
+      // itself has no equivalent cap (`runtime_eval_adstack_max_reduce_serial` at
+      // `runtime/llvm/runtime_module/adstack_runtime.cpp:622-636` iterates `total_length` unconditionally), so this
+      // evaluator must match - lift the cap to `UINT32_MAX` on CPU so legitimate above-cap workloads still complete.
+      const bool prog_is_cpu = (prog != nullptr) && arch_is_cpu(prog->compile_config().arch);
+      const int64_t kMaxOverRangeIterations = prog_is_cpu ? int64_t{UINT32_MAX} : (int64_t{1} << 24);
       QD_ERROR_IF(end > begin && end - begin > kMaxOverRangeIterations,
                   "SerializedSizeExpr MaxOverRange iteration count {} exceeds the {} guard; refusing to enumerate. "
                   "Shrink the enclosing reverse-mode loop or restructure the `SizeExpr` source kernel.",

tests/python/test_adstack.py

@@ -4735,21 +4735,25 @@ def compute(x: qd.types.NDArray, selector: qd.types.NDArray, out: qd.types.NDArr
         "above_cap_arith_combine",
     ],
 )
-@test_utils.test(require=qd.extension.adstack, cfg_optimization=False)
+@test_utils.test(arch=[qd.cuda, qd.amdgpu, qd.vulkan, qd.metal], require=qd.extension.adstack, cfg_optimization=False)
 def test_max_reducer_pins_stride_for_oversized_axis(shape, body_kind):
-    # A reverse-mode kernel with a parallel-for over an arbitrarily large ndarray axis and an inner range-for bound to a
-    # recognizer-accepted trip-count expression sizes its adstack at launch time and computes the right gradient,
-    # without the per-task sizer's `1<<24` cap firing.
+    # A reverse-mode kernel with a parallel-for over an arbitrarily large ndarray axis and an inner range-for bound to
+    # a recognizer-accepted trip-count expression sizes its adstack at launch time and computes the right gradient,
+    # without the per-task sizer's `1<<24` cap firing. GPU only: the max-reducer dispatch is GPU-specific - the host
+    # evaluator handles equivalent shapes on CPU.
     #
     # Internal details: the kernel lowers to `MaxOverRange(0, a.shape[0], <body>)` in the per-stack `SizeExpr`.
     # `recognize_adstack_max_reducer_specs` captures the spec; the launcher dispatches the parallel max-reducer before
-    # the per-task sizer walks the tree; `substitute_precomputed_max_over_range` rewrites the captured `MaxOverRange` to
-    # `Const`. The above-cap variants place the only non-zero cell at `arr_np[-1] = N_X` so heap-stride correctness
+    # the per-task sizer walks the tree; `substitute_precomputed_max_over_range` rewrites the captured `MaxOverRange`
+    # to `Const`. The above-cap variants place the only non-zero cell at `arr_np[-1] = N_X` so heap-stride correctness
     # depends on the dispatch walking every element of the axis rather than relying on a partial host-eval walk. The
     # `shape_in_body` / `field_in_body` variants additionally pin that closed leaves (`ExternalTensorShape`,
     # `FieldLoad`) host-fold to `kConst` at encode time and never reach the device interpreter; `arith_combine`
     # exercises every binary combinator (`Add`, `Sub`, `Mul`, `Max`) and `Const` leaf in a single body expression that
-    # algebraically reduces to `a[i_e]`.
+    # algebraically reduces to `a[i_e]`. The CPU codegen gate lives in
+    # `codegen/llvm/codegen_llvm.cpp::finalize_offloaded_task_function`; the lifted host-eval cap lives in
+    # `program/adstack/eval.cpp::evaluate_node`. On CPU `_get_max_reducer_dispatch_count` stays at 0 (no dispatch
+    # fires), which is why this test pins it on GPU arches only.
     N_X = 4
     arr_np = np.zeros(shape, dtype=np.int32)
     arr_np[-1] = N_X
@@ -4813,17 +4817,18 @@ def compute(a: qd.types.ndarray(dtype=qd.i32, ndim=1)):
         assert x.grad[k] == pytest.approx(2 * 0.1, rel=1e-5)
 
 
-@test_utils.test(require=qd.extension.adstack, cfg_optimization=False)
+@test_utils.test(arch=[qd.cuda, qd.amdgpu, qd.vulkan, qd.metal], require=qd.extension.adstack, cfg_optimization=False)
 def test_max_reducer_dispatch_counts_advance_on_input_mutation():
-    # Pins the dispatch + cache invalidation pipeline. The first launch must fire at least one max-reducer dispatch (the
-    # kernel's `MaxOverRange(0, a.shape[0], a[var])` matches the recognizer grammar so the recognizer captures the spec;
-    # the launcher dispatches once and bumps `Program.max_reducer_dispatch_count`). A subsequent host mutation of the
-    # gating ndarray must bump `ndarray_data_gen` and force the next launch to re-dispatch, advancing the counter beyond
-    # its post-first-launch value. Steady-state cache short-circuit on an unchanged ndarray is backend-dependent (the
-    # CPU launcher's `set_host_accessible_ndarray_ptrs` path converts qd.ndarray reads to `kNone` semantics and
-    # `bump_writes_for_kernel_llvm` then bumps the gen on every read; the SPIR-V launchers preserve the qd.ndarray
-    # dev-alloc-type and only bump on writes), so this test asserts only the mutation-triggers-redispatch contract that
-    # holds uniformly.
+    # Pins the dispatch + cache invalidation pipeline. The first launch must fire at least one max-reducer dispatch
+    # (the kernel's `MaxOverRange(0, a.shape[0], a[var])` matches the recognizer grammar so the recognizer captures
+    # the spec; the launcher dispatches once and bumps `Program.max_reducer_dispatch_count`). A subsequent host
+    # mutation of the gating ndarray must bump `ndarray_data_gen` and force the next launch to re-dispatch, advancing
+    # the counter beyond its post-first-launch value. Steady-state cache short-circuit on an unchanged ndarray is
+    # backend-dependent (the CPU launcher's `set_host_accessible_ndarray_ptrs` path converts qd.ndarray reads to
+    # `kNone` semantics and `bump_writes_for_kernel_llvm` then bumps the gen on every read; the SPIR-V launchers
+    # preserve the qd.ndarray dev-alloc-type and only bump on writes), so this test asserts only the
+    # mutation-triggers-redispatch contract that holds uniformly. GPU only: the max-reducer dispatch is GPU-specific -
+    # the host evaluator handles equivalent shapes on CPU.
     N = 4
 
     x = qd.field(qd.f32, shape=(N,), needs_grad=True)
@@ -4921,13 +4926,14 @@ def compute():
         "arr_bv_indexed_by_field_load",
     ],
 )
-@test_utils.test(require=qd.extension.adstack)
+@test_utils.test(arch=[qd.cuda, qd.amdgpu, qd.vulkan, qd.metal], require=qd.extension.adstack)
 def test_max_reducer_field_load_bound_var_dispatch(body_kind):
     # A reverse-mode kernel whose inner range-for trip count reads a `qd.field` indexed by the outer chain variable
     # captures via the parallel max-reducer dispatch and produces the analytical gradient. The body-shape
     # parametrization exercises every supported composition: bound-var FieldLoad on its own, mixed with bound-var ETR
     # via `Add` / `Max`, combined with `Const` / arithmetic, and the nested-load worst-case form (`field[field[i]]` /
-    # `arr[field[i]]`).
+    # `arr[field[i]]`). GPU only: the max-reducer dispatch is GPU-specific - the host evaluator handles equivalent
+    # shapes on CPU.
     #
     # Internal details: each variant lowers to `MaxOverRange(0, M, body)` where `body` is bound-var-indexed
     # `FieldLoad(field_a, [bound_var])` or a recognizer-accepted composition that includes one. The relaxed
@@ -5028,10 +5034,11 @@ def compute(a: qd.types.ndarray(dtype=qd.i32, ndim=1)):
         assert x.grad[k] == pytest.approx(expected, rel=1e-5)
 
 
-@test_utils.test(require=qd.extension.adstack)
+@test_utils.test(arch=[qd.cuda, qd.amdgpu, qd.vulkan, qd.metal], require=qd.extension.adstack)
 def test_max_reducer_field_load_bound_var_cache_invalidates_on_snode_mutation():
     # A reverse-mode kernel whose inner trip count reads a `qd.field` indexed by the outer chain variable redispatches
-    # the max-reducer when the gating field is mutated between launches.
+    # the max-reducer when the gating field is mutated between launches. GPU only: the max-reducer dispatch is
+    # GPU-specific - the host evaluator handles equivalent shapes on CPU.
     #
     # Internal details: the encoder emits a `kFieldLoad` device node and pushes a `FieldLoadObs` carrying the snode id
     # and the live `snode_write_gen` snapshot. On the second launch's `try_max_reducer_cache_hit`,
@@ -5087,10 +5094,11 @@ def compute():
     assert prog._get_max_reducer_dispatch_count() > pre_mutation
 
 
-@test_utils.test(require=qd.extension.adstack, cfg_optimization=False)
+@test_utils.test(arch=[qd.cuda, qd.amdgpu, qd.vulkan, qd.metal], require=qd.extension.adstack, cfg_optimization=False)
 def test_above_cap_out_of_grammar_kernel_raises():
     # A reverse-mode kernel whose inner `range(...)` trip count is bound to an out-of-grammar `MaxOverRange` body and
     # whose iteration count exceeds the `1<<24` adstack-sizer cap surfaces a `QuadrantsAssertionError` at `qd.sync()`.
+    # GPU only: on CPU the host-eval cap is lifted to UINT32_MAX, so a shape of `(1<<24)+1` resolves without raising.
     #
     # Internal details: the recognizer's body grammar accepts only `Const / ExternalTensorRead / Add / Sub / Mul / Max
     # / ExternalTensorShape / FieldLoad(literal-or-bound-var indices)`, and `max_reducer_body_is_recognizable` further