[Bug] Propagate the adstack size-expr onto scalarized tensor components.

quadrants/transforms/scalarize.cpp

@@ -719,6 +719,15 @@ class Scalarize : public BasicStmtVisitor {
         auto scalar_ad_stack = std::make_unique<AdStackAllocaStmt>(element_type, stmt->max_size);
         scalar_ad_stack->ret_type = element_type;
         scalar_ad_stack->ret_type.set_is_pointer(true);
+        // Carry the per-launch size bound onto each scalar component. `determine_ad_stack_size` runs before this
+        // pass and, for an adaptive tensor-typed stack, records a `size_expr` while leaving `max_size` at the seed
+        // value of 1 for the runtime to overwrite. Every component is pushed and popped in lockstep with the tensor
+        // stack, so they share its depth exactly; without cloning `size_expr` the component keeps only the seed
+        // `max_size`, the launcher finds no symbolic tree and sizes the heap to that seed, and a scalar stack pushed
+        // inside a runtime-bounded loop overflows once the loop runs more than once.
+        if (stmt->size_expr) {
+          scalar_ad_stack->size_expr = stmt->size_expr->clone();
+        }
 
         scalarized_ad_stack.push_back(scalar_ad_stack.get());
         delayed_modifier_.insert_before(stmt, std::move(scalar_ad_stack));

tests/python/test_adstack.py

@@ -5200,6 +5200,53 @@ def compute(a: qd.types.ndarray(dtype=qd.i64, ndim=1)):
         qd.sync()
 
 
+@test_utils.test(require=qd.extension.adstack, cfg_optimization=False)
+def test_adstack_scalarized_tensor_component_loop_trip_grad_correct():
+    # A tensor-valued loop-carried variable divided by a scalar (`q / q.norm()`) stays tensor-typed past
+    # `determine_ad_stack_size`, which records a `size_expr` on the tensor adstack while leaving `max_size` at the
+    # seed 1. The later scalarize pass splits the tensor stack into per-component scalar stacks; before the fix the
+    # split dropped `size_expr`, so each component kept only the seed and overflowed once the loop ran more than once
+    # (loud on SPIR-V, silently per-lane-replicated gradients on a `__debug__`-disabled build).
+    n_iter_val = 3
+    denom_val = 2.0
+    x_np = np.array([0.3, 0.7], dtype=np.float32)
+    n_env = x_np.size
+
+    n_iter = qd.field(qd.i32, shape=())
+    denom = qd.field(qd.f32, shape=())
+    x = qd.field(qd.f32, shape=n_env, needs_grad=True)
+    out = qd.field(qd.f32, shape=n_env, needs_grad=True)
+
+    @qd.kernel
+    def compute():
+        for i in range(n_env):
+            q = qd.Vector([x[i], x[i] * 0.5], dt=qd.f32)
+            s = denom[None]
+            for _ in range(n_iter[None]):
+                q = qd.Vector([q[0] + q[1] * 0.1, q[1] + q[0] * 0.2], dt=qd.f32)
+                q = q / s
+            out[i] = q[0] + q[1]
+
+    n_iter[None] = n_iter_val
+    denom[None] = denom_val
+    x.from_numpy(x_np)
+    compute()
+    out.grad.from_numpy(np.ones(n_env, dtype=np.float32))
+    x.grad.fill(0.0)
+    compute.grad()
+
+    # The recurrence q <- (M @ q) / s with M = [[1, 0.1], [0.2, 1]] is linear, so d(out)/d(x[i]) is the same for
+    # every env: sum of the columns of (M / s)^n_iter applied to the initial jacobian (1, 0.5).
+    mat = np.array([[1.0, 0.1], [0.2, 1.0]], dtype=np.float64) / denom_val
+    jac = np.array([1.0, 0.5], dtype=np.float64)
+    for _ in range(n_iter_val):
+        jac = mat @ jac
+    expected = float(jac.sum())
+    grad = x.grad.to_numpy()
+    for i in range(n_env):
+        assert grad[i] == pytest.approx(expected, rel=1e-5)
+
+
 @test_utils.test(require=qd.extension.adstack, cfg_optimization=False)
 def test_adstack_offset_array_difference_loop_trip_grad_correct():
     # Inner loop trip count is the difference of two adjacent offset-array reads (`starts[i + 1] - starts[i]`, the