CUDA 算子笔记 3

Stream Compaction

给一个包含正负数和 0 的数组，把所有的正数按照原来的顺序放到数组开头并输出

例子：

Input:  A = [1.0, -2.0, 3.0, 0.0, -1.0, 4.0]
Output: out = [1.0, 3.0, 4.0, 0.0, 0.0, 0.0]

这是 Radix Sort 的一个前置题目，用计数排序的思想，包含三个步骤：

Predicate：判断每个数是否是正数，得到 predicate[]
Scan：对 predicate[] 计算 Exclusive 前缀和，得到元素的新位置
Scatter：计算出来的新位置，把元素分散写回输出数组

最直接的方法

直接三个 Kernel 即可：

#include <cuda_runtime.h>

#define CEIL(x, y) (((x) + (y) - 1) / (y))
#define FLOAT4(x) (*(reinterpret_cast<float4*>(&(x))))
#define CONST_FLOAT4(x) (*(reinterpret_cast<const float4*>(&(x))))
#define INT4(x) (*(reinterpret_cast<int4*>(&(x))))
#define CONST_INT4(x) (*(reinterpret_cast<const int4*>(&(x))))

__device__ __forceinline__ int warp_inclusive_scan(int val) {
  const int lane = threadIdx.x & 31;
#pragma unroll
  for (int offset = 1; offset < 32; offset <<= 1) {
    const int other = __shfl_up_sync(0xffffffff, val, offset);
    if (lane >= offset) {
      val += other;
    }
  }
  return val;
}

__device__ __forceinline__ int warp_exclusive_scan(int val) {
  return warp_inclusive_scan(val) - val;
}

__device__ __forceinline__ int block_inclusive_scan(int val) {
  __shared__ int warp_offsets[32];

  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int num_warps = blockDim.x >> 5;

  const int prefix = warp_inclusive_scan(val);

  if (lane == 31) {
    warp_offsets[warp] = prefix;
  }

  __syncthreads();

  if (warp == 0) {
    const int warp_sum = lane < num_warps ? warp_offsets[lane] : 0;
    const int warp_offset = warp_exclusive_scan(warp_sum);
    if (lane < num_warps) {
      warp_offsets[lane] = warp_offset;
    }
  }

  __syncthreads();

  return prefix + warp_offsets[warp];
}

__device__ __forceinline__ int block_exclusive_scan(int val) {
  return block_inclusive_scan(val) - val;
}

__device__ int g_block_offsets[20000];

template <int k_items_per_thread>
__global__ void stream_compaction_block_part1_kernel(
    const float* __restrict__ A, int N) {
  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
  const int thread_idx = tid * k_items_per_thread;

  int thread_sum = 0;

  if (thread_idx + k_items_per_thread <= N) {
#pragma unroll
    for (int i = 0; i < k_items_per_thread; i += 4) {
      const float4 val = CONST_FLOAT4(A[thread_idx + i]);
      thread_sum += (val.x > 0) + (val.y > 0) + (val.z > 0) + (val.w > 0);
    }
  } else {
#pragma unroll
    for (int i = 0; i < k_items_per_thread; ++i) {
      const int idx = thread_idx + i;
      if (idx < N) {
        thread_sum += A[idx] > 0;
      }
    }
  }

  const int block_prefix = block_inclusive_scan(thread_sum);

  if (threadIdx.x == blockDim.x - 1) {
    g_block_offsets[blockIdx.x] = block_prefix;
  }
}

__global__ void stream_compaction_grid_kernel(int M) {
  const int items_per_thread = CEIL(M, blockDim.x);
  const int thread_idx = threadIdx.x * items_per_thread;

  int thread_offsets[128];
  int thread_sum = 0;

  for (int i = 0; i < items_per_thread; ++i) {
    const int idx = thread_idx + i;
    const int val = idx < M ? g_block_offsets[idx] : 0;
    thread_offsets[i] = thread_sum;
    thread_sum += val;
  }

  const int offset = block_exclusive_scan(thread_sum);

  for (int i = 0; i < items_per_thread; ++i) {
    const int idx = thread_idx + i;
    if (idx < M) {
      g_block_offsets[idx] = thread_offsets[i] + offset;
    }
  }
}

template <int k_items_per_thread>
__global__ void stream_compaction_block_part2_kernel(
    const float* __restrict__ A, float* __restrict__ out, int N) {
  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
  const int thread_idx = tid * k_items_per_thread;

  float thread_inputs[k_items_per_thread];

  int thread_sum = 0;

  if (thread_idx + k_items_per_thread <= N) {
#pragma unroll
    for (int i = 0; i < k_items_per_thread; i += 4) {
      const float4 val = CONST_FLOAT4(A[thread_idx + i]);
      FLOAT4(thread_inputs[i]) = val;
      thread_sum += (val.x > 0) + (val.y > 0) + (val.z > 0) + (val.w > 0);
    }
  } else {
#pragma unroll
    for (int i = 0; i < k_items_per_thread; ++i) {
      const int idx = thread_idx + i;
      const float val = idx < N ? A[idx] : 0.0f;
      thread_inputs[i] = val;
      if (idx < N) {
        thread_sum += val > 0;
      }
    }
  }

  const int thread_offset = block_exclusive_scan(thread_sum);
  int output_pos = g_block_offsets[blockIdx.x] + thread_offset;

#pragma unroll
  for (int i = 0; i < k_items_per_thread; ++i) {
    const int idx = thread_idx + i;
    const float val = thread_inputs[i];
    if (idx < N && val > 0) {
      out[output_pos++] = val;
    }
  }
}

// A, out are device pointers
extern "C" void solve(const float* A, int N, float* out) {
  // 512 / 32 = 16
  constexpr int block_dim = 512;
  constexpr int k_items_per_thread = 12;

  // 100,000,000 / (512 * 12) = 16276
  const int grid_dim = CEIL(N, block_dim * k_items_per_thread);

  stream_compaction_block_part1_kernel<k_items_per_thread>
      <<<grid_dim, block_dim>>>(A, N);
  stream_compaction_grid_kernel<<<1, block_dim>>>(grid_dim);
  stream_compaction_block_part2_kernel<k_items_per_thread>
      <<<grid_dim, block_dim>>>(A, out, N);
}

