Struct RadixSortPipeline 

pub struct RadixSortPipeline { /* private fields */ }

Introduction

This implementation of the radix-sort algorithm is based on the paper: Fast 4-way parallel radix sorting on GPUs

The radix sort algorithm can be divided into 3 steps:

1. Count the number of each radix in the block and generate a histogram (count_radix_pipeline).
2. Perform prefix sum operations on the histogram (scan_upsweep_pipeline, scan_dnsweep_pipeline, scan_last_block_pipeline).
3. Based on the histogram information, write the key values in the block to new ordered positions (scatter_pipeline).

Modification

I have made modifications to the Step 2: Compute prefix sum on the histogram in the paper to improve performance.

The memory layout of global_histogram_buffer in the paper is as follows:

              workgroup0      workgroup1        ...           workgroupN                                                                    
          +---------------+---------------+---------------+---------------+                                                                 
radix 0   |      V_0      |      V_1      |     ...       |     V_n-1     |                                                                 
          +---------------+---------------+---------------+---------------+                                                                 
radix 1   |      V_n      |      V_n+1    |     ...       |     V_2n-1    |                                                                 
          +---------------+---------------+---------------+---------------+                                                                 
 ...      |                              ...                              |                                                                 
          +---------------+---------------+---------------+---------------+                                                                 
radix 255 |     V_255n    |     V_255n+1  |     ...       |    V_256n-1   |                                                                 
          +---------------+---------------+---------------+---------------+                                                                 

This memory layout is beneficial for Step 2 prefix sum calculations, but not for Step 1 and Step 3, because the memory access to global_histogram_buffer in Step 1 and Step 3 is Non-Coalesced Memory Access.

Non-Coalesced Memory Access can cause significant delays, leading to a drastic decrease in the performance of radix_sort. To address this issue, I redesigned the memory layout of global_histogram_buffer:

                 radix 0         radix 1          ...            radix 255                                                                    
            +---------------+---------------+---------------+---------------+                                                                 
workgroup0  |      V_0      |      V_1      |     ...       |     V_255     |                                                                 
            +---------------+---------------+---------------+---------------+                                                                 
workgroup1  |     V_256     |     V_257     |     ...       |     V_511     |                                                                 
            +---------------+---------------+---------------+---------------+                                                                 
   ...      |                              ...                              |                                                                 
            +---------------+---------------+---------------+---------------+                                                                 
workgroupN  |     V_256n    |    V_256n+1   |     ...       |  V_256n+255   |                                                                 
            +---------------+---------------+---------------+---------------+                                                         

This solves the delay problem caused by Non-Coalesced Memory Access, but brings a new problem: how to perform prefix sum calculations on such a memory layout?

The idea is simple: change the unit of the prefix sum operation from radix to histogram, and perform the prefix sum operation on the histogram. This means accumulating histogram-workgroup0 to histogram-workgroup1, histogram-workgroup1 to histogram-workgroup2, and so on.

Of course, we cannot use a serial prefix sum as it would degrade performance. Therefore, we use the Blelloch parallel prefix sum algorithm, which consists of an up-sweep and a down-sweep step:

Up-Sweep

          ┌──┐ ┌──┐ ┌──┐ ┌──┐ ┌──┐ ┌──┐ ┌──┐ ┌──┐                                                       
          │H0│ │H1│ │H2│ │H3│ │H4│ │H5│ │H6│ │H7│                                                       
          └┬─┘ └┬─┘ └┬─┘ └┬─┘ └┬─┘ └┬─┘ └┬─┘ └┬─┘                                                       
           └───┬▼─┐  └───┬▼─┐  └───┬▼─┐  └───┬▼─┐                                                       
Round 0        │H1│      │H3│      │H5│      │H7│                                                       
               └┴─┴──────┴┬─┘      └──┴──────┴┬─┘                                                       
                         ┌▼─┐                ┌▼─┐                                                       
Round 1                  │H3│                │H7│                                                       
                         └──┴────────────────┴┬─┘                                                       
                                             ┌▼─┐                                                       
