Skip to content

warp_reduce runs a 10-instruction shuffle tree where Ampere's redux.sync does it in one: 4.74x measured, but no way to gate it on the target #811

Description

@vyncint

cooperative_groups::warp_reduce<T, Op, N> is a generic butterfly: log2(N) rounds of shfl_xor_via + Op::combine. For a full warp that is 5 shuffles + 5 combines = 10 instructions. redux.sync.add.u32 does the same job in one, cuda-oxide already exposes it as warp::redux_sync_add (#223, #258) — the generic API just never routes integers to it.

Measured

A10G (sm_86, driver 595.71.05), 1M threads x 64 full-warp u32 sum reductions each, 30 timed launches after 5 warmup. Same kernel either way; only the reduction call differs:

form us/launch correct
warp_reduce::<u32, Sum, _>(&tile, v) 159.4 yes
warp::redux_sync_add(FULL_MASK, v) 33.6 yes
4.74x

Both bit-identical against a host replay, which is expected: wrapping u32 addition is associative mod 2^32, and min/max/and/or/xor are associative and exact, so the reduction order cannot change the result.

PTX confirms the mechanism — 10 shfl.sync versus 8 redux.sync across the respective unroll factors, i.e. 10 instructions per reduction against 1.

Scope of that number: the probe kernel does 64 reductions and almost nothing else, so it is reduction-bound and this is the primitive's cost in isolation, not an application speedup. It sizes the prize; a memory-bound kernel with one reduction at the end would see little.

The participation contract already matches

This looked like it might be unsound, and it is not. #672 states the existing behaviour plainly: "The warp reductions shuffle with the full 32-lane member mask, so every lane must be launched and converged." #672 then added separate _partial entry points for short warps rather than relaxing warp_reduce.

So warp_reduce on a WarpTile<32> already requires a fully converged warp and the full mask — exactly redux.sync's requirement. Substituting it does not weaken anything. Sub-warp tiles (N < 32) would keep the shuffle path, since each tile carries its own mask and the wrapper's out-of-tile lane substitution is shuffle-specific.

What actually blocks it

redux.sync is sm_80+ and the PTX will not assemble below that. redux_sum's own comment says so, and that example handles it by checking ctx.compute_capability() at runtime and printing skipping: redux.sync.add requires sm_80+ (Ampere).

Device code has no way to ask what it is being compiled for. There is no arch-derived cfg supplied to device codegen — I looked; the sm_80+ requirement lives in doc comments and is enforced by the caller. --device-cfg exists but is user-supplied, not automatic.

So warp_reduce cannot simply emit redux.sync: every sm_75 build that uses a warp reduction would stop assembling. That is a hard regression, not a tuning question, which is why this is an issue and not a PR.

Options, all yours to pick

  1. An arch-derived cfg for device codegen (cfg(cuda_oxide_sm_80) or similar), so warp_reduce can specialize where the target allows. This is the general fix and would serve every other arch-gated API in cuda-device, not just this one.
  2. Recognise it in the compiler: match the butterfly in mir-lower and rewrite to redux.sync when the target is sm_80+. The backend does know the target. Pattern-matching a shuffle tree is fragile, though, and it would fight the convergent modelling.
  3. Do nothing to the API, document the gap. Callers who know they are on Ampere+ and reducing integers can already write warp::redux_sync_add directly, and a measured note in the book would tell them it is worth 4.74x on the primitive. Cheapest, and no risk.

Happy to implement whichever you prefer, including (3) on its own. I did not want to pick the mechanism for you when the enabling piece — target-aware device code — is a broader design decision than this one reduction.

(Related but distinct: #274 tracks adding SM80+ intrinsics. redux.sync is already added; this is about the high-level API not using it.)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    codegenDevice code-generation pipeline (Rust MIR to IR to PTX)device-apisUser-facing device-side and kernel-authoring APIs (cuda-device, cuda-macros)intrinsicsDevice intrinsics and libdevice math mappingsperfPerformance of generated code or of the compiler itself

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions