diff --git a/crates/cuda_std/src/warp.rs b/crates/cuda_std/src/warp.rs index 91f1937b..68283181 100644 --- a/crates/cuda_std/src/warp.rs +++ b/crates/cuda_std/src/warp.rs @@ -8,6 +8,9 @@ use crate::gpu_only; use core::arch::asm; use half::{bf16, f16}; +#[path = "warp_control.rs"] +mod control; + /// Synchronizes all of the threads inside of this warp according to `mask`. /// /// # Safety @@ -766,15 +769,7 @@ unsafe fn warp_shuffle_32( fn __nvvm_warp_shuffle(mask: u32, mode: u32, a: u32, b: u32, c: u32) -> u64; } - assert!( - !(width & (width - 1)) != 0 && width <= 32, - "width must be a power of 2 and less than or equal to 32" - ); - - // mimicking nvcc's behavior - let mut c = 0; - c |= 0b11111; - c |= (32 - width) << 8; + let c = control::shuffle_control(width, matches!(mode, WarpShuffleMode::Up)); let result = unsafe { __nvvm_warp_shuffle(mask, mode as u32, value, b, c) }; unpack_warp_result(result) diff --git a/crates/cuda_std/src/warp_control.rs b/crates/cuda_std/src/warp_control.rs new file mode 100644 index 00000000..1d1bb3f5 --- /dev/null +++ b/crates/cuda_std/src/warp_control.rs @@ -0,0 +1,57 @@ +//! PTX shuffle control encoding, shared by every shuffled value width. + +#[allow(dead_code)] // GPU-only callers are replaced with stubs on the host. +#[inline(always)] +pub(crate) fn shuffle_control(width: u32, is_up: bool) -> u32 { + assert!( + width.is_power_of_two() && width <= 32, + "width must be a power of 2 and less than or equal to 32" + ); + // UP compares against a segment's lower bound. Other directions compare + // against its upper bound. The upper bits partition the 32-lane warp. + let clamp = if is_up { 0 } else { 31 }; + ((32 - width) << 8) | clamp +} + +#[cfg(test)] +mod tests { + extern crate std; + use super::shuffle_control; + + #[test] + fn control_bounds_match_logical_warp_segments() { + for width in [1, 2, 4, 8, 16, 32] { + for lane in 0..32_u32 { + for is_up in [false, true] { + let control = shuffle_control(width, is_up); + // Apply the PTX ISA's maxLane expression independently. + let mask = (control >> 8) & 31; + let bound = (lane & mask) | ((control & 31) & !mask); + let segment_start = lane / width * width; + assert_eq!( + bound, + if is_up { + segment_start + } else { + segment_start + width - 1 + } + ); + if is_up { + assert_eq!(lane > bound, lane % width != 0); + } else { + assert_eq!(lane < bound, lane % width + 1 < width); + } + } + } + } + } + + #[test] + fn rejects_zero_non_power_of_two_and_oversized_widths() { + for width in [0, 3, 7, 15, 31, 33, 64, u32::MAX] { + for is_up in [false, true] { + assert!(std::panic::catch_unwind(|| shuffle_control(width, is_up)).is_err()); + } + } + } +} diff --git a/crates/rustc_codegen_nvvm/build.rs b/crates/rustc_codegen_nvvm/build.rs index 55ac3965..aed3a7c8 100644 --- a/crates/rustc_codegen_nvvm/build.rs +++ b/crates/rustc_codegen_nvvm/build.rs @@ -257,9 +257,20 @@ fn configure_libintrinsics(llvm_config: &Path, flavor: &LlvmFlavor) { build_helper::rerun_if_changed(Path::new("libintrinsics.ll")); - let input = manifest_dir.join("libintrinsics.ll"); - let output = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR was not set")) - .join(format!("libintrinsics_v{}.bc", flavor.major)); + let shuffle_version = if flavor.major >= 19 { 19 } else { 7 }; + let shuffle = format!("libintrinsics_shuffle_v{shuffle_version}.ll"); + build_helper::rerun_if_changed(Path::new(&shuffle)); + let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR was not set")); + let input = out_dir.join(format!("libintrinsics_v{}.ll", flavor.major)); + let output = out_dir.join(format!("libintrinsics_v{}.bc", flavor.major)); + // Modern NVVM encodes the shuffle operation in the intrinsic name. Keep + // the legacy wrapper separate so LLVM 7 retains its original interface. + let common = std::fs::read_to_string(manifest_dir.join("libintrinsics.ll")) + .expect("could not read common NVVM intrinsic wrappers"); + let shuffle = std::fs::read_to_string(manifest_dir.join(shuffle)) + .expect("could not read dialect-specific shuffle wrappers"); + std::fs::write(&input, format!("{common}\n{shuffle}")) + .expect("could not write assembled NVVM intrinsic source"); let llvm_as = find_llvm_as(llvm_config, flavor); let status = Command::new(&llvm_as) diff --git a/crates/rustc_codegen_nvvm/libintrinsics.ll b/crates/rustc_codegen_nvvm/libintrinsics.ll index 68ded105..709ad677 100644 --- a/crates/rustc_codegen_nvvm/libintrinsics.ll +++ b/crates/rustc_codegen_nvvm/libintrinsics.ll @@ -160,20 +160,6 @@ declare {i16, i1} @llvm.umul.with.overflow.i16(i16, i16) #0 ; pack into a plain i64 instead: low 32 bits = value, bit 32 = predicate. ; Primitive integer return ⇒ no struct ABI ⇒ no spurious return-attribute. -define i64 @__nvvm_warp_shuffle(i32, i32, i32, i32, i32) #1 { -start: - %r = call { i32, i1 } @llvm.nvvm.shfl.sync.i32(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4) - %val = extractvalue { i32, i1 } %r, 0 - %pred = extractvalue { i32, i1 } %r, 1 - %val64 = zext i32 %val to i64 - %pred64 = zext i1 %pred to i64 - %pred_hi = shl i64 %pred64, 32 - %packed = or i64 %val64, %pred_hi - ret i64 %packed -} - -declare { i32, i1 } @llvm.nvvm.shfl.sync.i32(i32, i32, i32, i32, i32) #1 - define i64 @__nvvm_warp_match_all_32(i32, i32) { start: %r = call { i32, i1 } @llvm.nvvm.match.all.sync.i32(i32 %0, i32 %1) diff --git a/crates/rustc_codegen_nvvm/libintrinsics_shuffle_v19.ll b/crates/rustc_codegen_nvvm/libintrinsics_shuffle_v19.ll new file mode 100644 index 00000000..37b3112f --- /dev/null +++ b/crates/rustc_codegen_nvvm/libintrinsics_shuffle_v19.ll @@ -0,0 +1,39 @@ +; LLVM 19 / modern NVVM dialect: the operation is encoded in the intrinsic +; name, not an extra mode operand. Preserve the Rust-facing packed i64 ABI. +define i64 @__nvvm_warp_shuffle(i32 %mask, i32 %mode, i32 %value, i32 %offset, i32 %clamp) convergent #1 { +start: + switch i32 %mode, label %invalid [ + i32 0, label %idx + i32 1, label %up + i32 2, label %down + i32 3, label %bfly + ] +idx: + %ri = call { i32, i1 } @llvm.nvvm.shfl.sync.idx.i32p(i32 %mask, i32 %value, i32 %offset, i32 %clamp) + br label %pack +up: + %ru = call { i32, i1 } @llvm.nvvm.shfl.sync.up.i32p(i32 %mask, i32 %value, i32 %offset, i32 %clamp) + br label %pack +down: + %rd = call { i32, i1 } @llvm.nvvm.shfl.sync.down.i32p(i32 %mask, i32 %value, i32 %offset, i32 %clamp) + br label %pack +bfly: + %rb = call { i32, i1 } @llvm.nvvm.shfl.sync.bfly.i32p(i32 %mask, i32 %value, i32 %offset, i32 %clamp) + br label %pack +invalid: + unreachable +pack: + %r = phi { i32, i1 } [ %ri, %idx ], [ %ru, %up ], [ %rd, %down ], [ %rb, %bfly ] + %val = extractvalue { i32, i1 } %r, 0 + %pred = extractvalue { i32, i1 } %r, 1 + %val64 = zext i32 %val to i64 + %pred64 = zext i1 %pred to i64 + %pred_hi = shl i64 %pred64, 32 + %packed = or i64 %val64, %pred_hi + ret i64 %packed +} + +declare { i32, i1 } @llvm.nvvm.shfl.sync.idx.i32p(i32, i32, i32, i32) convergent +declare { i32, i1 } @llvm.nvvm.shfl.sync.up.i32p(i32, i32, i32, i32) convergent +declare { i32, i1 } @llvm.nvvm.shfl.sync.down.i32p(i32, i32, i32, i32) convergent +declare { i32, i1 } @llvm.nvvm.shfl.sync.bfly.i32p(i32, i32, i32, i32) convergent diff --git a/crates/rustc_codegen_nvvm/libintrinsics_shuffle_v7.ll b/crates/rustc_codegen_nvvm/libintrinsics_shuffle_v7.ll new file mode 100644 index 00000000..49be3985 --- /dev/null +++ b/crates/rustc_codegen_nvvm/libintrinsics_shuffle_v7.ll @@ -0,0 +1,14 @@ +define i64 @__nvvm_warp_shuffle(i32, i32, i32, i32, i32) #1 { +start: + %r = call { i32, i1 } @llvm.nvvm.shfl.sync.i32(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4) + %val = extractvalue { i32, i1 } %r, 0 + %pred = extractvalue { i32, i1 } %r, 1 + %val64 = zext i32 %val to i64 + %pred64 = zext i1 %pred to i64 + %pred_hi = shl i64 %pred64, 32 + %packed = or i64 %val64, %pred_hi + ret i64 %packed +} + +declare { i32, i1 } @llvm.nvvm.shfl.sync.i32(i32, i32, i32, i32, i32) #1 + diff --git a/crates/rustc_codegen_nvvm/src/nvvm.rs b/crates/rustc_codegen_nvvm/src/nvvm.rs index 8ecff0cc..83f66fe5 100644 --- a/crates/rustc_codegen_nvvm/src/nvvm.rs +++ b/crates/rustc_codegen_nvvm/src/nvvm.rs @@ -341,12 +341,15 @@ unsafe fn internalize_pass(module: &Module, cx: &Context) { } let iter = GlobalIter::new(&module); - for func in iter { - let is_decl = LLVMIsDeclaration(func) == True; - - if !is_decl { - LLVMRustSetLinkage(func, Linkage::InternalLinkage); - LLVMRustSetVisibility(func, Visibility::Default); + for global in iter { + let is_decl = LLVMIsDeclaration(global) == True; + // llvm.used, llvm.compiler.used and other appending globals have + // special linker/optimizer semantics. Internalizing them produces + // invalid LLVM IR and prevents verified pre-NVVM optimization. + let is_appending = LLVMRustGetLinkage(global) == Linkage::AppendingLinkage; + if !is_decl && !is_appending { + LLVMRustSetLinkage(global, Linkage::InternalLinkage); + LLVMRustSetVisibility(global, Visibility::Default); } } } diff --git a/crates/rustc_codegen_nvvm/src/override_fns.rs b/crates/rustc_codegen_nvvm/src/override_fns.rs index 8534d790..38c405fb 100644 --- a/crates/rustc_codegen_nvvm/src/override_fns.rs +++ b/crates/rustc_codegen_nvvm/src/override_fns.rs @@ -44,7 +44,12 @@ fn should_override<'tcx>(func: Instance<'tcx>, cx: &CodegenCx<'_, 'tcx>) -> bool return false; } - let sym = cx.tcx.item_name(func.def_id()); + // Closures can remain as separate codegen items (for example libm's + // rint_round at opt-level=s). They have no item name and cannot name a + // libdevice intrinsic. item_name would ICE instead of compiling them. + let Some(sym) = cx.tcx.opt_item_name(func.def_id()) else { + return false; + }; let name = sym.as_str(); if is_unsupported_libdevice_fn(name) {