Skip to content

Chapter 09: Subgroups

Overview

Subgroups (also called waves or warps) are groups of invocations that execute in lockstep. This chapter covers:

  • Understanding the subgroup execution model
  • Subgroup operations (vote, ballot, arithmetic)
  • Efficient reductions and scans

What you'll learn:

  • Hardware-level parallelism
  • Avoiding shared memory for communication
  • Writing efficient reduction kernels

What is a Subgroup?

Within a workgroup, invocations are grouped into subgroups that execute together:

Workgroup (256 invocations)
├── Subgroup 0 (32 invocations) ─── Execute in lockstep
├── Subgroup 1 (32 invocations) ─── Execute in lockstep
├── Subgroup 2 (32 invocations) ─── Execute in lockstep
├── ...
└── Subgroup 7 (32 invocations) ─── Execute in lockstep

Subgroup size varies by GPU:

Vendor Subgroup Size
NVIDIA 32 (warp)
AMD 32 or 64 (wave)
Intel 8, 16, or 32 (SIMD width)
Apple 32

Querying Subgroup Properties

VkPhysicalDeviceSubgroupProperties subgroup_props = {
    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_PROPERTIES
};

VkPhysicalDeviceProperties2 props2 = {
    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2,
    .pNext = &subgroup_props
};

vkGetPhysicalDeviceProperties2(physical_device, &props2);

printf("Subgroup size: %u\n", subgroup_props.subgroupSize);
printf("Supported stages: %x\n", subgroup_props.supportedStages);
printf("Supported operations: %x\n", subgroup_props.supportedOperations);

Required Extensions

#version 450
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_arithmetic : require
#extension GL_KHR_shader_subgroup_ballot : require

Built-in Variables

// Your position within the subgroup (0 to subgroupSize-1)
uint gl_SubgroupInvocationID;

// Size of the subgroup
uint gl_SubgroupSize;

// Which subgroup you're in (within workgroup)
uint gl_SubgroupID;

// Total number of subgroups in workgroup
uint gl_NumSubgroups;

Subgroup Operations

Vote Operations

All invocations vote on a condition:

bool condition = (value > threshold);

// True if all invocations have condition = true
bool all_true = subgroupAll(condition);

// True if any invocation has condition = true
bool any_true = subgroupAny(condition);

// True if all invocations have same value
bool all_equal = subgroupAllEqual(value);

Broadcast

Share a value from one invocation to all:

// Get value from invocation 0
float shared = subgroupBroadcastFirst(my_value);

// Get value from specific invocation
float from_5 = subgroupBroadcast(my_value, 5);

Arithmetic Reductions

Reduce across the entire subgroup:

float sum = subgroupAdd(my_value);      // Sum all values
float product = subgroupMul(my_value);  // Multiply all
float minimum = subgroupMin(my_value);  // Find minimum
float maximum = subgroupMax(my_value);  // Find maximum

Inclusive/Exclusive Scans

Prefix operations:

// Inclusive: includes current invocation
float prefix_sum = subgroupInclusiveAdd(my_value);
// Invocation 0: v0
// Invocation 1: v0 + v1
// Invocation 2: v0 + v1 + v2
// ...

// Exclusive: excludes current invocation
float exclusive_sum = subgroupExclusiveAdd(my_value);
// Invocation 0: 0
// Invocation 1: v0
// Invocation 2: v0 + v1
// ...

Ballot

Get a bitmask of which invocations satisfy a condition:

bool condition = (my_value > 0);
uvec4 ballot = subgroupBallot(condition);

// Count how many invocations satisfy condition
uint count = subgroupBallotBitCount(ballot);

Shuffle

Exchange values between invocations:

// Get value from invocation at delta offset
float neighbor = subgroupShuffleDown(my_value, 1);  // From next invocation
float prev = subgroupShuffleUp(my_value, 1);        // From previous
float specific = subgroupShuffle(my_value, 5);       // From invocation 5

Example: Parallel Reduction

Sum 1 million values efficiently:

#version 460
#extension GL_KHR_shader_subgroup_arithmetic : require

layout(local_size_x = 256) in;

layout(set = 0, binding = 0) readonly buffer Input {
    float input_data[];
};

layout(set = 0, binding = 1) buffer Output {
    float output_data[];
};

layout(push_constant) uniform PushConstants {
    uint element_count;
};

shared float shared_data[256];

