Skip to content

Commit 2ffc19b

Browse files
vyncintnihalpasham
andauthored
feat(cuda-device): derive an SM floor cfg and use redux.sync in warp_reduce (#1305)
## What This PR Does Derives a shared compute-capability floor (`resolve_sm_floor`) across `mir-importer` and both backend target selectors, and uses single-instruction `redux.sync` for full-warp `u32`/`i32` reductions in `cooperative_groups::warp_reduce` when the target floor is `sm_80` or newer: 1. **Shared SM Floor (`cuda-target-spec`, `mir-importer`, `cuda-oxide-codegen`)**: `resolve_sm_floor` sets the floor from an explicit target (`--arch` / `CUDA_OXIDE_TARGET`), or `min(80, detected_gpu)` when unpinned. Both PTX and NVVM target selectors enforce `selected.capability() >= floor`, and `mir-importer` folds `cuda_device::warp::__target_has_integer_redux_sync()` at MIR-import time without needing a `cuda-device/build.rs` (keeping `cuda-device` and `cuda-host` cached across `--arch` switches). 2. **Full-Warp Integer `redux.sync` (`cuda-device`)**: Full-warp (`N == 32`) `u32` (`Sum`, `Min`, `Max`, `BitAnd`, `BitOr`, `BitXor`) and `i32` (`Sum`, `Min`, `Max`) `warp_reduce` calls emit `redux.sync` at `sm_80+`, while `f32`, sub-warp tiles (`N < 32`), and `sm_75` builds keep the shuffle butterfly. 3. **Legacy NVVM Legalization (`nvvm-transforms`)**: Rewrites the 8 integer `llvm.nvvm.redux.sync.*` intrinsics to convergent inline PTX on the LLVM 7 `libNVVM` path (`sm_80`..`sm_90`). ```text BEFORE (warp_reduce::<u32, Sum, 32> on sm_80+): 5-round butterfly: 5x shfl.sync.bfly.b32 + 5x add.s32 (10 instructions) AFTER (warp_reduce::<u32, Sum, 32> on sm_80+): sm_80+ floor: 1x `redux.sync.add.s32` instruction sm_75 floor: 5-round shuffle butterfly (unchanged) ``` --- ## Verification * **GPU Execution (`RTX 5090`, `sm_120a`)**: `cargo oxide run warp_reduce_redux` (default `sm_120a`, `--arch sm_75`, and `--emit-nvvm-ir --arch sm_80`), `verify-code-shape.sh`, and the `assert_only_expected_rebuilt` cache-freshness check across `--arch` switches all pass (`compute-sanitizer --tool memcheck`: `0 errors`). * **Tests & CI Gates**: `cargo test -p cuda-target-spec -p cuda-device -p nvvm-transforms -p cuda-oxide-codegen -p mir-importer --all-targets`, `just fmt-check`, `cargo clippy`, and all repository guards pass. Partially addresses #811 --------- Signed-off-by: Vyncint Ng <115854244+vyncint@users.noreply.github.com> Signed-off-by: nihalpasham <nihalp@nvidia.com> Co-authored-by: nihalpasham <nihalp@nvidia.com>
1 parent 8a04da1 commit 2ffc19b

21 files changed

Lines changed: 1802 additions & 22 deletions

File tree

‎cuda-oxide/crates/cuda-device/src/cooperative_groups.rs‎

Lines changed: 107 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -909,6 +909,44 @@ pub mod ops {
909909
fn identity() -> T;
910910
/// Binary combiner. Must be associative and commutative.
911911
fn combine(a: T, b: T) -> T;
912+
913+
/// Whether this operation has a full-warp instruction at the target floor.
914+
/// Other operation/type pairs use the shuffle implementation.
915+
const HAS_FULL_WARP_FORM: bool = false;
916+
917+
/// Reduce the participating warp. Called only when `HAS_FULL_WARP_FORM`
918+
/// is true and the tile contains all 32 lanes. An implementation that
919+
/// sets `HAS_FULL_WARP_FORM` must override it: this default returns
920+
/// `value` unreduced.
921+
#[inline(always)]
922+
fn reduce_full_warp(_mask: u32, value: T) -> T {
923+
value
924+
}
925+
}
926+
927+
// Exclude unsupported intrinsics before MIR and feature detection run.
928+
macro_rules! redux_full_warp {
929+
($ty:ty, |$mask:ident, $value:ident| $call:expr) => {
930+
const HAS_FULL_WARP_FORM: bool = true;
931+
932+
// Keep the intrinsic call as the return expression of the fast arm
933+
// and fall back to the shuffle butterfly when the resolved target
934+
// floor is below sm_80.
935+
#[inline(always)]
936+
fn reduce_full_warp($mask: u32, $value: $ty) -> $ty {
937+
if crate::warp::__target_has_integer_redux_sync() {
938+
return $call;
939+
}
940+
let mut acc = $value;
941+
let mut delta: u32 = 16;
942+
while delta > 0 {
943+
let other = crate::warp::shuffle_xor_sync($mask, acc as u32, delta) as $ty;
944+
acc = <Self as ReduceOp<$ty>>::combine(acc, other);
945+
delta >>= 1;
946+
}
947+
acc
948+
}
949+
};
912950
}
913951

914952
/// Sum: wrapping integer addition; IEEE-754 add for floats.
@@ -945,6 +983,7 @@ pub mod ops {
945983
fn combine(a: u32, b: u32) -> u32 {
946984
a.wrapping_add(b)
947985
}
986+
redux_full_warp!(u32, |m, v| crate::warp::redux_sync_add(m, v));
948987
}
949988
impl ReduceOp<u32> for Min {
950989
#[inline(always)]
@@ -955,6 +994,7 @@ pub mod ops {
955994
fn combine(a: u32, b: u32) -> u32 {
956995
if a < b { a } else { b }
957996
}
997+
redux_full_warp!(u32, |m, v| crate::warp::redux_sync_min_u32(m, v));
958998
}
959999
impl ReduceOp<u32> for Max {
9601000
#[inline(always)]
@@ -965,6 +1005,7 @@ pub mod ops {
9651005
fn combine(a: u32, b: u32) -> u32 {
9661006
if a > b { a } else { b }
9671007
}
1008+
redux_full_warp!(u32, |m, v| crate::warp::redux_sync_max_u32(m, v));
9681009
}
9691010
impl ReduceOp<u32> for BitAnd {
9701011
#[inline(always)]
@@ -975,6 +1016,7 @@ pub mod ops {
9751016
fn combine(a: u32, b: u32) -> u32 {
9761017
a & b
9771018
}
1019+
redux_full_warp!(u32, |m, v| crate::warp::redux_sync_and(m, v));
9781020
}
9791021
impl ReduceOp<u32> for BitOr {
9801022
#[inline(always)]
@@ -985,6 +1027,7 @@ pub mod ops {
9851027
fn combine(a: u32, b: u32) -> u32 {
9861028
a | b
9871029
}
1030+
redux_full_warp!(u32, |m, v| crate::warp::redux_sync_or(m, v));
9881031
}
9891032
impl ReduceOp<u32> for BitXor {
9901033
#[inline(always)]
@@ -995,6 +1038,7 @@ pub mod ops {
9951038
fn combine(a: u32, b: u32) -> u32 {
9961039
a ^ b
9971040
}
1041+
redux_full_warp!(u32, |m, v| crate::warp::redux_sync_xor(m, v));
9981042
}
9991043

10001044
// --- i32 ---
@@ -1007,6 +1051,7 @@ pub mod ops {
10071051
fn combine(a: i32, b: i32) -> i32 {
10081052
a.wrapping_add(b)
10091053
}
1054+
redux_full_warp!(i32, |m, v| crate::warp::redux_sync_add(m, v as u32) as i32);
10101055
}
10111056
impl ReduceOp<i32> for Min {
10121057
#[inline(always)]
@@ -1017,6 +1062,7 @@ pub mod ops {
10171062
fn combine(a: i32, b: i32) -> i32 {
10181063
if a < b { a } else { b }
10191064
}
1065+
redux_full_warp!(i32, |m, v| crate::warp::redux_sync_min_i32(m, v));
10201066
}
10211067
impl ReduceOp<i32> for Max {
10221068
#[inline(always)]
@@ -1027,6 +1073,7 @@ pub mod ops {
10271073
fn combine(a: i32, b: i32) -> i32 {
10281074
if a > b { a } else { b }
10291075
}
1076+
redux_full_warp!(i32, |m, v| crate::warp::redux_sync_max_i32(m, v));
10301077
}
10311078

10321079
// --- f32 ---
@@ -1256,13 +1303,13 @@ impl WarpShuffle for i64 {
12561303
/// bit-identical results. That is the primitive in isolation — a kernel that
12571304
/// reduces once after a memory-bound pass will see far less.
12581305
///
1259-
/// This function does not select that form for you: `redux.sync` will not
1260-
/// assemble below `sm_80`, and device code has no way to ask what target it is
1261-
/// being compiled for (see
1262-
/// [#811](https://github.com/NVlabs/cuda-oxide/issues/811)). Call it directly
1263-
/// when you know the target is Ampere+, gating on
1264-
/// `CudaContext::compute_capability` as the `redux_sum` example does. Floats
1265-
/// keep this butterfly — there is no `f32` `redux` before `sm_100`.
1306+
/// Full-warp integer reductions use `redux.sync` when the target floor is
1307+
/// `sm_80` or newer. Floats, smaller tiles and lower floors use shuffles.
1308+
///
1309+
/// `CUDA_OXIDE_TARGET` sets the floor. An advisory `CUDA_OXIDE_DEVICE_ARCH`
1310+
/// can lower the default `sm_80` floor, but cannot raise it: incompatible
1311+
/// kernels may target another architecture. The build script and backend
1312+
/// share this rule. Use `--arch sm_75` to select the shuffle implementation.
12661313
///
12671314
/// `N` is the tile size (1, 2, 4, 8, 16, or 32) — already validated by
12681315
/// [`ThreadBlock::tiled_partition`] at construction time.
@@ -1290,6 +1337,11 @@ where
12901337
T: WarpShuffle,
12911338
Op: ops::ReduceOp<T>,
12921339
{
1340+
// The const conditions remove this branch for floats and smaller tiles.
1341+
if N == 32 && Op::HAS_FULL_WARP_FORM {
1342+
return Op::reduce_full_warp(tile.mask(), value);
1343+
}
1344+
12931345
let mut acc = value;
12941346
let mut delta: u32 = N >> 1;
12951347
while delta > 0 {
@@ -1748,6 +1800,7 @@ where
17481800
#[cfg(test)]
17491801
mod tests {
17501802
use super::CoalescedThreads;
1803+
use super::ops::{BitAnd, BitOr, BitXor, Max, Min, ReduceOp, Sum};
17511804

17521805
#[test]
17531806
fn pack_lanes_uses_group_relative_bit_positions() {
@@ -1774,4 +1827,51 @@ mod tests {
17741827
);
17751828
}
17761829
}
1830+
1831+
/// Every integer pair `warp_reduce` can be asked for, and whether the
1832+
/// Ampere rung should give it a single-instruction form.
1833+
///
1834+
/// A `const` hook makes this checkable from the host with no device call,
1835+
/// which matters because the crate is `no_std` and the `redux_sync_*`
1836+
/// stubs are device functions.
1837+
const INTEGER_PAIRS: [(&str, bool); 9] = [
1838+
("u32 Sum", <Sum as ReduceOp<u32>>::HAS_FULL_WARP_FORM),
1839+
("u32 Min", <Min as ReduceOp<u32>>::HAS_FULL_WARP_FORM),
1840+
("u32 Max", <Max as ReduceOp<u32>>::HAS_FULL_WARP_FORM),
1841+
("u32 BitAnd", <BitAnd as ReduceOp<u32>>::HAS_FULL_WARP_FORM),
1842+
("u32 BitOr", <BitOr as ReduceOp<u32>>::HAS_FULL_WARP_FORM),
1843+
("u32 BitXor", <BitXor as ReduceOp<u32>>::HAS_FULL_WARP_FORM),
1844+
("i32 Sum", <Sum as ReduceOp<i32>>::HAS_FULL_WARP_FORM),
1845+
("i32 Min", <Min as ReduceOp<i32>>::HAS_FULL_WARP_FORM),
1846+
("i32 Max", <Max as ReduceOp<i32>>::HAS_FULL_WARP_FORM),
1847+
];
1848+
1849+
/// All nine are wired, not just the one that happened to get tested.
1850+
#[test]
1851+
fn every_integer_reduction_has_the_single_instruction_form_at_the_ampere_floor() {
1852+
for (name, wired) in INTEGER_PAIRS {
1853+
assert!(wired, "{name} lost its redux.sync form");
1854+
}
1855+
}
1856+
1857+
/// There is no `f32` `redux` before `sm_100`, so floats keep the butterfly
1858+
/// at every floor. The control a broadened hook would break: claiming the
1859+
/// form here would name an instruction that does not exist for the type.
1860+
#[test]
1861+
fn floats_never_claim_the_single_instruction_form() {
1862+
// `const` blocks: these are compile-time facts, so a regression is a
1863+
// build failure rather than a test failure.
1864+
const {
1865+
assert!(!<Sum as ReduceOp<f32>>::HAS_FULL_WARP_FORM);
1866+
assert!(!<Min as ReduceOp<f32>>::HAS_FULL_WARP_FORM);
1867+
assert!(!<Max as ReduceOp<f32>>::HAS_FULL_WARP_FORM);
1868+
}
1869+
}
1870+
1871+
/// The default hook must stay inert: it is reached only when the `const`
1872+
/// is false, and must not be mistaken for a reduction.
1873+
#[test]
1874+
fn the_default_hook_is_the_identity_on_its_input() {
1875+
assert_eq!(<Sum as ReduceOp<f32>>::reduce_full_warp(u32::MAX, 3.5), 3.5);
1876+
}
17771877
}

‎cuda-oxide/crates/cuda-device/src/warp.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -691,6 +691,14 @@ pub fn popc(predicate: bool) -> u32 {
691691
ballot(predicate).count_ones()
692692
}
693693

694+
/// Compiler-resolved predicate returning whether the target floor supports
695+
/// integer `redux.sync` (`sm_80` or newer).
696+
#[doc(hidden)]
697+
#[inline(never)]
698+
pub fn __target_has_integer_redux_sync() -> bool {
699+
unreachable!("__target_has_integer_redux_sync called outside CUDA kernel context")
700+
}
701+
694702
// =============================================================================
695703
// Warp Match Operations (sm_70+)
696704
// =============================================================================

‎cuda-oxide/crates/cuda-oxide-codegen/src/export.rs‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@ use crate::error::PipelineError;
77
use crate::generated::GeneratedModuleRequirements;
88
use crate::target::{
99
DetectedFeatures, arch_satisfies, generated_ptx_isa_requirement, generated_target_satisfied,
10-
select_target_with_generated, validate_generated_target, validate_target_features,
10+
honour_sm_floor, select_target_with_generated, validate_generated_target,
11+
validate_target_features,
1112
};
13+
use cuda_target_spec::resolve_sm_floor;
1214
use libnvvm_sys::CudaArch;
1315
use llvm_export::export::{
1416
DebugKind, DeviceExternType, ExportBackendConfig, FunctionLocalStaticPlacement, NvvmIrDialect,
@@ -82,6 +84,16 @@ impl llvm_export::export::AsDeviceExtern for DeviceExternDecl {
8284
}
8385
}
8486

87+
/// Resolve the shared device cfg floor when no target was explicitly selected.
88+
fn sm_floor_for_selection(device_arch_hint: Option<&CudaArch>) -> Result<u32, PipelineError> {
89+
resolve_sm_floor(None, device_arch_hint).map_err(|error| PipelineError::TargetSelection {
90+
target: String::new(),
91+
reason: format!(
92+
"cannot resolve the compute-capability floor device code is compiled against: {error}"
93+
),
94+
})
95+
}
96+
8597
// mir-importer pipeline plumbing; not part of the frontend contract.
8698
#[doc(hidden)]
8799
#[cfg(test)]
@@ -153,6 +165,12 @@ pub(crate) fn resolve_nvvm_target_with_generated(
153165
}
154166
let target =
155167
select_target_with_generated(features, generated).map_err(PipelineError::Export)?;
168+
// No explicit target reaches here, so this is the shared default
169+
// floor; it is resolved rather than restated so the ladder device
170+
// code was built against and the target selected here cannot drift.
171+
let floor = sm_floor_for_selection(device_arch_hint)?;
172+
let target =
173+
honour_sm_floor(target, floor, features, generated).map_err(PipelineError::Export)?;
156174
return Ok(target);
157175
}
158176

@@ -166,6 +184,9 @@ pub(crate) fn resolve_nvvm_target_with_generated(
166184
if !generated.is_empty() {
167185
let target = select_target_with_generated(DetectedFeatures::Basic, generated)
168186
.map_err(PipelineError::Export)?;
187+
let floor = sm_floor_for_selection(device_arch_hint)?;
188+
let target = honour_sm_floor(target, floor, DetectedFeatures::Basic, generated)
189+
.map_err(PipelineError::Export)?;
169190
return Ok(target);
170191
}
171192

‎cuda-oxide/crates/cuda-oxide-codegen/src/lib.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,17 @@ pub mod __private {
136136
#[doc(hidden)]
137137
pub use mir_transforms::unroll::{HintLocation, HintWarning};
138138

139+
/// Whether the resolved target floor supports Ampere integer `redux.sync` (`sm_80+`).
140+
#[doc(hidden)]
141+
pub fn target_has_integer_redux_sync(
142+
target_arch: Option<&str>,
143+
device_arch_hint: Option<&DeviceArchHint>,
144+
) -> bool {
145+
let device_arch = device_arch_hint.and_then(|hint| hint.as_device_arch().ok());
146+
cuda_target_spec::resolve_sm_floor(target_arch, device_arch.map(AsRef::as_ref))
147+
.is_ok_and(|floor| floor >= 80)
148+
}
149+
139150
/// Compiler-only attribute used to carry the Rust generated-intrinsic ABI marker.
140151
#[doc(hidden)]
141152
pub const GENERATED_INTRINSIC_MARKER_ATTR: &str =

‎cuda-oxide/crates/cuda-oxide-codegen/src/target/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,11 @@ pub(crate) use generated_requirements::{
2525
generated_ptx_isa_requirement, generated_target_satisfied, merge_generated_module_requirements,
2626
merge_generated_module_requirements_for_target, validate_generated_target,
2727
};
28-
pub use select::{required_ptx_feature, validate_target_features, validate_target_for_llvm_major};
2928
pub(crate) use select::{
30-
resolve_ptx_target_with_generated, select_target_with_generated,
29+
honour_sm_floor, resolve_ptx_target_with_generated, select_target_with_generated,
3130
validate_ptx_isa_for_llvm_major,
3231
};
32+
pub use select::{required_ptx_feature, validate_target_features, validate_target_for_llvm_major};
3333

3434
// The re-exports below preserve the module's pre-split surface. In-crate
3535
// consumers currently reach these names only from test code (or straight

‎cuda-oxide/crates/cuda-oxide-codegen/src/target/select.rs‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use super::generated_requirements::{
1111
};
1212
use crate::error::PipelineError;
1313
use crate::generated::GeneratedModuleRequirements;
14-
use cuda_target_spec::{CudaArch, RECORDED_PTX_FLOORS, recorded_ptx_floor};
14+
use cuda_target_spec::{CudaArch, RECORDED_PTX_FLOORS, recorded_ptx_floor, resolve_sm_floor};
1515

1616
/// Select a concrete architecture that satisfies every detected feature.
1717
///
@@ -167,6 +167,33 @@ pub fn resolve_ptx_target(
167167
)
168168
}
169169

170+
/// Raise an automatically selected target to the device cfg floor.
171+
/// Reject the combination if that floor cannot run the module's features.
172+
pub(crate) fn honour_sm_floor(
173+
selected: CudaArch,
174+
floor: u32,
175+
features: DetectedFeatures,
176+
generated: &GeneratedModuleRequirements,
177+
) -> Result<CudaArch, String> {
178+
if selected.capability() >= floor {
179+
return Ok(selected);
180+
}
181+
182+
let at_floor = CudaArch::new(floor, None).map_err(|error| {
183+
format!("compute capability {floor} is not a CUDA target this compiler knows: {error}")
184+
})?;
185+
if arch_satisfies(&at_floor, features) && generated_target_satisfied(&at_floor, generated) {
186+
return Ok(at_floor);
187+
}
188+
189+
Err(format!(
190+
"device code is compiled against compute capability {floor} or newer, but this module's \
191+
features run only on {}, which that floor cannot host. Pin the architecture you mean to \
192+
build for with `--arch` (or CUDA_OXIDE_TARGET), so device code and the emitted PTX agree.",
193+
selected.sm()
194+
))
195+
}
196+
170197
pub(crate) fn resolve_ptx_target_with_generated(
171198
explicit_override: Option<&str>,
172199
explicit_override_source: &'static str,
@@ -208,8 +235,21 @@ pub(crate) fn resolve_ptx_target_with_generated(
208235
return Ok((device.clone(), "detected GPU"));
209236
}
210237

238+
// Reached only with no explicit target, so only the detected-device hint
239+
// can move the floor -- downwards. Resolved through the same function the
240+
// build script uses rather than restated here.
241+
let floor = resolve_sm_floor(explicit_override, device_hint).map_err(|error| {
242+
PipelineError::TargetSelection {
243+
target: String::new(),
244+
reason: format!(
245+
"cannot resolve the compute-capability floor device code is compiled against: {error}"
246+
),
247+
}
248+
})?;
211249
let target =
212250
select_target_with_generated(detected, generated).map_err(PipelineError::PtxGeneration)?;
251+
let target = honour_sm_floor(target, floor, detected, generated)
252+
.map_err(PipelineError::PtxGeneration)?;
213253
Ok((target, "feature requirement"))
214254
}
215255

0 commit comments

Comments
 (0)