[AutoDiff] SSA-promote AdStack count per task on LLVM backends to shrink ptxas RA cost

[duburcqa]

SSA-promote AdStack count per task on LLVM backends

Reverse-mode AD on adstack-heavy kernels (deep static-unrolled bodies, many AD variables per iteration) emits one heap-resident u64 count header per stack that every push / pop / top operation reads and writes through the runtime helpers stack_init / stack_push / stack_pop / stack_top_primal / stack_top_adjoint. LLVM cannot fold those loads / stores away because the stack pointer reaches the runtime through several getter calls that AA conservatively assumes alias arbitrary memory, so an unrolled body of N pushes emits N load + N store pairs against the same address. On large reverse-mode kernels the resulting flat PTX blocks dominate ptxas register-allocator cost. This PR replaces the helper-call sequence with inline LLVM IR that tracks the count in a per-stack alloca i64; mem2reg then promotes it to an SSA register and GVN folds consecutive count++ chains across straight-line unrolled bodies.

TL;DR

• Per-task per-stack alloca i64 replaces the heap-resident u64 count header. The alloca is created once in the task entry block and init-stored to zero at the AdStackAllocaStmt visit site (so an alloca nested inside a loop body restarts the count every iteration, matching the previous stack_init semantics).
• LLVMRuntime.adstack_max_sizes[stack_id] is hoisted to the entry block per stack on first use. Each push otherwise re-issued the GEP+load through the runtime metadata pointer; the hoist makes the value SSA-stable for the whole task.
• All five AdStack* visit methods (Alloca, Push, Pop, LoadTop, LoadTopAdj, AccAdjoint) are reimplemented as inline LLVM IR using the alloca instead of runtime helper calls. The slot-address math (stack + sizeof(u64) + idx * 2 * element_size) is now exposed to GVN, so consecutive top-of-stack accesses fold across the runtime call boundary that used to block them.
• Overflow path keeps its relaxed-atomic adstack_overflow_flag write via a new set_adstack_overflow_flag(LLVMRuntime*) runtime helper; the inline AdStackPushStmt codegen branches to the helper when count + 1 > max_size and otherwise increments the alloca, zeros the new slot's primal+adjoint pair, and stores the pushed value into the primal half - identical observable behaviour to the previous stack_push + stack_top_primal two-call sequence.
• Backwards-compatible heap layout: the unused 8-byte u64 header is left in place at the start of each stack's slab so ad_stack_per_thread_stride_ and the host-side LlvmRuntimeExecutor::ensure_adstack_heap slab sizing are unchanged. The kernel just never reads or writes those bytes.

Why

