fix(nvvm): keep address-space-qualified pointers out of the generic space

Every access to an `#[address_space(shared)]` (or `constant`) static was
emitted as a *generic* access behind an `addrspacecast`, because `load`,
`volatile_load`, `gep`, `inbounds_gep` and `check_store` all cast their
pointer operand through `type_ptr_to`, which always yields an addrspace 0
pointer.

Correctness then rests entirely on libNVVM's `InferAddressSpaces` pass
folding the cast back into the access. It usually does, which is why this
mostly works. When it does not, the access is emitted as a generic `ld`/`st`
against the raw shared window offset -- an out-of-bounds access at runtime,
with no error or warning at compile time, and only on the warps whose data
reaches the affected branch.

Cast the pointer operand in the address space it already points into
instead. For the kernel added in this commit the emitted NVVM IR goes from

    %22 = getelementptr inbounds i32,
            i32* addrspacecast (i32 addrspace(3)* ... @sh ... to i32*), i64 %18
    store i32 %1, i32* %22, align 4

to

    %22 = getelementptr inbounds i32, i32 addrspace(3)* ... @sh ..., i64 %18
    store i32 %1, i32 addrspace(3)* %22, align 4

`raw_eq`'s integer-compare fast path had the same defect and is fixed with
it. Casts that are *semantically* required -- a pointer passed to a callee
whose parameter is generic, or to `memcmp` -- are left alone; those still go
through `bitcast`, which addrspacecasts to generic as before.

Add a `dis` compiletest pinning the shared-memory codegen shape: the kernel
must use `st.shared`/`ld.shared` on the raw symbol offset and needs no
`cvta.shared`.

Refs #402

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: nh13

crates/rustc_codegen_nvvm/src/builder.rs

@@ -509,7 +509,7 @@ impl<'ll, 'tcx, 'a> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 
     fn load(&mut self, ty: &'ll Type, ptr: &'ll Value, align: Align) -> &'ll Value {
         trace!("Load {ty:?} {:?}", ptr);
