[praveenperera]

Closes #1497.

The problem

warp_reduce, warp_scan, block_reduce, and block_scan accept only u32, i32, and f32. A u64 call such as warp_reduce::<u64, Sum, 32>(tile, value) fails with E0277 for both u64: WarpShuffle and Sum: ReduceOp<u64>. A downstream crate can't add the missing WarpShuffle impl for a primitive type, so the only workaround is a newtype with its own shuffles and combiners.

WarpShuffle is generic over WarpCollective, and that trait had no 64-bit shuffle. So the low-level warp::shuffle_*_u64_sync functions from #278 were not reachable with a group's mask and tile-boundary rules.

The fix

All library changes are in cuda-device/src/cooperative_groups.rs:

1. WarpCollective gains shfl_u64, shfl_xor_u64, shfl_down_u64, and shfl_up_u64, matching the existing _f32 set. WarpTile<N>, CoalescedThreads, and the private partial-warp prefix group implement them with the same masks and source-lane rules as their 32-bit shuffles, on top of warp::shuffle_*_u64_sync. Each one is a single inline-PTX block with two shfl.sync.*.b32, one per half. CoalescedThreads gets a private rank_to_lane helper, shared by the u32, f32, and u64 index shuffles, in place of three copies of the same loop.
2. WarpShuffle is implemented for u64 and for i64. i64 casts through u64, the same way i32 casts through u32.
3. ReduceOp<u64> is implemented for Sum (wrapping), Min, Max, BitAnd, BitOr, and BitXor. ReduceOp<i64> is implemented for Sum, Min, and Max.

I added i64 because it falls out of the same pattern: one WarpShuffle impl that casts through u64, and the three combiners i32 already has. Bitwise ops stay unsigned-only, as they are for 32 bits. f64 would work the same way, but I left it out of this PR.

The new trait methods have default bodies that move the two halves through the existing u32 methods, so a WarpCollective impl outside this crate still compiles and gets working u64 shuffles. The built-in groups override all four with the single-block form.

redux.sync has no 64-bit form, so the 64-bit types always use the shuffle path. The new impls sit after the f32 ones and use the default for the HAS_FULL_WARP_FORM item that #1305 adds, so they don't overlap its hunks. The book's supported-features table now lists the new types.

Testing

On an RTX 5090 (sm_120), CUDA 13.2, driver 595.99:

• New example coop_groups_u64 compares every result with a CPU reference:
  • WarpTile<N> for N = 1, 2, 4, 8, 16, and 32, in blocks of 128 threads, so each launch has several warps and, below 32 lanes, several tiles per warp. Each tile size runs the 18 reductions and inclusive scans (six u64 ops and three i64 ops) and the four u64 group shuffles. The shuffle offsets include ones that reach past the tile, where the caller must get its own value back. 28,672 results per tile size.
  • block_reduce and block_scan at 14 block sizes: 1, 7, 31, 32, 33, 45, 64, 100, 255, 256, 257, 640, 1000, and 1024 threads, three blocks each. Nine of the sizes leave a partial final warp. 202,230 results in total.
  • CoalescedThreads u64 shuffles on eight groups, from the full warp to a single lane, including sparse ones such as 0x8010_0089.
  • The same u64 shuffles on 8- and 32-lane groups through a wrapper that implements only the required WarpCollective methods, as a group from another crate would. This exercises the default bodies.
  • A total(&WarpTile<32>, u64) function written as in the issue.
  • The inputs include values whose low words alone would give the wrong Min and Max, low words whose Sum carries into the high word, values just below u64::MAX whose Sum wraps, the u64 and i64 extremes, and all-u64::MAX and all-negative blocks.
• Without this change the example does not compile. I also ran it against 18 single-line mutations of the new code. It caught 17: high word dropped from the u64 or i64 shuffle; Min comparing low words only; Sum without the carry between halves; 32-bit BitAnd and Min identities; 0 as the i64 Max identity; each WarpTile shuffle ignoring the tile boundary or base lane; CoalescedThreads treating a rank as a lane or skipping the member check; and, in the default bodies, a high half taken from the low half, shuffled the wrong way or read from the wrong rank, or reassembled with its halves swapped. The one it missed makes the partial-warp prefix use a full 32-lane mask. That is undefined behavior, but the 5090 returns correct results, and memcheck and synccheck report nothing. The existing u32 path has the same gap, and partial_warp_reduce documents the same behavior.
• compute-sanitizer 2026.1.1: memcheck 0 errors and synccheck 0 errors on the final code. Racecheck reported 0 hazards on the code before review, whose block kernels are unchanged.
• PTX: each u64 exchange in warp_total is one { mov.b64 {lo, hi}; shfl.sync.bfly.b32 lo; shfl.sync.bfly.b32 hi; mov.b64 } block followed by add.s64. That is five blocks and ten shfl.sync.bfly.b32 for a full warp, with no local memory. Before I added the kernels for the default bodies, which use two separate u32 shuffles by design, all 638 shuffle sites in the example module were that same two-shuffle block (308 up, 291 bfly, 29 down, 10 idx), and the module used no local memory.
• just check passes, along with the new example's cargo clippy --all-targets -D warnings, cargo fmt --check, and scripts/sync-example-locks.sh --check. Before review, scripts/smoketest.sh passed 241/241 in both GPU and --compile-only modes, and 235/241 in --full-debug, coop_groups_u64 included. The six full-debug failures fail the same way on unmodified main: constant_memory_simple, iket_trace, packed_aggregate_abi, partial_warp_reduce, warp_sums, and wgmma_mma_bf16. After the review changes (default bodies and docs), the 14 cooperative-groups, warp, reduce, scan and shuffle examples pass in GPU and --compile-only modes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.