PTX: avoid atomics on thread-private memory

[timesselens]

PTX has no atomic instructions for thread-private (local) memory. An atomicrmw or cmpxchg on a generic pointer that points to a stack slot becomes a generic atom, which on Pascal GPUs faults with CUDA_ERROR_INVALID_ADDRESS_SPACE (717), a sticky error that leaves the context unusable. Enzyme's reverse mode produces such atomics (EnzymeAD/Enzyme.jl#428). This PR proposes turning atomics on memory known to be local into a plain load, operation and store at the end of the PTX pipeline, and guarding atomics on generic pointers of unknown origin with a run-time isspacep.local check.

Before / after

On a GTX 1080 (Pascal), with CUDA.jl 6.4. That needs GPUCompiler 2, so Enzyme.jl comes from EnzymeAD/Enzyme.jl#3512, with the relocation-slot guards suggested there:

using CUDA, Enzyme
struct Pair2; a::Float64; b::Float64; end
@noinline pg(c::Pair2) = c.a * c.a + c.b
k!(y, v) = (i = threadIdx().x; @inbounds y[i] = pg(v[i]); nothing)
∇k!(y, ȳ, v, v̄) = (autodiff_deferred(Reverse, Const(k!), Const, Duplicated(y, ȳ), Duplicated(v, v̄)); nothing)
v̄ = CuArray([Pair2(0, 0)])
CUDA.@sync @cuda ∇k!(CUDA.zeros(1), CUDA.ones(1), CuArray([Pair2(0.5, 3)]), v̄)
Array(v̄)
# main:    CUDA error: operation not supported on global/shared address space (code 717, ERROR_INVALID_ADDRESS_SPACE),
#          then the context is unusable
# this PR: [Pair2(1.0, 1.0)]

(GTX 1080, CUDA.jl 6.4, Julia 1.12.5; with pg inlined, the gradient was already as expected.) In the IR, an atomic on a generic pointer becomes:

; before
%old = atomicrmw fadd ptr %p, float %v monotonic, align 4
; after (the helper is inlined)
%local = call i1 @llvm.nvvm.isspacep.local(ptr %p)
br i1 %local, label %plain, label %atomic      ; plain: load, fadd, store
                                               ; atomic: the original atomicrmw, marked !gpucompiler.local_checked

What seems to happen

• The derivative of a function that isn't inlined accumulates into the shadow of a by-reference argument with atomicrmw fadd, since the callee can't know whether the pointer is private to the thread. Here the caller's shadow is an alloca, so the pointer is generic and points to local memory.
• The PTX ISA on atom: "atom with scalar type may be used only with .global and .shared spaces and with generic addressing, where the address points to .global or .shared space." Pascal faults on such an address. Newer GPUs may tolerate it (only Pascal was tested here), but it is outside the ISA either way.
• LLVM's NVPTXAtomicLower (D98650) demotes an atomicrmw on local memory only when the pointer is already in the local address space after NVPTXLowerAlloca and InferAddressSpaces, which are both function-local. A pointer that arrives as an argument stays generic, and cmpxchg isn't handled.

Proposed change

ptx_local_atomics! (src/ptx.jl) is a module pass at the end of optimize_module!, after inlining has exposed what it can, followed by AlwaysInlinerPass. For every atomicrmw and cmpxchg:

• on a pointer known to be local (derived from an alloca, or in addrspace 5): a plain load, operation and store, keeping the atomic's alignment and volatility;
• on a generic pointer of unknown origin: a call to an alwaysinline helper that branches on llvm.nvvm.isspacep.local. It takes the plain path for local memory and the original atomic otherwise, keeping its ordering, scope, alignment, volatility and weakness. That atomic is marked, so a second run leaves it alone;
• on global, shared and constant memory: nothing changes. CuDeviceArray and shared-memory atomics carry their address space, so they are not touched.

The pass handles every operation that LLVM's C API names; uinc_wrap and udec_wrap stay atomic.

Tests

In test/ptx.jl, inside the "IR" set, so no GPU is needed. The IR is written in typed-pointer syntax, which parses under both pointer regimes.