基于 Decoupled Look-back Scan（DLBS）的方法

参考 Single-pass Parallel Prefix Scan with Decoupled Look-back

把 Scan 分成了几类：

方法	第一大阶段	第二大阶段	元素级 Global Memory Traffic	主要问题
Scan-then-propagate	Local Scan 并保存结果	加 Block Prefix	$\sim4n$	中间结果写一次、读一次
Reduce-then-scan	Local Reduce	重新读输入做 Seeded Scan	$\sim3n$	输入读两次
Chained-scan	Local Reduce/Scan	Prefix 在 Block 间链式传播	$\sim2n$	严格串行 Prefix Dependency
Decoupled Look-back	Local Reduce/Scan	向前查看 Aggregate/Prefix	$\sim2n$	少量冗余 Look-back 工作

上面的代码属于 Reduce-then-scan（第一个阶段换成 Reduce 也一样的，只需要每个 Block 的和），它的 Global Memory Traffic 是 $3n$：

Kernel 1 读取输入，一个 $n$
Kernel 3 读取输入，写入输出，两个 $n$

还有一种方法是 Scan-then-propagate，就是 Kernel 1 把中间结果写入 Global Memory，Kernel 3 就不需要再读了，这样是 $4n$

Chained-scan 是所有 Block 的结果间串行依赖，这样肯定慢，但是去掉了 Kernel 3 的读取操作，少了一个 $n$。Decoupled Look-back 则是优化了这个方法，每个 Block 在 Look-back 时一直向前遍历，每个 Block 有三个状态：未完成、计算内部和（Aggregate）完成、计算前缀和（Prefix）完成。如果遍历的 Block 前缀和计算完成了，直接加上然后 Break 即可，如果内部和计算完成那么加上继续向前遍历

代码：

#include <cuda_runtime.h>

#define CEIL(x, y) (((x) + (y) - 1) / (y))

#define FLOAT4(x) (*(reinterpret_cast<float4*>(&(x))))
#define CONST_FLOAT4(x) (*(reinterpret_cast<const float4*>(&(x))))

#define INT4(x) (*(reinterpret_cast<int4*>(&(x))))
#define CONST_INT4(x) (*(reinterpret_cast<const int4*>(&(x))))

__device__ __forceinline__ int warp_inclusive_scan(int val) {
  const int lane = threadIdx.x & 31;
#pragma unroll
  for (int offset = 1; offset < 32; offset <<= 1) {
    const int other = __shfl_up_sync(0xffffffff, val, offset);
    if (lane >= offset) {
      val += other;
    }
  }
  return val;
}

__device__ __forceinline__ int warp_exclusive_scan(int val) {
  return warp_inclusive_scan(val) - val;
}

__device__ __forceinline__ int block_inclusive_scan(int val) {
  __shared__ int warp_offsets[32];

  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int num_warps = blockDim.x >> 5;

  const int prefix = warp_inclusive_scan(val);

  if (lane == 31) {
    warp_offsets[warp] = prefix;
  }

  __syncthreads();

  if (warp == 0) {
    const int warp_sum = lane < num_warps ? warp_offsets[lane] : 0;
    const int warp_offset = warp_exclusive_scan(warp_sum);
    if (lane < num_warps) {
      warp_offsets[lane] = warp_offset;
    }
  }

  __syncthreads();

  return prefix + warp_offsets[warp];
}

__device__ __forceinline__ int block_exclusive_scan(int val) {
  return block_inclusive_scan(val) - val;
}

struct TileState {
  static constexpr unsigned long long VALUE_MASK = 0xffffffffull;
  static constexpr unsigned long long AGGREGATE_READY = 1ull << 32;
  static constexpr unsigned long long PREFIX_READY = 2ull << 32;

  unsigned long long data;

  // Publish:
  //   status = A
  //   value  = aggregate
  __device__ __forceinline__ void set_aggregate(int aggregate) {
    atomicExch(&data, AGGREGATE_READY | static_cast<unsigned int>(aggregate));
  }

  // Publish:
  //   status = P
  //   value  = inclusive prefix
  __device__ __forceinline__ void set_prefix(int prefix) {
    atomicExch(&data, PREFIX_READY | static_cast<unsigned int>(prefix));
  }

  // Atomic load.
  __device__ __forceinline__ unsigned long long load() {
    return atomicAdd(&data, 0ull);
  }

  __device__ __forceinline__ static bool is_ready(unsigned long long state) {
    return state != 0;
  }

  __device__ __forceinline__ static bool is_prefix_ready(
      unsigned long long state) {
    return (state & ~VALUE_MASK) == PREFIX_READY;
  }

  __device__ __forceinline__ static bool is_aggregate_ready(
      unsigned long long state) {
    return (state & ~VALUE_MASK) == AGGREGATE_READY;
  }

  __device__ __forceinline__ static int get_value(unsigned long long state) {
    return static_cast<int>(state & VALUE_MASK);
  }
};

__device__ __forceinline__ int decoupled_lookback(
    int tile_idx, int block_aggregate, TileState* __restrict__ tile_states) {
  // Tile 0: Publish prefix directly
  if (tile_idx == 0) {
    tile_states[0].set_prefix(block_aggregate);
    return 0;
  }

  // Step 1: Publish our aggregate
  tile_states[tile_idx].set_aggregate(block_aggregate);

  // Step 2: Look backward
  int exclusive_prefix = 0;

  for (int predecessor = tile_idx - 1;; --predecessor) {
    unsigned long long state;

    // Wait until predecessor reaches A or P
    do {
      state = tile_states[predecessor].load();
    } while (!TileState::is_ready(state));

    const int value = TileState::get_value(state);

    if (TileState::is_prefix_ready(state)) {
      // P: value = prefix[0 ... predecessor]
      exclusive_prefix += value;
      break;
    } else {
      // A: value = aggregate[predecessor]
      exclusive_prefix += value;
    }
  }

  // Step 3: Publish our complete inclusive prefix
  tile_states[tile_idx].set_prefix(exclusive_prefix + block_aggregate);
  return exclusive_prefix;
}

