[Lang] Inline AdStack ops on release LLVM codegen: dramatically reduces compile time for adstack-enabled reverse-mode kernels#584
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fc6350c64d
ℹ️ 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".
…d repeated-load-top kernel shapes
|
@claude review |
|
@claude review |
Coverage Report (
|
| File | Coverage | Missing |
|---|---|---|
🟢 tests/python/test_adstack.py |
83% | 2439-2444,2493-2498,2552-2557,2612-2619 |
Diff coverage: 83% · Overall: 52% · 152 lines, 26 missing
…d elides the bounds check
|
@claude review |
…d LLVM path: skip the count alloca entirely for AdStackAllocaStmts whose sizer-resolved capacity is compile-time max_size==1 with a Const SizeExpr - slot is provably slot 0 at every program point so push / pop / loadtop / loadtopadj / accadjoint reduce to constant-offset GEPs (stack + 8 for primal, stack + 8 + element_size for adjoint) with no count load / store, no mem2reg recurrence, no SCEV induction-variable analysis. On a representative reverse-mode AD cold compile this shape covers ~44 percent of the per-task adstack population, dominated by 1-byte u1 / i8 branch-outcome snapshots that the reverse pass replays without any adjoint accumulation; profile A/B against the multi-slot-only path drops the dominant ptxas symbol from 21.74s to 4.41s and shaves an additional ~35s off cold-compile wall-clock. Multi-slot stacks fall through to the existing alloca path; debug builds keep the runtime-helper path.
Coverage Report (
|
| File | Coverage | Missing |
|---|---|---|
🟢 tests/python/test_adstack.py |
83% | 2449-2454,2503-2508,2562-2567,2622-2629 |
Diff coverage: 83% · Overall: 73% · 155 lines, 26 missing
|
@claude review |
Coverage Report (
|
| File | Coverage | Missing |
|---|---|---|
🟢 tests/python/test_adstack.py |
83% | 2449-2454,2503-2508,2569-2574,2629-2636 |
Diff coverage: 83% · Overall: 61% · 155 lines, 26 missing
|
@claude review |
|
can you post please:
|
|
|
comparing Quadrants main vs this PR on Genesis main:
|
|
Nice! 🔥 checklist:
=> ok to merge |
…s single-slot specialization - SSA-promotes the count alloca and emits constant-offset GEPs for max_size==1 stacks, dramatically reducing compile time on adstack-heavy reverse-mode kernels
…s single-slot specialization - SSA-promotes the count alloca and emits constant-offset GEPs for max_size==1 stacks, dramatically reducing compile time on adstack-heavy reverse-mode kernels
…he four parametrized tests and the three debug=True overflow flags that landed on main via PR #584; only delta vs main is the new test_adstack_min_loop_carried_serial_range_for added by this PR (with the cpu / cuda / amdgpu arch restriction dropped now that the snap-stack fix lands all backends correctly)
…he four parametrized tests and the three debug=True overflow flags that landed on main via PR #584; only delta vs main is the new test_adstack_min_loop_carried_serial_range_for added by this PR (with the cpu / cuda / amdgpu arch restriction dropped now that the snap-stack fix lands all backends correctly)
…he four parametrized tests and the three debug=True overflow flags that landed on main via PR #584; only delta vs main is the new test_adstack_min_loop_carried_serial_range_for added by this PR (with the cpu / cuda / amdgpu arch restriction dropped now that the snap-stack fix lands all backends correctly)
* [Misc] Warn user to disable caching when print_ir/QD_DUMP_IR enabled (Genesis-Embodied-AI#425) Co-authored-by: v01dxyz <v01dxyz@v01d.xyz> * [Build] Pin torch version to CUDA 12.8 for CUDA tests (Genesis-Embodied-AI#428) * [Misc] Fixing up taichi-dev urls (Genesis-Embodied-AI#429) * [Perf] Rename cuda_graph to gpu_graph across the codebase (Genesis-Embodied-AI#430) * Misc: fix typo integeral -> integral (Genesis-Embodied-AI#434) Co-authored-by: v01dxyz <v01dxyz@v01d.xyz> * [Perf] CUDA graph 4: call from multiple locations (Genesis-Embodied-AI#420) * [Bug] Fix fastcache not restoring graph_do_while_arg (Genesis-Embodied-AI#435) * [Perf] Cache last-call result in perf_dispatch for single-compatible case (Genesis-Embodied-AI#438) * Fix gpu_graph fallback on old Nvidia GPU. (Genesis-Embodied-AI#443) * Fix shared memory offset not reset between CUDA kernels. (Genesis-Embodied-AI#442) * [Misc] Allow disabling GPU graph via QD_GPU_GRAPH=0 env var (Genesis-Embodied-AI#439) * [Misc] Add named top-level loops (Genesis-Embodied-AI#440) * [Misc] Rename gpu_graph to graph (Genesis-Embodied-AI#446) * [Misc] Add cross-platform shuffle (Genesis-Embodied-AI#447) * [Bug] Fix graph_do_while on Windows: search for cudadevrt.lib (Genesis-Embodied-AI#456) * [Bug] Also search default CUDA toolkit install location on Windows (Genesis-Embodied-AI#461) * [SPIRV] Feature Parity Atomics & Shared Array (Genesis-Embodied-AI#432) * [Misc] Change clang format to 120 characters (Genesis-Embodied-AI#463) * [Misc] CUDA graph 5 Add fatbin (Genesis-Embodied-AI#464) * [Bug] Reuse VkInstance across init/reset cycles (Genesis-Embodied-AI#465) * [Perf] Tiles 1: _load, _store, _eye_ (Genesis-Embodied-AI#466) * [Misc] Remove dead InternalFuncStmt type_check override (Genesis-Embodied-AI#471) * [Perf] Tiles 2: add cholesky and ger (Genesis-Embodied-AI#472) * [Perf] Tiles 2b: add triangular solve (Genesis-Embodied-AI#474) * [Misc] Refactor: use _get_col/_set_col in tiles load/store/init (Genesis-Embodied-AI#475) * [Build] Fix flaky test_clock_accuracy (Genesis-Embodied-AI#436) * Fix AARCH64 emitting invalid asm in CUDA kernels. (Genesis-Embodied-AI#473) Co-authored-by: Hugh Perkins <hughperkins@gmail.com> * [AMDGPU] Enable HIP memory pool and surface pool-exhaustion errors. (Genesis-Embodied-AI#485) * [AMDGPU] Scope hsaco tmp dir per-user to avoid collisions. (Genesis-Embodied-AI#484) * [Perf] Tiles 3: Add slice syntax, qd.outer() and initial doc (Genesis-Embodied-AI#477) * [AMDGPU] Fix gradient computation. (Genesis-Embodied-AI#486) * Enable all backends that are supported in unit tests. (Genesis-Embodied-AI#488) * Fix SPIRV ID overflow for large kernels due to autodiff. (Genesis-Embodied-AI#489) * [Misc] Fix purity checker to allow accessing constants from quadrants modules (Genesis-Embodied-AI#487) * [Misc] Increase tolerance for clock monotonic test (Genesis-Embodied-AI#492) * [CI] Serialize api doc workflow (Genesis-Embodied-AI#494) * [CI] Increase tolerance for clock test (Genesis-Embodied-AI#506) * [CI] Increase clock test tolerance to 20% (Genesis-Embodied-AI#509) * [Perf] Add tensor_type parametrization to tile16 tests (Genesis-Embodied-AI#504) * [Perf] Tiles 4b: Migrate tiles16 tests to enable fastcache (Genesis-Embodied-AI#505) * [Perf] Tiles 4c: add Tiles16x16 proxy (Genesis-Embodied-AI#507) * [Perf] Tiles 4d: Consolidate slice error tests using parametrize (Genesis-Embodied-AI#508) * [Perf] Tiles 4: add SharedArray slice support (Genesis-Embodied-AI#482) * [Perf] Tiles 5: add Cholesky benchmark demo (Genesis-Embodied-AI#483) * [Doc] Add user guide page for subgroup shuffle (Genesis-Embodied-AI#512) * [Perf] Implement cross-platform shuffle_down (Genesis-Embodied-AI#510) * [Perf] Add portable subgroup reduce_add and reduce_all_add (Genesis-Embodied-AI#511) * [Perf] Add first warmup config to perf dispatch (Genesis-Embodied-AI#422) * [AutoDiff] Autodiff 1: Add baseline adstack regression test for unary_collections (Genesis-Embodied-AI#500) * [AutoDiff] Autodiff 2: Implement derivative for tan (Genesis-Embodied-AI#501) * [AutoDiff] Autodiff 3: Recompute tanh/exp on the operand in the reverse pass (Genesis-Embodied-AI#502) * [AutoDiff] Autodiff 4: Mark rsqrt as non-linear for adstack promotion (Genesis-Embodied-AI#503) * [AutoDiff] Autodiff 5: Fix adjoint-alloca placement for GlobalLoads outside the current range-for (Genesis-Embodied-AI#496) * [AutoDiff] Autodiff 6: Adstack regression tests (Genesis-Embodied-AI#491) * [AutoDiff] Autodiff 7: Fix header size in AdStackAllocaStmt to match u64 runtime layout (Genesis-Embodied-AI#534) * [AutoDiff] Autodiff 8: Surface LLVM adstack push/pop overflow as a Python exception (Genesis-Embodied-AI#535) * [AutoDiff] Autodiff 9: Guard against LLVM worker-thread stack overflow from large per-task adstack budget (Genesis-Embodied-AI#495) * [AutoDiff] Autodiff 10: Implement adstack for SPIR-V (Genesis-Embodied-AI#490) * [AutoDiff] Autodiff 11: Latent adstack-adjacent fixes (AMDGPU hipFree, flush() keeps ctx_buffers_, always-preallocate) (Genesis-Embodied-AI#536) * [Doc] Add AGENTS.md with instructions for AI agents (Genesis-Embodied-AI#541) * [Bug] Abort kernel execution on assertion failure instead of segfaulting (Genesis-Embodied-AI#419) * [Type] ndarray typing 1: Add eval_str=True to inspect.signature() calls (Genesis-Embodied-AI#411) * [CI] Suppress reportPrivateImportUsage in torch-using files (Genesis-Embodied-AI#552) * [Misc] QD_DUMP_IR dumps to files with the task_id added to the filename (Genesis-Embodied-AI#441) * [Type] ndarray typing 2: Fix NDArray single-arg subscript crash (Genesis-Embodied-AI#412) * [Test] Flush xdist channel before worker exit so test failure reports are visible (Genesis-Embodied-AI#555) * [CI] Reduce test retries on CI from 3 to 1. (Genesis-Embodied-AI#554) * [AutoDiff] Autodiff 12: Heap-backed adstack on LLVM backends (CPU/CUDA/AMDGPU) (Genesis-Embodied-AI#537) * [AutoDiff] Autodiff 13: Heap-backed adstack on SPIR-V backends (Metal, Vulkan) (Genesis-Embodied-AI#493) * [AutoDiff] Autodiff 14: Resolve bounded-inner-loop adstacks without default_ad_stack_size fallback (Genesis-Embodied-AI#539) * [SPIRV] Vulkan SPIR-V correctness: atomic-view aliasing, PSB stride, narrow storage caps, u1 cast, per-init layer recheck (Genesis-Embodied-AI#513) * [Build] Autodiff 15: Replace 2022 MoltenVK pin with LunarG Vulkan SDK fetch and sanitise MoltenVK cap advertisement (Genesis-Embodied-AI#551) * [Test] Suppress stock pytest-timeout to avoid conflict with pytest_hardtle (Genesis-Embodied-AI#557) * [Vulkan] Use SDK validation layer for debugPrintf instead of apt package (Genesis-Embodied-AI#562) * [Test] Fix flaky perf_dispatch tests by increasing work amounts (Genesis-Embodied-AI#559) * [Test] Add --maxfail CLI option to run_tests.py (default 20) (Genesis-Embodied-AI#558) * [CI] Vulkan debug printf fix to address flaky tests (Genesis-Embodied-AI#563) * [Docs] Add a new page to help for first time contributors (Genesis-Embodied-AI#426) Authored-by: v01dxyz <v01dxyz@v01d.xyz> * [AutoDiff] Autodiff 16: Resolve reverse-mode adstack depths per-launch via runtime-evaluated SizeExpr (Genesis-Embodied-AI#543) * Fix: raise error if device memory allocation fails (Genesis-Embodied-AI#451) (Genesis-Embodied-AI#453) Co-authored-by: v01dxyz <v01dxyz@v01d.xyz> Co-authored-by: Hugh Perkins <hughperkins@gmail.com> * [CI] Add CI job to check line wrapping of comments and docs (Genesis-Embodied-AI#564) * [Misc] Add coverage report to PRs, including kernels (Genesis-Embodied-AI#470) * [CI] CI wrap check feeds only diffs to agent (Genesis-Embodied-AI#567) * Skip 'flaky' test on MacOS CI. (Genesis-Embodied-AI#573) * [Test] Fix missing `import sys` in test_fail_device_memory_allocation (Genesis-Embodied-AI#574) * [CI] Fix Vulkan debugPrintf flake with session-scoped warmup (Genesis-Embodied-AI#571) * [AutoDiff] determine_ad_stack_size: replace whole-CFG Bellman-Ford with SCC + DAG DP (Genesis-Embodied-AI#575) * [Test] Fix macOS OOM skip reason to describe actual root cause (Genesis-Embodied-AI#576) * [Lang] whole_kernel_cse: 2.5x compile time speedup on large kernels (Genesis-Embodied-AI#577) * [CI] Add CI check for unnecessarily deleted comments (Genesis-Embodied-AI#570) * [CI] Migrate coverage report to github Check page (Genesis-Embodied-AI#566) * [Lang] Skip IR verifier between passes unless debug=true (Genesis-Embodied-AI#579) * [Lang] Inline AdStack ops on release LLVM codegen: dramatically reduces compile time for adstack-enabled reverse-mode kernels (Genesis-Embodied-AI#584) * [CUDA] Honor offline_cache=False end-to-end so QD_OFFLINE_CACHE=0 actually gives a cold compile (Genesis-Embodied-AI#580) * [Type] Tensor 24 (Genesis-Embodied-AI#561) Co-authored-by: hugh <hugh@slurm-login-0.slurm-login.tenant-slurm.svc.cluster.local> * [Lang] auto_diff host-walk reductions: dramatically faster front-end compile time on adstack-enabled reverse-mode kernels (Genesis-Embodied-AI#587) * [AutoDiff] Speed up reverse-mode kernel launches on GPU backends (Genesis-Embodied-AI#578) * [Vulkan] Move adstack-sizer scratch out of Function-scope memory to fix SPIR-V pipeline build failures (Genesis-Embodied-AI#588) * [AutoDiff] Improve diagnosis of unsupported reverse-mode AD patterns (Genesis-Embodied-AI#590) * [Bug] Fix: promote Ndarray to AnyArray in build_Name for flattened struct fields (Genesis-Embodied-AI#592) * [SPIR-V] Shrink reverse-grad kernel MSL by ~50% (Genesis-Embodied-AI#591) * [CI] Add CI check that PR changes have test coverage (Genesis-Embodied-AI#596) * [Perf] Enable zero-copy in to_torch() and to_numpy() (Genesis-Embodied-AI#450) * Add BufferView: safe sub-range ndarray access for kernels (Genesis-Embodied-AI#585) Co-authored-by: alanray-tech <alanray-tech@users.noreply.github.com> Co-authored-by: Hugh Perkins <hughperkins@gmail.com> * [Doc] Add user-facing fastcache documentation (Genesis-Embodied-AI#597) Co-authored-by: hugh <hugh@slurm-login-0.slurm-login.tenant-slurm.svc.cluster.local> * [Misc] Upgrade to enable v1 dlpack so to_numpy(copy=False) writable (Genesis-Embodied-AI#598) Co-authored-by: root <root@rtx-209-201.slurm-compute.tenant-slurm.svc.cluster.local> * [AutoDiff] Cut reverse-mode adstack memory usage 10x on all backends (Genesis-Embodied-AI#599) * [Misc] Add CI check for feature file factorization (Genesis-Embodied-AI#606) * [Perf] Skip _recursive_set_args for all-Field frozen dataclass structs (Genesis-Embodied-AI#607) Co-authored-by: Cursor <cursoragent@cursor.com> * [AutoDiff] SNode-arm bound-expr capture rejects fold-attack gate indices (Genesis-Embodied-AI#610) * [Misc] Suppress field fastcache warning for qd.Tensor (Genesis-Embodied-AI#615) Co-authored-by: Cursor <cursoragent@cursor.com> * [AutoDiff] Adstack heap: clip reducer count by per-task loop trip count (compile-time and SizeExpr-evaluated) (Genesis-Embodied-AI#611) * [Misc] Forward copy= through qd.Tensor, add copy=None option (Genesis-Embodied-AI#616) Co-authored-by: Cursor <cursoragent@cursor.com> * [Doc] Update README (Genesis-Embodied-AI#617) Co-authored-by: Cursor <cursoragent@cursor.com> * [CI] Fix coverage report showing def lines as uncovered (Genesis-Embodied-AI#623) Co-authored-by: Cursor <cursoragent@cursor.com> * [Perf] Generic launcher: persistent context, JIT-pointer reuse, Metal compute encoder, LLVM-GPU async memory ops (Part 1/2) (Genesis-Embodied-AI#619) * [CI] Encode Python-first testing policy in coverage-check prompt (Genesis-Embodied-AI#622) Co-authored-by: Cursor <cursoragent@cursor.com> * [CI] Add PR Line change report (Genesis-Embodied-AI#624) Co-authored-by: Cursor <cursoragent@cursor.com> * [CI] Disable quadrants pytest plugin during quadrants internal coverage runs (Genesis-Embodied-AI#629) Co-authored-by: Cursor <cursoragent@cursor.com> * [AutoDiff] Adstack load+store eliminations: EliminateRecomputableAdStackPushes pass + leaf extensions (Genesis-Embodied-AI#621) * [CI] Simplify coverage PR comment to a single linked line (Genesis-Embodied-AI#630) * [CUDA] Add AGX Thor, SM_110 (Genesis-Embodied-AI#631) Co-authored-by: Johnny Nunez and Hugh Perkins * [CI] Lines changed report: collapse PR comment to a single linked totals line (Genesis-Embodied-AI#632) * [FEATURE] Support external Metal command queue via qd.init (Genesis-Embodied-AI#618) Co-authored-by: Cursor <cursoragent@cursor.com> * [Perf] Cache adstack-sizer metadata per task across SPIR-V + LLVM-GPU; per-snode / DeviceAllocation invalidation (Part 2/2) (Genesis-Embodied-AI#620) * [AutoDiff] Disable EliminateRecomputableAdStackPushes pending mutated-SNode chain-leaf fix (Genesis-Embodied-AI#633) * [AutoDiff] Adstack chain-clone safety: mutated-SNode leaf reject + load_top consumer-aware guard (Genesis-Embodied-AI#634) * [Docs] Add user-guide page for qd.simt.block.* primitives (Genesis-Embodied-AI#638) * [Docs] Expand qd.simt.subgroup user-guide page to cover every op (Genesis-Embodied-AI#639) * [Perf] Streams 1-4 (Genesis-Embodied-AI#410) * [Docs] Add user-guide page for matrix decompositions and solvers (Genesis-Embodied-AI#643) * [Bug] Revert "[Perf] Streams 1-4 (Genesis-Embodied-AI#410)" (Genesis-Embodied-AI#650) * [Docs] Add user-guide page for atomics and bit operations (Genesis-Embodied-AI#640) * [Docs] Add user-guide page for qd.simt.grid.* primitives (Genesis-Embodied-AI#641) * [AutoDiff] Adstack max-reducer: parallel multi-axis MaxOverRange dispatch (Genesis-Embodied-AI#635) * [AMDGPU] Fix amdgpu parallel rand init (Genesis-Embodied-AI#658) * [Perf] Adstack: skip max-reducer recognizer on CPU + lift host-eval cap (Genesis-Embodied-AI#655) * [Perf] Re-land Streams 1-4 with bug fixes (Genesis-Embodied-AI#653) * [AMDGPU] Apply device_memory_GB=0.3 cap to AMDGPU tests (Genesis-Embodied-AI#659) * [Perf] Per-launch host sync: drop wait_idle on SPIR-V, pin stream and drop stream_synchronize on CUDA/AMDGPU (Genesis-Embodied-AI#654) * [AMDGPU] Unload hipModule_t in JITModuleAMDGPU destructor (Genesis-Embodied-AI#660) * [AMDGPU] Trim default mempool on qd.reset() (Genesis-Embodied-AI#669) * [AMDGPU] Hoist rand-state buffer to process lifetime (Genesis-Embodied-AI#668) * [Streams] Use events for streams serialization on AMDGPU and CUDA (Genesis-Embodied-AI#667) * [Perf] Adstack max-reducer: launch cache + zero-copy result map; content-stable registry_id (Genesis-Embodied-AI#671) * [SPIR-V] dispatch_max_reducers: register each task with the real kernel name (Genesis-Embodied-AI#675) * [AutoDiff] Debug-mode field/grad/dual: dtype, layout, and access-time invariants (Genesis-Embodied-AI#677) * [Docs] Add user-guide page for qd.algorithms.* device-wide algorithms (Genesis-Embodied-AI#642) Co-authored-by: alanray-tech <alan.ray@genesis-ai.company> * [Docs] Doc for existing atomics: switch support table to per-backend columns (Genesis-Embodied-AI#657) Co-authored-by: alanray-tech <alan.ray@genesis-ai.company> * [GPU] Cross gpu atomics (Genesis-Embodied-AI#666) Co-authored-by: alanray-tech <alan.ray@genesis-ai.company> * [GPU] Make block operations portable cross-gpu (Genesis-Embodied-AI#664) * [Perf] CPU LLVM adstack-cache: skip per-launch bump-writes + ndarray_shapes capture on forward-only handles (Genesis-Embodied-AI#685) * [GPU] Cross-GPU for grid ops (Genesis-Embodied-AI#670) * [Math] Make bitop operations portable cross-gpu (Genesis-Embodied-AI#662) * [AMDGPU] Always use wave64, on both RDNA and CDNA (Genesis-Embodied-AI#687) * [AMDGPU] Use syncscope("agent") for atomix xor to avoid CAS livelock (Genesis-Embodied-AI#672) * [GPU] New bit ops for QIPC (Genesis-Embodied-AI#679) * [GPU] Subgroup ops cross-gpu (Genesis-Embodied-AI#665) * [Graph] Rename CUDA Graph to Graph in docs (Genesis-Embodied-AI#691) * [SPIR-V] Fix FIFO-queue ordering when sharing command queue. (Genesis-Embodied-AI#694) * [Atomics] New QIPC ops for atomics (Genesis-Embodied-AI#690) * Pass dataclass sub-structs into qd.func (Genesis-Embodied-AI#698) * [AMDGPU] HIP graph runtime support for @qd.kernel(graph=True) (Genesis-Embodied-AI#692) * [CI] Add per-file timing report to Mac Metal test job (Genesis-Embodied-AI#695) Co-authored-by: Cursor <cursoragent@cursor.com> * [CI] Enable kernel disk cache during tests (Genesis-Embodied-AI#696) * [Math] New QIPC ops for single-threaded linalg (Genesis-Embodied-AI#683) * [BREAKING][GPU] New QIPC ops for subgroups (Genesis-Embodied-AI#676) * [GPU] New QIPC ops for block (Genesis-Embodied-AI#684) * [GPU] New device-level ops for QIPC (Genesis-Embodied-AI#693) * [algorithms] PrefixSumExecutor: drop unused GRID_SZ local (Genesis-Embodied-AI#701) * [block] sync(): fix unsupported-arch error message (Genesis-Embodied-AI#700) * [volatile_load] add qd.volatile_load primitive (closes Genesis-Embodied-AI#648) (Genesis-Embodied-AI#702) * [AutoDiff] Reject recycled identity_key in AdStackCache::register_adstack_sizing_info (Genesis-Embodied-AI#708) * [Vulkan] Declare GroupNonUniform SPIR-V caps and enable shaderSubgroupExtendedTypes (Genesis-Embodied-AI#707) * Fix duplicate HIP graph driver-function declarations after v1.0.0 merge The amd-integration fork had cherry-picked the HIP graph driver functions (graph_create / graph_destroy / graph_add_kernel_node / graph_instantiate / graph_exec_destroy / graph_launch), and upstream v1.0.0 added the same set. The per-file 3-way merge appended both copies into amdgpu_driver_functions.inc.h, producing redeclaration errors that broke the AMDGPU RHI/runtime compile. Drop the upstream duplicate block; the signatures are identical to the fork's existing declarations. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix AMDGPU launcher coherence and num_instructions visibility after v1.0.0 merge - kernel_launcher.cpp: the 3-way merge spliced upstream v1.0.0's launch_llvm_kernel rewrite (ephemeral arg/context buffers, explicit-stream path, AmdgpuDefaultStream PinGuard) onto the AMD fork's kernarg-by-value + persistent-scratch design, leaving references to undefined `ephemeral_context_ptr`. Restore the fork's coherent launch_llvm_kernel verbatim; it calls the (already merged) enhanced launch_offloaded_tasks, which keeps the max-reducer dispatch and stream-parallel groups adapted onto the AMD launch path. - llvm_context.h: both the fork and upstream added `num_instructions`; the merge kept upstream's private placement, but the AMDGPU codegen force-inline heuristic calls it statically from outside the class. Move it back to the public section. Co-authored-by: Cursor <cursoragent@cursor.com> * Restore async result D2H and hoist kernarg vectors in AMDGPU launcher The v1.0.0 merge resolution regressed two amd-integration baseline optimizations in launch_llvm_kernel / launch_offloaded_tasks: - The per-launch result-buffer copy was a blocking memcpy_device_to_host, forcing a host stall on every value-returning launch and serializing the GPU pipeline. Restore the async D2H (the caller synchronizes lazily when it needs the value); external-array transfers still stream_synchronize once before reading back. - launch_task constructed the kernarg std::vectors from initializer lists ({kernarg_payload} / {kernarg_size}) on every dispatch (heap alloc + free per launch). Hoist arg_ptrs/arg_sizes out of the per-task launch and reuse. Co-authored-by: Cursor <cursoragent@cursor.com> * amdgpu: default to LDS permlane64 emulation; drop host-x86 barrier asm on retarget Two AMDGPU JIT-compile crashes surfaced after the v1.0.0 merge pulled in the QIPC subgroup ops (Genesis-Embodied-AI#676), which made the rigid constraint solver's wave-cooperative reductions route through `amdgpu_cross_half_shuffle_i32`. Both manifested as a SIGSEGV inside `llvm::SIInstrInfo::getInstSizeInBytes` during `JITSessionAMDGPU::compile_module_to_hsaco` (i.e. at first kernel launch), and reproduce on gfx942 / MI300X. Baseline 0.4.6 never emitted these constructs, which is why it was unaffected. 1. Native `llvm.amdgcn.permlane64` lowering crashes the bundled LLVM 22.1.0 AMDGPU backend. Default `amdgpu_permlane64` to the existing LDS-roundtrip software emulation on every target (it produces identical results). Add `QD_AMDGPU_USE_NATIVE_PERMLANE64=1` to opt back into the native instruction once the backend bug is fixed; the old `QD_AMDGPU_FORCE_PERMLANE64_FALLBACK` is now the default and still honored. This is the actual crash fix. 2. The runtime module is compiled by the host x86_64 clang and only retargeted to amdgcn here, so `amdgpu_cross_half_shuffle_i32`'s `__asm__ volatile("" : "+v"(byte))` optimization barrier carries x86 flag clobbers (`~{dirflag},~{fpsr},~{flags}`) that are meaningless on AMDGPU. The IR verifies but the empty-body INLINEASM is invalid on the amdgcn target. Neutralize empty-body barrier asm during retarget (forward the tied value, then erase) so no stale host asm reaches codegen. On the wave64 targets we ship `ds_bpermute` already addresses the full wave, so the hint is a no-op. Co-authored-by: Cursor <cursoragent@cursor.com> * style: apply clang-format (v19.1.7) to AMDGPU fn_attrs and launcher sources CI pre-commit's clang-format hook reformatted these files (long declarations/lambda signatures collapsed onto single lines per the repo's clang-format config). Apply the same formatting so the hook passes. No functional changes. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(amdgpu): use CreateNeg for branchless i32 sgn instead of CreateSub(0, input) clang-tidy (modernize-use-nullptr, -warnings-as-errors) flagged `builder->CreateSub(0, input)` in the i32 sgn path: the literal `0` binds to the `llvm::Value*` LHS parameter as a null pointer, not an integer zero. Replace with `builder->CreateNeg(input)`, which emits `0 - input` with a proper zero constant -- identical intended semantics, and clang-tidy clean. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Robert Dazi <14996868+v01dXYZ@users.noreply.github.com> Co-authored-by: v01dxyz <v01dxyz@v01d.xyz> Co-authored-by: Hugh Perkins <hughperkins@gmail.com> Co-authored-by: Alexis DUBURCQ <alexis.duburcq@gmail.com> Co-authored-by: hugh <hugh@slurm-login-0.slurm-login.tenant-slurm.svc.cluster.local> Co-authored-by: alanray-tech <alan.ray@genesis-ai.company> Co-authored-by: alanray-tech <alanray-tech@users.noreply.github.com> Co-authored-by: root <root@rtx-209-201.slurm-compute.tenant-slurm.svc.cluster.local> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Johnny <johnnynuca14@gmail.com>
TL;DR
alloca i64replaces the heap-resident u64 count header. The alloca is created once in the task entry block (somem2regpromotes it) and init-stored to zero at theAdStackAllocaStmtvisit site (so an alloca nested inside a loop body restarts the count every iteration, matching the previousstack_initsemantics).max_size==1with aConstSizeExpr, the count alloca is elided entirely. Slot is fixed at offset 8 from the stack base; push is a single store, loadtop is a single load, pop is a no-op. No mem2reg recurrence, no SCEV induction-variable analysis on this stack. On a representative reverse-mode AD cold compile this shape covers ~44 percent of the per-task adstack population, dominated by 1-byteu1/i8branch-outcome snapshots that the reverse pass replays without any adjoint accumulation.AdStack*visit methods (Alloca,Push,Pop,LoadTop,LoadTopAdj,AccAdjoint) are reimplemented as inline LLVM IR using the alloca (multi-slot) or constant-offset GEP (single-slot) instead of runtime helper calls whencompile_config.debug == false. 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.determine_ad_stack_sizeis mandatory and produces a valid upper bound on per-thread push count along every execution path - any unresolved stack is a hard compile error - so the runtimen + 1 > max_num_elementscompare and the relaxed-atomicadstack_overflow_flagwrite are dead code in correct compilations. Eliminating them removes one cond_br plus one overflow basic block from each push site, halving the basic-block count contributed by the unrolled push sequence and removing the runtime-pointer dependency that pulls the LLVMRuntime address into every push's live range.compile_config.debug == true) routes pushes throughstack_pushso any sizer bug surfaces as an overflow flag at sync. Same forstack_init / stack_pop / stack_top_primal / stack_top_adjointon the debug path.llvm.memsetthat zeroes the primal+adjoint slot pair now passes the true minimum slot alignment (min(8, 2 * element_size)) instead of an unconditionalMaybeAlign(8). For element sizes < 4 the slot pointer is genuinely 2- or 4-byte-aligned, not 8 - the previous over-stated claim was invalid IR metadata that future LLVM / NVPTX / AMDGCN lowering decisions could turn into a misaligned-store fault.ad_stack_per_thread_stride_and the host-sideLlvmRuntimeExecutor::ensure_adstack_heapslab sizing are unchanged. The kernel just never reads or writes those bytes on the release path.Why
Profiles of cold GPU compile of a reverse-mode AD kernel with adstack enabled showed
libnvidia-ptxjitcompileraccounting for ~20x more wall-clock than the same kernel with adstack disabled, with the dominant ptxas function (presumably register allocation or live-range analysis) at ~145s on a single function. Two independent escape hatches (CU_JIT_OPTIMIZATION_LEVEL=2exposed in #581 andCUDA_DISABLE_PTXAS_OPT=1env var) had no measurable effect, indicating the cost is in mandatory phases (RA / liveness) that scale superlinearly with PTX basic-block size and memory-dependence chain length, 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.Multi-slot SSA-promotion collapses that sequence to a chain of integer adds, shrinking the flat block ptxas's RA has to chew on and removing the per-push count's may-alias chain from
MemoryDependenceResults. Single-slot specialization eliminates the count alloca entirely for the dominant 1-byte branch-outcome population, removing ~4500 mem2reg recurrences from the LLVM IR, ~9000 LLVM instructions across all push / pop / loadtop sites, and the matching PTX body. Empirical A/B on a representative cold compile:ScalarEvolution::canReuseInstruction: 0s without adstack, 55s after multi-slot SSA-promotion (the alloca's loop-carried recurrence is what feeds SCEV), 47s after single-slot specialization (the easy recurrences come out, the harder multi-slot ones stay).Surface API
No public API change.
qd.initand theCompileConfigknobs 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 sixAdStack*visit methods inTaskCodeGenLLVMand three new private helpers (ensure_ad_stack_count_alloca_llvm,emit_ad_stack_top_slot_ptr,emit_ad_stack_single_slot_ptr).Mechanism
Before (one helper call per op)
GVN cannot fold consecutive
stack_pushcalls because they have function-call memory effects plus the relaxed-atomic overflow store acts as a memory barrier; it cannot fold consecutivestack_top_primalcalls either because the helper reads through%stackwhich AA cannot prove disjoint from the runtime metadata. Result: each unrolled iteration emits the full sequence again.After, multi-slot release path (inline IR, per-stack alloca)
After mem2reg promotes
count_0to SSA and GVN folds theadd 1chain 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.After, single-slot release path (constant-offset GEP, no count state)
Single-slot specialization gates on
max_size == 1ANDsize_expr->kind == Const && size_expr->const_value == 1(rejects placeholder cases where the sizer setmax_size = 1because the symbolic bound is non-Const and the host evaluates the actual capacity per launch).Heap layout (unchanged)
The 8-byte header gap is preserved so
ad_stack_per_thread_stride_(sum ofalign_up_8(size_in_bytes)per stack) and the host'sensure_adstack_heap(stride * num_threads)allocation are unchanged. Slot offsets stay atsizeof(u64) + idx * 2 * element_size.Per-backend matrix
Tests
The existing AD test suite already exercises every code path this PR touches:
tests/python/test_adstack.py,test_ad_basics.py,test_ad_for.py,test_ad_if.py,test_ad_atomic.py,test_ad_offload.py,test_ad_demote_dense.py,test_ad_grad_check.py,test_ad_dynamic_index.py- 1041 tests pass on the LLVM CPU backend, covering the full reverse-mode AD surface including unary loop-carried unrolled bodies (the dominant adstack shape on this PR's hot path), nested control-flow inside reverse pass, atomic adjoint accumulation, dynamic indexing, gradient checks against PyTorch autograd, and offload boundary handling. The top-of-stack saturating-count - 1underflow guard, slot-zeroing on push, and per-iteration count restart inside loop bodies all match the runtime helper semantics they replace.This PR also adds four new parametrized tests to
tests/python/test_adstack.pyto broaden coverage of adstack codegen kernel shapes that the existing suite did not pin tightly. The added tests cross-check reverse-mode gradients against PyTorch autograd:test_adstack_linear_only_accumulator_with_nonlinear_operand(parametrized overn_iterin{4, 32}) - kernel that mixes a loop-carried linear accumulator with a non-linear operandsin(a), weaving the accumulator's adjoint chain through the operand's adstack push / pop sites.test_adstack_repeated_load_top_of_outer_value_in_unrolled_inner_loop(parametrized overn_innerin{4, 32}) - kernel that consumes one outer-loop value via a non-linear unary op N times in a row in an unrolled inner loop, producing N consecutiveAdStackLoadTopStmtreads of the same stack inside one straight-line block.test_adstack_unrolled_many_pushes_across_multiple_stacks(parametrized overn_iterin{4, 16, 64}andn_stacksin{1, 3}) - kernel that fans out ton_stacksparallel adstacks, each receivingn_iterstraight-line pushes inside an unrolled inner loop.n_iter=64is well above the adstack capacity floor of 32, so any regression that miscalculates the slot offset under multiple sibling stacks, races a cross-stack increment, or short-circuits the inlinecount - 1saturation, surfaces as a wrong gradient long before it would surface as an overflow at runtime.test_adstack_f16_unrolled_pushes_alignment(parametrized overn_iterin{4, 16}, restricted to LLVM backends because Vulkan / Metal lack atomic-f16 support for the kernel's atomic-add to a scalar) - regression sentinel for the memset alignment fix. The kernel usesqd.f16so each adstack slot is 4-byte aligned (not 8), andn_iter=16ensures multiple non-zero slot offsets get exercised including odd indices where the alignment claim matters most. The current toolchain does not turn the over-stated alignment into a fault on this kernel shape, but a future LLVM / NVPTX / AMDGCN lowering decision that does start trusting the alignment metadata more aggressively trips this test in CI rather than at user runtime.These tests are independent of the codegen change in this PR (they pin gradient correctness, which the runtime-helper path on
mainalready produces correctly), so they pass on bothmainand this branch. They are bundled here to lock down regression coverage for the kernel shapes the inline AdStack codegen exercises hardest.Side-effect audit
mem2regrequires the alloca to be in the entry block;ensure_ad_stack_count_alloca_llvmuses anInsertPointGuardto emit there regardless of where the AdStackAllocaStmt visit site sits. Same pattern asensure_ad_stack_heap_base_llvm/ensure_ad_stack_metadata_llvm.AdStackAllocaStmtis nested inside a loop body the alloca is still created once at the entry block, but the init store runs every iteration, matching the previousstack_init(stack_ptr)semantics where the heap u64 header was zeroed on each entry.Const SizeExprso placeholdermax_size = 1cases (wheredetermine_ad_stack_sizeset 1 because the symbolic bound is non-Const and the host evaluates the actual capacity per launch) fall through to the multi-slot alloca path.emit_ad_stack_top_slot_ptr(multi-slot) andemit_ad_stack_single_slot_ptr(single-slot) both compute slot 0 when count == 0, matchingstack_top_primal'sn > 0 ? n - 1 : 0underflow path. The host raises onruntime->adstack_overflow_flagbefore any garbage value reaches user code.AdStackPushStmtslot index usesold_count(i.e.new_count - 1) directly on the multi-slot path, skipping the saturating-subtract because we just incremented andnew_countis provably >= 1.stack_init / stack_push / stack_pop / stack_top_primal / stack_top_adjointare still called, the bounds-check insidestack_pushstill fires, and the relaxed-atomicadstack_overflow_flagwrite is preserved.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.