Rust compiler generates spec-violating PTX for atomics on the nvptx64-nvidia-cuda target

This issue only affects the nvptx64-nvidia-cuda target, where the sematnics of atomics are slightly odd: they are restricted in what address spaces they can operate in.

Excerpt from the PTX spec:

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.

So, those instructions can be used on pointers to the global & shared address space, but not on the .local or .const or .param address spaces.

Currently, rustc generates PTX which performs such spec-violating operations.

#![no_std]
use core::sync::atomic::{AtomicUsize,Ordering};
#[unsafe(no_mangle)]
fn local_atomic(){
    // We create a local atomic, and operate on it. 
    // This is not very useful, but it is safe Rust code. 
    let local = AtomicUsize::new(0);
    core::hint::black_box(&local).fetch_add(1, Ordering::SeqCst);
}

This results in the following PTX:

{
        // Allocate space in the local address space
        mov.u64         %SPL, __local_depot0;
        // Cast this to a generic pointer 
        cvta.local.u64  %SP, %SPL;
        add.u64         %rd5, %SP, 0;
        add.u64         %rd6, %SPL, 0;
        add.u64         %rd7, %SP, 8;
        add.u64         %rd8, %SPL, 8;
        mov.b64         %rd9, 0;
        st.local.u64    [%rd8], %rd9;
        st.local.u64    [%rd6], %rd7;
        cvt.u32.u64     %r1, %rd5;
        ld.local.u64    %rd1, [%rd6];
        ld.u64  %rd11, [%rd1];
        add.s64         %rd10, %rd11, 1;
        // An atomic operation on a generic pointer to the local address space occurs here!
        // This violates the PTX spec. 
        atom.cas.b64    %rd4, [%rd1], %rd11, %rd10;
        setp.ne.s64     %p1, %rd4, %rd11;
        mov.u64         %rd11, %rd4;
        ret;
}

This is not valid according to the PTX spec. Some NVIDIA tools(e.g. libnvvm, the CUDA compiler) warn about such operations:

__global__ void test_kernel() {
    int local;
    atomicAdd(&local, 1);
}

Warning: Cannot do atomic on local memory

This strongly suggests this is not a valid thing to do, according to the manufacturer.

Furthermore, GCC requires software stack emulation to enable atomic operations on local variables in for their PTX target:

For -msoft-stack (implied by -mgomp), generate code that does not use .local memory or PTX alloca directly for stack storage. Instead, a per-warp stack pointer is maintained explicitly. This enables variable-length stack allocation (with variable-length arrays or alloca), and when global memory is used for underlying storage, makes it possible to access automatic variables from other threads, or with atomic instructions.

So, both official, and unofficial tooling avoid / warn about using the atom PTX instruction on locals.

Potential fixes

This issue can be mitigated checking the address space of a pointer before performing an atomic operation. This is what we do in Rust-CUDA(using atom for address spaces it supports, emulating it for thread-local memory, and traping on anything else).

This can be fixed on the LLVM side(make LLVM always generate spec-compilant PTX, even for this corner case), or on the rustc side(have a target-specifc workaround).

Start by reproducing the no_std local AtomicUsize example for the nvptx64-nvidia-cuda target and inspecting the generated PTX. Read the PTX atomic address-space restrictions and compare whether the fix belongs in rustc or LLVM; done means local or otherwise unsupported atomic operations no longer produce spec-violating PTX.