Radix Sort 效能優化:Memory Coalescing、Radix 大小與 Thread Coarsening

重點總覽 (Overview)

本筆記接續 13-Sorting/01-Sorting-Foundations-and-Parallel-Radix-Sort 的 one-thread-per-key 基礎核心,處理它的最大瓶頸:寫回 global memory 時無法 coalesce。三個優化彼此相關,目標都是讓「連續的 thread 寫到連續的記憶體位址」並降低 global scan 開銷。

為何 Global 寫入無法 Coalesce (The Coalescing Problem)

基礎核心(sibling note 01)讓每個 thread 算出 key 的 destination index 後直接 scatter 寫回 global output。問題在於:相鄰 thread 的 key 不一定落在同一個 bucket。

Block 0 input:   k0    k1    k2    k3
bit (LSB):        0     1     0     1
                  |     |     |     |
            +-----+     |     +--+  |
            | +---------+        |  +-----+
            v v                  v        v
Global out: | ... 0-bucket ... | ... 1-bucket ... |
                ^k0 ^k2            ^k1 ^k3
  t0->posA   t1->posX   t2->posA+1   t3->posX+1
  (一個 warp 內位址跳來跳去 => 多個 memory request => 未 coalesced)

• 相鄰 thread index 不對應相鄰記憶體位址 → 一個 warp 需發出多次 memory request。
• Ch.6 提供三種改善 coalescing 的方法:(1) 重排 threads、(2) 重排資料、(3) 把不可 coalesce 的存取放到 shared memory,再以 coalesced 方式在 shared/global 間搬運。本章採用 (3)。

Shared Memory 區域排序優化 (Optimizing for Memory Coalescing) — 13.4

做法:不再做 global sort;改為每個 block 在 shared memory 用同樣演算法做一次 block-level local sort(只需 block 內 local exclusive scan),把 0/1 桶在 shared memory 排成連續段,再整段 coalesced 寫回 global。

Block b shared memory (local sort 後,桶內連續):
   [ k k k | k k ]      <- 左段=local 0-bucket, 右段=local 1-bucket
     ^^^^^   ^^^
   寫 0-bucket: t0,t1,t2 -> 連續 global 位址 (coalesced)
   寫 1-bucket: t3,t4    -> 連續 global 位址 (coalesced)

核心挑戰:每個 block 要知道自己的 local bucket 在 global output 中的起始位置,而這取決於其他 block 的 local bucket 大小:

• block 的 local 0-bucket 位置 = 所有前面 block 的 local 0-bucket 之後。
• block 的 local 1-bucket 位置 = 所有 block 的 local 0-bucket + 前面 block 的 local 1-bucket 之後。

解法:對「各 block 的 local bucket 大小表」做一次 global exclusive scan。表以 row-major 排列(先所有 block 的 0-bucket size,再所有 block 的 1-bucket size):

              B0   B1   B2   B3
 row0 (#0):   s00  s01  s02  s03
 row1 (#1):   s10  s11  s12  s13
 linearized:  [s00 s01 s02 s03 | s10 s11 s12 s13]
                       |  exclusive scan
                       v
 result    :  各 block 各 local bucket 的「global 起始位移」

掃描結果即為各 local bucket 的 global 起始位移。寫回時每個 thread 依自己的 threadIdx 落在 local 0 或 1 桶,決定寫哪個 global 段。

基礎(未優化)核心回顧,優化版以此為「block 內」子程序(Fig. 13.4):

// 基礎 global 版 (sibling 01);優化版把它變成 block 內 local sort
__global__ void radix_sort_iter(unsigned int* in, unsigned int* out,
                                unsigned int* bits, unsigned int N,
                                unsigned int iter) {
    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;   // 03
    unsigned int key, bit;
    if (i < N) {                                              // 04 邊界檢查
        key = in[i];                                          // 06
        bit = (key >> iter) & 1;                              // 07 取出本回合的 bit
        bits[i] = bit;                                        // 08
    }
    exclusiveScan(bits, N);  // 10: 在 block 外 (全 thread 須活著做 barrier)
    if (i < N) {
        unsigned int onesBefore = bits[i];                    // 12 #ones before
        unsigned int onesTotal  = bits[N];                    // 13 #ones total
        unsigned int dst = (bit == 0)
            ? (i - onesBefore)                                // 14 destination of 0
            : (N - onesTotal + onesBefore);                   // 15 destination of 1
        out[dst] = key;                                       // 16 (未優化:scatter)
    }
}

優化版概念(本章把實作留作習題):

// 概念示意:block 在 shared memory 做 local sort,再 coalesced 寫回
__global__ void radix_sort_coalesced(unsigned int* in, unsigned int* out,
                                     unsigned int* blockBucketOffsets, /* 已 scan */
                                     unsigned int N, unsigned int iter) {
    __shared__ unsigned int s_keys[BLOCK_SIZE];
    // 1) 載入 + 取 bit + block 內 local exclusive scan -> 得 local destination
    // 2) 在 shared memory 內把 key 排成 [0-bucket | 1-bucket] 連續段
    // 3) 每 thread 依 threadIdx 落在 local 0 或 1 桶,從 blockBucketOffsets
    //    取得該桶 global 起始位移,連續 (coalesced) 寫回 out[]
}

Radix 大小的選擇 (Choice of Radix Value) — 13.5

用 r-bit radix:一次 iteration 處理 r 個 bits,分到 2^r 個 buckets。

• N-bit keys 所需 iteration 數:⌈N_bits / r⌉(1-bit 需 N 次;2-bit 4-bit keys 只需 2 次)。
• block 內 local sort:r-bit radix sort = 連續 r 次 1-bit local iteration,每次各需一次 local exclusive scan(皆 block 內,iteration 間無跨 block 協調)。
• global bucket-size 表:列數從 2 變成 2^r(2-bit 為 4 列)。

2-bit radix:一次分成 4 個 buckets
  bits: 00   01   10   11
        [..] [..] [..] [..]
local sort = 1-bit(LSB) 後再 1-bit(次低位),兩次 local scan

取捨表:小 radix vs 大 radix

Thread Coarsening 改善 Coalescing (Thread Coarsening) — 13.6

把 radix sort 切到很多 block 的代價是寫回 coalescing 差:block 越多 → 每 block key 越少 → local buckets 越小 → coalescing 機會越少。

做法:每個 thread 負責多個 key(而非一個)→ block 數減少、每 block key 增多。

無 coarsening:  t0  t1  t2  t3        (每 thread 1 key,local bucket 小)
有 coarsening:  t0  t0  t1  t1 ...    (每 thread 多 key,local bucket 大)
                local buckets 變大 => 相鄰 thread 更可能寫到相鄰位址 (coalesced)

兩個收益:

1. 改善 coalescing:local buckets 變大,連續 thread 更可能寫到連續位址(對比 Fig. 13.8 vs 13.10)。
2. 降低 global scan 開銷:scan table 大小 ∝ bucket 數 × block 數;coarsening 減少 block 數 → table 變小 → exclusive scan 開銷下降。

考試/面試重點 (Exam / Test Patterns)

Related Notes

• 13-Sorting/01-Sorting-Foundations-and-Parallel-Radix-Sort
• 13-Sorting/03-Parallel-Merge-Sort-and-Other-Methods
• 06-Performance-Considerations/01-Memory-Coalescing
• 06-Performance-Considerations/03-Thread-Coarsening
• 11-Prefix-Sum-Scan/03-Arbitrary-Length-and-Single-Pass-Scan
• 09-Parallel-Histogram/02-Histogram-Optimizations-Privatization-Coarsening-Aggregation