From 3a6cf270e920d094c842722134834bc8bdb8762d Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Tue, 21 Jul 2026 23:03:46 -0700 Subject: [PATCH] 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) --- crates/rustc_codegen_nvvm/src/builder.rs | 47 +++++++++++++++---- crates/rustc_codegen_nvvm/src/intrinsic.rs | 7 +-- tests/compiletests/ui/dis/shared_memory.rs | 27 +++++++++++ .../compiletests/ui/dis/shared_memory.stderr | 36 ++++++++++++++ 4 files changed, 105 insertions(+), 12 deletions(-) create mode 100644 tests/compiletests/ui/dis/shared_memory.rs create mode 100644 tests/compiletests/ui/dis/shared_memory.stderr diff --git a/crates/rustc_codegen_nvvm/src/builder.rs b/crates/rustc_codegen_nvvm/src/builder.rs index b864e0ab..28b58b89 100644 --- a/crates/rustc_codegen_nvvm/src/builder.rs +++ b/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)) } +} diff --git a/crates/rustc_codegen_nvvm/src/intrinsic.rs b/crates/rustc_codegen_nvvm/src/intrinsic.rs index f68dcc87..db1e6173 100644 --- a/crates/rustc_codegen_nvvm/src/intrinsic.rs +++ b/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 { diff --git a/tests/compiletests/ui/dis/shared_memory.rs b/tests/compiletests/ui/dis/shared_memory.rs new file mode 100644 index 00000000..2b0690a5 --- /dev/null +++ b/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; 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(); + } +} diff --git a/tests/compiletests/ui/dis/shared_memory.stderr b/tests/compiletests/ui/dis/shared_memory.stderr new file mode 100644 index 00000000..b385a6b8 --- /dev/null +++ b/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; + +} + +