template <const int k_items_per_thread>
__global__ void stream_compaction_kernel(const float* __restrict__ A,
                                         float* __restrict__ out, int N,
                                         TileState* __restrict__ tile_states,
                                         int* __restrict__ g_tile_counter) {
  __shared__ int tile_idx;
  __shared__ int block_aggregate;
  __shared__ int global_block_offset;

  if (threadIdx.x == 0) {
    tile_idx = atomicAdd(g_tile_counter, 1);
  }

  __syncthreads();

  const int thread_idx =
      (tile_idx * blockDim.x + threadIdx.x) * k_items_per_thread;

  float thread_inputs[k_items_per_thread];
  int thread_predicates[k_items_per_thread];

  int thread_sum = 0;

  if (thread_idx + k_items_per_thread <= N) {
    static_assert(k_items_per_thread % 4 == 0,
                  "k_items_per_thread must be a multiple of 4");
#pragma unroll
    for (int i = 0; i < k_items_per_thread; i += 4) {
      const float4 val = CONST_FLOAT4(A[thread_idx + i]);
      const int4 predict = {val.x > 0, val.y > 0, val.z > 0, val.w > 0};
      thread_sum += (val.x > 0) + (val.y > 0) + (val.z > 0) + (val.w > 0);
      INT4(thread_predicates[i]) = predict;
      FLOAT4(thread_inputs[i]) = val;
    }
  } else {
#pragma unroll
    for (int i = 0; i < k_items_per_thread; ++i) {
      const int idx = thread_idx + i;
      const float val = idx < N ? A[idx] : -0.0f;
      const int pred = val > 0;
      thread_sum += pred;
      thread_predicates[i] = pred;
      thread_inputs[i] = val;
    }
  }

  const int local_thread_prefix = block_inclusive_scan(thread_sum);
  const int local_thread_offset = local_thread_prefix - thread_sum;

  if (threadIdx.x == blockDim.x - 1) {
    block_aggregate = local_thread_prefix;
  }

  __syncthreads();

  if (threadIdx.x == 0) {
    global_block_offset =
        decoupled_lookback(tile_idx, block_aggregate, tile_states);
  }

  __syncthreads();

  int output_pos = global_block_offset + local_thread_offset;
#pragma unroll
  for (int i = 0; i < k_items_per_thread; ++i) {
    const int idx = thread_idx + i;
    if (idx < N && thread_predicates[i]) {
      out[output_pos++] = thread_inputs[i];
    }
  }
}

// A, out are device pointers
extern "C" void solve(const float* A, int N, float* out) {
  constexpr int block_dim = 512;
  constexpr int k_items_per_thread = 16;

  const int num_tiles = CEIL(N, block_dim * k_items_per_thread);
  const int grid_dim = num_tiles;

  TileState* d_tile_states;
  int* d_tile_counter;

  cudaMalloc(&d_tile_states, num_tiles * sizeof(TileState));
  cudaMalloc(&d_tile_counter, sizeof(int));

  cudaMemset(d_tile_states, 0, num_tiles * sizeof(TileState));
  cudaMemset(d_tile_counter, 0, sizeof(int));

  stream_compaction_kernel<k_items_per_thread>
      <<<grid_dim, block_dim>>>(A, out, N, d_tile_states, d_tile_counter);

  cudaFree(d_tile_states);
  cudaFree(d_tile_counter);
}

这个代码还可以继续优化，现在是一个 Block 向前 Look-back 的时候一个一个检查，还可以一次检查 32 个，即用第一个 Warp 的所有 Lane 向前 Look-back

代码：

#include <cuda_runtime.h>

#define CEIL(x, y) (((x) + (y) - 1) / (y))

#define FLOAT4(x) (*(reinterpret_cast<float4*>(&(x))))
#define CONST_FLOAT4(x) (*(reinterpret_cast<const float4*>(&(x))))

#define INT4(x) (*(reinterpret_cast<int4*>(&(x))))
#define CONST_INT4(x) (*(reinterpret_cast<const int4*>(&(x))))

__device__ __forceinline__ int warp_reduce(int val, int last_lane) {
  const int lane = threadIdx.x & 31;
#pragma unroll
  for (int offset = 1; offset < 32; offset <<= 1) {
    const int other = __shfl_down_sync(0xffffffff, val, offset);
    if (lane + offset <= last_lane) {
      val += other;
    }
  }
  return val;
}

__device__ __forceinline__ int warp_inclusive_scan(int val) {
  const int lane = threadIdx.x & 31;
#pragma unroll
  for (int offset = 1; offset < 32; offset <<= 1) {
    const int other = __shfl_up_sync(0xffffffff, val, offset);
    if (lane >= offset) {
      val += other;
    }
  }
  return val;
}

__device__ __forceinline__ int warp_exclusive_scan(int val) {
  return warp_inclusive_scan(val) - val;
}

__device__ __forceinline__ int block_inclusive_scan(int val) {
  __shared__ int warp_offsets[32];

  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int num_warps = blockDim.x >> 5;

  const int prefix = warp_inclusive_scan(val);

  if (lane == 31) {
    warp_offsets[warp] = prefix;
  }

  __syncthreads();

  if (warp == 0) {
    const int warp_sum = lane < num_warps ? warp_offsets[lane] : 0;
    const int warp_offset = warp_exclusive_scan(warp_sum);
    if (lane < num_warps) {
      warp_offsets[lane] = warp_offset;
    }
  }

  __syncthreads();

  return prefix + warp_offsets[warp];
}

__device__ __forceinline__ int block_exclusive_scan(int val) {
  return block_inclusive_scan(val) - val;
}

struct TileState {
  static constexpr unsigned long long VALUE_MASK = 0xffffffffull;
  static constexpr unsigned long long AGGREGATE_READY = 1ull << 32;
  static constexpr unsigned long long PREFIX_READY = 2ull << 32;

  unsigned long long data;

  // Publish:
  //   status = A
  //   value  = aggregate
  __device__ __forceinline__ void set_aggregate(int aggregate) {
    atomicExch(&data, AGGREGATE_READY | static_cast<unsigned int>(aggregate));
  }

  // Publish:
  //   status = P
  //   value  = inclusive prefix
  __device__ __forceinline__ void set_prefix(int prefix) {
    atomicExch(&data, PREFIX_READY | static_cast<unsigned int>(prefix));
  }

  // Atomic load.
  __device__ __forceinline__ unsigned long long load() {
    return atomicAdd(&data, 0ull);
  }

  __device__ __forceinline__ static bool is_ready(unsigned long long state) {
    return state != 0;
  }

  __device__ __forceinline__ static bool is_prefix_ready(
      unsigned long long state) {
    return (state & ~VALUE_MASK) == PREFIX_READY;
  }

  __device__ __forceinline__ static bool is_aggregate_ready(
      unsigned long long state) {
    return (state & ~VALUE_MASK) == AGGREGATE_READY;
  }

  __device__ __forceinline__ static int get_value(unsigned long long state) {
    return static_cast<int>(state & VALUE_MASK);
  }
};

__device__ __forceinline__ int decoupled_lookback_warp(
    int tile_idx, int block_aggregate, TileState* __restrict__ tile_states) {
  const int lane = threadIdx.x & 31;

  // Tile 0: Publish prefix directly
  if (tile_idx == 0) {
    if (lane == 0) {
      tile_states[0].set_prefix(block_aggregate);
    }
    return 0;
  }

  // Step 1: Publish our aggregate
  if (lane == 0) {
    tile_states[tile_idx].set_aggregate(block_aggregate);
  }

  // Step 2: Look backward
  int exclusive_prefix = 0;
  int predecessor = tile_idx - lane - 1;

  __nanosleep(450);

  while (true) {
    unsigned long long state;

    // Wait until all 32 predecessors reaches A or P

    // int __all_sync(unsigned mask, int predicate) 当且仅当所有线程的 predicate
    // 非零时返回 1，否则返回 0
    do {
      // 不要 busy 等待
      __nanosleep(350);
      if (predecessor >= 0) {
        state = tile_states[predecessor].load();
      } else {
        // [P | 0]
        state = TileState::PREFIX_READY;
      }
    } while (!__all_sync(0xffffffff, TileState::is_ready(state)));

    // 当前每个 lane 负责的 predecessors 为 [p, p-1, p-2, p-3, ...,
    // p-31]，只需要找到最低位的为 P 状态的 lane，然后求前缀即可

    // unsigned __ballot_sync(unsigned mask, int predicate) 线程 predicate
    // 为零的该函数返回值该位为 0，线程 predicate 非零的该函数返回值该位为 1
    const unsigned prefix_lanes =
        __ballot_sync(0xffffffff, TileState::is_prefix_ready(state));

    // __ffs(int active) 返回 active 最低有效位的位置。如果 active 为 0，__ffs()
    // 将返回 0
    const int last_lane = prefix_lanes == 0 ? 31 : (__ffs(prefix_lanes) - 1);

    const int value = TileState::get_value(state);

    // 只有 lane=0 的才有用
    int window_sum = warp_reduce(value, last_lane);
    exclusive_prefix += window_sum;

    if (prefix_lanes != 0) {
      break;
    }

    predecessor -= 32;
  }

  // Step 3: Publish our complete inclusive prefix
  if (lane == 0) {
    tile_states[tile_idx].set_prefix(exclusive_prefix + block_aggregate);
  }
  return exclusive_prefix;
}

template <const int k_items_per_thread, const int k_block_dim>
__global__ void stream_compaction_kernel(const float* __restrict__ A,
                                         float* __restrict__ out, int N,
                                         TileState* __restrict__ tile_states,
                                         int* __restrict__ g_tile_counter) {
  __shared__ int tile_idx;
  __shared__ int block_aggregate;
  __shared__ int global_block_offset;

  __shared__ float dense_output[k_items_per_thread * k_block_dim];

  if (threadIdx.x == 0) {
    tile_idx = atomicAdd(g_tile_counter, 1);
  }

  __syncthreads();

  const int thread_idx =
      (tile_idx * k_block_dim + threadIdx.x) * k_items_per_thread;

  float thread_inputs[k_items_per_thread];

  int thread_sum = 0;

#pragma unroll
  for (int i = 0; i < k_items_per_thread; ++i) {
    const int idx = thread_idx + i;
    const float val = idx < N ? A[idx] : -0.0f;
    thread_sum += (val > 0);
    thread_inputs[i] = val;
  }

  const int local_thread_prefix = block_inclusive_scan(thread_sum);
  const int local_thread_offset = local_thread_prefix - thread_sum;

  if (threadIdx.x == k_block_dim - 1) {
    block_aggregate = local_thread_prefix;
  }

  __syncthreads();

  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;

  if (warp == 0) {
    const int prefix =
        decoupled_lookback_warp(tile_idx, block_aggregate, tile_states);

    if (lane == 0) {
      global_block_offset = prefix;
    }
  }

  __syncthreads();

  int local_output_pos = local_thread_offset;

#pragma unroll
  for (int i = 0; i < k_items_per_thread; ++i) {
    const int idx = thread_idx + i;
    if (idx < N && thread_inputs[i] > 0) {
      dense_output[local_output_pos++] = thread_inputs[i];
    }
  }

  __syncthreads();

  for (int i = threadIdx.x; i < block_aggregate; i += k_block_dim) {
    out[global_block_offset + i] = dense_output[i];
  }
}

