From 11fe2c18daf43bc1ad3e02f6c3678b0d84579f12 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 20 Jul 2026 15:37:12 +0200 Subject: [PATCH 1/4] [Bug] Propagate the adstack size-expr onto scalarized tensor components. `determine_ad_stack_size` runs before the scalarize pass. For an adaptive tensor-typed adstack it records a symbolic `size_expr` and leaves `max_size` at the seed value of 1 for the runtime to overwrite from that tree. When scalarize splits the tensor stack into one scalar `AdStackAllocaStmt` per component it copied `max_size` but dropped `size_expr`, so the launcher found no tree and sized the component heap to the seed. A scalar component stack pushed inside a runtime-bounded loop then overflowed as soon as the loop ran more than once - a loud assertion on SPIR-V, silently corrupted gradients on a `__debug__`- disabled build. Clone `size_expr` onto every scalar component. The components are pushed and popped in lockstep with the tensor stack, so they share its depth exactly. Fixes #793. --- quadrants/transforms/scalarize.cpp | 9 ++++++ tests/python/test_adstack.py | 49 ++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/quadrants/transforms/scalarize.cpp b/quadrants/transforms/scalarize.cpp index 259c63b45c..3cf97e73a7 100644 --- a/quadrants/transforms/scalarize.cpp +++ b/quadrants/transforms/scalarize.cpp @@ -719,6 +719,15 @@ class Scalarize : public BasicStmtVisitor { auto scalar_ad_stack = std::make_unique(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)); diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index 2fd161af0c..3a5777130e 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -5198,3 +5198,52 @@ def compute(a: qd.types.ndarray(dtype=qd.i64, ndim=1)): x.grad[i] = 0.0 compute.grad(a) qd.sync() + + +@test_utils.test(require=qd.extension.adstack, cfg_optimization=False) +def test_adstack_scalarized_tensor_component_loop_trip_grad_correct(): + # A reverse-mode kernel with a tensor-valued loop-carried variable divided by a scalar inside a field-bounded + # loop must size each scalar component of the adstack for the full trip count. `determine_ad_stack_size` records + # a `size_expr` on the tensor-typed 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 + # replicated per-lane gradients on a `__debug__`-disabled build). The vector-by-scalar division keeps the value + # tensor-typed past the pre-pass so the split happens after sizing. + n_iter_np = np.array([3], dtype=np.int32) + denom_val = 2.0 + x_np = np.array([0.3, 0.7], dtype=np.float32) + n_env = x_np.size + + n_iter = qd.ndarray(qd.i32, shape=(1,)) + n_iter.from_numpy(n_iter_np) + denom = qd.ndarray(qd.f32, shape=(1,)) + denom.from_numpy(np.array([denom_val], dtype=np.float32)) + x = qd.ndarray(qd.f32, shape=(n_env,), needs_grad=True) + out = qd.ndarray(qd.f32, shape=(n_env,), needs_grad=True) + + @qd.kernel + def compute(x: qd.types.ndarray(), out: qd.types.ndarray(), n_iter: qd.types.ndarray(), denom: qd.types.ndarray()): + for i in range(n_env): + q = qd.Vector([x[i], x[i] * 0.5], dt=qd.f32) + s = denom[0] + for _ in range(n_iter[0]): + 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] + + x.from_numpy(x_np) + compute(x, out, n_iter, denom) + out.grad.from_numpy(np.ones(n_env, dtype=np.float32)) + x.grad.fill(0.0) + compute.grad(x, out, n_iter, denom) + + # 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(int(n_iter_np[0])): + 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) From fa19bf5536914b12985541d8dead835ab3271bd2 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 20 Jul 2026 16:11:50 +0200 Subject: [PATCH 2/4] fixup! [Bug] Propagate the adstack size-expr onto scalarized tensor components. --- tests/python/test_adstack.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index 3a5777130e..27ca376aa8 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -5202,13 +5202,11 @@ def compute(a: qd.types.ndarray(dtype=qd.i64, ndim=1)): @test_utils.test(require=qd.extension.adstack, cfg_optimization=False) def test_adstack_scalarized_tensor_component_loop_trip_grad_correct(): - # A reverse-mode kernel with a tensor-valued loop-carried variable divided by a scalar inside a field-bounded - # loop must size each scalar component of the adstack for the full trip count. `determine_ad_stack_size` records - # a `size_expr` on the tensor-typed 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 - # replicated per-lane gradients on a `__debug__`-disabled build). The vector-by-scalar division keeps the value - # tensor-typed past the pre-pass so the split happens after sizing. + # 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_np = np.array([3], dtype=np.int32) denom_val = 2.0 x_np = np.array([0.3, 0.7], dtype=np.float32) From 3af24e06eb6e75fffa06dc9d9b1703cfc4820506 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 20 Jul 2026 18:09:25 +0200 Subject: [PATCH 3/4] fixup! [Bug] Propagate the adstack size-expr onto scalarized tensor components. --- tests/python/test_adstack.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index 27ca376aa8..1fb5fd87ed 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -5207,39 +5207,39 @@ def test_adstack_scalarized_tensor_component_loop_trip_grad_correct(): # 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_np = np.array([3], dtype=np.int32) + 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.ndarray(qd.i32, shape=(1,)) - n_iter.from_numpy(n_iter_np) - denom = qd.ndarray(qd.f32, shape=(1,)) - denom.from_numpy(np.array([denom_val], dtype=np.float32)) - x = qd.ndarray(qd.f32, shape=(n_env,), needs_grad=True) - out = qd.ndarray(qd.f32, shape=(n_env,), needs_grad=True) + 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(x: qd.types.ndarray(), out: qd.types.ndarray(), n_iter: qd.types.ndarray(), denom: qd.types.ndarray()): + def compute(): for i in range(n_env): q = qd.Vector([x[i], x[i] * 0.5], dt=qd.f32) - s = denom[0] - for _ in range(n_iter[0]): + 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(x, out, n_iter, denom) + compute() out.grad.from_numpy(np.ones(n_env, dtype=np.float32)) x.grad.fill(0.0) - compute.grad(x, out, n_iter, denom) + 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(int(n_iter_np[0])): + for _ in range(n_iter_val): jac = mat @ jac expected = float(jac.sum()) grad = x.grad.to_numpy() From 0d067423f43194a775755ad0c69fdc699046b7d7 Mon Sep 17 00:00:00 2001 From: Alexis DUBURCQ Date: Tue, 21 Jul 2026 10:32:04 +0200 Subject: [PATCH 4/4] Update test_adstack.py --- tests/python/test_adstack.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index 464adf91a4..eb3eb5480f 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -5245,7 +5245,8 @@ def compute(): 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