diff --git a/cuda-oxide/crates/cuda-device/src/cooperative_groups.rs b/cuda-oxide/crates/cuda-device/src/cooperative_groups.rs index 07d50d592f..c6c1de84a3 100644 --- a/cuda-oxide/crates/cuda-device/src/cooperative_groups.rs +++ b/cuda-oxide/crates/cuda-device/src/cooperative_groups.rs @@ -47,8 +47,8 @@ //! On top of the typed handles this module also ships the four //! collective primitives `warp_reduce`, `warp_scan`, `block_reduce`, //! and `block_scan`, generic over a [`ReduceOp`](ops::ReduceOp) trait -//! with concrete impls for `Sum`/`Min`/`Max` (`u32`/`i32`/`f32`) and -//! `BitAnd`/`BitOr`/`BitXor` (`u32` only). Block-scoped variants +//! with concrete impls for `Sum`/`Min`/`Max` (`u32`/`i32`/`f32`/`u64`/`i64`) +//! and `BitAnd`/`BitOr`/`BitXor` (`u32`/`u64`). Block-scoped variants //! take a `*mut SharedArray` for scratch and obey the //! standard "block-sync between back-to-back uses of the same scratch" //! contract documented on each function. @@ -165,6 +165,51 @@ pub trait WarpCollective: ThreadGroup { /// `f32` variant of [`shfl_up`](Self::shfl_up). fn shfl_up_f32(&self, var: f32, delta: u32) -> f32; + /// `u64` variant of [`shfl`](Self::shfl). + /// + /// PTX has no 64-bit shuffle, so this moves the two 32-bit halves with + /// two `shfl.sync.idx.b32` under the same mask and source. Convert `i64` + /// with `as u64` / `as i64`, and `f64` with `to_bits` / `from_bits`. + /// + /// The default exchanges the halves through [`shfl`](Self::shfl). The + /// groups in this module override all four `u64` methods with + /// [`crate::warp::shuffle_u64_sync`] and its siblings, which emit both + /// halves from one inline-PTX block. + #[inline(always)] + fn shfl_u64(&self, var: u64, src_rank: u32) -> u64 { + join_halves( + self.shfl(var as u32, src_rank), + self.shfl((var >> 32) as u32, src_rank), + ) + } + + /// `u64` variant of [`shfl_xor`](Self::shfl_xor). + #[inline(always)] + fn shfl_xor_u64(&self, var: u64, lane_mask: u32) -> u64 { + join_halves( + self.shfl_xor(var as u32, lane_mask), + self.shfl_xor((var >> 32) as u32, lane_mask), + ) + } + + /// `u64` variant of [`shfl_down`](Self::shfl_down). + #[inline(always)] + fn shfl_down_u64(&self, var: u64, delta: u32) -> u64 { + join_halves( + self.shfl_down(var as u32, delta), + self.shfl_down((var >> 32) as u32, delta), + ) + } + + /// `u64` variant of [`shfl_up`](Self::shfl_up). + #[inline(always)] + fn shfl_up_u64(&self, var: u64, delta: u32) -> u64 { + join_halves( + self.shfl_up(var as u32, delta), + self.shfl_up((var >> 32) as u32, delta), + ) + } + /// Group-relative bitmask of lanes whose `value` equals mine. /// /// PTX `match.any.sync.b32` (sm_70+). Raw hardware lane bits are @@ -193,6 +238,12 @@ pub trait WarpCollective: ThreadGroup { // Helpers (not part of the public surface) // ============================================================================= +/// Reassembles a `u64` from the low and high 32-bit halves. +#[inline(always)] +fn join_halves(lo: u32, hi: u32) -> u64 { + ((hi as u64) << 32) | lo as u64 +} + /// Total threads per block. #[inline(always)] fn threads_per_block() -> u32 { @@ -537,6 +588,27 @@ impl WarpCollective for WarpTile { warp::shuffle_up_f32_sync(self.mask(), var, self.up_operand_or_self(delta)) } + #[inline(always)] + fn shfl_u64(&self, var: u64, src_rank: u32) -> u64 { + let abs_src = self.tile_base_lane() | (src_rank & (N - 1)); + warp::shuffle_u64_sync(self.mask(), var, abs_src) + } + + #[inline(always)] + fn shfl_xor_u64(&self, var: u64, lane_mask: u32) -> u64 { + warp::shuffle_xor_u64_sync(self.mask(), var, self.xor_operand_or_self(lane_mask)) + } + + #[inline(always)] + fn shfl_down_u64(&self, var: u64, delta: u32) -> u64 { + warp::shuffle_down_u64_sync(self.mask(), var, self.down_operand_or_self(delta)) + } + + #[inline(always)] + fn shfl_up_u64(&self, var: u64, delta: u32) -> u64 { + warp::shuffle_up_u64_sync(self.mask(), var, self.up_operand_or_self(delta)) + } + #[inline(always)] fn match_any(&self, value: u32) -> u32 { let raw = warp::match_any_sync(self.mask(), value); @@ -634,6 +706,23 @@ impl CoalescedThreads { packed } + /// Absolute warp lane of the member whose group rank is `src_rank`, or + /// the calling lane when `src_rank` is out of range. + #[inline(always)] + fn rank_to_lane(&self, src_rank: u32) -> u32 { + let mut m = self.mask; + let mut k = src_rank; + while k > 0 && m != 0 { + m &= m - 1; + k -= 1; + } + if m == 0 { + warp::lane_id() + } else { + m.trailing_zeros() + } + } + #[inline(always)] fn xor_operand_or_self(&self, lane_mask: u32) -> u32 { let source = warp::lane_id() ^ (lane_mask & 31); @@ -702,9 +791,9 @@ impl WarpCollective for CoalescedThreads { warp::any_sync(self.mask, predicate) } - // `shfl` translates the group-relative `src_rank` to an absolute - // warp lane (the lane whose own `thread_rank` inside the coalesced - // group equals `src_rank`) by popping the `src_rank` lowest set + // The index shuffles translate the group-relative `src_rank` to an + // absolute warp lane (the lane whose own `thread_rank` inside the + // coalesced group equals `src_rank`) by popping the `src_rank` lowest set // bits of the participation mask. When `src_rank` is a runtime // value the loop survives as a small `popc`/branch sequence; when // it folds to a constant the loop unrolls. @@ -715,18 +804,7 @@ impl WarpCollective for CoalescedThreads { #[inline(always)] fn shfl(&self, var: u32, src_rank: u32) -> u32 { - let mut m = self.mask; - let mut k = src_rank; - while k > 0 && m != 0 { - m &= m - 1; - k -= 1; - } - let abs_lane = if m == 0 { - warp::lane_id() - } else { - m.trailing_zeros() - }; - warp::shuffle_sync(self.mask, var, abs_lane) + warp::shuffle_sync(self.mask, var, self.rank_to_lane(src_rank)) } #[inline(always)] @@ -746,18 +824,7 @@ impl WarpCollective for CoalescedThreads { #[inline(always)] fn shfl_f32(&self, var: f32, src_rank: u32) -> f32 { - let mut m = self.mask; - let mut k = src_rank; - while k > 0 && m != 0 { - m &= m - 1; - k -= 1; - } - let abs_lane = if m == 0 { - warp::lane_id() - } else { - m.trailing_zeros() - }; - warp::shuffle_f32_sync(self.mask, var, abs_lane) + warp::shuffle_f32_sync(self.mask, var, self.rank_to_lane(src_rank)) } #[inline(always)] @@ -775,6 +842,26 @@ impl WarpCollective for CoalescedThreads { warp::shuffle_up_f32_sync(self.mask, var, self.up_operand_or_self(delta)) } + #[inline(always)] + fn shfl_u64(&self, var: u64, src_rank: u32) -> u64 { + warp::shuffle_u64_sync(self.mask, var, self.rank_to_lane(src_rank)) + } + + #[inline(always)] + fn shfl_xor_u64(&self, var: u64, lane_mask: u32) -> u64 { + warp::shuffle_xor_u64_sync(self.mask, var, self.xor_operand_or_self(lane_mask)) + } + + #[inline(always)] + fn shfl_down_u64(&self, var: u64, delta: u32) -> u64 { + warp::shuffle_down_u64_sync(self.mask, var, self.down_operand_or_self(delta)) + } + + #[inline(always)] + fn shfl_up_u64(&self, var: u64, delta: u32) -> u64 { + warp::shuffle_up_u64_sync(self.mask, var, self.up_operand_or_self(delta)) + } + #[inline(always)] fn match_any(&self, value: u32) -> u32 { self.pack_lanes(warp::match_any_sync(self.mask, value)) @@ -836,15 +923,15 @@ pub mod ops { #[derive(Copy, Clone, Debug)] pub struct Max; - /// Bitwise AND (`u32` only — meaningless for signed/float). + /// Bitwise AND (`u32` and `u64` only — meaningless for signed/float). #[derive(Copy, Clone, Debug)] pub struct BitAnd; - /// Bitwise OR (`u32` only). + /// Bitwise OR (`u32` and `u64` only). #[derive(Copy, Clone, Debug)] pub struct BitOr; - /// Bitwise XOR (`u32` only). + /// Bitwise XOR (`u32` and `u64` only). #[derive(Copy, Clone, Debug)] pub struct BitXor; @@ -976,14 +1063,111 @@ pub mod ops { if a > b { a } else { b } } } + + // --- u64 --- + impl ReduceOp for Sum { + #[inline(always)] + fn identity() -> u64 { + 0 + } + #[inline(always)] + fn combine(a: u64, b: u64) -> u64 { + a.wrapping_add(b) + } + } + impl ReduceOp for Min { + #[inline(always)] + fn identity() -> u64 { + u64::MAX + } + #[inline(always)] + fn combine(a: u64, b: u64) -> u64 { + if a < b { a } else { b } + } + } + impl ReduceOp for Max { + #[inline(always)] + fn identity() -> u64 { + 0 + } + #[inline(always)] + fn combine(a: u64, b: u64) -> u64 { + if a > b { a } else { b } + } + } + impl ReduceOp for BitAnd { + #[inline(always)] + fn identity() -> u64 { + u64::MAX + } + #[inline(always)] + fn combine(a: u64, b: u64) -> u64 { + a & b + } + } + impl ReduceOp for BitOr { + #[inline(always)] + fn identity() -> u64 { + 0 + } + #[inline(always)] + fn combine(a: u64, b: u64) -> u64 { + a | b + } + } + impl ReduceOp for BitXor { + #[inline(always)] + fn identity() -> u64 { + 0 + } + #[inline(always)] + fn combine(a: u64, b: u64) -> u64 { + a ^ b + } + } + + // --- i64 --- + impl ReduceOp for Sum { + #[inline(always)] + fn identity() -> i64 { + 0 + } + #[inline(always)] + fn combine(a: i64, b: i64) -> i64 { + a.wrapping_add(b) + } + } + impl ReduceOp for Min { + #[inline(always)] + fn identity() -> i64 { + i64::MAX + } + #[inline(always)] + fn combine(a: i64, b: i64) -> i64 { + if a < b { a } else { b } + } + } + impl ReduceOp for Max { + #[inline(always)] + fn identity() -> i64 { + i64::MIN + } + #[inline(always)] + fn combine(a: i64, b: i64) -> i64 { + if a > b { a } else { b } + } + } } /// Scalars that ride a `WarpCollective` shuffle. /// /// Reductions and scans need to shuffle the value being combined, but the -/// underlying PTX `shfl.sync.{bfly,up,down,idx}` only ships in `b32` and -/// `f32` flavors. This trait routes each supported `T` (`u32`, `i32`, -/// `f32`) to the right one and bitcasts where required (`i32`). +/// underlying PTX `shfl.sync.{bfly,up,down,idx}` only moves 32-bit +/// registers. This trait routes each supported `T` (`u32`, `i32`, `f32`, +/// `u64`, `i64`) to the matching [`WarpCollective`] shuffle and bitcasts +/// where required (`i32` through `u32`, `i64` through `u64`). The 64-bit +/// types cost two `b32` shuffles per exchange, and `redux.sync` has no +/// 64-bit form, so they always take the shuffle path. pub trait WarpShuffle: Copy + Sized { /// Butterfly exchange under the group's mask. See /// [`WarpCollective::shfl_xor`]. @@ -1029,6 +1213,29 @@ impl WarpShuffle for f32 { } } +impl WarpShuffle for u64 { + #[inline(always)] + fn shfl_xor_via(group: &G, value: Self, lane_mask: u32) -> Self { + group.shfl_xor_u64(value, lane_mask) + } + #[inline(always)] + fn shfl_up_via(group: &G, value: Self, delta: u32) -> Self { + group.shfl_up_u64(value, delta) + } +} + +impl WarpShuffle for i64 { + // same two's-complement round-trip as `i32`, through the `u64` shuffle + #[inline(always)] + fn shfl_xor_via(group: &G, value: Self, lane_mask: u32) -> Self { + group.shfl_xor_u64(value as u64, lane_mask) as i64 + } + #[inline(always)] + fn shfl_up_via(group: &G, value: Self, delta: u32) -> Self { + group.shfl_up_u64(value as u64, delta) as i64 + } +} + /// Reduce one value per lane down to a single value across a `WarpTile`. /// /// Implements the standard butterfly reduction via `shfl.sync.bfly` — @@ -1039,8 +1246,9 @@ impl WarpShuffle for f32 { /// /// # Performance: integers on `sm_80`+ have a one-instruction form /// -/// This is generic over `T`, so a full warp always costs five shuffles plus -/// five combines. For **integer** reductions on Ampere and newer, +/// This is generic over `T`, so a full warp always costs five exchanges plus +/// five combines; each exchange is one `shfl.sync` for 32-bit types and two +/// for `u64`/`i64`. For **32-bit integer** reductions on Ampere and newer, /// [`crate::warp::redux_sync_add`] and its `min`/`max`/`and`/`or`/`xor` /// siblings do the whole reduction in a single `redux.sync` instruction. /// Measured on an A10G (`sm_86`) over 1M threads reducing 64 times each: @@ -1213,6 +1421,26 @@ impl WarpCollective for ContiguousWarpPrefix { warp::shuffle_up_f32_sync(self.0.mask, var, delta) } + #[inline(always)] + fn shfl_u64(&self, var: u64, src_rank: u32) -> u64 { + self.0.shfl_u64(var, src_rank) + } + + #[inline(always)] + fn shfl_xor_u64(&self, var: u64, lane_mask: u32) -> u64 { + self.0.shfl_xor_u64(var, lane_mask) + } + + #[inline(always)] + fn shfl_down_u64(&self, var: u64, delta: u32) -> u64 { + self.0.shfl_down_u64(var, delta) + } + + #[inline(always)] + fn shfl_up_u64(&self, var: u64, delta: u32) -> u64 { + warp::shuffle_up_u64_sync(self.0.mask, var, delta) + } + #[inline(always)] fn match_any(&self, value: u32) -> u32 { self.0.match_any(value) diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/Cargo.lock new file mode 100644 index 0000000000..81e45daa59 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/Cargo.lock @@ -0,0 +1,753 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[package]] +name = "bindgen" +version = "0.69.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "lazy_static", + "lazycell", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn", + "which", +] + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clang-sys" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "coop_groups_u64" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cuda-artifact-finalizer" +version = "0.2.1" +dependencies = [ + "libnvvm-sys", + "nvjitlink-sys", + "serde", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-bindings" +version = "0.4.0" +dependencies = [ + "bindgen", + "libloading", + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-core" +version = "0.4.0" +dependencies = [ + "anyhow", + "cuda-bindings", + "cuda-core-derive", + "half", + "oxide-artifacts", +] + +[[package]] +name = "cuda-core-derive" +version = "0.4.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cuda-device" +version = "0.2.1" +dependencies = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-macros", + "half", + "libc", + "oxide-artifacts", + "ptx-parse", + "sha2", + "thiserror", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[package]] +name = "cuda-target-spec" +version = "0.1.0" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnvvm-sys" +version = "0.2.1" +dependencies = [ + "cuda-target-spec", + "libloading", + "thiserror", +] + +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nvjitlink-sys" +version = "0.2.1" +dependencies = [ + "libloading", + "thiserror", +] + +[[package]] +name = "object" +version = "0.36.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +dependencies = [ + "crc32fast", + "hashbrown 0.15.5", + "indexmap", + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "oxide-artifacts" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d9a941c0c43d2dceaf25f5c9ca4c9f9cceb28dfa5e4e3cfd14e476f003e7c9" +dependencies = [ + "object", +] + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "ptx-parse" +version = "0.1.0" + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reserved-oxide-symbols" +version = "0.2.1" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.59.0", +] + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.154" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/Cargo.toml new file mode 100644 index 0000000000..8fe81365f5 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "coop_groups_u64" +version = "0.1.0" +edition = "2024" +authors = ["NVIDIA Corporation"] +license = "Apache-2.0" + +[workspace] + +[dependencies] +cuda-core = { path = "../../../../../cuda-core" } +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/src/main.rs new file mode 100644 index 0000000000..569ea15fe2 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_u64/src/main.rs @@ -0,0 +1,889 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! 64-bit reductions, scans and shuffles through `cooperative_groups`. +//! +//! `warp_reduce`, `warp_scan`, `block_reduce` and `block_scan` accept `u64` +//! with every combiner in `ops`, and `i64` with `Sum`, `Min` and `Max`. PTX +//! has no 64-bit shuffle, so each exchange moves the two 32-bit halves. A +//! collective that loses the high word, drops a carry between the halves, +//! orders values by their low words, or reads across a tile boundary gives +//! a different answer from the host reference here. +//! +//! The kernels cover: +//! +//! - `tile_1` .. `tile_32`: every reduction and inclusive scan on +//! `WarpTile` for each tile size, plus the four `u64` group shuffles, +//! including offsets that leave the tile and must return the caller's own +//! value. Each block of 128 threads holds several warps and, below 32 +//! lanes, several tiles per warp. +//! - `block_collectives`: every block reduction and inclusive scan, launched +//! with block sizes from 1 to 1024 threads, most of them leaving a partial +//! final warp. +//! - `coalesced_shuffles`: the `u64` shuffles on sparse `CoalescedThreads` +//! groups, whose ranks and lanes differ. +//! - `defaults_8` and `defaults_32`: the same shuffles on a group that +//! implements only the required `WarpCollective` methods, so they run the +//! trait's default bodies. +//! - `warp_total`: a caller that sums `u64` values over a full warp through +//! the public API alone. +//! +//! The inputs follow fixed patterns per warp or per block: full-width random +//! values, values whose low words alone would pick the wrong `Min` and +//! `Max`, values just below `u64::MAX` whose `Sum` wraps, the `u64` and +//! `i64` extremes, and low words whose `Sum` carries into the high word. +//! +//! Build and run with: +//! cargo oxide run coop_groups_u64 + +use cuda_core::simt::LaunchConfig; +use cuda_core::{CudaContext, DeviceBuffer}; +use cuda_device::cooperative_groups::{ + ThreadGroup, WarpCollective, WarpTile, block_reduce, block_scan, coalesced_threads, + ops::{BitAnd, BitOr, BitXor, Max, Min, Sum}, + this_thread_block, warp_reduce, warp_scan, +}; +use cuda_device::{DisjointSlice, SharedArray, device, kernel, thread, warp}; +use cuda_host::cuda_module; + +/// Reductions and scans per thread: six `u64` combiners and three `i64` +/// combiners, each reduced and scanned. +const COLLECTIVE_OUTPUTS: usize = 18; +/// Group shuffles per thread: one index, three butterfly, three down and +/// three up shuffles. +const SHUFFLE_OUTPUTS: usize = 10; +/// Results per thread from a tile kernel: the collectives, then the shuffles. +const TILE_OUTPUTS: usize = COLLECTIVE_OUTPUTS + SHUFFLE_OUTPUTS; +/// Results per member thread from `coalesced_shuffles`. +const COALESCED_OUTPUTS: usize = 7; + +/// Pre-filled into every output so a thread that never writes is visible. +const SENTINEL: u64 = 0xA5A5_A5A5_A5A5_A5A5; + +#[cuda_module] +mod kernels { + use super::*; + + /// Stores `value` at column `col` of the row that starts at `row`. + #[device] + fn put(out: &mut DisjointSlice, row: usize, col: usize, value: u64) { + // SAFETY: the host sizes `out` to one full row per launched thread, + // and each thread passes the start of its own row + unsafe { *out.get_unchecked_mut(row + col) = value }; + } + + /// Every reduction, scan and `u64` shuffle a `WarpTile` offers, in the + /// column order the host reference uses. + #[device] + fn tile_row(input: &[u64], mut out: DisjointSlice) { + let gid = thread::index_1d().get(); + let tile = this_thread_block().tiled_partition::(); + let v = input[gid]; + let s = v as i64; + let row = gid * TILE_OUTPUTS; + + put(&mut out, row, 0, warp_reduce::(&tile, v)); + put(&mut out, row, 1, warp_reduce::(&tile, v)); + put(&mut out, row, 2, warp_reduce::(&tile, v)); + put(&mut out, row, 3, warp_reduce::(&tile, v)); + put(&mut out, row, 4, warp_reduce::(&tile, v)); + put(&mut out, row, 5, warp_reduce::(&tile, v)); + put(&mut out, row, 6, warp_scan::(&tile, v)); + put(&mut out, row, 7, warp_scan::(&tile, v)); + put(&mut out, row, 8, warp_scan::(&tile, v)); + put(&mut out, row, 9, warp_scan::(&tile, v)); + put(&mut out, row, 10, warp_scan::(&tile, v)); + put(&mut out, row, 11, warp_scan::(&tile, v)); + put( + &mut out, + row, + 12, + warp_reduce::(&tile, s) as u64, + ); + put( + &mut out, + row, + 13, + warp_reduce::(&tile, s) as u64, + ); + put( + &mut out, + row, + 14, + warp_reduce::(&tile, s) as u64, + ); + put(&mut out, row, 15, warp_scan::(&tile, s) as u64); + put(&mut out, row, 16, warp_scan::(&tile, s) as u64); + put(&mut out, row, 17, warp_scan::(&tile, s) as u64); + + shuffle_cols::<_, N>(&tile, v, &mut out, row + COLLECTIVE_OUTPUTS); + } + + /// The four `u64` group shuffles of an `N`-lane group, written to the + /// `SHUFFLE_OUTPUTS` columns that start at `start`. + #[device] + fn shuffle_cols( + group: &G, + v: u64, + out: &mut DisjointSlice, + start: usize, + ) { + let rank = group.thread_rank(); + + // the third offset of each kind reaches past a sub-warp tile, where + // the group must hand back the caller's own value + let wide_xor = if N < 32 { N } else { 16 }; + let wide_delta = if N < 32 { N } else { 17 }; + put(out, start, 0, group.shfl_u64(v, rank * 7 + 3)); + put(out, start, 1, group.shfl_xor_u64(v, 1)); + put(out, start, 2, group.shfl_xor_u64(v, N - 1)); + put(out, start, 3, group.shfl_xor_u64(v, wide_xor)); + put(out, start, 4, group.shfl_down_u64(v, 1)); + put(out, start, 5, group.shfl_down_u64(v, N - 1)); + put(out, start, 6, group.shfl_down_u64(v, wide_delta)); + put(out, start, 7, group.shfl_up_u64(v, 1)); + put(out, start, 8, group.shfl_up_u64(v, N - 1)); + put(out, start, 9, group.shfl_up_u64(v, wide_delta)); + } + + /// A group that implements only the required `WarpCollective` methods, + /// as one defined outside `cuda-device` would, so its `u64` shuffles run + /// the trait's default bodies over two 32-bit shuffles. + #[derive(Clone, Copy)] + pub struct HalvesOnly(WarpTile); + + impl ThreadGroup for HalvesOnly { + fn size(&self) -> u32 { + self.0.size() + } + + fn thread_rank(&self) -> u32 { + self.0.thread_rank() + } + + fn sync(&self) { + self.0.sync() + } + } + + impl WarpCollective for HalvesOnly { + fn ballot(&self, predicate: bool) -> u32 { + self.0.ballot(predicate) + } + + fn all(&self, predicate: bool) -> bool { + self.0.all(predicate) + } + + fn any(&self, predicate: bool) -> bool { + self.0.any(predicate) + } + + fn shfl(&self, var: u32, src_rank: u32) -> u32 { + self.0.shfl(var, src_rank) + } + + fn shfl_xor(&self, var: u32, lane_mask: u32) -> u32 { + self.0.shfl_xor(var, lane_mask) + } + + fn shfl_down(&self, var: u32, delta: u32) -> u32 { + self.0.shfl_down(var, delta) + } + + fn shfl_up(&self, var: u32, delta: u32) -> u32 { + self.0.shfl_up(var, delta) + } + + fn shfl_f32(&self, var: f32, src_rank: u32) -> f32 { + self.0.shfl_f32(var, src_rank) + } + + fn shfl_xor_f32(&self, var: f32, lane_mask: u32) -> f32 { + self.0.shfl_xor_f32(var, lane_mask) + } + + fn shfl_down_f32(&self, var: f32, delta: u32) -> f32 { + self.0.shfl_down_f32(var, delta) + } + + fn shfl_up_f32(&self, var: f32, delta: u32) -> f32 { + self.0.shfl_up_f32(var, delta) + } + + fn match_any(&self, value: u32) -> u32 { + self.0.match_any(value) + } + + fn match_any_i64(&self, value: u64) -> u32 { + self.0.match_any_i64(value) + } + + fn match_all(&self, value: u32) -> u32 { + self.0.match_all(value) + } + + fn match_all_i64(&self, value: u64) -> u32 { + self.0.match_all_i64(value) + } + } + + /// The `u64` shuffles through the trait's default bodies. + #[device] + fn default_row(input: &[u64], mut out: DisjointSlice) { + let gid = thread::index_1d().get(); + let group = HalvesOnly(this_thread_block().tiled_partition::()); + shuffle_cols::<_, N>(&group, input[gid], &mut out, gid * SHUFFLE_OUTPUTS); + } + + #[kernel] + pub fn defaults_8(input: &[u64], out: DisjointSlice) { + default_row::<8>(input, out); + } + + #[kernel] + pub fn defaults_32(input: &[u64], out: DisjointSlice) { + default_row::<32>(input, out); + } + + #[kernel] + pub fn tile_1(input: &[u64], out: DisjointSlice) { + tile_row::<1>(input, out); + } + + #[kernel] + pub fn tile_2(input: &[u64], out: DisjointSlice) { + tile_row::<2>(input, out); + } + + #[kernel] + pub fn tile_4(input: &[u64], out: DisjointSlice) { + tile_row::<4>(input, out); + } + + #[kernel] + pub fn tile_8(input: &[u64], out: DisjointSlice) { + tile_row::<8>(input, out); + } + + #[kernel] + pub fn tile_16(input: &[u64], out: DisjointSlice) { + tile_row::<16>(input, out); + } + + #[kernel] + pub fn tile_32(input: &[u64], out: DisjointSlice) { + tile_row::<32>(input, out); + } + + /// Every block reduction and inclusive scan, in the same column order as + /// the first `COLLECTIVE_OUTPUTS` columns of a tile row. + #[kernel] + pub fn block_collectives(input: &[u64], mut out: DisjointSlice) { + static mut SMEM: SharedArray = SharedArray::UNINIT; + static mut SMEM_I64: SharedArray = SharedArray::UNINIT; + + let block = this_thread_block(); + let gid = thread::index_1d().get(); + let v = input[gid]; + let s = v as i64; + let row = gid * COLLECTIVE_OUTPUTS; + + // each call leaves its scratch in an arbitrary state, so a barrier + // separates it from the next call on the same allocation + macro_rules! block_op { + ($f:ident, $t:ty, $op:ty, $x:expr, $smem:ident) => {{ + // SAFETY: every thread of the block makes this call with the + // same scratch, and the barrier below orders its reuse + let r = unsafe { $f::<$t, $op, 32>(&block, $x, &raw mut $smem) }; + block.sync(); + r + }}; + } + + put(&mut out, row, 0, block_op!(block_reduce, u64, Sum, v, SMEM)); + put(&mut out, row, 1, block_op!(block_reduce, u64, Min, v, SMEM)); + put(&mut out, row, 2, block_op!(block_reduce, u64, Max, v, SMEM)); + put( + &mut out, + row, + 3, + block_op!(block_reduce, u64, BitAnd, v, SMEM), + ); + put( + &mut out, + row, + 4, + block_op!(block_reduce, u64, BitOr, v, SMEM), + ); + put( + &mut out, + row, + 5, + block_op!(block_reduce, u64, BitXor, v, SMEM), + ); + put(&mut out, row, 6, block_op!(block_scan, u64, Sum, v, SMEM)); + put(&mut out, row, 7, block_op!(block_scan, u64, Min, v, SMEM)); + put(&mut out, row, 8, block_op!(block_scan, u64, Max, v, SMEM)); + put( + &mut out, + row, + 9, + block_op!(block_scan, u64, BitAnd, v, SMEM), + ); + put( + &mut out, + row, + 10, + block_op!(block_scan, u64, BitOr, v, SMEM), + ); + put( + &mut out, + row, + 11, + block_op!(block_scan, u64, BitXor, v, SMEM), + ); + put( + &mut out, + row, + 12, + block_op!(block_reduce, i64, Sum, s, SMEM_I64) as u64, + ); + put( + &mut out, + row, + 13, + block_op!(block_reduce, i64, Min, s, SMEM_I64) as u64, + ); + put( + &mut out, + row, + 14, + block_op!(block_reduce, i64, Max, s, SMEM_I64) as u64, + ); + put( + &mut out, + row, + 15, + block_op!(block_scan, i64, Sum, s, SMEM_I64) as u64, + ); + put( + &mut out, + row, + 16, + block_op!(block_scan, i64, Min, s, SMEM_I64) as u64, + ); + put( + &mut out, + row, + 17, + block_op!(block_scan, i64, Max, s, SMEM_I64) as u64, + ); + } + + /// The `u64` shuffles on a sparse coalesced group. Lanes named in + /// `masks[warp]` enter the branch together and form the group; the + /// others write nothing. + #[kernel] + pub fn coalesced_shuffles( + input: &[u64], + masks: &[u32], + mut out: DisjointSlice, + mut seen: DisjointSlice, + ) { + let gid = thread::index_1d().get(); + let members = masks[gid / 32]; + + if members & (1u32 << warp::lane_id()) != 0 { + let group = coalesced_threads(); + let v = input[gid]; + let row = gid * COALESCED_OUTPUTS; + + put( + &mut out, + row, + 0, + group.shfl_u64(v, group.thread_rank() * 7 + 3), + ); + put(&mut out, row, 1, group.shfl_xor_u64(v, 1)); + put(&mut out, row, 2, group.shfl_xor_u64(v, 6)); + put(&mut out, row, 3, group.shfl_down_u64(v, 1)); + put(&mut out, row, 4, group.shfl_down_u64(v, 5)); + put(&mut out, row, 5, group.shfl_up_u64(v, 1)); + put(&mut out, row, 6, group.shfl_up_u64(v, 5)); + + // SAFETY: the host sizes `seen` to one entry per launched thread + unsafe { *seen.get_unchecked_mut(gid) = group.raw_mask() }; + } + } + + /// Sums one `u64` per lane over a full warp, written against the public + /// API alone. + #[device] + pub fn total(tile: &WarpTile<32>, value: u64) -> u64 { + warp_reduce::(tile, value) + } + + #[kernel] + pub fn warp_total(input: &[u64], mut out: DisjointSlice) { + let gid = thread::index_1d().get(); + let tile = this_thread_block().tiled_partition::<32>(); + put(&mut out, gid, 0, total(&tile, input[gid])); + } +} + +// ============================================================================= +// HOST REFERENCE +// ============================================================================= + +#[derive(Clone, Copy)] +enum Op { + Sum, + Min, + Max, + And, + Or, + Xor, +} + +const U64_OPS: [Op; 6] = [Op::Sum, Op::Min, Op::Max, Op::And, Op::Or, Op::Xor]; +const I64_OPS: [Op; 3] = [Op::Sum, Op::Min, Op::Max]; + +const COLLECTIVE_LABELS: [&str; COLLECTIVE_OUTPUTS] = [ + "reduce u64 Sum", + "reduce u64 Min", + "reduce u64 Max", + "reduce u64 BitAnd", + "reduce u64 BitOr", + "reduce u64 BitXor", + "scan u64 Sum", + "scan u64 Min", + "scan u64 Max", + "scan u64 BitAnd", + "scan u64 BitOr", + "scan u64 BitXor", + "reduce i64 Sum", + "reduce i64 Min", + "reduce i64 Max", + "scan i64 Sum", + "scan i64 Min", + "scan i64 Max", +]; + +const SHUFFLE_LABELS: [&str; SHUFFLE_OUTPUTS] = [ + "shfl_u64 rank*7+3", + "shfl_xor_u64 1", + "shfl_xor_u64 N-1", + "shfl_xor_u64 wide", + "shfl_down_u64 1", + "shfl_down_u64 N-1", + "shfl_down_u64 wide", + "shfl_up_u64 1", + "shfl_up_u64 N-1", + "shfl_up_u64 wide", +]; + +const COALESCED_LABELS: [&str; COALESCED_OUTPUTS] = [ + "shfl_u64 rank*7+3", + "shfl_xor_u64 1", + "shfl_xor_u64 6", + "shfl_down_u64 1", + "shfl_down_u64 5", + "shfl_up_u64 1", + "shfl_up_u64 5", +]; + +fn combine_u64(op: Op, a: u64, b: u64) -> u64 { + match op { + Op::Sum => a.wrapping_add(b), + Op::Min => a.min(b), + Op::Max => a.max(b), + Op::And => a & b, + Op::Or => a | b, + Op::Xor => a ^ b, + } +} + +fn combine_i64(op: Op, a: i64, b: i64) -> i64 { + match op { + Op::Sum => a.wrapping_add(b), + Op::Min => a.min(b), + Op::Max => a.max(b), + Op::And | Op::Or | Op::Xor => unreachable!("i64 has no bitwise combiners"), + } +} + +/// The reductions and inclusive scans one member of `group` at `rank` +/// should see, in kernel column order. +fn collective_row(group: &[u64], rank: usize) -> Vec { + let fold_u64 = |op, values: &[u64]| { + values[1..] + .iter() + .fold(values[0], |acc, &x| combine_u64(op, acc, x)) + }; + let fold_i64 = |op, values: &[u64]| { + values[1..] + .iter() + .fold(values[0] as i64, |acc, &x| combine_i64(op, acc, x as i64)) as u64 + }; + let prefix = &group[..=rank]; + + let mut row = Vec::with_capacity(COLLECTIVE_OUTPUTS); + row.extend(U64_OPS.iter().map(|&op| fold_u64(op, group))); + row.extend(U64_OPS.iter().map(|&op| fold_u64(op, prefix))); + row.extend(I64_OPS.iter().map(|&op| fold_i64(op, group))); + row.extend(I64_OPS.iter().map(|&op| fold_i64(op, prefix))); + row +} + +/// Expected rows for a tile kernel with `n`-lane tiles. Launches use whole +/// warps, so consecutive chunks of `n` threads are exactly the tiles. +fn tile_expected(input: &[u64], n: usize) -> Vec { + let wide_xor = if n < 32 { n } else { 16 }; + let wide_delta = if n < 32 { n } else { 17 }; + + let mut expected = Vec::with_capacity(input.len() * TILE_OUTPUTS); + for tile in input.chunks(n) { + for rank in 0..n { + let own = tile[rank]; + let xor = |mask: usize| tile.get(rank ^ mask).copied().unwrap_or(own); + let down = |delta: usize| tile.get(rank + delta).copied().unwrap_or(own); + let up = |delta: usize| rank.checked_sub(delta).map_or(own, |src| tile[src]); + + expected.extend(collective_row(tile, rank)); + expected.push(tile[(rank * 7 + 3) & (n - 1)]); + expected.extend([xor(1), xor(n - 1), xor(wide_xor)]); + expected.extend([down(1), down(n - 1), down(wide_delta)]); + expected.extend([up(1), up(n - 1), up(wide_delta)]); + } + } + expected +} + +/// Expected rows for `block_collectives` launched with `block`-thread blocks. +fn block_expected(input: &[u64], block: usize) -> Vec { + let mut expected = Vec::with_capacity(input.len() * COLLECTIVE_OUTPUTS); + for group in input.chunks(block) { + for rank in 0..group.len() { + expected.extend(collective_row(group, rank)); + } + } + expected +} + +/// Expected rows for `coalesced_shuffles`: a member at `lane` reads from the +/// lane its offset names when that lane is in the group, and otherwise +/// keeps its own value. Non-members keep the sentinel. +fn coalesced_expected(input: &[u64], masks: &[u32]) -> Vec { + let mut expected = vec![SENTINEL; input.len() * COALESCED_OUTPUTS]; + for (w, &members) in masks.iter().enumerate() { + let values = &input[w * 32..(w + 1) * 32]; + let member = |lane: usize| lane < 32 && members & (1u32 << lane) != 0; + + for lane in (0..32).filter(|&lane| member(lane)) { + let own = values[lane]; + let from = |src: Option| match src { + Some(src) if member(src) => values[src], + _ => own, + }; + let rank = (members & ((1u32 << lane) - 1)).count_ones() as usize; + let rank_lane = (0..32).filter(|&l| member(l)).nth(rank * 7 + 3); + let row = [ + from(rank_lane), + from(Some(lane ^ 1)), + from(Some(lane ^ 6)), + from(Some(lane + 1)), + from(Some(lane + 5)), + from(lane.checked_sub(1)), + from(lane.checked_sub(5)), + ]; + let start = (w * 32 + lane) * COALESCED_OUTPUTS; + expected[start..start + COALESCED_OUTPUTS].copy_from_slice(&row); + } + } + expected +} + +// ============================================================================= +// INPUTS +// ============================================================================= + +/// Number of value patterns; warp `w` (or block `b`) uses pattern `w % PATTERNS`. +const PATTERNS: usize = 8; + +fn splitmix64(seed: u64) -> u64 { + let mut z = seed.wrapping_add(0x9E37_79B9_7F4A_7C15); + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) +} + +/// Value `i` of `pattern`. Each pattern targets one way a 64-bit collective +/// built from 32-bit shuffles can go wrong. +fn pattern_value(pattern: usize, i: u64) -> u64 { + let r = splitmix64(i ^ ((pattern as u64) << 40)); + match pattern { + // full-width random values + 0 => r, + // the high word decides Min and Max, and the low words alone would + // pick the opposite value; the low words also carry when summed + 1 if i.is_multiple_of(2) => (1 << 32) | (r & 0xF), + 1 => 0xFFFF_FFF0 | (r & 0xF), + // just below u64::MAX, so Sum wraps + 2 => u64::MAX - (r & 0xFFFF), + // the u64 and i64 extremes + 3 => [0, 1, u64::MAX, u64::MAX - 1, 1 << 63, (1 << 63) - 1][(r % 6) as usize], + // all u64::MAX, the Min and BitAnd identity + 4 => u64::MAX, + // only the high words differ + 5 => (r & 0xFFFF_FFFF_0000_0000) | 0xDEAD_BEEF, + // negative as i64, with low words that carry into the high word + 6 => 0xFFFF_FFFF_8000_0000 | (r & 0x7FFF_FFFF), + // one bit each, spread over both words + _ => 1 << (r % 64), + } +} + +/// `len` values where every consecutive run of `run` values shares a +/// pattern, starting from pattern `first`. +fn patterned_input(len: usize, run: usize, first: usize) -> Vec { + (0..len) + .map(|i| pattern_value((first + i / run) % PATTERNS, i as u64)) + .collect() +} + +// ============================================================================= +// HOST CODE +// ============================================================================= + +const TILE_BLOCK: u32 = 128; +const TILE_BLOCKS: u32 = 8; +const TILE_THREADS: usize = (TILE_BLOCK * TILE_BLOCKS) as usize; + +/// Block sizes for `block_collectives`, most with a partial final warp. +const BLOCK_SIZES: [u32; 14] = [ + 1, 7, 31, 32, 33, 45, 64, 100, 255, 256, 257, 640, 1000, 1024, +]; +const BLOCKS_PER_SIZE: u32 = 3; + +/// Coalesced group per warp of `coalesced_shuffles`. +const COALESCED_MASKS: [u32; 8] = [ + u32::MAX, + 0x8010_0089, + 0x5555_5555, + 0xAAAA_AAAA, + 0x0000_0001, + 0xFFFF_0000, + 0x7FFF_FFFE, + 0xF0F0_0F0F, +]; + +/// Compares `got` with `expected` row by row and prints the first few +/// mismatches. Returns how many values differ. +fn compare(case: &str, got: &[u64], expected: &[u64], labels: &[&str]) -> usize { + let row_len = labels.len(); + let mut mismatches = 0; + for (i, (&g, &e)) in got.iter().zip(expected).enumerate() { + if g != e { + if mismatches < 5 { + println!( + " {case}: thread {} {}: got {g:#018x}, expected {e:#018x}", + i / row_len, + labels[i % row_len], + ); + } + mismatches += 1; + } + } + if got.len() != expected.len() { + println!( + " {case}: {} results, expected {}", + got.len(), + expected.len() + ); + mismatches += 1; + } + if mismatches == 0 { + println!(" {case}: {} results match", got.len()); + } else { + println!(" {case}: {mismatches} mismatches"); + } + mismatches +} + +fn main() -> Result<(), Box> { + let ctx = CudaContext::new(0)?; + let stream = ctx.default_stream(); + let (major, minor) = ctx.compute_capability()?; + println!("=== cooperative_groups u64 / i64 collectives on sm_{major}{minor} ==="); + + let module = kernels::load(&ctx)?; + let mut failures = 0; + + // ----- warp tiles ----- + println!("\n--- WarpTile reductions, scans and u64 shuffles ---"); + let tile_cfg = LaunchConfig { + block_dim: (TILE_BLOCK, 1, 1), + grid_dim: (TILE_BLOCKS, 1, 1), + shared_mem_bytes: 0, + }; + let host_input = patterned_input(TILE_THREADS, 32, 0); + let input = DeviceBuffer::from_host(&stream, &host_input)?; + let tile_labels: Vec<&str> = COLLECTIVE_LABELS + .iter() + .chain(SHUFFLE_LABELS.iter()) + .copied() + .collect(); + + for n in [1usize, 2, 4, 8, 16, 32] { + let mut out = + DeviceBuffer::from_host(&stream, &vec![SENTINEL; TILE_THREADS * TILE_OUTPUTS])?; + let s = stream.as_ref(); + // SAFETY: each launch covers TILE_THREADS threads in blocks of 128, + // a multiple of every tile size, and `input` and `out` hold one value + // and one row per thread + unsafe { + match n { + 1 => module.tile_1(s, tile_cfg, &input, &mut out), + 2 => module.tile_2(s, tile_cfg, &input, &mut out), + 4 => module.tile_4(s, tile_cfg, &input, &mut out), + 8 => module.tile_8(s, tile_cfg, &input, &mut out), + 16 => module.tile_16(s, tile_cfg, &input, &mut out), + _ => module.tile_32(s, tile_cfg, &input, &mut out), + } + }?; + stream.synchronize()?; + failures += compare( + &format!("WarpTile<{n}>"), + &out.to_host_vec(&stream)?, + &tile_expected(&host_input, n), + &tile_labels, + ); + } + + // ----- blocks ----- + println!("\n--- block_reduce / block_scan ---"); + for (launch, block) in BLOCK_SIZES.into_iter().enumerate() { + let threads = (block * BLOCKS_PER_SIZE) as usize; + let host_input = + patterned_input(threads, block as usize, launch * BLOCKS_PER_SIZE as usize); + let input = DeviceBuffer::from_host(&stream, &host_input)?; + let mut out = + DeviceBuffer::from_host(&stream, &vec![SENTINEL; threads * COLLECTIVE_OUTPUTS])?; + let cfg = LaunchConfig { + block_dim: (block, 1, 1), + grid_dim: (BLOCKS_PER_SIZE, 1, 1), + shared_mem_bytes: 0, + }; + // SAFETY: blocks hold at most 1024 threads, the 32 warp totals the + // scratch holds, and `input` and `out` hold one value and one row per + // thread + unsafe { module.block_collectives(stream.as_ref(), cfg, &input, &mut out) }?; + stream.synchronize()?; + failures += compare( + &format!("{BLOCKS_PER_SIZE} blocks of {block}"), + &out.to_host_vec(&stream)?, + &block_expected(&host_input, block as usize), + &COLLECTIVE_LABELS, + ); + } + + // ----- default trait bodies ----- + println!("\n--- WarpCollective default u64 shuffles ---"); + let host_input = patterned_input(TILE_THREADS, 32, 0); + let input = DeviceBuffer::from_host(&stream, &host_input)?; + for n in [8usize, 32] { + let mut out = + DeviceBuffer::from_host(&stream, &vec![SENTINEL; TILE_THREADS * SHUFFLE_OUTPUTS])?; + let s = stream.as_ref(); + // SAFETY: as for the tile kernels, with one shuffle row per thread + unsafe { + match n { + 8 => module.defaults_8(s, tile_cfg, &input, &mut out), + _ => module.defaults_32(s, tile_cfg, &input, &mut out), + } + }?; + stream.synchronize()?; + // the shuffle columns of the tile reference, which used the same calls + let expected: Vec = tile_expected(&host_input, n) + .chunks(TILE_OUTPUTS) + .flat_map(|row| row[COLLECTIVE_OUTPUTS..].to_vec()) + .collect(); + failures += compare( + &format!("HalvesOnly<{n}>"), + &out.to_host_vec(&stream)?, + &expected, + &SHUFFLE_LABELS, + ); + } + + // ----- coalesced groups ----- + println!("\n--- CoalescedThreads u64 shuffles ---"); + let threads = COALESCED_MASKS.len() * 32; + let host_input = patterned_input(threads, threads, 0); + let input = DeviceBuffer::from_host(&stream, &host_input)?; + let masks = DeviceBuffer::from_host(&stream, &COALESCED_MASKS)?; + let mut out = DeviceBuffer::from_host(&stream, &vec![SENTINEL; threads * COALESCED_OUTPUTS])?; + let mut seen = DeviceBuffer::from_host(&stream, &vec![0u32; threads])?; + let cfg = LaunchConfig { + block_dim: (threads as u32, 1, 1), + grid_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + // SAFETY: one block of whole warps, with one mask per warp and one value, + // row and mask slot per thread + unsafe { + module.coalesced_shuffles(stream.as_ref(), cfg, &input, &masks, &mut out, &mut seen) + }?; + stream.synchronize()?; + let seen = seen.to_host_vec(&stream)?; + for (gid, &mask) in seen.iter().enumerate() { + let want = COALESCED_MASKS[gid / 32]; + let observed = if want & (1 << (gid % 32)) != 0 { + want + } else { + 0 + }; + if mask != observed { + println!(" thread {gid} saw group mask {mask:#010x}, expected {observed:#010x}"); + failures += 1; + } + } + failures += compare( + "sparse groups", + &out.to_host_vec(&stream)?, + &coalesced_expected(&host_input, &COALESCED_MASKS), + &COALESCED_LABELS, + ); + + // ----- public-API caller ----- + println!("\n--- total(&WarpTile<32>, u64) ---"); + let host_input = patterned_input(TILE_THREADS, 32, 0); + let input = DeviceBuffer::from_host(&stream, &host_input)?; + let mut out = DeviceBuffer::from_host(&stream, &vec![SENTINEL; TILE_THREADS])?; + // SAFETY: TILE_THREADS threads in whole warps, one value and one result + // per thread + unsafe { module.warp_total(stream.as_ref(), tile_cfg, &input, &mut out) }?; + stream.synchronize()?; + let expected: Vec = host_input + .chunks(32) + .flat_map(|warp| { + let sum = warp.iter().fold(0u64, |acc, &x| acc.wrapping_add(x)); + [sum; 32] + }) + .collect(); + failures += compare( + "warp sums", + &out.to_host_vec(&stream)?, + &expected, + &["total"], + ); + + if failures != 0 { + return Err(format!("{failures} results disagree with the host reference").into()); + } + println!("\nSUCCESS: u64 and i64 collectives match the host reference"); + Ok(()) +} diff --git a/cuda-oxide/cuda-oxide-book/appendix/supported-features.md b/cuda-oxide/cuda-oxide-book/appendix/supported-features.md index 71a5df7b38..4e2e30327a 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/supported-features.md +++ b/cuda-oxide/cuda-oxide-book/appendix/supported-features.md @@ -256,8 +256,8 @@ Anonymous promoted allocations remain unsupported. | Typed Group Handles | **Full** | `Grid`, `Cluster`, `ThreadBlock`, `WarpTile` (N ∈ {1,2,4,8,16,32}), `CoalescedThreads`. | | Group Universal API | **Full** | `size()`, `thread_rank()`, `sync()` on every group handle. | | Warp Tile Partitioning | **Full** | `ThreadBlock::tiled_partition::()` partitions a block into fixed-size `WarpTile` groups; the block thread count must be evenly divisible by `N`. `coalesced_threads()` materialises the active-lane group. | -| Warp Collectives | **Full** | `ballot`, `all`, `any`, `shfl`, `shfl_xor`, `shfl_down`, `shfl_up` (`u32` and `f32`); `match_any` / `match_all` (`i32` and `i64`); `active_mask`. | -| Warp Reductions / Scans | **Full** | `warp_reduce`, `warp_scan` (inclusive). `Sum`/`Min`/`Max` for `u32`/`i32`/`f32`; `BitAnd`/`BitOr`/`BitXor` for `u32`. | +| Warp Collectives | **Full** | `ballot`, `all`, `any`, `shfl`, `shfl_xor`, `shfl_down`, `shfl_up` (`u32`, `f32`, and `u64` as two 32-bit shuffles); `match_any` / `match_all` (`i32` and `i64`); `active_mask`. | +| Warp Reductions / Scans | **Full** | `warp_reduce`, `warp_scan` (inclusive). `Sum`/`Min`/`Max` for `u32`/`i32`/`f32`/`u64`/`i64`; `BitAnd`/`BitOr`/`BitXor` for `u32`/`u64`. | | Block Reductions / Scans | **Full** | `block_reduce`, `block_scan` (inclusive). `NUM_WARPS` is the shared-memory capacity in warp totals and must be at least `ceil(block_threads / 32)`; partial final warps are supported. Same op/type matrix as warp variants. | | Cooperative Kernel Launch | **Full** | `#[cooperative_launch]` on a `#[cuda_module]` kernel (or `unsafe { cuda_launch! { cooperative: true, ... } }`) enables `Grid::sync()` for grid-wide barriers. |