[AutoDiff] Autodiff 12: Heap-backed adstack on LLVM backends (CPU/CUDA/AMDGPU)

[duburcqa]

Heap-backed adstack on LLVM backends (CPU / CUDA / AMDGPU)

Replaces the per-thread worker-stack allocation of AdStackAllocaStmt with a host-grown shared heap, lifting the 256 KB CPU thread-stack budget that the prior create_entry_block_alloca path imposed on deep reverse-mode kernels. SPIR-V side is the following PR.

TL;DR

Prior behaviour on CPU / CUDA / AMDGPU: every AdStackAllocaStmt lowered to a function-scope alloca at the task's entry block, so every adstack lived on the LLVM stack frame (= worker-thread stack on CPU, per-thread local memory on GPU). A kernel with many loop-carried values at default_ad_stack_size=256 crossed the worker-thread limit and silently corrupted adjacent stack memory; the previous PR in the stack added a 256 KB codegen-time guard that hard-aborted those kernels.

This PR moves the storage off the stack:

• Codegen: a pre-scan of each offloaded task body computes per-task {ad_stack_offsets_, ad_stack_per_thread_stride_}, and visit(AdStackAllocaStmt) emits base = runtime->adstack_heap_buffer + linear_thread_idx * stride + offset instead of an alloca. Base is loaded once in entry_block and reused.
• Runtime: LlvmRuntimeExecutor::ensure_adstack_heap(needed_bytes) grows the per-runtime slab via amortised doubling, publishes the new pointer/size into runtime->{adstack_heap_buffer, adstack_heap_size} by caching the two device field-pointer addresses on first grow and writing through them on every subsequent grow.
• Launchers: CPU / CUDA / AMDGPU kernel launchers call ensure_adstack_heap(per_thread_stride * num_threads) before each task launch. Dynamic-bound range-for tasks resolve num_threads by reading begin / end from runtime->temporaries via a host-side DtoH memcpy.
• CUDA graphs: rejected at launch when any task has per_thread_stride > 0, because graph baking precludes the host-side ensure_adstack_heap step between dispatches.
• default_ad_stack_size exposed via qd.init(); raised from 32 → 256 now that the per-thread on-chip / worker-stack budget no longer caps it.
• The codegen budget guard from the previous PR is removed; the stack frame no longer carries adstack storage, so the 256 KB ceiling is obsolete.

Nothing changes for kernels that don't enable the adstack extension.

Why

Prior to this PR, a kernel like a reverse-mode articulated-body dynamics step in Genesis hit the 256 KB CPU-stack budget at modest capacities (4 loop-carried f64 variables × 4096 entries × 16 bytes each already crosses it). The two alternatives — ship with default_ad_stack_size capped at a value small enough to fit on every worker stack, or ask users to lower ad_stack_size per-kernel — either regress correctness on large kernels or force tuning noise on the user. Moving the storage off-stack removes the constraint entirely: per-thread slice size is bounded only by num_threads * per_thread_stride and the driver's allocator.

Changes

Codegen (quadrants/codegen/llvm/codegen_llvm.{h,cpp}, llvm_compiled_data.h)

TaskCodeGenLLVM grows three new per-task fields:

• ad_stack_per_thread_stride_ — sum of AdStackAllocaStmt::size_in_bytes() (aligned up to 8) for every adstack in the task.
• ad_stack_offsets_ — map from each alloca stmt to its offset within the per-thread slice.
• ad_stack_heap_base_llvm_ — cached SSA value of the heap base pointer, emitted once in entry_block.

init_offloaded_task_function pre-scans the task body before any codegen runs and populates the first two, so that later sibling allocas never shift an earlier alloca's offset out from under a cached SSA pointer.

visit(AdStackAllocaStmt) now emits:

base   = LLVMRuntime_get_adstack_heap_buffer(runtime)  // cached in entry_block
tid_64 = zext(linear_thread_idx(context))              // i32 → i64
slice  = tid_64 * stride                               // widened mul to avoid i32 overflow
ptr    = base + slice + offset

linear_thread_idx is the arch-appropriate invocation id (RuntimeContext::cpu_thread_id on CPU; block_idx * block_dim + thread_idx on CUDA / AMDGPU), matching how rand_states is indexed.

The old 256 KB function-scope budget guard (introduced in the previous PR) is deleted; its ad_stack_fn_scope_bytes_ accumulator is gone too. Heap-backed storage makes the ceiling irrelevant.

OffloadedTask gains an AdStackSizingInfo ad_stack sub-struct that propagates sizing to the host launcher: per_thread_stride, static_num_threads, dynamic_gpu_range_for, plus const values and gtmps byte offsets for range-for begin / end.