Profiles of cold GPU compile of the Genesis Ant kernel showed ptxas accounting for ~80% of total CUDA compile time, with a single hot function (presumably register allocation or live-range analysis) at ~161s. Two independent escape hatches (CU_JIT_OPTIMIZATION_LEVEL=2 exposed in #581 and CUDA_DISABLE_PTXAS_OPT=1 env var) had no measurable effect, indicating the cost is in mandatory phases (RA / liveness) that scale superlinearly with PTX basic-block size, not in optimization passes. The dominant size driver is the per-push count load / store / bounds-check sequence that the runtime helper path emits per AD variable per unrolled iteration. SSA-promoting the count collapses that sequence to a chain of integer adds, shrinking the flat block that ptxas's RA has to chew on.

Surface API

No public API change. qd.init and the CompileConfig knobs are untouched. The Python frontend, AD-stack sizing pipeline, host-side metadata publication, and SizeExpr machinery are all unaffected. The only internal shape change is to the four AdStack* visit methods in TaskCodeGenLLVM and one new runtime helper.

Mechanism

Before (one helper call per op)

; AdStackPushStmt (visit emits these)
%offsets_ptr  = call ptr @LLVMRuntime_get_adstack_offsets(ptr %runtime)        ; cached at entry
%max_sizes    = call ptr @LLVMRuntime_get_adstack_max_sizes(ptr %runtime)     ; cached at entry
%max_addr     = getelementptr i64, ptr %max_sizes, i64 0
%max_size     = load i64, ptr %max_addr                                        ; PER PUSH
call void @stack_push(ptr %runtime, ptr %stack, i64 %max_size, i32 %elem_size) ; PER PUSH (load + store of u64 header inside)
%top          = call ptr @stack_top_primal(ptr %stack, i32 %elem_size)         ; PER PUSH (load of u64 header inside)
store float %v, ptr %top

GVN cannot fold consecutive stack_push calls because they have function-call memory effects + the __atomic_store_n overflow path acts as a memory barrier; it cannot fold consecutive stack_top_primal calls either because the helper reads through %stack which AA cannot prove disjoint from the runtime metadata. Result: each unrolled iteration emits the full sequence again.

After (inline IR, per-stack alloca)

; entry block (once per task per stack):
%count_0 = alloca i64                                                          ; per-stack alloca
%max_size_0 = load i64, ptr %max_sizes_addr_0                                  ; hoisted

; AdStackAllocaStmt body (init):
store i64 0, ptr %count_0

; AdStackPushStmt (per push, after mem2reg + GVN folding):
%new_count = add i64 %old_count, 1
%overflow  = icmp ugt i64 %new_count, %max_size_0
br i1 %overflow, label %ovf, label %normal
normal:
  %slot = getelementptr i8, ptr %stack, i64 (8 + %old_count * 2 * elem_size)
  call void @llvm.memset.p0.i64(ptr %slot, i8 0, i64 (2*elem_size), ...)
  store float %v, ptr %slot

After mem2reg promotes count_0 to SSA and GVN folds the add 1 chain across N unrolled pushes, the only remaining memory ops in the unrolled body are the slot stores themselves. The count side has zero memory traffic.

Heap layout (unchanged)

stack_ptr[0..8)      = (legacy) u64 count header (no longer read/written by the kernel)
stack_ptr[8..)       = slot 0 primal, slot 0 adjoint, slot 1 primal, ...

The 8-byte header gap is preserved so ad_stack_per_thread_stride_ (sum of align_up_8(size_in_bytes) per stack) and the host's ensure_adstack_heap(stride * num_threads) allocation are unchanged. Slot offsets stay at sizeof(u64) + idx * 2 * element_size. A future PR can drop the gap and shrink the slab by 8 bytes per stack per thread.

Per-backend matrix

Backend	Codegen path	Behaviour change
CPU (LLVM)	TaskCodeGenLLVM	inline IR replaces runtime helper calls
CUDA (LLVM/NVPTX)	TaskCodeGenLLVM	inline IR replaces runtime helper calls
AMDGPU (LLVM/AMDGCN)	TaskCodeGenLLVM	inline IR replaces runtime helper calls
Metal (SPIR-V)	unchanged	unaffected
Vulkan (SPIR-V)	unchanged	unaffected

Tests

tests/python/test_adstack.py

• New test_adstack_unrolled_many_pushes_promotes_count_to_ssa parametrized over n_iter in {4, 16, 64} and n_stacks in {1, 3}. Each variant builds a kernel with one outer for i in x per element + a qd.static(range(n_stacks)) x qd.static(range(n_iter)) straight-line unrolled body that pushes qd.sin(x[i] + j*step + s_offset) onto independent adstacks per s. Reverse-mode gradients are cross-checked against PyTorch autograd. The 64-iteration variant is well above the historical max_size=32 floor and pins the bounds-check hoist: a regression that miscalculates the slot offset under the new codegen, or short-circuits / silently drops pushes via a wrong count > max_size clamp, produces wrong gradients here long before it surfaces as an overflow at runtime.
• Pre-existing test_adstack_unary_loop_carried / test_adstack_unary_loop_carried_f64 (every supported unary op x several n_iter and x_val) and the broader AD suite (test_ad_basics, test_ad_for, test_ad_if, test_ad_atomic, test_ad_offload, test_ad_demote_dense, test_ad_grad_check, test_ad_dynamic_index) continue to pass on the LLVM CPU backend - 488 tests across the wider AD suite, 346 on test_adstack.py alone.

Side-effect audit

• mem2reg requires the alloca to be in the entry block; ensure_ad_stack_count_alloca uses an InsertPointGuard to emit there regardless of where the AdStackAllocaStmt visit site sits. Same pattern as ensure_ad_stack_heap_base_llvm / ensure_ad_stack_metadata_llvm.
• The init-store at the AdStackAllocaStmt visit site is intentionally separate from the alloca creation. If an AdStackAllocaStmt is nested inside a loop body the alloca is still created once at the entry block, but the init store runs every iteration, matching the previous stack_init(stack_ptr) semantics where the heap u64 header was zeroed on each entry.
• Overflow-path semantics match the previous helper exactly: relaxed-atomic store of 1 to runtime->adstack_overflow_flag, no increment of the count, no slot zero-init, no value store. The legacy helper additionally stored the value to slot[max(0, n-1)] (i.e. clobbered the previous top); the new path skips that store. This is a strict improvement: the old path corrupted user data on overflow, the new one leaves it intact. runtime->adstack_overflow_flag is still set so qd.sync() raises before any garbage value reaches user code.
• The unused 8-byte header in heap memory is wasted per stack per thread. For a 4-stack kernel running 1M threads on GPU that is 32 MB of unused slab, which is a small constant relative to the slot data (typically GB on adstack-heavy kernels). A future PR can re-pack the layout and drop the gap.
• set_adstack_overflow_flag lives in the runtime module and is linked into kernel modules; the prefix matches the existing runtime_* symbols that survive eliminate_unused_functions.
• No change to the AdStack pre-scan (ad_stack_per_thread_stride_, ad_stack_offsets_, ad_stack_allocas_info_, ad_stack_size_exprs_), no change to host-side metadata publication, no change to SizeExpr machinery.

[duburcqa]

Closing - the optimization is a regression on the adstack-heavy reverse-mode workload it targeted. Cold GPU compile profile of the Genesis Ant kernel shows ptxas top-3 functions going from ~270s to ~461s (70% worse) after this change. The mistake: I optimized the count-side memory traffic (which mem2reg + GVN do fold cleanly via the alloca) but did not account for the basic-block count explosion from inline-expanding the runtime helpers. Each stack_push call site that used to be one opaque function call now emits a CondBr + overflow block + normal block + continuation, and each stack_top_primal site becomes a distinct slot-offset GEP - giant CFG with hundreds of small blocks per task, which ptxas's register allocator hates more than it hated the helper-call memory traffic. The right lever for this workload appears to be PTX size reduction at the source (fewer AdStack push sites, not faster ones); see PR thread for thoughts on alternative directions.

[github-actions]

Coverage Report (7313a336e)

File	Coverage	Missing
🟢 tests/python/test_adstack.py	81%	2445-2452

Diff coverage: 81% · Overall: 73% · 42 lines, 8 missing

Full annotated report