• "atomics on local memory" (114 tests):
  • stack slot, addrspace(5) and cmpxchg on a stack slot: no atomic left;
  • generic pointer: the check, with ordering, scope, alignment and weakness kept;
  • global, cast-from-global and shared: unchanged;
  • every integer operation; idempotence;
  • the plain form against the atomic, executed on the host through llvmcall, for every operation and six input pairs.
• "atomics on local memory, in the pipeline" (2 tests): PTX.code_llvm of a kernel with an atomic on its own stack slot has no atomicrmw left, and one through a pointer argument gets the isspacep.local check.

Run with julia --project -e 'using Pkg; Pkg.test(test_args=["ptx"])'. With this test file on main, the first set gives 31 failures and 3 errors, and both pipeline tests fail.

Open questions

• The check is unconditional. It could be limited to targets before sm_70, which is where the fault showed up (on Pascal); I kept it for all targets because the ISA doesn't allow these atomics anywhere. Which would you prefer?
• Metal: lower LLVM atomics to AIR atomic intrinsics #942 makes atomics on thread-private memory plain accesses on Metal, with the same test (every object the pointer can be derived from is an alloca). This PR is the PTX side of that; PTX also needs the run-time check for generic pointers, because an atom on local memory faults there. Once Metal: lower LLVM atomics to AIR atomic intrinsics #942 lands, the demotion could share its helper, and I'd be happy to rebase onto it or fold this into your rework.

Related

• CUDA Error 717: Invalid Address Space -- Atomic operation not supported on stack memory. EnzymeAD/Enzyme.jl#428: this crash, reported by jgreener64, with reproducer work by leios. vchuravy and wsmoses diagnosed it there (an alloca passed to a function that isn't inlined; InferAddressSpaces is function-local) and kept the issue open for a later fix. They judged a fix inside Enzyme (knowing across calls that a pointer is thread-local) a large change; this PR lowers the atomics in the back-end instead, as NVPTXAtomicLower does for the known-local case.
• LLVM D98650 (NVPTXAtomicLower): the known-local case in the back-end.
• [2.7.0] Metal: atomic RMW on thread-private memory is silently dropped #934 and Metal: lower LLVM atomics to AIR atomic intrinsics #942: atomics on thread-private memory on Metal. Lower LLVM intrinsics to platform-specific ops #479 and Native LLVM atomics Metal.jl#478 (native LLVM atomics), WIP: Support LLVM atomic intrinsics. #652 and Work towards atomic legalization #701: where atomic legalization is planned.
• Hook the Enzyme kernel rules into the @cuda launch pipeline CUDA.jl#3264 hooks Enzyme's rules into @cuda launches, so raw @cuda reverse mode now reaches this path directly.

Verification

• Pkg.test(test_args=["ptx"]) on macOS aarch64 with Julia 1.10.12, 1.11.9, 1.12.5 and 1.13.0: the two new sets pass (114/114 and 2/2). "Julia value global names" errors on every version, with or without this PR, because the NVPTX back-end can't be loaded on macOS.
• GTX 1080 (driver 580, CUDA runtime 12.9, CUDA.jl 6.4, Julia 1.12.5): the reproducer gives [Pair2(1.0, 1.0)]. WaterLily.jl's reverse-mode test set, whose kernels take Active closures, crashed at its first GPU test before; with this change (and fixes to Enzyme's CUDA math functions) it got past it, and the set as it was then passed on CPU + CUDA. Linux x86_64, on this head: ptx 188/188 (with ptx/precompile). The Before/after block, pasted as is on this head: [Pair2(1.0, 1.0)]; with main, the error shown, followed at exit by cuModuleUnload finalizer errors from the sticky context.

[codecov]

Codecov Report

❌ Patch coverage is 95.10490% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.93%. Comparing base (d36ab74) to head (d8cfd5c).

Files with missing lines	Patch %	Lines
src/ptx.jl	95.10%	7 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #963      +/-   ##
==========================================
+ Coverage   86.65%   86.93%   +0.27%     
==========================================
  Files          29       29              
  Lines        5786     5924     +138     
==========================================
+ Hits         5014     5150     +136     
- Misses        772      774       +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.