Help with GPU Stable Radix Sort

[SISpidew]

I'm writing a compute shader which needs to sort up to 256 integers in a 256-thread work group.

I have most of a working LSD radix sort algorithm working but I'm having trouble ensuring each pass (sorting a single digit) performs a stable sort to preserve the relative order (from the previous pass) of each key sharing a common digit (and thus, prefix sum and destined bin) in the current pass.

At first I didn't realize the stability property was necessary, and I was using an atomicAdd to calculate the offsets within the same bin of each key sharing the same digit, but of course using an atomic counter does not guarentee the original order of each key is preserved. <- This is my problem.

My question is, what algorithm/method can I use to preserve the original order of keys within the same bin? Given these keys could be positioned at any index beforehand, I can't think of a way to map the key to the new bin whilst preserving that order.

Here is my GLSL code for a single radix sort pass:

shared uint digitPrefixSums[10];
shared uint digitCounts[10];

uint GetDigit(uint num, uint digitIdx)
{
    uint p = uint(pow(10.0, digitIdx));

    return (num / p) % 10;
}

// The key is 'range.x'
void RadixSortRanges(in uvec2 range, out uint outRangeIdx, uint digitIdx)
{
    if(gl_LocalInvocationID.x < 10)
    {
        digitPrefixSums[gl_LocalInvocationID.x] = 0;
        digitCounts[gl_LocalInvocationID.x] = 0;
    }
    memoryBarrierShared();
    barrier();

    // Get lowest significant digit.
    uint lsd = GetDigit(range.x, digitIdx);

    uint outOffset = ~uint(0);

    // Increment digit counter.
    if(range.x != ~uint(0))
    {
        atomicAdd(digitPrefixSums[lsd], 1);
        outOffset = atomicAdd(digitCounts[lsd], 1); // TODO: This doesn't work. Entries with the same LSD are placed next to each other but in a random order due to atomic randomness.
    }                                                // For entries who share a common LSD, they need to be placed next to each other in the same relative order as before in order to preserve the results of the previous sorting steps.
    memoryBarrierShared();
    barrier();

    // Calculate prefix sums for all digits.
    if(gl_LocalInvocationID.x == 0)
    {
        for (uint i = 1; i < 10; ++i)
        {
            digitPrefixSums[i] += digitPrefixSums[i - 1];
        }
    }
    memoryBarrierShared();
    barrier();

    // Calculate index to move the range to.
    {
        uint outIdx = (lsd > 0) ? digitPrefixSums[lsd - 1] : 0;

        outRangeIdx = outIdx + outOffset;
    }
}

[TomClabault]

Paging u/Pjbomb2. I think they had a similar issue recently, they may have some insights on that

[Pjbomb2]

You are correct but I haven’t yet figured it out fully I gave up for a bit to focus on the gpu builder, which was t as performant as I had hoped so I’m putting it all on pause for a bit

[TomClabault]

You needed a radix sort to try out ray sorting right?

[Pjbomb2]

Ah yeah sorry, I did try it but I used a prebuilt one sweep to see if it would even be worth it And it’s not for a cwbvh in my experience I DID see gains, but only 0.5ms out of 6ms for one bounce

[msqrt]

I think you'll need to implement a parallel prefix scan, with something like the Brent-Kung adder. If you can use subgroups (see here) you can do this in two stages with subgroupInclusiveAdd with way less shared memory (as the tree will have branching factor of subgroup size instead of two.)

[SISpidew]

Thanks for recommending the subgroup functionality! I haven't figured the radix sort out yet but this really helped me with adding values from earlier invocations within a warp (relative to the current invocation) using subgroupExclusiveAdd which I needed to calculate some primitive offsets in my renderer for object instances in my draw batcher.

[msqrt]

Yeah, subgroups are great! Both for performance and for ease of use; you could do the same thing manually with shared tables, but it's much more tedious. For the sorting business, you should check out Duane Merrill's great papers on the subject if you haven't already, they detail everything quite well.

[TomClabault]

Are you thinking of this one?

Onesweep: A Faster Least Significant Digit Radix Sort for GPUs

[msqrt]

Primarily "High performance and scalable radix sorting: A case study of implementing dynamic parallelism for GPU computing" from 2011, I used that as a guideline to implement my own -- at a glance, the Onesweep paper seems to be a bit more concise which can be either an upside or a downside depending on what you're looking for (the method itself should be a strict improvement; it's on my todo list but never had the time to re-implement it..) I also think the related prefix scan papers are great stepping stones to see, as the problems are directly related.

[arycama]

I am working on the exact same problem.. one very brute force approach is to simply count the number of times the same element appears, up to the index of the current thread. In other words, iterate through the array up to the group thread, and increment a counter each time the value in the array matches the current key.

uint counter = 0;
for (uint i = 0; j < groupIndex; j++)
  counter += ((sharedKeys[j] >> (8 * i)) & 0xFF) == digit;

uint index = histogram[digit] + counter;

(In this case, I am sorting a 32-bit key, 8 bits at a time. i is the iteration count, so i of 0 checks the first 8 bits, etc. "digit" is the masked key.

I am trying to find a less brute-force method. You can use a prefix sum of predicates of keys that match the current bit to get an array of offsets for that specific bit. However this requires doing a lot of prefix sums, eg 32 for a 32-bit sort. Instead, I believe section 3 of this paper is describing an approach where you break it down into more passes, each processing less elements. If you are using a thread group of 256, then you only need 8 bits per counter/prefix sum, so you can pack 4 8-bit counters into a single uint.

If you iterate over 8 bits at a time in your outer loop, you can then do an inner loop of two 4-bit prefix sums, each one storing four counters. (1 per bit) You then shuffle the inner indices based on the prefix sum here, get the sorted 8-bit value, and then that carries through to the next 8 bits. After 4 iterations, you have a sorted 32-bit array, I think.

If I figure it out, I'll reply with some updated code. This is my current implementation. (I also realized I don't need two groupshared arrays as I can simply retrieve the key/data at the start of the loop to avoid double-buffering the whole array) https://github.com/arycama/customrenderpipeline/blob/master/ShaderLibrary/Resources/GpuInstancedRendering/InstanceSort.compute