Skip to content

feat(cuda-device): add f32 and f64 warp reduce utilities#475

Open
honeyspoon wants to merge 1 commit into
NVlabs:mainfrom
honeyspoon:up/warp-reduce-f32
Open

feat(cuda-device): add f32 and f64 warp reduce utilities#475
honeyspoon wants to merge 1 commit into
NVlabs:mainfrom
honeyspoon:up/warp-reduce-f32

Conversation

@honeyspoon

Copy link
Copy Markdown
Contributor

warp::redux_sync_* is integer-only. The surface is redux_sync_add/and/or/xor plus redux_sync_{min,max}_{i32,u32}, with no float form — which mirrors PTX, where redux.sync defines only integer operand types. It is a hardware limit, not a gap in the bindings.

The consequence is that anyone reducing floats across a warp writes the same butterfly shuffle by hand. This adds it once, for both precisions:

reduce_sum_f32   reduce_max_f32   reduce_min_f32
reduce_sum_f64   reduce_max_f64   reduce_min_f64

Each is five XOR shuffles at offsets 16, 8, 4, 2, 1, built on the existing shuffle_xor_f32 / shuffle_xor_f64.

Notes for review

Every lane ends with the reduced value, not just lane 0. That is deliberate — it lets a caller feed the result straight into a block-level reduction without a second barrier, and it is the property a block reduce built on top of these depends on.

Safe functions, not unsafe. These are warp-synchronous but take no pointers and require no convergence beyond what the underlying shuffle intrinsics already document, so there is no additional obligation to put behind unsafe.

#[must_use] on all six: the return value is the entire point, so discarding one is always a mistake.

Uses f32::max / f32::min for the max and min variants rather than an open-coded compare, so the NaN behaviour follows whatever those lower to rather than being independently specified here.

fmt --check clean, no warnings, no new dependencies.

`warp::redux_sync_*` is integer-only - the verified surface is
`redux_sync_add/and/or/xor` plus `redux_sync_{min,max}_{i32,u32}`, with no
float form. That mirrors PTX, where `redux.sync` defines only integer operand
types, so it is a hardware limit rather than a gap in the bindings.

The consequence is that every user reimplementing a float reduction writes the
same butterfly shuffle by hand. This adds it once, for both precisions:

    reduce_sum_f32   reduce_max_f32   reduce_min_f32
    reduce_sum_f64   reduce_max_f64   reduce_min_f64

Each is five XOR shuffles at offsets 16, 8, 4, 2, 1, built on the existing
`shuffle_xor_f32`/`shuffle_xor_f64`. Every lane ends with the reduced value,
not just lane 0, which is what lets a caller feed the result straight into a
block-level reduction without a second barrier.

Safe functions: these are warp-synchronous but take no pointers and require no
convergence the shuffle intrinsics do not already document, so there is no
additional obligation to put behind `unsafe`.

`#[must_use]` on all six - the return value is the entire point, so discarding
one is always a mistake.

Signed-off-by: abder <bobmatt911@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant