[AutoDiff] Autodiff 9: Guard against LLVM worker-thread stack overflow from large per-task adstack budget

[duburcqa]

Guard against LLVM worker-thread stack overflow from large per-task adstack budget

CPU-only codegen guard. Rejects compilation when the cumulative AdStackAllocaStmt::size_in_bytes() in a single LLVM task crosses the ~256 KB secondary-thread stack budget; without it, the frame silently clobbers adjacent stack memory and the reverse pass returns zero / garbage gradients.

TL;DR

void TaskCodeGenLLVM::visit(AdStackAllocaStmt *stmt) {
  QD_ASSERT_INFO(stmt->max_size > 0, "...");
  auto type = llvm::ArrayType::get(llvm::Type::getInt8Ty(*llvm_context), stmt->size_in_bytes());

  if (arch_is_cpu(current_arch())) {
    constexpr std::size_t kFnScopeAdStackBudgetBytes = 256 * 1024;
    ad_stack_fn_scope_bytes_ += stmt->size_in_bytes();
    QD_ERROR_IF(ad_stack_fn_scope_bytes_ > kFnScopeAdStackBudgetBytes,
                "LLVM autodiff-stack budget exceeded: cumulative `AdStackAllocaStmt` size {} bytes in task "
                "'{}' crosses the {} byte function-scope budget. ...",
                ad_stack_fn_scope_bytes_, kernel_name, kFnScopeAdStackBudgetBytes);
  }

  auto alloca = create_entry_block_alloca(type, sizeof(int64));
  // ...
}

The check is CPU-only by design: on CUDA / AMDGPU the same LLVM allocas lower to per-thread GPU local memory (a separate address space sized by the driver, not shared with the CPU call stack), so the 256 KB CPU-stack budget is not meaningful there. A non-gated version of the check would falsely reject valid GPU kernels with f64 loop-carried variables (4 adstacks at ad_stack_size=4096 already cross 256 KB).

Why 256 KB

macOS secondary threads default to a ~512 KB stack. The worker-thread pool used by the LLVM JIT runs on those. A function-scope alloca for an adstack sits on the LLVM stack frame; if the sum of its sizes across a task crosses the thread's stack limit, the frame corrupts adjacent stack memory (typically the next page, sometimes guard pages). Downstream reverse-mode accumulators read zero, producing silently-wrong gradients with no crash.

256 KB is a conservative upper bound that leaves ~256 KB for other locals and nested call frames. Linux defaults to ~8 MB per secondary thread, so the limit is strictly conservative there — the codegen is more protective on Linux than strictly necessary, which is fine.

Why QD_ERROR_IF and not throw

QD_ERROR_IF logs a descriptive message and then calls QD_UNREACHABLE (compiler-assisted abort). The codegen runs inside the LLVM compilation worker thread pool, where a C++ exception thrown from a worker doesn't cleanly propagate back to the Python-level caller (pybind11's exception translation only works at pybind binding boundaries, which are on the main thread). Throwing QuadrantsRuntimeError from here results in std::terminate() and a bare SIGABRT — a worse user experience than the guard's current behaviour, which at least logs the message before aborting.

The trade-off is that the abort shows up as a SIGABRT in the child process rather than a catchable Python exception. The testing approach below handles that.

Obsoleted by subsequent PRs

Autodiff 12 (LLVM heap-backed adstack) moves the storage off the worker thread stack entirely, making this guard unnecessary. When that PR lands, the guard — along with the ad_stack_fn_scope_bytes_ accumulator field — is removed. This PR is still useful in the interim to make the silent-corruption failure loud, and to provide coverage for any future path that might re-introduce function-scope allocas for adstack storage.

Changes

quadrants/codegen/llvm/codegen_llvm.{h,cpp}

• TaskCodeGenLLVM gains an ad_stack_fn_scope_bytes_ accumulator field, reset to 0 at the top of each offloaded task.
• init_offloaded_task_function resets the accumulator.
• visit(AdStackAllocaStmt) tallies per-alloca size and raises via QD_ERROR_IF when the per-task total crosses 256 KB on CPU.

Tests

test_adstack_codegen_budget_guard_runs_in_child_process

Runs the overflowing kernel in a child process (since QD_ERROR_IF aborts the process rather than raising a catchable Python exception) and asserts:

1. The child exits with a non-zero returncode (the guard fired and terminated the process).
2. The guard's message ("autodiff-stack budget exceeded") appears in the child's combined stdout/stderr.