void main() {
    uint global_idx = gl_GlobalInvocationID.x;
    uint workgroup_idx = gl_WorkGroupID.x;
    uint stride = gl_NumWorkGroups.x * gl_WorkGroupSize.x;

    // Grid-stride load: each invocation accumulates every element it can
    // reach, so the dispatch need not cover one invocation per element.
    float value = 0.0;
    for (uint i = global_idx; i < element_count; i += stride) {
        value += input_data[i];
    }

    // Step 1: Reduce within subgroup (no shared memory needed!)
    float subgroup_sum = subgroupAdd(value);

    // Step 2: First invocation of each subgroup writes to shared
    uint lane = gl_SubgroupInvocationID;
    uint num_subgroups = gl_NumSubgroups;

    if (lane == 0) {
        shared_data[gl_SubgroupID] = subgroup_sum;
    }

    barrier();

    // Step 3: First subgroup folds the per-subgroup sums, looping in case
    // there are more subgroups than the subgroup is wide.
    if (gl_SubgroupID == 0) {
        float workgroup_sum = 0.0;
        for (uint base = 0; base < num_subgroups; base += gl_SubgroupSize) {
            uint slot = base + lane;
            float partial = (slot < num_subgroups) ? shared_data[slot] : 0.0;
            workgroup_sum += subgroupAdd(partial);
        }

        if (lane == 0) {
            output_data[workgroup_idx] = workgroup_sum;
        }
    }
}

Size the dispatch, or stride the load

A reduction pass that launches fewer invocations than there are elements will silently sum only the part it covers. The grid-stride loop makes the kernel correct for any dispatch size — the second pass here reduces 1024 partial sums with a single 256-thread workgroup.

Indexing the second stage

Step 3 must index shared_data by gl_SubgroupInvocationID, not gl_LocalInvocationID. They coincide only for subgroup 0 on the first iteration, so the wrong one produces subtly wrong sums.

Running the Example

./build/bin/ch09_subgroups

Output from an Apple M3 Pro:

=== Subgroup Properties ===
  Subgroup Size: 32
  Supported Stages: COMPUTE FRAGMENT 
  Supported Operations:
    - Basic (elect, barrier)
    - Vote (all, any, equal)
    - Arithmetic (add, mul, min, max)
    - Ballot
    - Shuffle
    - Shuffle Relative
    - Clustered
    - Quad

This GPU uses 32-wide subgroups (like CUDA warps of 32 threads)

=== Test Data ===
Array size: 262144 elements
Expected sum: 262144

=== Running Subgroup Reduction ===
Pass 1: 262144 elements -> 1024 partial sums
Pass 2: 1024 partial sums -> 1 final result

=== Results ===
Computed sum: 262144
Expected sum: 262144
Match: YES
Time: 3.758 ms

=== Benchmark ===
Iterations: 1000
Time per reduction: 0.2212 ms
Throughput: 4.42 GB/s
Elements/second: 1.19 billion

The input is 262144 copies of 1.0, so the expected sum is just the element count — an easy value to eyeball when the reduction goes wrong.

Performance Comparison

The sample measures only the subgroup path (0.22 ms per reduction of 262144 elements on an M3 Pro); it does not implement shared-memory or CPU baselines to compare against. Published comparisons generally put subgroup reductions ahead of the classic shared-memory tree, but the margin is hardware-specific — benchmark your own target before relying on a number.

Subgroups are faster because: - No shared memory writes/reads - No barriers within subgroup - Hardware-optimized operations

Subgroup Barriers

Sometimes you need explicit synchronization:

// Wait for all subgroup invocations
subgroupBarrier();

// Memory-specific barriers
subgroupMemoryBarrier();           // All memory
subgroupMemoryBarrierBuffer();     // Buffer memory
subgroupMemoryBarrierShared();     // Shared memory
subgroupMemoryBarrierImage();      // Image memory

Elect Operation

Choose one invocation to do special work:

if (subgroupElect()) {
    // Only one invocation per subgroup executes this
    atomicAdd(counter, 1);
}

Exercises

  1. Parallel Max: Implement finding the maximum value using subgroups.

  2. Prefix Sum: Implement an exclusive prefix sum (scan) for a large array.

  3. Histogram: Use subgroup ballot to count values in ranges.

Common Errors

Extension Not Enabled

// Error: subgroupAdd not found
#version 450
// Missing: #extension GL_KHR_shader_subgroup_arithmetic : require

Assuming Subgroup Size

// Wrong: Assumes 32
if (gl_SubgroupInvocationID < 32) { ... }

// Right: Use actual size
if (gl_SubgroupInvocationID < gl_SubgroupSize) { ... }

Cross-Subgroup Communication

Subgroup operations only work within a subgroup: 

// This won't share data between subgroups!
float wrong = subgroupBroadcast(my_value, 40);  // Invalid if size is 32

Feature Support

Check what your GPU supports:

VkSubgroupFeatureFlags ops = subgroup_props.supportedOperations;

if (ops & VK_SUBGROUP_FEATURE_BASIC_BIT)      printf("Basic\n");
if (ops & VK_SUBGROUP_FEATURE_VOTE_BIT)       printf("Vote\n");
if (ops & VK_SUBGROUP_FEATURE_ARITHMETIC_BIT) printf("Arithmetic\n");
if (ops & VK_SUBGROUP_FEATURE_BALLOT_BIT)     printf("Ballot\n");
if (ops & VK_SUBGROUP_FEATURE_SHUFFLE_BIT)    printf("Shuffle\n");

What's Next?

We've covered the major compute concepts. In Chapter 10, we'll learn how to debug and profile Vulkan applications effectively.