[AutoDiff] Improve diagnosis of unsupported reverse-mode AD patterns

[duburcqa]

Improve diagnosis of unsupported reverse-mode AD patterns

Two diagnostic fixes for previously-silent classes of wrong-gradient bugs in reverse-mode autodiff. (1) An AST-level fix that propagates the caller's loop_depth across @qd.func boundaries so the existing dynamic-range guard fires on @qd.func-routed shapes the same way it already does on inline shapes. (2) A new IR-level guard in auto_diff.cpp that catches cross-iteration read-after-write on a needs_grad global field at the offload-level loop, a shape neither the AST guard nor the adstack handles.

TL;DR

The kernel below silently produced x.grad == 1 instead of the correct x.grad == n in both adstack settings:

@qd.func
def body():
    for i in range(n_iter[None]):
        if i == 0:
            out[i] = x[None]
        else:
            out[i] = out[i - 1] + x[None]

@qd.kernel
def fwd():
    body()

It now raises a clear RuntimeError from auto_diff.cpp pointing at the two valid workarounds (wrap the call in an outer for-loop, or restructure to remove the cross-iteration RAW). The companion @qd.func-bypass class - a non-static range(...) nested inside an outer for-loop in the kernel but routed through a @qd.func body - now also raises the same QuadrantsCompilationError the inline shape already raised.

Why

Two distinct silent-wrong-gradient classes that previously slipped through every guard.

Class 1 - @qd.func bypassed the existing AST guard

The upstream loop_depth > 0 check at ASTTransformer.build_For fires when a non-static range(...) is AST-nested inside another for-loop. The check uses ctx.loop_depth, which is local to a single AST transformer context and resets to zero every time the transformer enters a @qd.func body, so:

# Inline: AST loop_depth == 1 at the inner range -> existing check fires (without adstack).
@qd.kernel
def k():
    for i in x:
        for _ in range(n_iter):
            v = v * 0.95 + 0.01

# Routed through @qd.func: func ctx loop_depth resets to 0 -> existing check did NOT fire,
# kernel compiled and silently produced wrong gradients.
@qd.func
def inner(...):
    for _ in range(n_iter):
        v = v * 0.95 + 0.01

@qd.kernel
def k():
    for i in x:
        y[None] += inner(x[i])

This shape is pinned by the new wrap_inner_in_func=True parametrization on test_adstack_basic_gradient_negative; before this PR it would have passed with no diagnostic and a wrong adjoint.

Class 2 - cross-iteration RAW on a needs_grad global field at the offload-level loop

