Bitonic Merge Sort in GLSL Compute Shaders

This article explains how to implement the bitonic merge sort algorithm across workgroup threads in a GLSL compute shader. It covers the underlying mechanics of parallel bitonic networks, the mapping of invocation IDs to paired data elements, the use of shared memory for high-throughput data exchanges, and the mandatory synchronization barriers needed to preserve execution order on the GPU.

Understanding the Bitonic Sort Network

Bitonic merge sort is a comparison-based sorting algorithm designed specifically for parallel hardware. Unlike QuickSort or MergeSort, bitonic sort performs comparisons in a predetermined sequence that does not depend on the data values themselves. This deterministic execution pattern eliminates thread divergence on SIMD/SIMT architectures like GPUs.

A sequence is "bitonic" if it first monotonically increases and then monotonically decreases, or can be cyclically shifted to do so. The sorting process consists of two primary loops:

  1. Stages (\(k\)): The size of the monotonic sub-sequences being created, doubling each iteration (\(2, 4, 8, \dots, N\)).
  2. Steps (\(j\)): The comparison distance within the current stage, halving each iteration (\(k/2, k/4, \dots, 1\)).

Workgroup Memory and Thread Mapping

On the GPU, global memory bandwidth is often the primary performance bottleneck. To optimize bitonic sort within a single workgroup, the unsorted data is transferred from a global buffer (such as a Shader Storage Buffer Object, or SSBO) into high-speed shared memory.

Each thread within the workgroup is assigned to process one or two elements. When sorting an array of size \(N\) with \(N/2\) threads:

  • Each thread handles a distinct pair of indices.
  • gl_LocalInvocationIndex identifies the calling thread within the workgroup.
  • Workgroup barriers ensure all threads finish a comparison step before the next step begins.

The Indexing and Swapping Logic

To determine which two indices must be compared at a given stage (\(k\)) and step (\(j\)), bitwise operations are employed:

  • Partner Index: For a given element index \(i\), its comparison partner is calculated using a bitwise XOR operation: partner = i ^ j.
  • Direction: The overall sort direction (ascending vs. descending) for the local sub-array is determined by whether the bit at position \(k\) is set: bool ascending = (i & k) == 0.
  • Execution Guard: Because each comparison involves two indices, only the thread with the lower index (i < partner) performs the compare-and-swap operation to avoid redundant or conflicting writes.

GLSL Compute Shader Implementation

Below is a complete implementation of an in-workgroup bitonic merge sort designed for a power-of-two dataset (e.g., 1024 elements using 512 threads).

#version 450

// Define the number of elements to sort (must be a power of two)
#define NUM_ELEMENTS 1024
#define WORKGROUP_SIZE (NUM_ELEMENTS / 2)

layout(local_size_x = WORKGROUP_SIZE, local_size_y = 1, local_size_z = 1) in;

// Input/Output storage buffer
layout(std430, binding = 0) buffer DataBuffer {
    float elements[];
};

// Workgroup shared memory for fast access
shared float shared_data[NUM_ELEMENTS];

void main() {
    uint local_id = gl_LocalInvocationIndex;

    // Step 1: Load global data into workgroup shared memory
    uint idx1 = local_id;
    uint idx2 = local_id + WORKGROUP_SIZE;

    shared_data[idx1] = elements[idx1];
    shared_data[idx2] = elements[idx2];

    // Ensure all loads are complete before sorting begins
    barrier();
    memoryBarrierShared();

    // Step 2: Perform the bitonic sort passes
    for (uint k = 2; k <= NUM_ELEMENTS; k <<= 1) {
        for (uint j = k >> 1; j > 0; j >>= 1) {
            // Each thread processes one pair of indices
            // Map thread ID to target element index
            uint i = 2 * local_id - (local_id & (j - 1));
            uint partner = i ^ j;

            if (partner > i) {
                // Determine direction: true = ascending, false = descending
                bool ascending = ((i & k) == 0);

                float val_i = shared_data[i];
                float val_partner = shared_data[partner];

                // Swap values if they violate the expected order
                if ((val_i > val_partner) == ascending) {
                    shared_data[i] = val_partner;
                    shared_data[partner] = val_i;
                }
            }

            // Synchronize threads before the next comparison step
            barrier();
            memoryBarrierShared();
        }
    }

    // Step 3: Write sorted results back to global memory
    elements[idx1] = shared_data[idx1];
    elements[idx2] = shared_data[idx2];
}

Synchronization Constraints

Correct execution hinges on two synchronization primitives:

  1. barrier(): Stalls the execution of the calling thread until all other threads in the workgroup reach this point.
  2. memoryBarrierShared(): Guarantees that all writes to shared memory variables made by the calling thread are committed and visible to all other threads in the workgroup.

Omitting either primitive results in race conditions, where a thread compares and overwrites data that has not yet been updated by its partner thread from the preceding step.

Scaling Beyond a Single Workgroup

The shared memory approach is constrained by the maximum workgroup size supported by hardware (typically 1024 threads). For datasets exceeding the local workgroup limit:

  • The initial stages (up to the workgroup capacity) are executed inside local shared memory as demonstrated above.
  • Subsequent stages (where comparison distances span across different workgroups) cannot use workgroup barriers. These must be executed via multiple separate compute shader dispatches, reading from and writing to global buffers directly, using the host command buffer to enforce pipeline barriers between dispatches.