[timesselens]

PTX only defines atomic instructions (atom, red, and relaxed, acquire or release loads and stores) on global and shared memory, and on generic addresses that point there. On the stack they are undefined, and two things go wrong in practice:

• An atomic on a generic pointer that points to a stack slot fails at run time on Pascal GPUs with CUDA_ERROR_INVALID_ADDRESS_SPACE, which leaves the context unusable (V100 reportedly fails sometimes; newer GPUs have run it fine so far). Enzyme emits these: when a stack slot of the caller is passed to a function that isn't inlined, the adjoint accumulates into its shadow with an atomicrmw fadd, because Enzyme can't see where the pointer points (CUDA Error 717: Invalid Address Space -- Atomic operation not supported on stack memory. EnzymeAD/Enzyme.jl#428). WaterLily.jl's reverse-mode kernels crash this way.
• When NVPTX can see that the pointer is local, it demotes atomicrmw to plain accesses, but only when optimizing, and never cmpxchg, for which it selects an atom.local.cas that ptxas rejects (NVPTX selects atom.local for atomics on local memory, which ptxas rejects CUDA.jl#3402):

  ptxas jl_FK5s9YlihS.ptx, line 141; error   : State space incorrect for instruction 'atom'
  

This PR handles both from the PTX finish_ir!, after optimization and before validation, the way the Metal and SPIR-V back-ends already deal with atomics on thread memory:

• Atomics on pointers that are known to point to the stack (derived from an alloca, or in the local address space) become plain accesses, since no other thread can access that memory. They also no longer fail validation, e.g. for system scope before sm_60.
• Atomics on generic pointers that may point to the stack check isspacep.local at run time, and use plain accesses on local memory and the original atomic otherwise. That is what NVIDIA's libcu++ does for every cuda::atomic_ref operation, on every architecture (Fix UB in atomics with automatic storage NVIDIA/cccl#2586), and what AMDGPU does in its back-end for flat atomics that may hit private memory. The kept atomic gets !noalias.addrspace metadata, so it isn't checked again.

Atomics on global or shared pointers, which is what CUDA.jl's arrays use, are left alone, and NVPTX folds the check when it can infer the address space, e.g. for kernel arguments. In the worst case, a generic atomic in a callee that isn't inlined, a kernel doing 256 atomic additions per thread took 0.60 ms with and without the check on an RTX 5080.

The check covers atomicrmw, cmpxchg, and atomic loads and stores that are at least monotonic. Julia's unordered accesses of heap references aren't checked, as they are too common to pay for it: NVPTX only lowers those to the relaxed instructions that the stack doesn't support from sm_70 on, where they haven't been seen to fail.

The new tests in test/ptx.jl work on IR and PTX, and pass on Julia 1.10 (typed pointers) and 1.13. On an RTX 5080, a kernel doing a cmpxchg on a stack slot fails to compile on main and runs correctly with this PR, as do atomics in a non-inlined callee on its caller's stack slot. I could not test on a Pascal GPU.

[codecov]

Codecov Report

❌ Patch coverage is 95.09804% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.72%. Comparing base (5419b7b) to head (e0b534d).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
src/ptx.jl	95.09%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #963      +/-   ##
==========================================
+ Coverage   89.54%   89.72%   +0.18%     
==========================================
  Files          31       31              
  Lines        6935     7029      +94     
==========================================
+ Hits         6210     6307      +97     
+ Misses        725      722       -3     

☔ 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.

[maleadt]

I'd rather this would be fixed in NVPTX, but a workaround is fine for now. Rebased, and slightly reworked; will merge when it passes CI.

[timesselens]

Thank you! Will report back with Pascal tests

[timesselens]

Confirmed on Pascal (GTX 1080, Linux, Julia 1.12.7): the ptx tests and CUDA.jl's atomics tests pass, and the Enzyme reverse-mode kernel that failed with CUDA_ERROR_INVALID_ADDRESS_SPACE (717) before this PR now gives the right result.