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

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;
+
+}
+
+