Subgroup primitives#

Subgroup operations let threads within the same subgroup (warp on NVIDIA, wave on AMD, subgroup on Vulkan / Metal) cooperate directly - exchanging register values, voting on predicates, scanning, and electing a leader - without going through shared memory or block barriers. They are the building block for fast in-warp data exchange and are used internally by Tile16x16 and Tile32x32 (see tile).

Subgroup ops live under qd.simt.subgroup and are written so the same Python source compiles to the right vendor primitive on each backend. Calling a backend that has not implemented an op fails when the kernel is compiled, not silently at runtime.

This page is a high-level tour: what each op does, how the pieces fit together, and what to expect performance-wise. The complete per-op reference - exact signatures, dtype support, per-backend behavior, and caller contracts - is generated from the docstrings in the API reference: quadrants.lang.simt.subgroup (data movement, identification / control, voting, and lane masks), quadrants.lang.simt.reductions (reductions and scans), and quadrants.lang.simt.sorting (sorting).

What’s available#

Every op is listed below, grouped by category, with the backends that currently support it.

Data movement#

Op	CUDA	AMDGPU	SPIR-V (Vulkan / Metal)	dtypes
`subgroup.shuffle(value, index)`	yes	yes	yes	i32, u32, f32, f64, i64, u64
`subgroup.shuffle_down(value, offset)`	yes	yes*	yes	i32, u32, f32, f64, i64, u64
`subgroup.shuffle_up(value, offset)`	yes	yes*	yes	i32, u32, f32, f64, i64, u64
`subgroup.shuffle_xor(value, mask)`	yes	yes	yes	i32, u32, f32, f64, i64, u64
`subgroup.broadcast(value, index)`	yes	yes	yes	i32, u32, f32, f64, i64, u64
`subgroup.broadcast_first(value)`	yes	yes	yes	i32, u32, f32, f64, i64, u64

* On AMDGPU, shuffle_down / shuffle_up use a slower cross-lane path rather than a single-cycle register move - see Performance notes for cycle counts.

shuffle_xor and broadcast_first are portable wrappers over shuffle / broadcast (shuffle_xor(value, mask) ≡ shuffle(value, lane ^ mask); broadcast_first(value) ≡ broadcast(value, qd.u32(0))). They inline at compile time and run wherever the underlying op runs.

SPIR-V caveat: broadcast(value, index) requires index to be dynamically uniform (the same on every lane). Use shuffle if you need a per-lane-varying source lane. See the broadcast docstring for the full contract.

Identification and control#

Op	CUDA	AMDGPU	SPIR-V (Vulkan / Metal)
`subgroup.invocation_id()`	yes	yes	yes
`subgroup.group_size()`	yes	yes	yes
`subgroup.log2_group_size()`	yes	yes	yes
`subgroup.elect()`	yes	yes	yes
`subgroup.sync()`	yes	yes	yes
`subgroup.mem_fence()`	yes**	yes**	yes

Default scope: every op in this module operates over the full active subgroup (32 lanes on wave32, 64 on wave64) unless its name carries a _tiled suffix. The _tiled variants take an extra log2_size template argument and split the subgroup into independent 2**log2_size-aligned tiles - see Tiled variants.

Caller contract for every op on this page: call from uniform control flow with all lanes active - the whole subgroup must execute the call together. Calling from divergent control flow is implementation-defined on most backends.

** mem_fence() compiles to a workgroup-scope fence on CUDA and AMDGPU. Both are over-strict for the subgroup-scope ask but correct: a workgroup-scope fence orders memory as observed by the whole workgroup, of which the subgroup is a strict subset.

Renames relative to the previous qd.simt.subgroup API:

subgroup.barrier() → subgroup.sync() (matching block.sync()).
subgroup.memory_barrier() → subgroup.mem_fence() (matching the planned block.mem_fence() and grid.mem_fence()).
subgroup.ballot_full_subgroup(predicate) → subgroup.ballot(predicate).
Every <op>(value, log2_size) reduction / scan / vote that previously took log2_size directly is now <op>_tiled(value, log2_size); the bare <op>(value) form is the full-subgroup convenience wrapper. This is a breaking change - call sites that hard-coded log2_size = 5 (the old “full warp” idiom) need to either drop the argument or add the _tiled suffix to keep the old behavior on wave64.

The barrier() / memory_barrier() / ballot_full_subgroup() names remain as deprecated aliases that emit a DeprecationWarning on first use and forward to the new ones; they will be removed in a future release. The rest of this page uses the new names.

Voting and predicate ops#

