Implement Prefix Sum Using DirectX Wave Intrinsics

This article provides an overview and practical guide on implementing parallel prefix sums (scan operations) using HLSL Shader Model 6.0 DirectX wave intrinsics. By leveraging hardware-level SIMD lane communication, wave intrinsics allow compute shaders to perform both inclusive and exclusive prefix scans across an active wave without routing data through slow groupshared memory.

Understanding Wave-Level Prefix Sums

A prefix sum calculates a sequence of running totals across an array of values. In GPU execution, an active collection of threads executing in lockstep is called a wave (or warp). Prefix sums across a wave are divided into two main categories:

  • Exclusive Scan: The element at index i contains the sum of all elements prior to i (excluding element i itself).
  • Inclusive Scan: The element at index i contains the sum of all elements up to and including i.

Native Wave Intrinsics

DirectX 12 (Shader Model 6.0 and above) provides built-in intrinsic functions optimized by hardware vendors specifically for exclusive prefix sums.

The primary function is WavePrefixSum():

// Computes the exclusive prefix sum of 'value' across all active lower-index lanes
float WavePrefixSum(float value);
uint WavePrefixSum(uint value);

Implementing Exclusive and Inclusive Scans

Because the native HLSL intrinsic executes an exclusive scan, an inclusive scan is achieved by adding the calling thread's own value to the result:

[numthreads(64, 1, 1)]
void CSMain(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    uint localValue = dispatchThreadID.x;

    // Exclusive scan: sum of lanes [0, WaveGetLaneIndex() - 1]
    uint exclusiveScan = WavePrefixSum(localValue);

    // Inclusive scan: sum of lanes [0, WaveGetLaneIndex()]
    uint inclusiveScan = exclusiveScan + localValue;

    // Output or use results
}

Manual Prefix Sum Implementation (Hillis-Steele)

When implementing custom reduction operators (such as custom data types or non-standard associative operations) where WavePrefixSum cannot be used directly, a logarithmic scan algorithm can be constructed using WaveReadLaneAt:

uint CustomWavePrefixSum(uint value)
{
    uint laneIndex = WaveGetLaneIndex();
    uint waveSize = WaveGetLaneCount();

    // Hillis-Steele logarithmic scan
    for (uint offset = 1; offset < waveSize; offset <<= 1)
    {
        uint neighbor = WaveReadLaneAt(value, laneIndex - offset);
        if (laneIndex >= offset)
        {
            value += neighbor;
        }
    }

    return value; // Returns the inclusive sum
}

Extending to Threadgroup Scans

A single wave typically consists of 32 or 64 lanes. When a threadgroup exceeds the native wave size, prefix sums must be performed hierarchically:

  1. Intra-Wave Scan: Each wave computes its local prefix sum using WavePrefixSum().
  2. Wave Total Export: The last lane of each wave (WaveGetLaneIndex() == WaveGetLaneCount() - 1) writes the total wave sum to a small groupshared memory array.
  3. Inter-Wave Scan: A single wave scans the groupshared array containing the sums of each wave.
  4. Final Accumulation: Each thread adds the scanned wave offset from shared memory back to its local wave sum.

This hierarchical approach provides maximal memory bandwidth efficiency and avoids heavy memory barriers across large compute workloads.