// A, out are device pointers
extern "C" void solve(const float* A, int N, float* out) {
  constexpr int block_dim = 512;
  constexpr int k_items_per_thread = 10;

  const int num_tiles = CEIL(N, block_dim * k_items_per_thread);
  const int grid_dim = num_tiles;

  TileState* d_tile_states;
  int* d_tile_counter;

  cudaMalloc(&d_tile_states, num_tiles * sizeof(TileState));
  cudaMalloc(&d_tile_counter, sizeof(int));

  cudaMemset(d_tile_states, 0, num_tiles * sizeof(TileState));
  cudaMemset(d_tile_counter, 0, sizeof(int));

  stream_compaction_kernel<k_items_per_thread, block_dim>
      <<<grid_dim, block_dim>>>(A, out, N, d_tile_states, d_tile_counter);

  cudaFree(d_tile_states);
  cudaFree(d_tile_counter);
}

这种方法并不稳定，有时候快有时候慢

Radix Sort

输入一个 32 位无符号整数数组，实现从小到大的基数排序并输出

例子：

Input:  [170, 45, 75, 90, 2, 802, 24, 66]
Output: [2, 24, 45, 66, 75, 90, 170, 802]

基数排序将待排序的元素拆分为 $k$ 个关键字，逐一对各个关键字排序后完成对所有元素的排序

分为从第一关键字到最后一个关键字（Most Significant Digit first，MSD）和最后一个关键字到第一个关键字（Least Significant Digit first，LSD）两种遍历顺序，对每一层内部排序的时候，一般使用一种叫做计数排序的稳定排序（并不改变不关心元素的顺序，LSD 的正确性来源于此）来完成

MSD 通常只能使用递归的形式完成，时间常数比较大，且不适用于 GPU，所以一般使用 LSD：

一个 LSD 基数排序全流程的例子

CPU 代码长这样：

#include <algorithm>
#include <iostream>
#include <utility>

void radix_sort(int n, int a[]) {
  int *b = new int[n];  // 临时空间
  int *cnt = new int[1 << 8];
  int mask = (1 << 8) - 1;
  int *x = a, *y = b;
  for (int i = 0; i < 32; i += 8) {
    for (int j = 0; j != (1 << 8); ++j) cnt[j] = 0;
    for (int j = 0; j != n; ++j) ++cnt[x[j] >> i & mask];
    for (int sum = 0, j = 0; j != (1 << 8); ++j) {
      // 等价于 std::exclusive_scan(cnt, cnt + (1 << 8), cnt, 0);
      sum += cnt[j], cnt[j] = sum - cnt[j];
    }
    for (int j = 0; j != n; ++j) y[cnt[x[j] >> i & mask]++] = x[j];
    std::swap(x, y);
  }
  delete[] cnt;
  delete[] b;
}

int main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(nullptr);
  int n;
  std::cin >> n;
  int *a = new int[n];
  for (int i = 0; i < n; ++i) std::cin >> a[i];
  radix_sort(n, a);
  for (int i = 0; i < n; ++i) std::cout << a[i] << ' ';
  delete[] a;
  return 0;
}

可以发现，它的 $\log$ 次循环内部包含这三个操作：

Predicate：判断每个数的当前二进制位置是 0 还是 1，得到 predicate[]（0 的话 predicate[i]=1，否则设置为 0）
Scan：对 predicate[] 计算 Exclusive 前缀和，得到元素的新位置
Scatter：计算出来的新位置，把元素分散写回输出数组

最简单的代码

整体代码长这样：

#include <cuda_runtime.h>

#define CEIL(x, y) (((x) + (y) - 1) / (y))

__device__ __forceinline__ int warp_reduce(int val) {
#pragma unroll
  for (int offset = 16; offset > 0; offset >>= 1) {
    val += __shfl_down_sync(0xffffffff, val, offset);
  }
  return val;
}

__device__ __forceinline__ int warp_inclusive_scan(int val) {
  const int lane = threadIdx.x & 31;
#pragma unroll
  for (int offset = 1; offset < 32; offset <<= 1) {
    const int other = __shfl_up_sync(0xffffffff, val, offset);
    if (lane >= offset) {
      val += other;
    }
  }
  return val;
}

__device__ __forceinline__ int warp_exclusive_scan(int val) {
  return warp_inclusive_scan(val) - val;
}

__device__ __forceinline__ int block_inclusive_scan(int val) {
  __shared__ int warp_offsets[32];

  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int num_warps = blockDim.x >> 5;

  const int prefix = warp_inclusive_scan(val);

  if (lane == 31) {
    warp_offsets[warp] = prefix;
  }

  __syncthreads();

  if (warp == 0) {
    const int warp_sum = lane < num_warps ? warp_offsets[lane] : 0;
    const int warp_offset = warp_exclusive_scan(warp_sum);
    if (lane < num_warps) {
      warp_offsets[lane] = warp_offset;
    }
  }

  __syncthreads();

  return prefix + warp_offsets[warp];
}

__device__ __forceinline__ int block_exclusive_scan(int val) {
  return block_inclusive_scan(val) - val;
}

__device__ __forceinline__ int block_reduce(int val) {
  __shared__ int warp_sums[32];

  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int num_warps = blockDim.x >> 5;

  val = warp_reduce(val);

  if (lane == 0) {
    warp_sums[warp] = val;
  }

  __syncthreads();

  if (warp == 0) {
    val = lane < num_warps ? warp_sums[lane] : 0;
    val = warp_reduce(val);
  }

  return val;
}

struct BlockState {
  int aggregate;
  int offset;
};

template <int k_items_per_thread, int k_block_dim>
__global__ void counting_sort_k1(const unsigned int* __restrict__ input,
                                 BlockState* __restrict__ block_states, int N,
                                 int bit) {
  const int block_idx = blockIdx.x * k_block_dim * k_items_per_thread;

  int sum = 0;
  for (int i = threadIdx.x; i < k_block_dim * k_items_per_thread;
       i += k_block_dim) {
    const int idx = block_idx + i;
    if (idx < N) {
      sum += ((input[idx] & (1u << bit)) == 0);
    }
  }

  sum = block_reduce(sum);

  if (threadIdx.x == 0) {
    block_states[blockIdx.x].aggregate = sum;
  }
}