all_true(p) / any_true(p) / all_equal(v) vote across the entire subgroup and broadcast the i32 (0 or 1) result to every lane. The two ballot variants return a lane-per-bit mask: ballot_first_n(predicate, n) returns a u32 covering lanes [0, n) (with n a compile-time constant <= 32), and ballot(predicate) returns a u64 covering every lane in the subgroup (32 lanes on wave32, 64 on wave64). lanemask_{lt,le,eq,gt,ge}(lane_id) are closed-form u32 lane-mask constants.

Op	CUDA	AMDGPU	SPIR-V (Vulkan / Metal)
`subgroup.ballot_first_n(predicate, n)`	yes	yes	yes
`subgroup.ballot(predicate)`	yes	yes	yes
`subgroup.{all,any}_true(predicate)`	yes	yes	yes
`subgroup.all_equal(value)`	yes	yes	yes
`subgroup.lanemask_{lt,le,eq,gt,ge}(lane_id)`	yes	yes	yes

all_equal uses the backend’s native ==, so for floats NaN != NaN and +0.0 == -0.0. Bit-cast to a same-width integer dtype first if you need bit-equality on floats.

Reductions and scans#

reduce_* returns the reduction in lane 0 (other lanes are undefined); reduce_all_* broadcasts the reduction to every lane in the subgroup. inclusive_* / exclusive_* produce the per-lane prefix; lane i ends up with the scan of value[0..i] (inclusive) or value[0..i-1] (exclusive). segmented_reduce_* resets the scan at every head_flag != 0 and returns the per-segment inclusive reduction in every lane of that segment.

Op	CUDA	AMDGPU	SPIR-V (Vulkan / Metal)	dtypes
`subgroup.reduce_add(v)`	yes	yes*	yes	any type supporting `+`
`subgroup.reduce_all_add(v)`	yes	yes	yes	any type supporting `+`
`subgroup.segmented_reduce_add(v, head_flag)`	yes	yes*	yes	any type supporting `+`
`subgroup.reduce_{min,max}(v)`	yes	yes*	yes	integer + float
`subgroup.reduce_all_{min,max}(v)`	yes	yes	yes	integer + float
`subgroup.segmented_reduce_{min,max}(v, head_flag)`	yes	yes*	yes	integer + float
`subgroup.inclusive_{add,mul,min,max,and,or,xor}(v)`	yes	yes*	yes	integer + float (`_and` / `_or` / `_xor`: integer only)
`subgroup.exclusive_{add,mul,min,max,and,or,xor}(v)`	yes	yes*	yes	integer + float (`_and` / `_or` / `_xor`: integer only)

* On AMDGPU these use a slower cross-lane path (~tens of cycles per step) - see Performance notes.

Every op above has a paired _tiled form that takes an extra log2_size template parameter and operates on independent 2**log2_size-aligned tiles within the subgroup - see Tiled variants. For reductions other than the ones listed above, build a sized helper on top of shuffle_down / shuffle following the same pattern as reduce_add_tiled / reduce_all_add_tiled.

Float NaN handling for min / max reductions is implementation-defined (CUDA and AMDGPU suppress NaN; some SPIR-V drivers propagate it). Avoid NaN inputs if you need a portable result.

Sorting#

In-register key/value sort across the subgroup, one (key, value) pair per lane. Pure shuffle - no shared memory, no barriers - fully unrolled at compile time.

Op	CUDA	AMDGPU	SPIR-V (Vulkan / Metal)	dtypes
`subgroup.bitonic_sort_kv(key, value)`	yes	yes	yes	key & value: i32, u32, f32, f64, i64, u64 (independently typed)

Returns (key, value) - assign with key, value = subgroup.bitonic_sort_kv(key, value). Sorts ascending on the (key, value) lex tuple; ties on key break on ascending value (not a textbook-stable sort - equal-keyed lanes come back in ascending-value order, not in original-lane order). Tiled variant: bitonic_sort_kv_tiled(key, value, log2_size) runs the same sort independently on each 2**log2_size-aligned tile - see Tiled variants. See Bitonic key/value sort example for the short-input (sentinel-padding) pattern; the docstring covers the float-NaN and stability caveats in full.

Subgroup size#

group_size() and log2_group_size() return the subgroup size (and its base-2 log) in effect for the current program, as plain Python ints. They are usable both inside @qd.kernel / @qd.func bodies - where the value is folded in as a compile-time constant - and from host scope (handy for setting block_dim / grid shapes that match the subgroup width). Because the value is a plain int, it can be passed as a qd.template() argument; that is how every full-subgroup op picks the right log2_size per backend without a runtime branch (each one is <op>_tiled(v, log2_group_size()) under the hood).