Per-arch codegen tweaks

• codegen_cpu.cpp — fills current_task->ad_stack with the pre-scanned stride and sets static_num_threads = cpu_thread_id_range (CPU thread count is known at compile time).
• codegen_cuda.cpp — fills current_task->ad_stack.static_num_threads = grid_dim * block_dim for const-bound tasks, and marks dynamic_gpu_range_for = true + records begin_offset_bytes / end_offset_bytes / begin_const_value / end_const_value for dynamic range-for tasks so the launcher can resolve the actual iteration count at launch time.
• codegen_amdgpu.cpp — same as CUDA.

Runtime (llvm_runtime_executor.{h,cpp}, runtime.cpp)

LLVMRuntime gains two new fields: Ptr adstack_heap_buffer = nullptr; u64 adstack_heap_size = 0;. These are read by every adstack-backed task on the device side; the host writes to them through the cached field-pointer addresses.

LlvmRuntimeExecutor::ensure_adstack_heap(needed_bytes):

1. No-op if needed_bytes == 0 || needed_bytes <= adstack_heap_size_.
2. Otherwise new_size = max(needed_bytes, 2 * adstack_heap_size_) (amortised doubling).
3. Allocates through the per-arch driver (llvm_device()->allocate_memory), wraps in a DeviceAllocationGuard.
4. On first grow, calls a one-shot runtime query to fetch the device addresses of runtime->adstack_heap_buffer and runtime->adstack_heap_size, caches them.
5. Publishes via memcpy_host_to_device (CUDA / AMDGPU) or plain pointer stores (CPU) against the cached addresses — no per-grow kernel launch.
6. Releases the previous DeviceAllocationGuard via move-assignment. Safety of the release (see the detailed block comment in the .cpp and the matching field comment in the .h): CPU uses std::free (trivially safe); CUDA cuMemFree_v2 synchronises before returning; AMDGPU dealloc_memory pools through CachingAllocator::release without sync, and cross-launch safety on AMDGPU is provided by the synchronous hipFree(context_pointer) at the tail of amdgpu::KernelLauncher::launch_llvm_kernel (the latent-fix in [AutoDiff] Autodiff 11: Latent adstack-adjacent fixes (AMDGPU hipFree, flush() keeps ctx_buffers_, always-preallocate) #536).

get_runtime_temporaries_device_ptr() — cached lookup of runtime->temporaries, used by the GPU launchers to read back dynamic range-for bounds.

Per-arch launchers

• runtime/cpu/kernel_launcher.{h,cpp} — Context gains a parallel ad_stack_needed_bytes vector, precomputed at register time (CPU sizing is static). launch_offloaded_tasks calls ensure_adstack_heap per task.
• runtime/cuda/kernel_launcher.cpp — adds resolve_num_threads(task) which DtoH-memcpys begin / end from runtime->temporaries for dynamic range-for tasks; calls ensure_adstack_heap per task.
• runtime/amdgpu/kernel_launcher.cpp — same as CUDA.
• runtime/cuda/graph_manager.cpp — hard-errors graph=True on kernels where any task has per_thread_stride > 0. Graph baking precludes host-side intervention between dispatches.

Capacity knob (compile_config.h, python/export_lang.cpp)

• default_ad_stack_size raised from 32 to 256.
• Exposed as a qd.init() kwarg. The comment block is rewritten to reflect the new heap-backing reality.
• ad_stack_size (per-stack explicit capacity) is unchanged.

Docs (docs/source/user_guide/autodiff.md)

Drops the "SPIR-V on-chip cap" limitation from the known-limitations list (that was about the prior Function-scope SPIR-V path; with the SPIR-V heap landing in the next PR, it's gone too). Adds a "Tuning the capacity" section explaining default_ad_stack_size vs ad_stack_size and the K+2-pushes-per-iteration rule for picking N.

Tests (tests/python/test_adstack.py)

Heap-specific additions:

• test_adstack_heap_grow_on_demand — two launches at increasing capacity pinpoint that the amortised-doubling grow path fires and the second launch reuses the bigger slab.
• test_adstack_heap_backed_exceeds_old_threadstack_budget — a kernel whose per-thread adstack bytes exceed the pre-PR 256 KB ceiling now compiles and runs correctly.
• test_adstack_cuda_graph_rejected_with_adstack — graph=True on an adstack kernel raises.

Side-effect audit

Concern	Where checked	Verdict
Offline cache key	AdStackSizingInfo fields serialised via the existing OffloadedTask key hash	Auto-covered, kernels without adstack get zero stride and hash identically
Stmt clone / serialization	No new Stmt fields (all sizing lives in codegen / OffloadedTask)	N/A
CUDA graphs	Hard rejection with clear error message when any task has per_thread_stride > 0	Fails loudly, not silently
AMDGPU release lifecycle	Relies on #536's synchronous hipFree(context_pointer) tail	Cross-launch invariant spelled out in both .cpp and .h comments
Non-adstack kernels	per_thread_stride == 0 path short-circuits the launcher-side ensure_adstack_heap	Zero-cost for kernels that don't enable the extension

Stack

Split 2/3 of the former "heap-backed adstack" PR. Based on #536 (latent fixes). Followed by #493 (SPIR-V heap).

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a619832bc4

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

[claude]

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

• 🔴 quadrants/program/compile_config.h:54-68 — The comment added to compile_config.h falsely claims SPIR-V now uses heap-backed StorageBuffers, but SPIR-V heap-backing is explicitly deferred to PR #493; SPIR-V still allocates adstacks as Function-scope per-thread on-chip memory bounded by the driver shader-compiler limit. The 32 to 256 default increase therefore multiplies per-thread private memory 8x on Metal/Vulkan, which can cause pipeline-creation failures for kernels that previously compiled correctly. The companion docs change compounds this by removing the only warning about the SPIR-V on-chip cap and replacing it with advice to bump default_ad_stack_size, which on SPIR-V causes shader-compiler rejection rather than heap growth.

  Extended reasoning...

  What is wrong

  compile_config.h lines 57-64 add a comment justifying the 32 to 256 default increase: "Both backends now heap-back the primal/adjoint slots: SPIR-V uses per-dispatch StorageBuffers (BufferType::AdStackHeapFloat + AdStackHeapInt, sliced by invocation)". This claim is factually incorrect for the current PR. The PR description itself explicitly states "SPIR-V side is the following PR", meaning PR #493 has not landed yet. After this PR merges, SPIR-V (Metal and Vulkan) still allocates every AdStackAllocaStmt using ir_->alloca_variable(arr_type) with spv::StorageClassFunction - Function-scope per-thread on-chip private memory - exactly as before.

  Concrete code path

  spirv_codegen.cpp:2221-2223 (unchanged in this diff) visits AdStackAllocaStmt and calls alloca_variable() for count_var, primal_arr, and adjoint_arr. There is no BufferType::AdStackHeapFloat or BufferType::AdStackHeapInt anywhere in the SPIR-V codegen, and no SPIR-V files appear in the list of changed files. The 8x raise is therefore applied to the SPIR-V path unconditionally by ControlFlowGraph::determine_ad_stack_size() in transforms/determine_ad_stack_size.cpp, which is arch-agnostic and falls back to default_ad_stack_size for any stack whose worst-case trip count cannot be statically proven.

  Why existing code does not prevent it

  The SPIR-V codegen has no guard that caps AdStackAllocaStmt::max_size against a per-thread on-chip budget - that responsibility fell on the deliberately-conservative 32-slot default. Removing that conservatism by raising the default 8x while the SPIR-V heap path is absent eliminates the only protection.

  Impact

  For a kernel with 4 f32 loop-carried variables and a dynamic loop whose trip count the compiler cannot prove:

  • Old default (32): per-thread Function-scope demand = 4 * (8 + 3224) bytes ~= 1 KB
  • New default (256): per-thread Function-scope demand = 4 * (8 + 25624) bytes ~= 8 KB

  Apple Metal's MSL shader compiler rejects pipelines whose per-thread private-variable footprint exceeds its budget (typically 4-16 KB on consumer hardware), and the existing test test_adstack_shader_compile_failure_raises with ad_stack_size=65536 confirms this failure mode exists. A kernel that compiled and ran correctly at the 32-slot default may now fail at kernel-launch time with a pipeline-creation error whose message does not mention adstack size.

  Step-by-step proof

  1. User calls qd.init(arch=qd.metal, ad_stack_experimental_enabled=True) - default_ad_stack_size is now 256 after this PR.
  2. A kernel has 4 f32 loop-carried variables under a range(n[None]) whose bound comes from a field.
  3. determine_ad_stack_size cannot prove the worst-case trip count, sets max_size = 256 for each adstack.
  4. SPIR-V codegen emits 4 x OpVariable StorageClassFunction arrays of 8x256 bytes each = 8192 bytes of per-thread on-chip private memory.
  5. Metal MSL compiler rejects the pipeline; compute.grad() raises RuntimeError: Failed to create pipeline with no hint that adstack sizing is the cause.
  6. The user consults the updated docs, which now say to bump default_ad_stack_size when they see a SPIR-V overflow. Doing so makes the pipeline failure worse, not better.

  Docs regression

  The removed Known Limitations bullet was the only explicit warning that on SPIR-V backends the adstack is allocated as per-thread on-chip memory capped by the driver. The new Tuning the capacity section and the Memory cost statement that "The buffer grows on demand to match the largest size any launch has needed so far" are presented without qualification and are false for SPIR-V - on SPIR-V the capacity is baked into the compiled shader at compile time, there is no grow-on-demand behavior, and bumping the value risks compile-time pipeline rejection rather than graceful heap growth.

  Fix

  Keep default_ad_stack_size at 32 until PR #493 lands. Update the comment in compile_config.h to remove the false SPIR-V StorageBuffer claim. Restore the Known Limitations bullet about the SPIR-V on-chip cap so Metal/Vulkan users receive the correct guidance.

[hughperkins]

lets have the units please
default_ad_stack_size_mb

[hughperkins]

if this is in units, or int32s, or something then maybe something like default_ad_stack_size_count or default_ad_stack_size_units or default_ad_stack_size_i32s`?

[hughperkins]

what happens if there is an i64 in the loop?

[hughperkins]

units ad_stack_size_mb

[hughperkins]

default_ad_stack_size_mb

[hughperkins]

dont we need to multiply by 4?

[hughperkins]

by the way, why + 2?

[hughperkins]

x 4?

[hughperkins]

• 4?

[hughperkins]

why do we need to qualify with "that the reverse pass has to remember."? Are there loop-carried variables tha the reverse pass does not have to remember?

[hughperkins]

can you clarify where num_threads suddenly springs from? I'm guessing it's from the top level for loop, but you don't introduce t his I think. Or at least, I dont remember your introducing this.

[hughperkins]

lets ditch the where, otherwsie no room to breathe. Seems lik a bunch new concepts here, so lets give the reader time to breathe. its a new paragraph.

[hughperkins]

this all seesm like way too much detail. Do we really need to know this to use autodiff? Move it to an 'advanced' or 'under the hood' section.

[hughperkins]

keep 'The buffer grows on demand to match the largest size any launch has needed so far and is then reused across subsequent launches, so you do not need to reserve memory up front.'

[hughperkins]

break the first sentence into two. so each setnence just states a single concept.

[hughperkins]

" the autodiff stack (adstack) " => "adstack". We only ever refer to it as adstack, so let's say adstack is its name. We can put the long form in brakcets i youf want "called the adstack (short for "(a)uto(d)iff (stack)")"

[hughperkins]

I think remove "and therefore changes from one iteration to the next"

[hughperkins]

" It is opt-in" => "adstack is opt-in"

[hughperkins]

"Running without it when it is needed raises a QuadrantsCompilationError in most cases (the autodiff pass rejects a non-static range that would otherwise lose its primal);" => nice 🙌

[hughperkins]

"in the narrow cases where the kernel compiles anyway, the reverse pass reads a stale value for every iteration and the gradients come out wrong but non-zero." => would be nice to get rid of such exceptional cases. Do we know what they are? Can we document them?

[hughperkins]

I know it's not part of your changes in this pr, but it's part of your set of prs overall, and anyway, I think it would be good to address:

I couldnt understand this sentence. I mean, I probably could if I thought about it a long time, bu I think it would be good to explain step by step, so people can just read fluidly, and undersatnd it, without grinding to a halt, and having to work through stuff in their head.

[hughperkins]

can we have a counter-example of a non-linear op that doesnt need adstack, in the 'do not need it' section below please.

[hughperkins]

mixtue of fonts is visually distracting. either put backticks around 16 x 16 too, or remove from = 16

[hughperkins]

whats a 'piece'? per stack carried variable?

[hughperkins]

oh you define it next 🤔

[hughperkins]

Maybe we just avoid the issue by avoiding having to use this noun at all? for example

"The pipeline allocates a scratch buffer for each loop-carried variable, and also for any loop counter, and any boolean branch flags." ?

Still, I'm unclear about this 'integer counter' and 'boolean branch flag'. You havent defined them before. Could you define these in a previous paragraph please.

[hughperkins]

Nice 🙌

[hughperkins]

have we used the term 'element' thus far? Have we defined it?

[hughperkins]

Checklist:

• Doc updated, and looks good to me (readable, makes sense, a few variable names that seem questionable, but they are pre-existing, and should be addressed seprately from this PR, to avoid inflating the PR)

=> ok to merge