Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 18 additions & 17 deletions crates/cuda_std/src/warp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -313,20 +313,20 @@ unsafe fn match_any_64(mask: u32, value: u64) -> u32 {
#[inline(always)]
unsafe fn match_all_32(mask: u32, value: u32) -> (u32, bool) {
unsafe extern "C" {
#[allow(improper_ctypes)]
fn __nvvm_warp_match_all_32(mask: u32, value: u32) -> (u32, bool);
// see libintrinsics.ll — packs (value, predicate) into i64
fn __nvvm_warp_match_all_32(mask: u32, value: u32) -> u64;
}
unsafe { __nvvm_warp_match_all_32(mask, value) }
unpack_warp_result(unsafe { __nvvm_warp_match_all_32(mask, value) })
}

#[gpu_only]
#[inline(always)]
unsafe fn match_all_64(mask: u32, value: u64) -> (u32, bool) {
unsafe extern "C" {
#[allow(improper_ctypes)]
fn __nvvm_warp_match_all_64(mask: u32, value: u64) -> (u32, bool);
// see libintrinsics.ll — packs (value, predicate) into i64
fn __nvvm_warp_match_all_64(mask: u32, value: u64) -> u64;
}
unsafe { __nvvm_warp_match_all_64(mask, value) }
unpack_warp_result(unsafe { __nvvm_warp_match_all_64(mask, value) })
}

/// Synchronizes a subset of threads in a warp then performs a reduce-and-broadcast
Expand Down Expand Up @@ -741,14 +741,16 @@ pub enum WarpShuffleMode {
Xor = 3,
}

// C-compatible struct to match LLVM IR's {i32, i8} return type
// This fixes an ABI mismatch where Rust would represent (u32, bool) as [2 x i32]
// but the LLVM intrinsic returns {i32, i8} (a struct, not an array)
#[doc(hidden)]
#[repr(C)]
pub struct WarpShuffleResult {
value: u32,
predicate: u8,
// The libintrinsics.ll wrappers pack their (value, predicate) result into a
// single i64: low 32 bits = value, bit 32 = predicate. Returning a primitive
// integer avoids the small-aggregate ABI path where rustc attaches `align N`
// to the call's return value — an attribute LLVM 19's verifier rejects on
// non-pointer returns.
// Unused on host targets — every caller is `#[gpu_only]`.
#[allow(dead_code)]
#[inline(always)]
fn unpack_warp_result(packed: u64) -> (u32, bool) {
(packed as u32, (packed >> 32) & 1 != 0)
}

#[gpu_only]
Expand All @@ -761,8 +763,7 @@ unsafe fn warp_shuffle_32(
) -> (u32, bool) {
unsafe extern "C" {
// see libintrinsics.ll
// Returns {i32, i8} in LLVM IR, which maps to our WarpShuffleResult struct
fn __nvvm_warp_shuffle(mask: u32, mode: u32, a: u32, b: u32, c: u32) -> WarpShuffleResult;
fn __nvvm_warp_shuffle(mask: u32, mode: u32, a: u32, b: u32, c: u32) -> u64;
}

assert!(
Expand All @@ -776,7 +777,7 @@ unsafe fn warp_shuffle_32(
c |= (32 - width) << 8;

let result = unsafe { __nvvm_warp_shuffle(mask, mode as u32, value, b, c) };
(result.value, result.predicate != 0)
unpack_warp_result(result)
}

unsafe fn warp_shuffle_128(
Expand Down
62 changes: 35 additions & 27 deletions crates/rustc_codegen_nvvm/libintrinsics.ll
Original file line number Diff line number Diff line change
Expand Up @@ -152,44 +152,52 @@ start:
}
declare {i16, i1} @llvm.umul.with.overflow.i16(i16, i16) #0

; Required because we need to explicitly generate { i32, i1 } for the following intrinsics
; except rustc will not generate them (it will make { i32, i8 }) which libnvvm rejects.

define { i32, i8 } @__nvvm_warp_shuffle(i32, i32, i32, i32, i32) #1 {
; NVVM intrinsics return { i32, i1 }, but rustc lowering of (u32, bool) — or any
; small two-field aggregate — produces { i32, i8 }, which libnvvm rejects. We
; used to bridge by re-packing into { i32, i8 } here, but that aggregate return
; causes rustc's call-site ABI to attach `align N` to the return value, which
; LLVM 19's verifier rejects (align is only valid on pointer returns). So we
; 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:
%5 = call { i32, i1 } @llvm.nvvm.shfl.sync.i32(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4)
%6 = extractvalue { i32, i1 } %5, 1
%7 = zext i1 %6 to i8
%8 = extractvalue { i32, i1 } %5, 0
%9 = insertvalue { i32, i8 } undef, i32 %8, 0
%10 = insertvalue { i32, i8 } %9, i8 %7, 1
ret { i32, i8 } %10
%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 { i32, i8 } @__nvvm_warp_match_all_32(i32, i32) {
define i64 @__nvvm_warp_match_all_32(i32, i32) {
start:
%2 = call { i32, i1 } @llvm.nvvm.match.all.sync.i32(i32 %0, i32 %1)
%3 = extractvalue { i32, i1 } %2, 1
%4 = zext i1 %3 to i8
%5 = extractvalue { i32, i1 } %2, 0
%6 = insertvalue { i32, i8 } undef, i32 %5, 0
%7 = insertvalue { i32, i8 } %6, i8 %4, 1
ret { i32, i8 } %7
%r = call { i32, i1 } @llvm.nvvm.match.all.sync.i32(i32 %0, i32 %1)
%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.match.all.sync.i32(i32, i32) #1

define { i32, i8 } @__nvvm_warp_match_all_64(i32, i64) {
define i64 @__nvvm_warp_match_all_64(i32, i64) {
start:
%2 = call { i32, i1 } @llvm.nvvm.match.all.sync.i64(i32 %0, i64 %1)
%3 = extractvalue { i32, i1 } %2, 1
%4 = zext i1 %3 to i8
%5 = extractvalue { i32, i1 } %2, 0
%6 = insertvalue { i32, i8 } undef, i32 %5, 0
%7 = insertvalue { i32, i8 } %6, i8 %4, 1
ret { i32, i8 } %7
%r = call { i32, i1 } @llvm.nvvm.match.all.sync.i64(i32 %0, i64 %1)
%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.match.all.sync.i64(i32, i64) #1
Expand Down
Loading