__global__ void counting_sort_k2(BlockState* __restrict__ block_states,
                                 int* __restrict__ grid_sum, int M) {
  const int items_per_thread = CEIL(M, blockDim.x);
  const int thread_idx = threadIdx.x * items_per_thread;

  int thread_offsets[128];
  int thread_sum = 0;

  for (int i = 0; i < items_per_thread; ++i) {
    const int idx = thread_idx + i;
    const int val = idx < M ? block_states[idx].aggregate : 0;
    thread_offsets[i] = thread_sum;
    thread_sum += val;
  }

  const int offset = block_exclusive_scan(thread_sum);

  for (int i = 0; i < items_per_thread; ++i) {
    const int idx = thread_idx + i;
    if (idx < M) {
      block_states[idx].offset = thread_offsets[i] + offset;
      if (idx == M - 1) {
        *grid_sum = block_states[idx].offset + block_states[idx].aggregate;
      }
    }
  }
}

template <int k_items_per_thread, int k_block_dim>
__global__ void counting_sort_k3(const unsigned int* __restrict__ input,
                                 BlockState* __restrict__ block_states,
                                 const int* __restrict__ grid_sum,
                                 unsigned int* __restrict__ output, int N,
                                 int bit) {
  unsigned int thread_inputs[k_items_per_thread];
  __shared__ unsigned int dense_output[2][k_items_per_thread * k_block_dim];

  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
  const int thread_idx = tid * k_items_per_thread;

  int thread_sum_0 = 0;

#pragma unroll
  for (int i = 0; i < k_items_per_thread; ++i) {
    const int idx = thread_idx + i;
    const unsigned int val = idx < N ? input[idx] : 0;
    thread_inputs[i] = val;
    if (idx < N) {
      thread_sum_0 += ((input[idx] & (1u << bit)) == 0);
    }
  }

  const int thread_sum_1 =
      min(k_items_per_thread, N - thread_idx) - thread_sum_0;

  const int thread_offset_0 = block_exclusive_scan(thread_sum_0);
  const int thread_offset_1 = block_exclusive_scan(thread_sum_1);

  const int block_offset_0 = block_states[blockIdx.x].offset;
  const int block_offset_1 =
      blockIdx.x * k_block_dim * k_items_per_thread - block_offset_0;

  const int block_aggregate_0 = block_states[blockIdx.x].aggregate;
  const int block_aggregate_1 =
      min(k_block_dim * k_items_per_thread,
          N - blockIdx.x * k_block_dim * k_items_per_thread)
      - block_aggregate_0;

  int local_pos_0 = thread_offset_0;
  int local_pos_1 = thread_offset_1;

#pragma unroll
  for (int i = 0; i < k_items_per_thread; ++i) {
    const int idx = thread_idx + i;
    if (idx < N) {
      const unsigned int x = thread_inputs[i];
      if (thread_inputs[i] & (1u << bit)) {
        dense_output[1][local_pos_1++] = x;
      } else {
        dense_output[0][local_pos_0++] = x;
      }
    }
  }

  __syncthreads();

  for (int i = threadIdx.x; i < block_aggregate_0; i += k_block_dim) {
    output[block_offset_0 + i] = dense_output[0][i];
  }

  const int total_zeros = *grid_sum;

  for (int i = threadIdx.x; i < block_aggregate_1; i += k_block_dim) {
    output[total_zeros + block_offset_1 + i] = dense_output[1][i];
  }
}

// input, output are device pointers
extern "C" void solve(const unsigned int* input, unsigned int* output, int N) {
  constexpr int block_dim = 256;
  constexpr int k_items_per_thread = 12;

  const int grid_dim = CEIL(N, block_dim * k_items_per_thread);

  BlockState* d_block_states;
  int* d_grid_sum;

  cudaMalloc(&d_block_states, grid_dim * sizeof(BlockState));
  cudaMalloc(&d_grid_sum, sizeof(int));

  unsigned int* temp;
  cudaMalloc(&temp, N * sizeof(unsigned int));

  const unsigned int* src = input;
  unsigned int* dst = temp;

  for (int bit = 0; bit < 32; ++bit) {
    counting_sort_k1<k_items_per_thread, block_dim>
        <<<grid_dim, block_dim>>>(src, d_block_states, N, bit);
    counting_sort_k2<<<1, block_dim>>>(d_block_states, d_grid_sum, grid_dim);
    counting_sort_k3<k_items_per_thread, block_dim>
        <<<grid_dim, block_dim>>>(src, d_block_states, d_grid_sum, dst, N, bit);

    if (dst == temp) {
      src = temp;
      dst = output;
    } else {
      src = output;
      dst = temp;
    }
  }

  cudaFree(temp);

  cudaFree(d_block_states);
  cudaFree(d_grid_sum);
}

一次处理多位

由于每一个 Pass 都需要读取所有的元素，一次处理多位可以减少 Global Memory Traffic，用三个 Kernel：

Kernel 1：Histogram，统计每个 Pass 的每个状态的元素个数
Kernel 2：把 Kernel 1 得到的结果经过 Exclusive Scan 转化为每个 Pass 的每个状态的所有元素在输出数组中的起点。这里只用 Block Scan 就可以了
Kernel 3：计算新位置并且 Scatter，目标位置 = 桶的全局起点 + 前面 Tile 的同桶元素数 + 本 Tile 内同桶的稳定排名。这里需要用到全局的 Scan 方法（在 CUB 库里面就用的 DLBS，当然也可以用其他的方法，用其他的方法的话可以把代码分散到前两个 Kernel 中）

Kernel 1 和 Kernel 2 一次处理全部的 Pass

代码：

#include <cuda_runtime.h>

#include <algorithm>

#define CEIL(x, y) (((x) + (y) - 1) / (y))

constexpr int k_bits = 32;

constexpr int k_radix_bits = 8;
constexpr int k_radix_digits = 1 << k_radix_bits;
constexpr int k_radix_mask = k_radix_digits - 1;
constexpr int k_radix_passes = k_bits / k_radix_bits;

constexpr int k_items_per_thread = 12;
constexpr int k_threads_per_block = 256;
constexpr int k_items_per_block = k_threads_per_block * k_items_per_thread;

constexpr int k_hist_parts = 4;
constexpr int k_bin_count = k_radix_passes * k_radix_digits;

