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