Backend	Subgroup size
CUDA	`32` (every sm_30+ NVIDIA arch)
AMDGPU	`64` (every AMDGPU target is pinned to wave64, see supported_systems)
SPIR-V (Vulkan / Metal)	probed from the device at `qd.init()` (`32` on Metal and most Vulkan GPUs)

The value is fixed for the lifetime of a program, so two kernels launched under the same qd.init() always see the same subgroup size.

Tiled variants#

Every reduce / scan / vote op in this module has a paired _tiled form that takes an extra log2_size template parameter and runs group_size() / (2**log2_size) independent reductions / scans / votes in parallel - one per 2**log2_size-aligned tile within the subgroup. The base op is the special case where the tile is the whole subgroup (i.e. log2_size = log2_group_size()).

With log2_size = k, the subgroup splits into tiles of 2**k consecutive lanes each, and each tile does its own reduction completely independently of every other tile. The caller arranges 2**k <= group_size() so every tile is full; a smaller k simply gives more, narrower tiles. It does not mean “only the first tile is active”.

`log2_size`	wave32 (CUDA / most Vulkan / Metal)	wave64 (AMDGPU - see supported_systems)
5 (tile = 32)	1 tile: lanes 0-31 (= base op)	2 tiles: 0-31, 32-63
4 (tile = 16)	2 tiles: 0-15, 16-31	4 tiles: 0-15, 16-31, 32-47, 48-63
3 (tile = 8)	4 tiles of 8	8 tiles of 8
0 (tile = 1)	every lane is its own tile (no-op)	same

Supported `_tiled` ops#

Tiled op	Result placement
`subgroup.{all,any}_true_tiled(p, log2_size)`	broadcast-to-tile
`subgroup.all_equal_tiled(v, log2_size)`	broadcast-to-tile
`subgroup.reduce_add_tiled(v, log2_size)`	tile-local lane 0
`subgroup.reduce_all_add_tiled(v, log2_size)`	broadcast-to-tile
`subgroup.segmented_reduce_add_tiled(v, head_flag, log2_size)`	broadcast-to-tile
`subgroup.reduce_{min,max}_tiled(v, log2_size)`	tile-local lane 0
`subgroup.reduce_all_{min,max}_tiled(v, log2_size)`	broadcast-to-tile
`subgroup.segmented_reduce_{min,max}_tiled(v, head_flag, log2_size)`	broadcast-to-tile
`subgroup.inclusive_{add,mul,min,max,and,or,xor}_tiled(v, log2_size)`	broadcast-to-tile
`subgroup.exclusive_{add,mul,min,max,and,or,xor}_tiled(v, log2_size)`	broadcast-to-tile
`subgroup.bitonic_sort_kv_tiled(key, value, log2_size)`	broadcast-to-tile (every lane in the tile holds its sorted-position pair)

Broadcast-to-tile forms: every lane in each tile holds that tile’s result. Lanes in different tiles hold different results (their own tile’s).
Tile-local lane-0 forms: only the tile-local lane 0 holds the reduction. That’s lane 0 alone with log2_size=5 on wave32, lanes 0 and 32 with log2_size=5 on wave64, and so on. Other lanes hold partial reductions and should be treated as undefined. Use the reduce_all_*_tiled counterparts if you want every lane to see its tile’s result.

log2_size is a qd.template() compile-time constant in [0, log2_group_size()] - [0, 5] on wave32 backends and [0, 6] on AMDGPU wave64. The caller must ensure 2**log2_size <= group_size(); passing a larger value silently computes the wrong result on most backends, and there is no runtime check. Ops a backend does not support raise a qd.static_assert at compile time.

Examples#

Broadcast lane 0 to all lanes example#

import quadrants as qd
from quadrants.lang.simt import subgroup

@qd.kernel
def broadcast(a: qd.types.ndarray(dtype=qd.f32, ndim=1)):
    qd.loop_config(block_dim=64)
    for i in range(a.shape[0]):
        a[i] = subgroup.shuffle(a[i], qd.u32(0))

After the kernel, every lane in a subgroup holds the original value of its lane 0. subgroup.broadcast(a[i], qd.u32(0)) is interchangeable here.

Ballot: count how many lanes satisfy a condition example#