Reverse-mode AD relies on either (a) the adstack to replay loop-carried local primals, or (b) the kernel-internal IB analysis at quadrants/transforms/auto_diff.cpp to chain adjoints through AST-nested loops. Neither mechanism applies when the cross-iteration dependency is on a global field at the offload-level loop (i.e., the kernel's outermost loop, not nested under another loop):

• AdStackAllocaJudger only promotes cross-iteration dependencies on local AllocaStmts. Cross-iter through GlobalStoreStmt / GlobalLoadStmt on a has_adjoint() SNode is never promoted.
• MakeAdjoint's reverse pass for the offload-level loop body emits per-iteration adjoints; the cross-iter chain out[i-1] -> out[i] is not threaded into that pass.

So a kernel that writes out[i] = out[i-1] + x[None] at the offload-level loop body silently produced wrong gradients in both adstack settings, with no diagnostic. The TL;DR repro is the minimal expression of this shape.

Surface API

No new user-visible API. Two diagnostics fire at compile time when the kernel hits an unsupported reverse-mode shape:

1. Cannot use non static range in Backwards mode (raised as QuadrantsCompilationError) - the upstream AST guard at ASTTransformer.build_For. Now also fires for @qd.func-routed shapes thanks to the caller-loop_depth propagation, so inline and func-routed forms of the diagnostic behave identically.
2. Cross-iteration read-after-write on the needs_gradglobal field at the offload-level loop body is not supported in reverse-mode autodiff: ... (raised as RuntimeError from QD_ERROR) - new IR-level diagnostic. Fires in both adstack settings because adstack does not handle this shape regardless of the flag. Multi-axis sibling-component access (out[i, 0] = out[i, 1] * 2) and in-place accumulation (out[i] = out[i] + x) are exempted via the per-axis rule described below.

test_adstack_basic_gradient_negative accepts both exception types since the same shape is now reachable through either guard depending on which fires first; everything else is unchanged.

Mechanism

1. AST-level: caller-loop_depth propagation

Three Python files - _func_base.py, ast_transformer_utils.py, call_transformer.py - thread the caller's ctx.loop_depth across @qd.func invocations:

• ASTTransformerGlobalContext carries a caller_loop_depth: int field (the snapshot site).
• CallTransformer.build_Call writes ctx.loop_depth into the global context's caller_loop_depth before invoking the callable, and restores the previous value in a finally (so chained calls and exceptions both restore correctly).
• _func_base.py seeds the freshly-created ASTTransformerFuncContext with ctx.loop_depth = global_context.caller_loop_depth for non-kernel contexts, so a non-static range(...) inside the func body sees the same loop_depth > 0 the caller is already inside. Kernel contexts always begin at depth 0.

The upstream loop_depth > 0 check at ASTTransformer.build_For is otherwise untouched.

2. IR-level: OffloadLevelGlobalCrossIterRAWChecker in auto_diff.cpp

A BasicStmtVisitor gated on autodiff_mode == kReverse, run before RegulateTensorTypedStatements so the diagnostic still points at the user-shaped IR (post-frontend-lowering, but before any AD rewrites).

Walks the body of every direct-child RangeForStmt of the kernel root once. Inside that body collects, per has_adjoint() SNode:

• every GlobalStoreStmt destination GlobalPtrStmt (also drilling through MatrixPtrStmt to its origin),
• every GlobalLoadStmt source GlobalPtrStmt (same).

inside_nested_ is set on entry into any nested RangeForStmt / StructForStmt / MeshForStmt / WhileStmt body (the four overrides go through a shared walk_nested_loop_body helper) so loads / stores inside nested loops are excluded - the AD pass's IB analysis handles cross-iter through global fields correctly when the cross-iter range is AST-nested under another loop, and the adstack handles cross-iter through local allocas.

After the walk, for each (snode, store, load) triple where the SNode is has_adjoint():

• skip if the store index does not transitively reference a LoopIndexStmt (the store writes the same slot every iteration; nothing iterates),
• skip if the load index does not transitively reference a LoopIndexStmt (the load reads from a disjoint iteration slice the loop never writes to),
• otherwise compare the indices per-axis via indices_have_no_cross_iter_dependency:
  • axes where neither index references LoopIndexStmt are iter-independent and ignored regardless of value (allows sibling-component access on a constant axis like out[i, 0] vs out[i, 1]),
  • axes where both indices reference LoopIndexStmt must be the same SSA Stmt* (allows in-place out[i] = out[i] + x; flags out[i] = out[i-1] + x because the load axis is BinaryOpStmt(LoopIndexStmt, Sub, ConstStmt) and the store axis is the bare LoopIndexStmt),
  • axes where exactly one references LoopIndexStmt are conservatively flagged (one iter-dependent, one constant on the same axis - cannot prove statically whether the constant aliases an iteration).

StructFor offloads are out of scope for run - their loop variable is an SNode index produced by the runtime, not a user-controlled integer arithmeticed against itself, so the out[i-1] shape this guard targets does not arise.

Tests

Existing test patches

• tests/python/test_adstack.py::test_adstack_basic_gradient_negative - new wrap_inner_in_func parametrization axis pins the @qd.func-bypass class. The pytest.raises accepts both qd.QuadrantsCompilationError (from the upstream AST guard) and RuntimeError (from QD_ERROR via auto_diff.cpp) since either may surface depending on the shape.

New tests in tests/python/test_adstack.py

A single _run_reverse_global_cross_iter(has_outer_for, in_func, load_kind, expect_raise) factory using the range(1) if qd.static(has_outer_for) else qd.static(range(1)) inline-if pattern (already used by _run_unary_loop_carried in this file) covers eight reverse-mode shapes that read and write the same needs_grad global field from inside a non-static range:

has_outer_for	in_func	load_kind	Shape
False	False	cross_iter	non-static range at the kernel top; cross-iter inline
False	True	cross_iter	non-static range at the top of a @qd.func body called from kernel
True	False	cross_iter	non-static range nested under outer range(1); cross-iter inline
True	True	cross_iter	non-static range nested under outer range(1); cross-iter via @qd.func
False	False	constant_outside_iter_range	kernel-top range; load reads out[n] (slot strictly outside the iteration range)
False	True	constant_outside_iter_range	as above, body in @qd.func
True	False	constant_outside_iter_range	nested under outer range(1); load reads out[n]
True	True	constant_outside_iter_range	as above, body in @qd.func

Two parametrized test functions (_with_adstack and _without_adstack) drive all eight shapes. The cross_iter variants pin the offload-level guard's positive cases; the constant_outside_iter_range variants pin the references_loop_index gate so loads on a disjoint iteration slice do not trigger a false positive (with adstack on, the constant-outside variants compile cleanly across all four (has_outer_for, in_func) combinations, regardless of nesting).

A separate _run_reverse_sibling_component_correct_gradient factory plus test_reverse_sibling_component_correct_gradient_with_adstack (parametrized over the same (has_outer_for, in_func) axes) pins the multi-axis sibling-component case: out[i, 1] = x; out[i, 0] = out[i, 1] * 2 where the per-axis rule must allow the differing constant on axis 1.

Three sanity tests pin behavior the new checker must not break:

• test_reverse_kernel_top_indep_iters_{with,without}_adstack - kernel-top range with no cross-iter RAW on any needs_grad field; gradient must be analytical in both adstack settings.
• test_reverse_kernel_top_in_place_accum_with_adstack - out[i] = out[i] + x[None]; store and load on out reference the same LoopIndexStmt so the per-axis rule treats the indices as matching.

Side-effect audit

Concern	Where checked	Verdict
qd.ndrange / qd.grouped / mesh-for nested cases	the upstream AST guard at ASTTransformer.build_For only fires on node.iter.func.id == "range", never on the other lowering paths, and the new IR-level checker does not regress this	Existing behavior preserved
AdStackAllocaJudger semantics	not touched	Existing behavior unchanged; only invoked from ReplaceLocalVarWithStacks as before
Adstack-on path	OffloadLevelGlobalCrossIterRAWChecker runs before RegulateTensorTypedStatements and the IB / adstack flow is unchanged	No new check on the adstack-on path's per-iteration adjoint emission
Independent-iteration kernel-top range-fors	OffloadLevelGlobalCrossIterRAWChecker only fires on same-SNode (store, load) pairs that BOTH transitively reference LoopIndexStmt and differ structurally per axis	out[i] = x[None] * 2 (no load on out) compiles unchanged in both settings
In-place accumulation out[i] = out[i] + x	per-axis rule: store and load reference the same LoopIndexStmt on every axis	Compiles unchanged; pinned by test_reverse_kernel_top_in_place_accum_with_adstack
Multi-axis sibling-component access out[i, 0] = out[i, 1] * 2	per-axis rule: axes where neither index references LoopIndexStmt are ignored	Compiles unchanged; pinned by test_reverse_sibling_component_correct_gradient_with_adstack
Loads on disjoint iteration slices out[i] = out[external_const]	the load index does not reference LoopIndexStmt and is skipped	Compiles unchanged; pinned by the load_kind=constant_outside_iter_range variants
Forward-mode AD	both fixes gated on autodiff_mode == kReverse	No effect
Non-AD kernels	both fixes gated on autodiff_mode == kReverse	No effect
Nested mesh-for / while bodies	walk_nested_loop_body overrides cover MeshForStmt and WhileStmt so loads / stores inside them are excluded from the offload-level pairing	Symmetric with RangeForStmt / StructForStmt

[hughperkins]

INteresting. Lets make sure to run genesis unit tests and benchmarks please. For loops and ast transformer is clearly used all over the place.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 772b3b20fc

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[duburcqa]

Lets make sure to run genesis unit tests and benchmarks please.

I'm afraid the unit tests are failing. Exactly as it should:

>       loss.backward()

tests/test_grad.py:102:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
genesis/grad/tensor.py:78: in backward
    self.scene._backward()
genesis/engine/scene.py:1450: in _backward
    self._step_grad()
genesis/engine/scene.py:1043: in _step_grad
    self._sim._step_grad()
genesis/engine/simulator.py:304: in _step_grad
    self.sub_step_grad(self.cur_substep_local)
genesis/engine/simulator.py:328: in sub_step_grad
    self._coupler.couple_grad(f)
genesis/engine/couplers/legacy_coupler.py:1010: in couple_grad
    self.mpm_grid_op.grad(
../quadrants/python/quadrants/lang/_quadrants_callable.py:126: in grad
    return self.quadrants_callable._adjoint(self.instance, *args, **kwargs)
../quadrants/python/quadrants/lang/kernel.py:699: in __call__
    key = self.ensure_compiled(*py_args)  # type: ignore[arg-type]
../quadrants/python/quadrants/lang/kernel.py:635: in ensure_compiled
    self.materialize(key=key, py_args=py_args, arg_features=arg_features)
../quadrants/python/quadrants/lang/kernel.py:429: in materialize
    quadrants_kernel = impl.get_runtime().prog.create_kernel(
../quadrants/python/quadrants/lang/kernel.py:233: in __call__
    transform_tree(tree, ctx)
../quadrants/python/quadrants/lang/ast/transform.py:8: in transform_tree
    ASTTransformer()(ctx, tree)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <quadrants.lang.ast.ast_transformer.ASTTransformer object at 0x7c99cc4cab90>
ctx = <quadrants.lang.ast.ast_transformer_utils.ASTTransformerFuncContext object at 0x7c99cc4c9db0>
node = <ast.Module object at 0x7c99dafa2fb0>

    def __call__(self, ctx: "ASTTransformerFuncContext", node: ast.AST):
        method_name = "build_" + node.__class__.__name__
        method = getattr(self, method_name, None)
        try:
            if method is None:
                error_msg = f'Unsupported node "{node.__class__.__name__}"'
                raise QuadrantsSyntaxError(error_msg)
            info = ctx.get_pos_info(node) if isinstance(node, (ast.stmt, ast.expr)) else ""
            with impl.get_runtime().src_info_guard(info):
                res = method(ctx, node)
                if not hasattr(node, "violates_pure"):
                    # assume False until proven otherwise
                    node.violates_pure = False
                    node.violates_pure_reason = None
                return res
        except Exception as e:
            stack_trace = traceback.format_exc()
            if impl.get_runtime().print_full_traceback:
                raise e
            if ctx.raised or not isinstance(node, (ast.stmt, ast.expr)):
>               raise e.with_traceback(None)
E               quadrants.lang.exception.QuadrantsCompilationError: quadrants stack trace:
E               ===
E               Traceback (most recent call last):
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/ast_transformer_utils.py", line 40, in __call__
E                   res = method(ctx, node)
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/ast_transformer.py", line 480, in build_Call
E                   return CallTransformer.build_Call(ctx, node, build_stmt, build_stmts)
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/ast_transformers/call_transformer.py", line 377, in build_Call
E                   node.ptr = func(*py_args, **py_kwargs)
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/_quadrants_callable.py", line 111, in __call__
E                   return self.wrapper(self.instance, *args, **kwargs)
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/func.py", line 89, in __call__
E                   ret = transform_tree(tree, ctx)
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/transform.py", line 8, in transform_tree
E                   ASTTransformer()(ctx, tree)
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/ast_transformer_utils.py", line 51, in __call__
E                   raise e.with_traceback(None)
E               quadrants.lang.exception.QuadrantsCompilationError: quadrants stack trace:
E               ===
E               Traceback (most recent call last):
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/ast_transformer_utils.py", line 40, in __call__
E                   res = method(ctx, node)
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/ast_transformer.py", line 1258, in build_For
E                   raise exception.QuadrantsCompilationError("Cannot use non static range in Backwards mode")
E               quadrants.lang.exception.QuadrantsCompilationError: Cannot use non static range in Backwards mode
E
E               ===
E
E               Your code:
E               File "/mnt/home/duburcqa/workspace/src/genesis/genesis/engine/couplers/legacy_coupler.py", line 149, in _func_collide_with_rigid:
E                       for i_g in range(self.rigid_solver.n_geoms):
E                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E               Cannot use non static range in Backwards mode
E
E
E               ===
E
E               Your code:
E               File "/mnt/home/duburcqa/workspace/src/genesis/genesis/engine/couplers/legacy_coupler.py", line 391, in mpm_grid_op:
E                                   vel_mpm = self._func_collide_with_rigid(
E                                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E                                       f,
E                                       ^^
E                                       pos,
E                                       ^^^^
E                                       vel_mpm,
E                                       ^^^^^^^^
E                                       mass_mpm,
E                                       ^^^^^^^^^
E                                       i_b,
E                                       ^^^^
E                                       geoms_state=geoms_state,
E                                       ^^^^^^^^^^^^^^^^^^^^^^^^
E                                       geoms_info=geoms_info,
E                                       ^^^^^^^^^^^^^^^^^^^^^^
E                                       links_state=links_state,
E                                       ^^^^^^^^^^^^^^^^^^^^^^^^
E                                       rigid_global_info=rigid_global_info,
E                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E                                       sdf_info=sdf_info,
E                                       ^^^^^^^^^^^^^^^^^^
E                                       collider_static_config=collider_static_config,
E                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E                                   )
E                                   ^
E               quadrants stack trace:
E               ===
E               Traceback (most recent call last):
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/ast_transformer_utils.py", line 40, in __call__
E                   res = method(ctx, node)
E                 File "/mnt/home/duburcqa/workspace/src/quadrants/python/quadrants/lang/ast/ast_transformer.py", line 1258, in build_For
E                   raise exception.QuadrantsCompilationError("Cannot use non static range in Backwards mode")
E               quadrants.lang.exception.QuadrantsCompilationError: Cannot use non static range in Backwards mode
E
E               ===
E
E               Your code:
E               File "/mnt/home/duburcqa/workspace/src/genesis/genesis/engine/couplers/legacy_coupler.py", line 149, in _func_collide_with_rigid:
E                       for i_g in range(self.rigid_solver.n_geoms):
E                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E               Cannot use non static range in Backwards mode

../quadrants/python/quadrants/lang/ast/ast_transformer_utils.py:51: QuadrantsCompilationError
======================================================================================================= 2 failed, 645 passed, 3 skipped, 2 xfailed in 1109.45s (0:18:29) =======================================================================================================

[duburcqa]

Comparing Quadrants main vs this PR on Genesis main:

env	batch_size	backend	gjk_collision	constraint_solver	runtime_fps_588	runtime_fps_590	runtime_fps_delta_pct
anymal_random	30000	cuda	-	-	9288483	9314890	+0.28
anymal_uniform	30000	cuda	-	-	12277746	12361607	+0.68
anymal_uniform_kinematic	0	cpu	-	-	2000	2028	+1.40
anymal_uniform_kinematic	30000	cuda	-	-	10382797	10266957	-1.12
anymal_zero	0	cpu	-	-	7186	7073	-1.57
anymal_zero	30000	cuda	-	-	18857814	19096316	+1.26
box_pyramid_3	4096	cuda	-	-	978981	968342	-1.09
box_pyramid_4	4096	cuda	-	-	391086	395824	+1.21
box_pyramid_5	4096	cuda	-	-	139273	141715	+1.75
box_pyramid_6	4096	cuda	False	-	58248	59547	+2.23
box_pyramid_6	4096	cuda	True	-	59916	61918	+3.34
dex_hand	4096	cuda	-	-	17364	17081	-1.63
duck_in_box_easy	30000	cuda	False	-	26593384	26469341	-0.47
duck_in_box_easy	30000	cuda	True	-	9530599	9681964	+1.59
duck_in_box_hard	0	cpu	-	-	5144	5204	+1.17
duck_in_box_hard	30000	cuda	False	-	10162732	10077402	-0.84
duck_in_box_hard	30000	cuda	True	-	3368778	3493847	+3.71
franka	30000	cuda	-	-	21630563	21764102	+0.62
franka_accessors	0	cpu	-	-	1203	1182	-1.75
franka_accessors	30000	cuda	-	-	15764165	15853954	+0.57
franka_free	30000	cuda	-	-	32238005	32626193	+1.20
franka_random	0	cpu	-	-	6298	6301	+0.05
franka_random	30000	cuda	-	CG	16696014	16793645	+0.58
franka_random	30000	cuda	-	Newton	16450966	16535074	+0.51
franka_random	30000	cuda	False	-	16345136	16786850	+2.70
franka_random	30000	cuda	True	-	11410466	11327811	-0.72
g1_fall	4096	cuda	-	Newton	922187	920318	-0.20
go2	4096	cuda	False	CG	3673338	3668476	-0.13
go2	4096	cuda	False	Newton	4464312	4386696	-1.74
go2	4096	cuda	True	-	3243521	3230236	-0.41
shadow_hand_cubes	0	cpu	-	-	41	41	+0.00
shadow_hand_cubes_sparse	0	cpu	-	-	64	66	+3.12

speed_test_590.txt

[hughperkins]

checklist:

• genesis benchmarks neutral
• genesis unit test failures explained by the purpose of this pr

=> ok to merge

[github-actions]

Coverage Report (887264c2c)

File	Coverage	Missing
🟢 python/quadrants/lang/_func_base.py	100%
🟢 python/quadrants/lang/ast/ast_transformer_utils.py	100%
🔴 python/quadrants/lang/ast/ast_transformers/call_transformer.py	68%	387-393,395
🟢 tests/python/test_adstack.py	100%

Diff coverage: 95% · Overall: 73% · 177 lines, 8 missing

Full annotated report