-        let ptr = self.pointercast(ptr, self.cx.type_ptr_to(ty));
+        let ptr = self.pointercast_preserving_addrspace(ptr, ty);
         unsafe {
             #[cfg(feature = "llvm19")]
             let load = llvm::LLVMBuildLoad2(self.llbuilder, ty, ptr, UNNAMED);
@@ -522,7 +522,7 @@ impl<'ll, 'tcx, 'a> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 
     fn volatile_load(&mut self, ty: &'ll Type, ptr: &'ll Value) -> &'ll Value {
         trace!("Volatile load `{:?}`", ptr);
-        let ptr = self.pointercast(ptr, self.cx.type_ptr_to(ty));
+        let ptr = self.pointercast_preserving_addrspace(ptr, ty);
         unsafe {
             #[cfg(feature = "llvm19")]
             let load = llvm::LLVMBuildLoad2(self.llbuilder, ty, ptr, UNNAMED);
@@ -794,7 +794,7 @@ impl<'ll, 'tcx, 'a> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 
     fn gep(&mut self, ty: &'ll Type, ptr: &'ll Value, indices: &[&'ll Value]) -> &'ll Value {
         trace!("gep: {ty:?} {:?} with indices {:?}", ptr, indices);
-        let ptr = self.pointercast(ptr, self.cx().type_ptr_to(ty));
+        let ptr = self.pointercast_preserving_addrspace(ptr, ty);
         unsafe {
             llvm::LLVMBuildGEP2(
                 self.llbuilder,
@@ -814,7 +814,7 @@ impl<'ll, 'tcx, 'a> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         indices: &[&'ll Value],
     ) -> &'ll Value {
         trace!("gep inbounds: {ty:?} {:?} with indices {:?}", ptr, indices);
-        let ptr = self.pointercast(ptr, self.cx().type_ptr_to(ty));
+        let ptr = self.pointercast_preserving_addrspace(ptr, ty);
         unsafe {
             llvm::LLVMBuildInBoundsGEP2(
                 self.llbuilder,
@@ -1497,17 +1497,35 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
 
     fn noundef_metadata(&mut self, _load: &'ll Value) {}
 
+    /// Casts `ptr` to a pointer to `ty`, leaving it in the address space it already points into.
+    ///
+    /// Use this rather than casting to `type_ptr_to(ty)` for the pointer operand of a GEP or of
+    /// a memory access. `type_ptr_to` always yields a generic (addrspace 0) pointer, so casting
+    /// through it turns every access to an `#[address_space(shared)]` (or `constant`) static
+    /// into a generic access behind an `addrspacecast`. libNVVM's `InferAddressSpaces` pass
+    /// usually folds those back into `ld.shared`/`st.shared`, but it is not obliged to, and
+    /// when it misses one the access is emitted as a generic `ld`/`st` against the raw shared
+    /// window offset -- an out-of-bounds access at runtime, with no diagnostic at compile time.
+    pub(crate) fn pointercast_preserving_addrspace(
+        &mut self,
+        ptr: &'ll Value,
+        ty: &'ll Type,
+    ) -> &'ll Value {
+        let addrspace = pointer_addrspace(self.cx, ptr);
+        self.pointercast(ptr, self.cx.type_ptr_to_ext(ty, addrspace))
+    }
+
     fn check_store(&mut self, val: &'ll Value, ptr: &'ll Value) -> &'ll Value {
         let dest_ptr_ty = self.cx.val_ty(ptr);
         let stored_ty = self.cx.val_ty(val);
-        let stored_ptr_ty = self.cx.type_ptr_to(stored_ty);
-
-        assert_eq!(self.cx.type_kind(dest_ptr_ty), TypeKind::Pointer);
+        let stored_ptr_ty = self
+            .cx
+            .type_ptr_to_ext(stored_ty, pointer_addrspace(self.cx, ptr));
 
         if dest_ptr_ty == stored_ptr_ty {
             ptr
         } else {
-            self.bitcast(ptr, stored_ptr_ty)
+            self.pointercast(ptr, stored_ptr_ty)
         }
     }
 
@@ -1601,7 +1619,11 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
 impl<'ll, 'tcx, 'a> Builder<'a, 'll, 'tcx> {
     /// Implements a standard atomic, using LLVM intrinsics(in `atomic_supported`, if `dst` is in a supported address space)
     /// or emulation(with `emulate_local`, if `dst` points to a thread-local address space).
-    /// FIXME(FractalFir): this code assumess all pointers are generic. Adjust it once we support address spaces.
+    /// `dst` may already be in a specific address space, in which case the `isspacep` checks
+    /// below are redundant: `check_call` casts them to generic for the predicate calls, and the
+    /// atomic itself stays on the original pointer.
+    /// FIXME(FractalFir): skip the runtime predicate dance entirely when the address space of
+    /// `dst` is statically known.
     fn atomic_op(
         &mut self,
         dst: &'ll Value,
@@ -1694,3 +1716,10 @@ impl<'ll, 'tcx, 'a> Builder<'a, 'll, 'tcx> {
         )
     }
 }
+
+/// The address space `ptr` points into.
+fn pointer_addrspace<'ll>(cx: &CodegenCx<'ll, '_>, ptr: &'ll Value) -> AddressSpace {
+    let ty = cx.val_ty(ptr);
+    assert_eq!(cx.type_kind(ty), TypeKind::Pointer);
+    unsafe { AddressSpace(llvm::LLVMGetPointerAddressSpace(ty)) }
+}

crates/rustc_codegen_nvvm/src/intrinsic.rs

@@ -705,10 +705,11 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     self.const_bool(true)
                 } else if use_integer_compare {
                     let integer_ty = self.type_ix(layout.size.bits());
-                    let ptr_ty = self.type_ptr_to(integer_ty);
-                    let a_ptr = self.bitcast(a, ptr_ty);
+                    // Keep both operands in their own address space; casting to a generic
+                    // pointer here would make the loads below generic accesses.
+                    let a_ptr = self.pointercast_preserving_addrspace(a, integer_ty);
                     let a_val = self.load(integer_ty, a_ptr, layout.align.abi);
-                    let b_ptr = self.bitcast(b, ptr_ty);
+                    let b_ptr = self.pointercast_preserving_addrspace(b, integer_ty);
                     let b_val = self.load(integer_ty, b_ptr, layout.align.abi);
                     self.icmp(IntPredicate::IntEQ, a_val, b_val)
                 } else {

tests/compiletests/ui/dis/shared_memory.rs

@@ -0,0 +1,27 @@
+// build-pass
+// compile-flags: -Cllvm-args=--disassemble-entry=shared_broadcast --error-format=human -Cdebuginfo=0
+
+// Regression test: accesses to an `#[address_space(shared)]` static must be emitted in the
+// shared address space, not as generic accesses behind an `addrspacecast` to addrspace 0.
+//
+// The kernel below must use `st.shared`/`ld.shared` on the raw symbol offset, and must contain
+// no `cvta.shared` -- there is nothing here that needs a generic pointer.
+
+use core::mem::MaybeUninit;
+use cuda_std::{address_space, kernel, thread};
+
+#[address_space(shared)]
+static mut SH: [MaybeUninit<u32>; 32] = [MaybeUninit::uninit(); 32];
+
+#[kernel]
+pub unsafe fn shared_broadcast(out: *mut u32, n: u32) {
+    let lane = thread::thread_idx_x() & 31;
+    let warp = ((thread::thread_idx_x() >> 5) & 31) as usize;
+    unsafe {
+        if lane == 0 {
+            SH[warp] = MaybeUninit::new(n);
+        }
+        thread::sync_threads();
+        *out = SH[warp].assume_init();
+    }
+}

tests/compiletests/ui/dis/shared_memory.stderr

@@ -0,0 +1,36 @@
+.visible .entry shared_broadcast(
+    .param .u64 shared_broadcast_param_0,
+    .param .u32 shared_broadcast_param_1
+)
+{
+    .reg .pred     %p<2>;
+    .reg .b32     %r<6>;
+    .reg .b64     %rd<6>;
+    // demoted variable
+    .shared .align 4 .b8 _RNvCs6I4kclYupDO_13shared_memory2SH[128];
+
+    ld.param.u64     %rd2, [shared_broadcast_param_0];
+    ld.param.u32     %r1, [shared_broadcast_param_1];
+    mov.u32     %r2, %tid.x;
+    and.b32      %r3, %r2, 31;
+    shr.u32     %r4, %r2, 5;
+    setp.eq.s32     %p1, %r3, 0;
+    mul.wide.u32     %rd3, %r4, 4;
+    mov.u64     %rd4, _RNvCs6I4kclYupDO_13shared_memory2SH;
+    add.s64     %rd1, %rd4, %rd3;
+    @%p1 bra     $L__BB0_1;
+    bra.uni     $L__BB0_2;
+
+$L__BB0_1:
+    st.shared.u32     [%rd1], %r1;
+
+$L__BB0_2:
+    bar.sync     0;
+    ld.shared.u32     %r5, [%rd1];
+    cvta.to.global.u64     %rd5, %rd2;
+    st.global.u32     [%rd5], %r5;
+    ret;
+
+}
+
+