Round 2                                      │H7│                                                       
                                             └──┘   

Down-Sweep

    Round 1   Round 0   Round 1                                                                          
    ┌────┐    ┌────┐    ┌────┐                                                                           
┌──┐│┌──┐│┌──┐│┌──┐│┌──┐│┌──┐│┌──┐ ┌──┐                                                                  
│H0│││H1│││H2│││H3│││H4│││H5│││H6│ │H7│                                                                  
└──┘│└──┘│└▲─┘│└──┘│└▲─┘│└──┘│└▲─┘ └──┘                                                                  
    │┌──┐│ │  │┌──┐│ │  │┌──┐│ │   ┌──┐                                                                  
    ││H1├┼─┘  ││H3├┼─┘  ││H5├┼─┘   │H7│                                                                  
    │└──┘│    │└──┘│    │└▲─┘│     └──┘                                                                  
    └────┘    │┌──┐│    └─┼──┘     ┌──┐                                                                  
              ││H3├┼──────┘        │H7│                                                                  
              │└──┘│               └──┘                                                                  
              └────┘               ┌──┐                                                                  
                                   │H7│                                                                  
                                   └──┘  

The performance of parallel algorithms on GPUs is often bandwidth-sensitive. Therefore, the number of memory accesses is a crucial indicator of the algorithm’s performance.

The memory access count of the Blelloch algorithm is quite high: 2R1W + 2R1W = 4R2W, which is a rather poor situation. Moreover, since it executes on device-memory, it cannot hide latency through shared memory. (General parallel prefix sum algorithms can achieve 2R1W memory access count, so the performance of device-scope Blelloch is not good, only half of the pre-modified performance)

However, the performance degradation brought by Blelloch is not a big issue here, because we can reduce the overall cost of Step 2 by reducing the number of histograms. To reduce the number of histograms, the method is very straightforward: increase the number of keys processed during the statistics phase (Step 1).

Previously, one workgroup processed 256 keys and output a histogram of length 256. It can be modified to one workgroup processing 16*256 keys and outputting a histogram of length 256. This reduces the number of histograms by 16x, thereby reducing the memory access overhead of Step 2.

Advanced?

In fact, the algorithm can be further optimized!

We can use the decoupled look-back algorithm to halve memory access overhead, achieving nearly 100% performance improvement. However, the problem is that decoupled look-back relies on communication between workgroups, and there are two ways to achieve workgroup communication:

1. device-scope storage-buffer memory barrier.
2. forward progress guarantees.

The rendering layer of bevy relies on wgpu, which is an implementation of the webgpu standard. Unfortunately, webgpu does not support either of these methods. This is mainly due to Apple’s Metal graphics API not supporting these two features, which leads to webgpu, aiming for full compatibility, also not supporting them.

For more details, see atomic concerns.

But fortunately, for the vast majority of desktop GPUs (Nvidia/AMD), forward progress guarantees are implicitly supported. So if your target platform does not include Apple and mobile devices, you can still use webgpu to implement the decoupled look-back radix sort algorithm.

However, for compatibility reasons, a more traditional radix sort algorithm is used here.

Tips

The number of pass these 3 steps are executed depends on the number of bits in the keys and the number of bits processed per pass. For keys of type u32 and RADIX_BITS_PER_PASS set to 8, 4 passes are required.

If the effective number of bits is less than 32 (but the key type is still u32), for example, if all keys are less than 256, then only 1 pass is needed.

You can customize the number of pass and positions for processing the keys according to your specific requirements to improve performance.

Implementations

impl RadixSortPipeline

pub fn bind_group_layout(&self) -> &BindGroupLayout

Trait Implementations

impl Clone for RadixSortPipeline

fn clone(&self) -> RadixSortPipeline

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for RadixSortPipeline

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl FromWorld for RadixSortPipeline

fn from_world(world: &mut World) -> Self

Creates Self using data from the given World.

impl Resource for RadixSortPipeline

where Self: Send + Sync + 'static,