GLSL Compute Radix Sort Local Scan Algorithm

This article explains the algorithms for digit extraction and local prefix scanning within a GPU-accelerated radix sort implemented in GLSL compute shaders. It details how keys are decomposed into radix digits using bitwise manipulation, how thread groups construct local histograms in shared memory, and how parallel prefix scans (such as Blelloch or Hillis-Steele) compute the local offsets required to reorder keys efficiently across compute workgroups.

Digit Extraction

In a GPU radix sort, multi-byte keys (typically 32-bit unsigned integers) are sorted in multiple passes. In each pass, a specific slice of bits—known as the radix digit—is processed. A standard choice is an 8-bit radix (\(b = 8\)), which partitions the sort into 4 passes over 256 possible digit bins (\(2^8 = 256\)). A 4-bit radix (\(b = 4\), 16 bins) is also common for low-register-pressure implementations.

The digit extraction algorithm isolates these bits using a bitwise right-shift followed by a bitwise AND with a bitmask.

uint extract_digit(uint key, uint pass_index, uint radix_bits) {
    uint bit_shift = pass_index * radix_bits;
    uint bit_mask = (1u << radix_bits) - 1u;
    return (key >> bit_shift) & bit_mask;
}

For an 8-bit radix, the mask is 0xFFu (255). On pass 0, the shift is 0; on pass 1, 8; on pass 2, 16; and on pass 3, 24.

Local Prefix Scanning

Once each thread in a workgroup extracts the digit for its assigned key, the workgroup must determine where each key will land relative to others with the same digit. This requires two steps: constructing a frequency histogram of digits within the workgroup and running a prefix sum (scan) over that histogram.

Depending on the architecture, local prefix scanning operates either at the thread level (local ranking) or at the bin level (histogram prefix sum). The standard local prefix scan pipeline consists of three phases:

1. Local Histogram Accumulation

Threads load keys from global memory, extract the digits, and tally the count of each digit within shared memory (shared uint).

#define RADIX_SIZE 256
#define WORKGROUP_SIZE 256

shared uint s_histogram[RADIX_SIZE];

void build_local_histogram(uint thread_id, uint key, uint pass) {
    // Zero out shared memory
    s_histogram[thread_id] = 0u;
    barrier();

    // Extract digit and increment count atomically
    uint digit = (key >> (pass * 8u)) & 0xFFu;
    atomicAdd(s_histogram[digit], 1u);
    barrier();
}

2. Multi-Pass or Hillis-Steele Workgroup Scan

To convert bin totals into base write offsets, an exclusive prefix scan is executed over s_histogram. A Hillis-Steele scan is simple and effective for execution inside GLSL shared memory:

  1. Step-based Addition: For offsets \(d = 1, 2, 4, 8, \dots < N\), each element at index \(i\) adds the value at index \(i - d\) to itself.
  2. Synchronization: A barrier() must be placed between every step to prevent data hazards.
  3. Inclusive to Exclusive Conversion: Shift the final array right by one element and pad index 0 with a 0.
shared uint s_scan[RADIX_SIZE];

void workgroup_exclusive_scan(uint thread_id) {
    // Load histogram to scan array
    s_scan[thread_id] = s_histogram[thread_id];
    barrier();

    // Hillis-Steele inclusive scan
    for (uint offset = 1u; offset < RADIX_SIZE; offset <<= 1u) {
        uint val = 0u;
        if (thread_id >= offset) {
            val = s_scan[thread_id - offset];
        }
        barrier();
        s_scan[thread_id] += val;
        barrier();
    }

    // Convert inclusive to exclusive scan
    uint inclusive_sum = s_scan[thread_id];
    barrier();
    
    if (thread_id == 0u) {
        s_histogram[0] = 0u;
    } else {
        s_histogram[thread_id] = s_scan[thread_id - 1u];
    }
    barrier();
}

3. Intra-Workgroup Thread Ranking

In addition to bin offsets, when multiple threads hold keys with identical digits, each thread needs a unique index within that bin. This is resolved via intra-workgroup ranking:

  1. Subgroup Intrinsics: Modern GLSL compute shaders utilize subgroup (wavefront/warp) operations. A ballot instruction (subgroupBallot) matches all threads in the subgroup holding the exact same digit.
  2. Bit Counting: subgroupBallotExclusiveBitCount() counts how many active threads with the same digit precede the current thread.
  3. Offset Calculation: The final local index for a key is: \[\text{Local Index} = \text{s\_histogram}[\text{digit}] + \text{subgroup\_rank}\]

After computing the local indices, keys are rearranged in shared memory or directly scattered to their designated positions in a globally partitioned prefix sum scheme (such as OneSweep or Chained Decoupled Lookback).