Understanding GL_NV_shader_subgroup_partitioned in GLSL

The GL_NV_shader_subgroup_partitioned GLSL extension enables invocations within a compute subgroup or warp to be dynamically partitioned into disjoint subsets based on arbitrary data values. By allowing operations to execute independently across these dynamically defined groups, the extension bypasses standard uniform subgroup restrictions. It significantly boosts performance in scenarios involving divergent data—such as binning, clustering, and local atomics—by replacing costly memory atomics and serialized control flow with efficient, hardware-level intra-subgroup reductions.

The Limitation of Standard Subgroups

Standard cross-invocation subgroup operations, defined in core specifications like GL_KHR_shader_subgroup_arithmetic and GL_KHR_shader_subgroup_ballot, evaluate arithmetic and reduction operations either across the entire active mask of a subgroup or within a fixed ballot mask. While powerful for uniform workloads, this model breaks down when invocations within the same subgroup process different tasks or categories simultaneously.

For instance, if different threads in a single warp are processing different material IDs, spatial bins, or hash keys, standard subgroup operations cannot perform reductions per category in a single step. Developers traditionally had to serialize execution by iterating through distinct values, use shared memory with complex synchronization, or fall back to high-contention global atomic operations.

How Partitioned Subgroups Function

The GL_NV_shader_subgroup_partitioned extension introduces functions that partition active invocations into non-overlapping subsets where every invocation in a subset shares the same user-specified key.

The mechanism operates in two primary phases:

  1. Partition Identification: The function subgroupPartitionNV(genType value) samples the provided value across all active invocations in the subgroup. It returns a 4-component integer vector representing a ballot mask of all invocations that evaluated to the exact same value. Invocations that provided matching keys receive identical ballot masks.
  2. Partitioned Operations: The extension provides a suite of partitioned arithmetic, scan, and ballot functions (such as subgroupPartitionedAddNV, subgroupPartitionedMinNV, and subgroupPartitionedInclusiveAddNV). These functions take the payload data alongside the partition ballot mask obtained from subgroupPartitionNV, executing the operation exclusively among the invocations identified in that mask.

Practical Roles in Non-Uniform Shader Workflows

Intra-Warp Reduction and Aggregation

In workloads like path tracing or clustered forward rendering, threads often need to update global histograms or accumulate energy buffers based on dynamic criteria. Using partitioned subgroups, threads sharing an ID compute their collective sum within the subgroup first. Only a single invocation per partition then emits a single atomic operation to global memory, dramatically cutting down atomic lock contention.

Handling Divergent Data Without Serialization

When control flow diverges based on dynamic properties, traditional subgroup reductions require branching loops that process one property at a time across the warp. Partitioned operations allow all subsets to compute their respective reductions in parallel without loop unrolling or serial execution overhead.

Stream Compaction and Binning

Building acceleration structures, particle systems, and tile lists requires fast compaction. By combining partitioned scans with ballot operations, non-uniform groups of data can calculate their relative output offsets and target locations entirely within hardware registers, avoiding intermediate storage in shared memory (LDS).

Through GL_NV_shader_subgroup_partitioned, GLSL provides access to low-level hardware ballot network features, allowing non-uniform data structures to achieve near-peak register bandwidth during intra-subgroup communication.