@qd.kernel
def count_positive(a: qd.types.ndarray(dtype=qd.f32, ndim=1),
                   counts: qd.types.ndarray(dtype=qd.u32, ndim=1)):
    qd.loop_config(block_dim=32)
    for i in range(a.shape[0]):
        mask = subgroup.ballot_first_n(qd.i32(a[i] > 0.0), 32)
        if subgroup.invocation_id() == 0:
            counts[i // 32] = mask

After the kernel, counts[g] contains a bitmask of which lanes in group g had positive values. Use popcount(mask) on the host to get the count.

Identity shuffle (each lane reads its own id) example#

Useful as a sanity check:

@qd.kernel
def identity(src: qd.types.ndarray(dtype=qd.f32, ndim=1),
             dst: qd.types.ndarray(dtype=qd.f32, ndim=1)):
    qd.loop_config(block_dim=64)
    for i in range(src.shape[0]):
        lane = subgroup.invocation_id()
        dst[i] = subgroup.shuffle(src[i], qd.cast(lane, qd.u32))

dst[i] equals src[i] on every lane.

Swap neighbors (xor pattern via explicit lane) example#

@qd.kernel
def swap_pairs(src: qd.types.ndarray(dtype=qd.f32, ndim=1),
               dst: qd.types.ndarray(dtype=qd.f32, ndim=1)):
    qd.loop_config(block_dim=64)
    for i in range(src.shape[0]):
        lane = subgroup.invocation_id()
        dst[i] = subgroup.shuffle(src[i], qd.cast(lane ^ 1, qd.u32))

Pairs (0,1), (2,3), … swap their values.

Arbitrary per-lane gather example#

@qd.kernel
def reverse4(src: qd.types.ndarray(dtype=qd.f32, ndim=1),
             dst: qd.types.ndarray(dtype=qd.f32, ndim=1)):
    qd.loop_config(block_dim=64)
    for i in range(src.shape[0]):
        lane = subgroup.invocation_id()
        group_base = (lane // 4) * 4
        src_lane = group_base + 3 - lane % 4
        dst[i] = subgroup.shuffle(src[i], qd.cast(src_lane, qd.u32))

Within each group of 4 contiguous lanes the values are reversed.

Tree reduction with `shuffle_down` example#

Classic warp-level sum of 4 values - after the second step, lane 0 of each group of 4 holds the total:

@qd.kernel
def reduce4(src: qd.types.ndarray(dtype=qd.f32, ndim=1),
            dst: qd.types.ndarray(dtype=qd.f32, ndim=1)):
    qd.loop_config(block_dim=64)
    for i in range(src.shape[0]):
        val = src[i]
        val = val + subgroup.shuffle_down(val, qd.u32(2))
        val = val + subgroup.shuffle_down(val, qd.u32(1))
        dst[i] = val

Extend the pattern (offsets 16, 8, 4, 2, 1, …) to reduce a full subgroup; only lane 0’s final value is meaningful, because the lanes near the top read past the end of the subgroup.

Sum 32 lanes with `reduce_add_tiled` example#

The same tree, packaged as a one-liner. Lane 0 of each group of 32 holds the total; other lanes hold partial sums:

@qd.kernel
def sum32(src: qd.types.ndarray(dtype=qd.f32, ndim=1),
          dst: qd.types.ndarray(dtype=qd.f32, ndim=1)):
    qd.loop_config(block_dim=32)
    for i in range(src.shape[0]):
        total = subgroup.reduce_add_tiled(src[i], 5)
        if subgroup.invocation_id() == 0:
            dst[i // 32] = total

5 is log2_size; 2**5 == 32 matches the block dim. The body of reduce_add_tiled unrolls at compile time into five shuffle_down + add pairs, so the generated code is identical to a hand-written tree reduction.

Broadcast the sum to all lanes with `reduce_all_add_tiled` example#

When every lane needs the reduction result - e.g. to normalize by the sum - use reduce_all_add_tiled, which leaves the sum in every lane. No follow-up broadcast needed:

@qd.kernel
def normalize32(a: qd.types.ndarray(dtype=qd.f32, ndim=1)):
    qd.loop_config(block_dim=32)
    for i in range(a.shape[0]):
        total = subgroup.reduce_all_add_tiled(a[i], 5)
        a[i] = a[i] / total

Every lane in each group of 32 sees the same total.

Partial-subgroup reductions example#

log2_size does not have to match the full subgroup. Sum groups of 8 with reduce_add_tiled(v, 3) or groups of 16 with reduce_all_add_tiled(v, 4); the caller just ensures 2**log2_size <= group_size() (so log2_size <= 5 on CUDA / Metal / Vulkan-wave32, <= 6 on AMDGPU wave64). Use the bare reduce_add(v) / reduce_all_add(v) form when you want “the whole subgroup” without hard-coding the limit.

Bitonic key/value sort example#

Sort up to 32 (key, value) pairs in registers, one per lane, with bitonic_sort_kv. The pattern below is the contact-pruning sort used by Genesis: each lane carries a packed link-pair id (key) and a contact index (value); after the sort, lane i holds the i-th smallest pair under the lex order (key, value). Lanes past the real data (lane >= n_con) carry a sentinel key (+inf here) that the sort moves to the tail:

@qd.kernel
def sort_contacts(keys: qd.types.ndarray(dtype=qd.f32, ndim=1),
                  idxs: qd.types.ndarray(dtype=qd.i32, ndim=1),
                  n_con: qd.i32):
    qd.loop_config(block_dim=32)
    for tid in range(32):
        # Load real data into the low n_con lanes; sentinel-pad the rest.  +inf compares greater than every real
        # key, so the sentinels drift to the high end of the sort.
        my_key = qd.f32(1.0e30)
        my_idx = qd.i32(-1)
        if tid < n_con:
            my_key = keys[tid]
            my_idx = idxs[tid]

        my_key, my_idx = subgroup.bitonic_sort_kv(my_key, my_idx)

        if tid < n_con:
            keys[tid] = my_key
            idxs[tid] = my_idx

After the kernel, keys[0..n_con] is sorted ascending and idxs is the matching permutation. The body unrolls at compile time into 30 shuffles + lex compares (for log2_size = 5, the wave32 default); no shared memory, no barriers.

Use bitonic_sort_kv_tiled(k, v, log2_size) directly to run multiple independent sorts per subgroup - e.g. bitonic_sort_kv_tiled(k, v, 3) runs group_size() / 8 independent 8-element sorts in parallel. The tiles are 2**log2_size-aligned within the subgroup and do not interact.

Inclusive scan with `inclusive_add_tiled` example#

@qd.kernel
def cumsum(a: qd.types.ndarray(dtype=qd.i32, ndim=1)):
    qd.loop_config(block_dim=32)
    for i in range(a.shape[0]):
        a[i] = subgroup.inclusive_add_tiled(a[i], 5)

After the call, lane k (within each group of 32) holds a[group_start] + a[group_start+1] + ... + a[k]. The 5 is log2_size; 2**5 == 32 matches the block dim. The body unrolls at compile time into five shuffle_up + add pairs. Use a smaller log2_size to scan over partial-subgroup groups (e.g. inclusive_add_tiled(v, 3) produces independent prefix sums in groups of 8).

Performance notes#

Shuffles are register-to-register on CUDA, and on SPIR-V where the GPU has hardware support - typically a handful of cycles, no memory traffic.
On AMDGPU, shuffle / shuffle_down / shuffle_up use a slower cross-lane path (roughly tens of cycles), so they cost more than the register-to-register shuffles on CUDA. On wave64 AMD GPUs, shuffles that cross the 32-lane half-boundary can cost a few extra cycles on some GPU architectures, but the underlying reads issue in parallel so latency stays close to a single shuffle.
shuffle_xor and broadcast_first are wrappers over shuffle / broadcast and inline at compile time, so on every backend they cost exactly the same as the underlying op.
Both ballot_first_n and ballot compile to a single hardware instruction on every backend. At n == 32, ballot_first_n skips the predicate-masking step entirely; at n < 32 it adds one multiply on the predicate.
reduce_add and reduce_all_add each issue exactly log2_group_size() shuffles and log2_group_size() adds per call (5 on wave32, 6 on AMDGPU wave64). No barriers, no shared memory, no launch overhead (they inline). The same holds for the _tiled form at any log2_size.
Prefer reduce_all_add over reduce_add + broadcast when you need the result in every lane - same cost, one fewer shuffle.
64-bit dtypes (i64, u64, f64) are emulated as two 32-bit shuffles on AMDGPU. Prefer 32-bit values when you have a choice.
The inclusive_* / exclusive_* scans cost log2_group_size() shuffle+op pairs (exclusive adds one extra shuffle + select), the same as a hand-rolled warp scan, on every backend. Quadrants uses a portable shuffle tree rather than a hardware scan intrinsic even where one is available, trading a small cost difference for predictable, uniform behavior across backends.
bitonic_sort_kv runs log2_size * (log2_size + 1) shuffles - 30 for log2_size = 5 (wave32), 42 for log2_size = 6 (wave64) - all compile-time unrolled into the calling kernel; no shared memory, no barriers within the sort.