Skip-gated on qd.cpu having the adstack and f64 extensions. Kernel shape: five qd.f64 loop-carried variables applying qd.sin inside a dynamic range(n_iter[None]) at ad_stack_size=4096. Each adstack is 8 + 4096 * 16 = 65,544 bytes ≈ 64 KB; five × 64 KB = 320 KB, comfortably past the 256 KB guard.

The "dynamic n_iter from a field" shape is load-bearing: with a Python-literal range(3) the compile-time trip count is known and the determine-ad-stack-size pass sizes each stack to 3 slots, not default_ad_stack_size=4096. Only a runtime-bound range defeats that pass and leaves each stack at the full default_ad_stack_size.

Side-effect audit

Concern	Verdict
Existing CPU kernels with small adstacks	Unaffected — 256 KB budget is generous for default-sized adstacks (32 slots × entry bytes).
CUDA / AMDGPU	Guard is gated on arch_is_cpu(current_arch()); never fires on GPU.
Linux vs macOS	256 KB is conservative on both; Linux stacks are larger by default, so the guard is strictly more protective than necessary there.
Valid-but-large GPU kernels	Not rejected (CPU-only gate).
Silent corruption path	Surfaced with a descriptive abort message instead of wrong gradients.
Future removal	Explicitly obsoleted by the LLVM heap-backed PR (Autodiff 12), which removes the guard and its accumulator.

Stack

Autodiff 9 of 13. Third commit of the "LLVM adstack safety" triplet split and the top-most of that split (this PR is #495). Based on #535 (runtime overflow). Followed by #490 (SPIR-V adstack).

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a5d3009268

ℹ️ 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".

[hughperkins]

Opus review of test factoriztaion:

Looking at the four tests side by side:

Test	Loop type	Body	Expected grad
test_adstack_sum_fixed_coeff	dynamic range(3)	y += v	3.0
test_adstack_sum_constant_coeffs	dynamic range(3)	y += v * (a+1)	6.0
test_adstack_sum_static_loop_correct	static qd.static(range(3))	y += v	3.0
test_adstack_basic_gradient	dynamic range(10)	v = v * 0.95 + 0.01	0.95^10

Tests 2, 3, and 4 share identical setup/teardown/assertion structure and only differ in two dimensions: dynamic vs static loop, and fixed vs varying coefficient. These three could collapse into a single parametrized test using if qd.static(...) to switch at compile time:

@pytest.mark.parametrize("use_static_loop", [True, False])
@pytest.mark.parametrize("use_varying_coeff", [True, False])
@test_utils.test(require=qd.extension.adstack, ad_stack_experimental_enabled=True)
def test_adstack_sum_linear(use_static_loop, use_varying_coeff):
    x = qd.field(qd.f32)
    y = qd.field(qd.f32)
    qd.root.dense(qd.i, 1).place(x, x.grad)
    qd.root.place(y, y.grad)

    @qd.kernel
    def compute():
        for i in x:
            v = x[i]
            if qd.static(use_static_loop):
                r = qd.static(range(3))
            else:
                r = range(3)
            for a in r:
                if qd.static(use_varying_coeff):
                    y[None] += v * qd.cast(a + 1, qd.f32)
                else:
                    y[None] += v

    x[0] = 1.0
    y[None] = 0.0
    compute()
    y.grad[None] = 1.0
    x.grad[0] = 0.0
    compute.grad()

    expected = sum((a + 1) for a in range(3)) if use_varying_coeff else 3.0
    assert x.grad[0] == test_utils.approx(expected, rel=1e-4)

This gives you 4 combinations (including static + varying coeff, which is new free coverage) in one test, and eliminates ~50 lines of duplication.

test_adstack_basic_gradient should stay separate — it's the only one with loop-carried state mutation (v = v * 0.95 + 0.01), which is a structurally different pattern. Though it could benefit from a brief comment explaining that distinction.

[hughperkins]

Please could you address Opus's suggetsions on test factroization.

[hughperkins]

nice

[hughperkins]

Asking Opus about any runtime impact:

Hot-path branches in stack operations. stack_push, stack_pop, and stack_top_primal are called on every iteration of every AD loop that uses the adstack. This PR adds a new
branch to each:

• stack_push: bounds check if (n + 1 > max_num_elements) before every push, plus an extra LLVMRuntime *runtime pointer argument threaded through every call site.
• stack_pop: underflow guard if (n > 0) before every pop.
• stack_top_primal: conditional n > 0 ? n - 1 : 0 instead of a bare n - 1.

In normal (non-overflowing) operation the branches are always not-taken, so CPU branch prediction should handle them well. On GPU, all threads in a warp would take the same
path (not-taken), so no divergence cost. But they do add instruction count and register pressure to what is essentially the innermost loop of the backward pass.

synchronize() poll. Every qd.sync() now calls check_adstack_overflow(), which makes a JIT call into the runtime to read and reset the flag, plus a result buffer fetch. This is
once per sync (not per iteration), so it's unlikely to be measurable unless the user syncs at very high frequency.

In practice, the cost is probably negligible — a couple of predictable branches per stack operation. But if someone wanted to benchmark it, the stack push/pop overhead in a
tight AD loop with many iterations would be the place to look.

[hughperkins]

Opus response for test coverage:

Code changes in this PR:

Change	Exercised by test?
stack_push overflow check + early return	Yes — the 64-iteration kernel triggers it
stack_push sets runtime->adstack_overflow_flag = 1	Yes — the test expects the resulting exception
stack_push new LLVMRuntime *runtime parameter	Yes — indirectly, the kernel wouldn't compile otherwise
stack_pop underflow guard (if (n > 0))	No — only reachable after overflow, when the reversed loop pops more times than were successfully pushed. The test raises before observing this path.
stack_top_primal clamping (n > 0 ? n - 1 : 0)	No — same situation, only reachable post-overflow.
LLVMRuntime::adstack_overflow_flag field	Yes
runtime_retrieve_and_reset_adstack_overflow	Yes — called by check_adstack_overflow
LlvmRuntimeExecutor::check_adstack_overflow()	Yes
check_adstack_overflow early return when llvm_runtime_ == nullptr	No — only reachable during early teardown before materialize_runtime
result_buffer_cache_ caching in materialize_runtime	Yes — indirectly, needed for check_adstack_overflow to work
LlvmProgramImpl::synchronize() calls check_adstack_overflow	Yes
LlvmProgramImpl::finalize() sets finalizing_ = true	No — no test verifies that overflow during teardown doesn't crash
codegen_llvm.cpp passes get_runtime() to stack_push	Yes — indirectly
internal_functions.h test_stack updated	No — no test calls test_internal_func_args / test_stack

Untested functionality implied by the design:

1. Flag reset after catch. The flag is reset in runtime_retrieve_and_reset_adstack_overflow, but no test verifies that after catching the overflow exception, a subsequent qd.sync() does not raise again. A second sync after the pytest.raises block would cover this.

2. Increasing ad_stack_size resolves the overflow. The error message tells users to pass ad_stack_size=N to qd.init(). No test verifies that doing so actually makes the same kernel succeed.

3. Overflow on SPIR-V. The test accepts both AssertionError and RuntimeError, and the docstring describes the SPIR-V path, but the SPIR-V code path isn't in this diff. If SPIR-V doesn't have equivalent overflow detection, the test would silently pass on LLVM-only CI without ever validating the SPIR-V claim.

4. Multi-threaded overflow. The comment in stack_push explains the race is benign (all threads write the same sentinel). No test uses a multi-element field where multiple threads overflow simultaneously.

5. Gradients are actually wrong without the safety check. The test only checks that an exception is raised, not that the gradients would have been silently wrong without it. A companion test at a just-under-capacity iteration count showing correct gradients would strengthen the argument.

6. Teardown safety. The finalizing_ flag exists specifically so that an overflow during ~Program() → finalize() → synchronize() doesn't throw into a destructor and terminate the process. No test covers this path — e.g., triggering overflow and then letting the qd.init() scope exit without an explicit sync.

[hughperkins]

Please could you add tests for untested fucntionality implied by the design:

1. Flag reset after catch. The flag is reset in runtime_retrieve_and_reset_adstack_overflow, but no test verifies that after catching the overflow exception, a subsequent qd.sync() does not raise again. A second sync after the pytest.raises block would cover this.

2. Increasing ad_stack_size resolves the overflow. The error message tells users to pass ad_stack_size=N to qd.init(). No test verifies that doing so actually makes the same kernel succeed.

3. Overflow on SPIR-V. The test accepts both AssertionError and RuntimeError, and the docstring describes the SPIR-V path, but the SPIR-V code path isn't in this diff. If SPIR-V doesn't have equivalent overflow detection, the test would silently pass on LLVM-only CI without ever validating the SPIR-V claim.

4. Multi-threaded overflow. The comment in stack_push explains the race is benign (all threads write the same sentinel). No test uses a multi-element field where multiple threads overflow simultaneously.

5. Gradients are actually wrong without the safety check. The test only checks that an exception is raised, not that the gradients would have been silently wrong without it. A companion test at a just-under-capacity iteration count showing correct gradients would strengthen the argument.

6. Teardown safety. The finalizing_ flag exists specifically so that an overflow during ~Program() → finalize() → synchronize() doesn't throw into a destructor and terminate the process. No test covers this path — e.g., triggering overflow and then letting the qd.init() scope exit without an explicit sync.

[hughperkins]

Oh, also please add tests for the rows in the table above with 'no', ie:

• stack_pop underflow guard (if (n > 0))
• stack_top_primal clamping (n > 0 ? n - 1 : 0)
• check_adstack_overflow early return when llvm_runtime_ == nullptr
• LlvmProgramImpl::finalize() sets finalizing_ = true
• internal_functions.h test_stack updated

[claude]

Additional findings (outside current diff — PR may have been updated during review):

• 🔴 quadrants/runtime/llvm/runtime_module/runtime.cpp:1107-1112 — The PR deletes all eight shuffle_down runtime helpers (amdgpu_shuffle_down_i32/f32/i64/f64 and cuda_shuffle_down_i32/f32/i64/f64) from runtime.cpp as incidental cleanup, but codegen_cuda.cpp (~lines 746-752) and codegen_amdgpu.cpp (~lines 449-455) still emit JIT calls to these exact symbol names when lowering subgroup shuffle_down operations. Any CUDA or AMDGPU kernel that uses qd.subgroup.shuffle_down() will fail at JIT link time with an undefined symbol error. Restore the eight deleted functions or update the codegen sites to inline the shuffle_down logic.

  Extended reasoning...

  What the bug is and how it manifests

  The PR diff shows two deletion hunks in runtime.cpp: one removing amdgpu_shuffle_down_i32/f32/i64/f64 (around line 1107) and one removing cuda_shuffle_down_i32/f32/i64/f64 (around line 1143). These are described as incidental cleanup but they are not dead code — they are JIT-linked runtime symbols that GPU codegen actively resolves at kernel compilation time.

  The specific code path that triggers it

  codegen_cuda.cpp contains an emit_cuda_shuffle_down() function (approximately lines 744–755) that calls:
  call("cuda_shuffle_down_i32", offset, value)
  call("cuda_shuffle_down_f32", offset, value)
  call("cuda_shuffle_down_f64", offset, value)
  call("cuda_shuffle_down_i64", offset, value)

  codegen_amdgpu.cpp contains an emit_amdgpu_shuffle_down() function (approximately lines 447–458) that calls:
  call("amdgpu_shuffle_down_i32", offset, value)
  call("amdgpu_shuffle_down_f32", offset, value)
  call("amdgpu_shuffle_down_f64", offset, value)
  call("amdgpu_shuffle_down_i64", offset, value)

  Neither codegen file is modified by this PR, so these eight call() sites remain unchanged and still reference the deleted symbols.

  Why existing code does not prevent it

  The call() invocations in LLVM codegen resolve symbol names from the JIT-linked runtime bitcode module at kernel compilation time. There is no compile-time check that the symbol exists in the module — the linker error only surfaces when a CUDA or AMDGPU kernel that exercises a subgroup.shuffle_down() operation is compiled. The PR's own test suite does not include CUDA/AMDGPU subgroup tests (the tests added are all CPU adstack tests), so the breakage escapes CI.

  What the impact would be

  Any user calling qd.subgroup.shuffle_down() on a CUDA or AMDGPU backend will receive a JIT linker error ('undefined symbol: cuda_shuffle_down_i32' or equivalent) the first time the kernel is compiled, turning previously working code into a hard failure. This is a complete regression for GPU subgroup reduction patterns.

  How to fix it

  Either (a) restore the eight deleted functions — they were not related to the adstack work and their deletion was purely cosmetic — or (b) update codegen_cuda.cpp and codegen_amdgpu.cpp to inline the shuffle_down logic directly instead of calling the runtime helpers.

  Step-by-step proof

  1. PR diff, runtime.cpp hunk 1 (line ~1107): deletes amdgpu_shuffle_down_i32, amdgpu_shuffle_down_f32, amdgpu_shuffle_down_i64, amdgpu_shuffle_down_f64.
  2. PR diff, runtime.cpp hunk 2 (line ~1143): deletes cuda_shuffle_down_i32, cuda_shuffle_down_f32, cuda_shuffle_down_i64, cuda_shuffle_down_f64.
  3. grep for 'shuffle_down' in the post-PR runtime.cpp returns zero matches — the functions are gone.
  4. codegen_cuda.cpp (not modified by PR): emit_cuda_shuffle_down() at lines 746/748/750/752 calls call("cuda_shuffle_down_i32", ...), call("cuda_shuffle_down_f32", ...), call("cuda_shuffle_down_f64", ...), call("cuda_shuffle_down_i64", ...).
  5. codegen_amdgpu.cpp (not modified by PR): emit_amdgpu_shuffle_down() at lines 449/451/453/455 calls the amdgpu variants.
  6. When any CUDA/AMDGPU kernel containing a subgroup.shuffle_down is compiled, the JIT linker cannot resolve the missing symbols and raises an undefined symbol error.

[claude]

LGTM — focused CPU-only budget guard with correct arch_is_cpu gate, properly reset per task, and a well-constructed child-process test.

Extended reasoning...

Overview

Touches three files: codegen_llvm.{cpp,h} and test_adstack.py. The change adds a single new field (ad_stack_fn_scope_bytes_) and a per-task accumulator reset in init_offloaded_task_function, then enforces a 256 KB budget via QD_ERROR_IF inside visit(AdStackAllocaStmt*). The test runs the budget-crossing kernel in a subprocess and asserts a non-zero exit code plus the expected guard message in output.

Security risks

No security-sensitive paths are touched. The change adds a compile-time guard that aborts the process on overflow — the same category of safety-check that already exists elsewhere in the codebase. No auth, crypto, or permissions code is involved.

Level of scrutiny

Low-to-moderate. The diff is small and self-contained. The two most important correctness properties — (a) the arch_is_cpu(current_arch()) gate preventing false rejections of valid GPU kernels, and (b) using current_task->name rather than the kernel-level kernel_name in the error message — are both correctly implemented in the current diff. The field-level comment in the header is already properly scoped with "On CPU arches only".

Other factors

All issues surfaced in the review thread (comment inaccuracies, teardown safety, flag-reset, ad_stack_size header mismatch, GPU guard) were addressed in earlier PRs or earlier commits of this PR. The one inline comment from today that lacks a formal "resolved" marker (codegen_llvm.h field comment) reflects code that is already correct in the current diff. The test is gated on both qd.extension.adstack and qd.extension.data64, and the subprocess approach correctly handles the QD_ERROR_IF → std::terminate termination path.

[claude]

LGTM — the CPU-only budget guard is correctly implemented and the inline comments I raised today were false positives (the code already used current_task->name and the header comment was already CPU-scoped).

Extended reasoning...

Overview

PR #495 adds a CPU-only per-task adstack budget guard to TaskCodeGenLLVM::visit(AdStackAllocaStmt*). The change touches three files: codegen_llvm.cpp (guard logic + per-task accumulator reset), codegen_llvm.h (new ad_stack_fn_scope_bytes_ field), and tests/python/test_adstack.py (child-process test that verifies the guard fires). This is the focused 'codegen budget guard' slice of the autodiff safety series, matching the split that was requested.

Security risks

No security-sensitive code paths touched. The guard is a compile-time check that produces a loud abort rather than silent memory corruption; it cannot be reached from user-supplied data at runtime, only from kernel compilation. No auth, crypto, or permission logic is involved.

Level of scrutiny

Moderate. The guard is CPU-only (correctly gated on arch_is_cpu(current_arch()) in the .cpp; the header comment already notes the gate and the GPU-local-memory distinction), uses current_task->name for per-task identification in the error message, and resets the accumulator at the start of each offloaded task. Earlier rounds of review caught and fixed a real ordering bug (pre_finalize timing), a header-size mismatch (sizeof(int32) vs sizeof(int64)), a missing arch guard, and several comment inaccuracies — all addressed by the author. No bugs were found by the automated system on the current revision.

Other factors

My two inline comments posted today (2026-04-21 07:35 and 08:48) were false positives: the diff already uses current_task->name and the header already carries the 'On CPU arches only' scoping language. The subprocess.run timeout nit in the budget-guard test is a minor style concern; the teardown test's explicit comment explains the team's rationale for relying on pytest's per-test timeout instead, which applies equally here.

[hughperkins]

Checklist:

• tightening of error checking => doesn't change existing API or usage
  • no need for doc changes
• doesn't add significantly more code around diverse files

=> ok to merge