__global__ void histogram_kernel(const unsigned int* __restrict__ input, int n,
                                 unsigned int* __restrict__ d_bins) {
  // hist[passes][digits][parts]
  // 再分出来 parts 是为了减少 atomic 冲突
  __shared__ unsigned int hist[k_bin_count * k_hist_parts];

  const int part = (threadIdx.x & 31) % k_hist_parts;

  for (int i = threadIdx.x; i < k_bin_count * k_hist_parts;
       i += k_threads_per_block) {
    hist[i] = 0;
  }
  __syncthreads();

  for (int base = blockIdx.x * k_items_per_block; base < n;
       base += gridDim.x * k_items_per_block) {
    unsigned int values[k_items_per_thread];
    bool valid[k_items_per_thread];

#pragma unroll
    for (int item = 0; item < k_items_per_thread; ++item) {
      const int idx = base + item * k_threads_per_block + threadIdx.x;

      valid[item] = idx < n;
      values[item] = valid[item] ? input[idx] : 0u;
    }

#pragma unroll
    for (int pass = 0; pass < k_radix_passes; ++pass) {
#pragma unroll
      for (int item = 0; item < k_items_per_thread; ++item) {
        if (valid[item]) {
          const unsigned int digit =
              (values[item] >> (pass * k_radix_bits)) & k_radix_mask;

          const int slot =
              (pass * k_radix_digits + digit) * k_hist_parts + part;

          atomicAdd(&hist[slot], 1u);
        }
      }
    }
  }

  __syncthreads();

  for (int bin = threadIdx.x; bin < k_bin_count; bin += k_threads_per_block) {
    unsigned int count = 0;

#pragma unroll
    for (int part_id = 0; part_id < k_hist_parts; ++part_id) {
      count += hist[bin * k_hist_parts + part_id];
    }

    if (count != 0) {
      atomicAdd(&d_bins[bin], count);
    }
  }
}

__device__ __forceinline__ unsigned int warp_inclusive_scan(unsigned int val) {
  const int lane = threadIdx.x & 31;
#pragma unroll
  for (int offset = 1; offset < 32; offset <<= 1) {
    const unsigned int other = __shfl_up_sync(0xffffffffu, val, offset);
    if (lane >= offset) {
      val += other;
    }
  }
  return val;
}

__device__ __forceinline__ unsigned int warp_exclusive_scan(unsigned int val) {
  return warp_inclusive_scan(val) - val;
}

template <int k_block_dim>
__device__ __forceinline__ int block_inclusive_scan(int val) {
  const int warp = threadIdx.x >> 5;
  const int lane = threadIdx.x & 31;
  const int k_num_warps = k_block_dim >> 5;

  __shared__ int warp_offsets[k_num_warps];

  const int prefix = warp_inclusive_scan(val);

  if (lane == 31) {
    warp_offsets[warp] = prefix;
  }

  __syncthreads();

  if (warp == 0) {
    const int warp_sum = lane < k_num_warps ? warp_offsets[lane] : 0;
    const int warp_offset = warp_exclusive_scan(warp_sum);
    if (lane < k_num_warps) {
      warp_offsets[lane] = warp_offset;
    }
  }

  __syncthreads();

  return prefix + warp_offsets[warp];
}

template <int k_block_dim>
__device__ __forceinline__ unsigned int block_exclusive_scan(unsigned int val) {
  return block_inclusive_scan<k_block_dim>(val) - val;
}

__global__ void bins_scan_kernel(unsigned int* __restrict__ d_bins) {
  static_assert(
      k_radix_digits >= 32 && k_radix_digits <= 1024
          && k_radix_digits % 32 == 0,
      "Requires 32..1024 bins, complete warps, and one thread per bin.");

  const int idx = blockIdx.x * k_radix_digits + threadIdx.x;
  d_bins[idx] = block_exclusive_scan<k_radix_digits>(d_bins[idx]);
}

struct TileState {
  static constexpr unsigned int k_value_mask = (1u << 30) - 1u; // 0x3fffffff
  static constexpr unsigned int k_status_mask = 3u << 30;       // 0xc0000000
  static constexpr unsigned int k_aggregate_ready = 1u << 30;   // 0x40000000
  static constexpr unsigned int k_prefix_ready = 2u << 30;      // 0x80000000

  unsigned int data;

  // Publish:
  //   status = A
  //   value  = aggregate
  __device__ __forceinline__ void set_aggregate(unsigned int aggregate) {
    atomicExch(&data, k_aggregate_ready | aggregate);
  }

  // Publish:
  //   status = P
  //   value  = inclusive prefix
  __device__ __forceinline__ void set_prefix(unsigned int prefix) {
    atomicExch(&data, k_prefix_ready | prefix);
  }

  // Atomic load
  __device__ __forceinline__ unsigned int load() {
    return atomicAdd(&data, 0u);
  }

  __device__ __forceinline__ static bool is_ready(unsigned int state) {
    return (state & k_status_mask) != 0u;
  }

  __device__ __forceinline__ static bool is_prefix_ready(unsigned int state) {
    return (state & k_status_mask) == k_prefix_ready;
  }

  __device__ __forceinline__ static bool is_aggregate_ready(
      unsigned int state) {
    return (state & k_status_mask) == k_aggregate_ready;
  }

  __device__ __forceinline__ static unsigned int get_value(unsigned int state) {
    return state & k_value_mask;
  }
};

