无线程组/Warp共享内存前缀和的16路基数排序实现优化问询

我为SPH模拟实现基数排序已有一段时间，最初改成了4位基数的16路基数排序，但这类实现的相关文档比较少。近期我尝试完全舍弃线程间通过专用共享内存计算前缀和的方式，转而依赖Metal的SIMD命令（对应CUDA中的Warp命令），在不完善的测试环境下，相比原16路实现，弃用线程组/Warp共享内存前缀和后性能提升显著。

现在有两个疑问：

1. 是否存在我没找到的同类实现？
2. 当前实现还有哪些可进一步优化的空间？

当前实现代码 #

void quantumsort_iter(const device uint32_t* keys_sorted_in,
                      device uint32_t* keys_sorted_out,
                      const uint32_t thread_id,
                      const uint32_t warp_sz,
                      const int radix,
                      const int start_idx,
                      const int end_idx)
{
    int prev_first = 0;
    int prefixsum[16] = { 0 };
    // count prefixes per thread
    for (int index = start_idx; index < end_idx; ++index)
        ++prefixsum[((keys_sorted_in[index] >> radix) & 0x0F)];
    // calculate prefix sums
    for (int index = 0; index < 16; ++index)
        prefixsum[index] = simd_prefix_inclusive_sum(prefixsum[index]);
    // cumulative sums
    for (int index = 0; index < 15; ++index)
        prefixsum[index + 1] += simd_broadcast(prefixsum[index], warp_sz - 1);
    // move prefixes right
    for (int index = 0; index < 16; ++index) {
        const int next_prev_first = simd_broadcast(prefixsum[index], warp_sz - 1);
        prefixsum[index] = simd_shuffle_and_fill_up(prefixsum[index], prev_first, 1);
        prev_first = next_prev_first;
    }
    // shuffle
    for (int index = start_idx; index < end_idx; ++index) {
        const uint32_t key_value = keys_sorted_in[index];
        const uint32_t key_masked = (key_value >> radix) & 0x0F;
        keys_sorted_out[prefixsum[key_masked]++] = key_value;
    }
}

kernel void quantumsort(const device uint32_t* keys_to_sort [[ buffer(0) ]],
                        device uint32_t* keys_sorted [[ buffer(1) ]],
                        device uint32_t* keys_sorted_swap [[ buffer(2) ]],
                        const device int& keys_sz [[ buffer(3) ]],
                        const uint32_t thread_id [[ thread_index_in_simdgroup ]],
                        const uint32_t warp_sz [[ threads_per_simdgroup ]])
{
    // constants
    const int size = (keys_sz + warp_sz - 1) / warp_sz;
    const int start = thread_id * size;
    const int end = min(start + size, keys_sz);
    // main loop
    for (int radix = 0; radix < 32; radix += 8) {
        quantumsort_iter(keys_to_sort, keys_sorted, thread_id, warp_sz, radix + 0, start, end);
        quantumsort_iter(keys_sorted, keys_to_sort, thread_id, warp_sz, radix + 4, start, end);
    }
}

我认为类似方法也可以通过CPU SIMD命令实现简单高效的基数排序。上述代码针对MacBook Pro M5开发，也适用于其他Mac设备。

内容的提问来源于stack exchange，提问作者harism