__global__ void onsweep_pass_kernel(const unsigned int* __restrict__ src,
                                    unsigned int* __restrict__ dst, int n,
                                    unsigned int bit,
                                    const unsigned int* __restrict__ d_bins,
                                    TileState* __restrict__ d_lookback,
                                    unsigned int* __restrict__ d_ctrs) {
  static_assert(k_threads_per_block == k_radix_digits,
                "Requires k_radix_digits threads per block");

  __shared__ unsigned int s_tile_id;
  __shared__ unsigned int s_values[k_items_per_block];
  __shared__ unsigned int s_counts[k_radix_digits];
  __shared__ unsigned int s_total_zeros;
  __shared__ unsigned int s_global_adjust[k_radix_digits];

  if (threadIdx.x == 0) {
    s_tile_id = atomicAdd(d_ctrs, 1u);
  }
  s_counts[threadIdx.x] = 0;

  __syncthreads();

  const int tile = s_tile_id;
  const int base = tile * k_items_per_block;

  const int remaining = n - base;
  const int valid_items =
      remaining < k_items_per_block ? remaining : k_items_per_block;

#pragma unroll
  for (int item = 0; item < k_items_per_thread; ++item) {
    const int i = item * k_threads_per_block + threadIdx.x;
    const unsigned int value = i < valid_items ? src[base + i] : 0xffffffffu;

    s_values[i] = value;

    if (i < valid_items) {
      const unsigned int digit = (value >> bit) & k_radix_mask;
      atomicAdd(&s_counts[digit], 1u);
    }
  }

  __syncthreads();

  TileState& own = d_lookback[tile * k_radix_digits + threadIdx.x];

  const unsigned int local_count = s_counts[threadIdx.x];
  own.set_aggregate(local_count);

  __syncthreads();

  const unsigned int local_start =
      block_exclusive_scan<k_threads_per_block>(local_count);
  unsigned int values[k_items_per_thread];

#pragma unroll
  for (int item = 0; item < k_items_per_thread; ++item) {
    values[item] = s_values[threadIdx.x * k_items_per_thread + item];
  }

  __syncthreads();

#pragma unroll
  for (int b = 0; b < k_radix_bits; ++b) {
    unsigned int zero_count = 0;

#pragma unroll
    for (int item = 0; item < k_items_per_thread; ++item) {
      zero_count += ((values[item] >> (bit + b)) & 1u) == 0;
    }

    const unsigned int zero_start =
        block_exclusive_scan<k_threads_per_block>(zero_count);

    if (threadIdx.x == k_threads_per_block - 1) {
      s_total_zeros = zero_start + zero_count;
    }

    __syncthreads();

    unsigned int seen_zeros = 0;

#pragma unroll
    for (int item = 0; item < k_items_per_thread; ++item) {
      const unsigned int position = threadIdx.x * k_items_per_thread + item;
      const unsigned int zeros_before = zero_start + seen_zeros;
      const bool is_zero = ((values[item] >> (bit + b)) & 1u) == 0;

      const unsigned int target =
          is_zero ? zeros_before : s_total_zeros + position - zeros_before;

      s_values[target] = values[item];
      seen_zeros += is_zero;
    }

    __syncthreads();

#pragma unroll
    for (int item = 0; item < k_items_per_thread; ++item) {
      values[item] = s_values[threadIdx.x * k_items_per_thread + item];
    }

    __syncthreads();
  }

  unsigned int preceding = 0;

  for (int previous = tile - 1; previous >= 0; --previous) {
    TileState& prior = d_lookback[previous * k_radix_digits + threadIdx.x];
    unsigned int state;

    do {
      state = prior.load();
    } while (!TileState::is_ready(state));

    preceding += TileState::get_value(state);

    if (TileState::is_prefix_ready(state)) {
      break;
    }
  }

  own.set_prefix(preceding + local_count);
  s_global_adjust[threadIdx.x] = d_bins[threadIdx.x] + preceding - local_start;

  __syncthreads();

#pragma unroll
  for (int item = 0; item < k_items_per_thread; ++item) {
    const int i = item * k_threads_per_block + threadIdx.x;

    if (i < valid_items) {
      const unsigned int value = s_values[i];
      const unsigned int digit = (value >> bit) & k_radix_mask;

      const int target = i + s_global_adjust[digit];

      dst[target] = value;
    }
  }
}

extern "C" void solve(const unsigned int* input, unsigned int* output, int N) {
  const int num_blocks = CEIL(N, k_items_per_block);

  int device = 0;
  int sm_count = 0;
  int active_blocks_per_sm = 0;

  cudaGetDevice(&device);
  cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device);
  cudaOccupancyMaxActiveBlocksPerMultiprocessor(
      &active_blocks_per_sm, histogram_kernel, k_threads_per_block, 0);

  int histogram_blocks = std::min(sm_count * active_blocks_per_sm, num_blocks);

  // d_bins[passes][radix_digits]
  // d_lookback[blocks][radix_digits]
  // d_tmp[N]
  // d_ctrs[passes]
  unsigned int *d_bins, *d_temp, *d_ctrs;
  TileState* d_lookback;

  const int bins_bytes = k_bin_count * sizeof(unsigned int);
  const int lookback_bytes = num_blocks * k_radix_digits * sizeof(unsigned int);
  const int ctrs_bytes = k_radix_passes * sizeof(unsigned int);

  cudaMalloc(&d_bins, bins_bytes);
  cudaMalloc(&d_lookback, num_blocks * k_radix_digits * sizeof(TileState));
  cudaMalloc(&d_temp, N * sizeof(unsigned int));
  cudaMalloc(&d_ctrs, ctrs_bytes);

  cudaMemsetAsync(d_bins, 0, bins_bytes);
  cudaMemsetAsync(d_ctrs, 0, ctrs_bytes);

  histogram_kernel<<<histogram_blocks, k_threads_per_block>>>(input, N, d_bins);
  bins_scan_kernel<<<k_radix_passes, k_radix_digits>>>(d_bins);

  const unsigned int* src = input;

#pragma unroll
  for (int pass = 0; pass < k_radix_passes; ++pass) {
    unsigned int* dst = (pass & 1) ? output : d_temp;
    cudaMemsetAsync(d_lookback, 0, lookback_bytes);
    onsweep_pass_kernel<<<num_blocks, k_threads_per_block>>>(
        src, dst, N, pass * k_radix_bits, d_bins + pass * k_radix_digits,
        d_lookback, d_ctrs + pass);
    src = dst;
  }

  cudaDeviceSynchronize();

  cudaFree(d_ctrs);
  cudaFree(d_temp);
  cudaFree(d_lookback);
  cudaFree(d_bins);
}