キャッシュ効率型基数ソートを使用する

キャッシュ効率型基数ソート (CRadix sort) は、最上位桁優先の基数ソートをキャッシュミスが少なくなるよう改めた文字列整列向けのアルゴリズムである。

通常の最上位桁優先基数ソートでは、文字列そのものではなく「各文字列へのポインタ」の配列を並べ替えることが多い。ポインタ配列は連続に読めても、各ポインタの先にある文字列はメモリ上の別々の場所にある。桁を調べるたびにその先へ辿ると、アクセス先が飛び飛びになりキャッシュミスが増えやすい。

キャッシュ効率型基数ソートは、各キーに短いキーバッファを割り当て、先に使う数桁をそこへコピーしてから区分する。並べ替えではポインタ(本記事では整数キーそのもの)と対応するバッファを同じ順番で動かす。次の桁はばらばらの文字列を追い直さず、並んだバッファを順に読めばよいので、キャッシュに載りやすい。

  1. キーバッファの確保: キーごとに長さ bs のバッファを用意する。理論上の目安はアルファベットサイズ m・件数 n に対しおよそ log n / log m だが、実装では小さな定数(本記事では bs = 2)で足りることが多い。
  2. バッファへの読込み: まだ見ていない桁のうち先頭 bs 個を各バッファへコピーする。キー本体へのアクセスはこのタイミングに寄せる。
  3. バッファ先頭桁での区分: 最上位桁優先と同様、バッファの先頭文字(桁)0..r-1 で安定にバケット分けする。キーポインタ(本記事では整数キーそのもの)とバッファブロックを同じ順で入れ替える。
  4. 使用済み桁の廃棄: 調べた桁をバッファ先頭から捨て、残りを前へ詰める。次のパスでも常に先頭だけ見ればよい。
  5. 再充填と再帰: バッファが空になったら次の bs 桁を読み直す。要素が 2 個以上残る各バケットについて、下位桁で 3〜5 を繰り返す。
procedure cradix_sort(A)
  if length(A) = 0 then return
  width = digit_width(maximum(A))
  B[i] = fill_buffer(A[i], start=0, width) for each i
  cradix(A, B, digit_pos=0, width)

procedure cradix(A, B, digit_pos, width)
  if length(A) <= 1 or digit_pos >= width then return
  stable_partition A and B by B[i][0]
  for each non-empty bucket S
    discard_front_digit(B in S)
    next = digit_pos + 1
    if next >= width then continue
    if next mod bs = 0 then
      refill B[i] from A[i] at digit next
    cradix(S, B in S, next, width)

桁数を w、基数を r とすると時間計算量は通常の最上位桁優先と同様に \(O(w \cdot (n + r))\) 程度である。追加でキーバッファ \(O(bs \cdot n)\) と区分用の作業領域を要する。各パスが安定ソートなら全体も安定ソートである。

本記事と計測コードでは、サイト共通の整数配列向けに十進桁へ写した簡略版を用いる。文字列ポインタ版と同じく「キーとバッファを同じ順で動かす」点が中心で、素朴な最下位桁優先の基数ソートとは設計目標が異なる。

以下のデモでは 3 桁の整数を bs = 2 で扱う。棒の上の括弧内がキーバッファの中身である。

類似アルゴリズムとの相違点

基数ソートの記事は最下位桁優先(LSD; Least Significant Digit)の素朴なカウンティング繰り返しが中心である。

キャッシュ効率型基数ソートは最上位桁優先(MSD; Most Significant Digit)側に立ち、キーバッファでキャッシュ線上の参照をまとめる点が異なる。

アメリカ国旗ソートも最上位桁優先だが、インプレース交換でバケットを作ることに主眼があり、キーバッファは用いない。

バーストソートはキャッシュ効率をトライの遅延展開で稼ぐ別系統である。

時間計算量および空間計算量を計測する

Size Average time (s) Maximum time (s) Average memory (KiB) Maximum memory (KiB)
256 0.000004 0.000041 4 4
512 0.000008 0.000042 7 7
1024 0.000022 0.000174 22 22
2048 0.000040 0.000118 34 34
4096 0.000078 0.000484 58 58
8192 0.000152 0.000438 106 106
16384 0.000378 0.001945 300 300
32768 0.000681 0.001442 492 492
65536 0.001301 0.004840 876 876
131072 0.003345 0.007392 2621 2621
262144 0.006073 0.011526 4157 4157
計測に使用したコードを表示する

set -euo pipefail

WORKDIR="$(mktemp -d)"
trap 'rm -rf "$WORKDIR"' EXIT

cat > "$WORKDIR/Dockerfile" <<'EOF'
FROM swift:6.0

WORKDIR /app

RUN cat > alloc_track.c <<'ALLOC'
#define _GNU_SOURCE
#include <dlfcn.h>
#include <malloc.h>
#include <stdatomic.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

static atomic_size_t live_bytes = 0;
static atomic_size_t peak_bytes = 0;

static void *(*real_malloc)(size_t) = NULL;
static void *(*real_calloc)(size_t, size_t) = NULL;
static void *(*real_realloc)(void *, size_t) = NULL;
static void (*real_free)(void *) = NULL;

static void init_reals(void) {
    if (real_malloc) {
        return;
    }
    real_malloc = (void *(*)(size_t))dlsym(RTLD_NEXT, "malloc");
    real_calloc = (void *(*)(size_t, size_t))dlsym(RTLD_NEXT, "calloc");
    real_realloc = (void *(*)(void *, size_t))dlsym(RTLD_NEXT, "realloc");
    real_free = (void (*)(void *))dlsym(RTLD_NEXT, "free");
}

static void record_alloc(size_t size) {
    size_t live = atomic_fetch_add(&live_bytes, size) + size;
    size_t peak = atomic_load(&peak_bytes);
    while (live > peak) {
        if (atomic_compare_exchange_weak(&peak_bytes, &peak, live)) {
            break;
        }
    }
}

void alloc_track_reset_peak(void) {
    atomic_store(&peak_bytes, atomic_load(&live_bytes));
}

size_t alloc_track_live(void) { return atomic_load(&live_bytes); }
size_t alloc_track_peak(void) { return atomic_load(&peak_bytes); }

void *malloc(size_t size) {
    init_reals();
    void *p = real_malloc(size);
    if (p) {
        record_alloc(malloc_usable_size(p));
    }
    return p;
}

void *calloc(size_t nmemb, size_t size) {
    init_reals();
    void *p = real_calloc(nmemb, size);
    if (p) {
        record_alloc(malloc_usable_size(p));
    }
    return p;
}

void *realloc(void *ptr, size_t size) {
    init_reals();
    size_t old_size = 0;
    if (ptr) {
        old_size = malloc_usable_size(ptr);
    }
    void *p = real_realloc(ptr, size);
    if (p) {
        atomic_fetch_sub(&live_bytes, old_size);
        record_alloc(malloc_usable_size(p));
    } else if (size == 0) {
        atomic_fetch_sub(&live_bytes, old_size);
    }
    return p;
}

void free(void *ptr) {
    init_reals();
    if (ptr) {
        atomic_fetch_sub(&live_bytes, malloc_usable_size(ptr));
        real_free(ptr);
    }
}

ALLOC

RUN cat > main.swift <<'SWIFT'
import Foundation
#if canImport(Glibc)
import Glibc
#elseif canImport(Darwin)
import Darwin
#endif

@_silgen_name("alloc_track_live") func alloc_track_live() -> Int
@_silgen_name("alloc_track_peak") func alloc_track_peak() -> Int
@_silgen_name("alloc_track_reset_peak") func alloc_track_reset_peak()

extension UnsafeMutableBufferPointer where Element == Int {
    func swapAt(_ i: Int, _ j: Int) {
        let t = self[i]; self[i] = self[j]; self[j] = t
    }
}

let MIN_POWER: Int = 8
let MAX_POWER: Int = 18
let RUNS: Int = 8192


private let CRADIX_RADIX = 10
private let CRADIX_BS = 2

private func cradix_digit_width(_ max: Int) -> Int {
    if max == 0 {
        return 1
    }
    var w = 0
    var v = max
    while v > 0 {
        w += 1
        v /= CRADIX_RADIX
    }
    return w
}

private func cradix_digit_at(_ value: Int, _ pos: Int, _ width: Int) -> UInt8 {
    let power = width - 1 - pos
    var div = 1
    for _ in 0..<power {
        if div > Int.max / CRADIX_RADIX {
            div = Int.max
            break
        }
        div *= CRADIX_RADIX
    }
    return UInt8((value / div) % CRADIX_RADIX)
}

private func cradix_fill_buffer(_ value: Int, _ start: Int, _ width: Int) -> [UInt8] {
    var buf = [UInt8](repeating: 0, count: CRADIX_BS)
    for i in 0..<CRADIX_BS {
        if start + i < width {
            buf[i] = cradix_digit_at(value, start + i, width)
        }
    }
    return buf
}

private func cradix_rec(
    _ a: UnsafeMutableBufferPointer<Int>,
    _ buffers: inout [[UInt8]],
    _ digit_pos: Int,
    _ width: Int
) {
    let n = a.count
    if n <= 1 || digit_pos >= width {
        return
    }

    var count = [Int](repeating: 0, count: CRADIX_RADIX)
    for b in buffers {
        count[Int(b[0])] += 1
    }

    var offset = [Int](repeating: 0, count: CRADIX_RADIX)
    for i in 1..<CRADIX_RADIX {
        offset[i] = offset[i - 1] + count[i - 1]
    }

    var out_a = [Int](repeating: 0, count: n)
    var out_b = [[UInt8]](repeating: [UInt8](repeating: 0, count: CRADIX_BS), count: n)
    var cursor = offset
    for i in 0..<n {
        let d = Int(buffers[i][0])
        out_a[cursor[d]] = a[i]
        out_b[cursor[d]] = buffers[i]
        cursor[d] += 1
    }
    for i in 0..<n {
        a[i] = out_a[i]
    }
    buffers = out_b

    for r in 0..<CRADIX_RADIX {
        let start = offset[r]
        let len = count[r]
        if len <= 1 {
            continue
        }
        let next_pos = digit_pos + 1
        if next_pos >= width {
            continue
        }

        let end = start + len
        for i in start..<end {
            for j in 0..<(CRADIX_BS - 1) {
                buffers[i][j] = buffers[i][j + 1]
            }
            buffers[i][CRADIX_BS - 1] = 0
        }

        if next_pos % CRADIX_BS == 0 {
            for i in start..<end {
                buffers[i] = cradix_fill_buffer(a[i], next_pos, width)
            }
        }

        var subBuffers = Array(buffers[start..<end])
        cradix_rec(
            UnsafeMutableBufferPointer(rebasing: a[start..<end]),
            &subBuffers,
            next_pos,
            width
        )
        for i in 0..<len {
            buffers[start + i] = subBuffers[i]
        }
    }
}

func cradix_sort(_ a: inout [Int]) {
    a.withUnsafeMutableBufferPointer { cradix_sort($0) }
}

func cradix_sort(_ a: UnsafeMutableBufferPointer<Int>) {
    if a.isEmpty {
        return
    }

    var max = a[0]
    for i in 1..<a.count {
        if a[i] > max { max = a[i] }
    }
    let width = cradix_digit_width(max)
    let n = a.count
    var buffers = [[UInt8]](repeating: [UInt8](repeating: 0, count: CRADIX_BS), count: n)
    for i in 0..<n {
        buffers[i] = cradix_fill_buffer(a[i], 0, width)
    }
    cradix_rec(a, &buffers, 0, width)
}


func benchmark_sort(_ array: inout [Int]) {

    cradix_sort(&array)

}

func is_non_decreasing(_ a: [Int]) -> Bool {
    guard a.count >= 2 else { return true }
    for i in 1..<a.count {
        if a[i - 1] > a[i] { return false }
    }
    return true
}

func same_multiset(_ a: [Int], _ b: [Int]) -> Bool {
    if a.count != b.count {
        return false
    }

    var left = a
    var right = b
    left.sort()
    right.sort()
    return left == right
}

func check_correctness_case(_ label: String, _ input: [Int]) {
    var input = input
    let original = input

    benchmark_sort(&input)

    if !is_non_decreasing(input) {
        fatalError("correctness case \(label): output is not sorted")
    }

    if !same_multiset(input, original) {
        fatalError("correctness case \(label): elements were lost or added")
    }
}

// Skip cases larger than the algorithm's measured size cap (MAX_POWER). That
// cap exists because larger inputs are impractically slow; forcing them here
// would stall the published measurement script before any table rows print.
func check_correctness_case_within_limit(_ label: String, _ input: [Int]) {
    if input.count > (1 << MAX_POWER) {
        return
    }
    check_correctness_case(label, input)
}

func few_unique_values(_ size: Int, _ unique: Int, _ seed: UInt64) -> [Int] {
    var state = seed
    var result = [Int]()
    result.reserveCapacity(size)
    for _ in 0..<size {
        state ^= state << 13
        state ^= state >> 7
        state ^= state << 17
        result.append(Int(state % UInt64(unique)) + 1)
    }
    return result
}

func run_correctness_checks() {
    check_correctness_case("empty", [])
    check_correctness_case("single", [42])
    check_correctness_case("duplicates", [3, 1, 3, 2, 1, 2])
    check_correctness_case("sorted", [1, 2, 3, 4, 5])
    check_correctness_case("reverse", [5, 4, 3, 2, 1])
    check_correctness_case("all_equal", [7, 7, 7, 7])
    check_correctness_case("skewed_range", [1_000_000, 2, 1_000_001, 1, 999_999])
    // Static-buffer Grail skips the in-buffer build when key collection is sparse
    // (ideal_buffer = false). Exercising that path catches regressions in buffer gating.
    check_correctness_case(
        "few_keys_len16",
        [2, 2, 2, 2, 2, 2, 2, 2, 4, 3, 1, 2, 3, 4, 1, 4]
    )
    // Seed 0 is a fixed point of the xorshift below, so it would degenerate into
    // yet another all-equal case instead of a 4-value mix. Start at 1.
    for seed in 1...32 {
        check_correctness_case(
            "few_keys_len32_seed_\(seed)",
            few_unique_values(32, 4, UInt64(seed))
        )
    }
    // Small-input cutoffs (insertion sort below 32 elements, etc.) hide duplicate-key
    // bugs in the recursive path, so repeat the duplicate cases at the smallest
    // benchmark size, which every algorithm must handle within reasonable time.
    check_correctness_case("all_equal_len256", [Int](repeating: 7, count: 256))
    for seed in 1...4 {
        check_correctness_case(
            "few_keys_len256_seed_\(seed)",
            few_unique_values(256, 4, UInt64(seed))
        )
    }
    // Blit's equal-key second sweep used to copy the whole range into a fixed
    // 512-element swap; lengths above that must still sort without panicking.
    // Respect MAX_POWER so algorithms with a low measured-size cap (slow,
    // sleep) do not hang here for minutes or months.
    check_correctness_case_within_limit("all_equal_len600", [Int](repeating: 7, count: 600))
    for seed in 1...4 {
        check_correctness_case_within_limit(
            "few_keys_len2048_seed_\(seed)",
            few_unique_values(2048, 4, UInt64(seed))
        )
    }
}


func shuffled(_ size: Int, seed: UInt64) -> [Int] {
    guard size > 0 else { return [] }

    var v = Array(1...size)
    var state = seed

    if size > 1 {
        for i in stride(from: size - 1, through: 1, by: -1) {
            state ^= state << 13
            state ^= state >> 7
            state ^= state << 17

            let j = Int(state % UInt64(i + 1))
            v.swapAt(i, j)
        }
    }

    return v
}

func micros(_ d: Duration) -> UInt64 {
    let c = d.components
    let fromSeconds = UInt64(c.seconds) * 1_000_000
    let fromAttos = UInt64(max(0, c.attoseconds / 1_000_000_000_000))
    return fromSeconds + fromAttos
}

func padLeft(_ value: String, _ width: Int) -> String {
    if value.count >= width {
        return value
    }
    return String(repeating: " ", count: width - value.count) + value
}

func formatSeconds(_ micros: UInt64) -> String {
    let whole = micros / 1_000_000
    let frac = micros % 1_000_000
    let fracStr = padLeft(String(frac), 6).replacingOccurrences(of: " ", with: "0")
    return "\(whole).\(fracStr)"
}

func input_array(_ size: Int, seed: UInt64) -> [Int] {
    shuffled(size, seed: seed)
}

/// Peak heap growth during `benchmark_sort`, in bytes (explicit buffers such as swap).
/// Kept in bytes so the parent can average before rounding; converting to KiB here
/// would truncate sub-KiB buffers to 0 in every run and hide them from the average.
func run_once(size: Int, seed: Int) -> (UInt64, Int) {
    var array = input_array(size, seed: UInt64(seed))

    let baseBytes = alloc_track_live()
    alloc_track_reset_peak()

    let start = ContinuousClock.now

    benchmark_sort(&array)

    let elapsed = ContinuousClock.now - start
    let peakBytes = alloc_track_peak()
    let auxBytes = max(0, peakBytes - baseBytes)

    let expected: [Int] = size > 0 ? Array(1...size) : []
    if array != expected {
        fatalError("sort failed with seed \(seed) for size \(size)")
    }

    return (micros(elapsed), auxBytes)
}

func run_child(_ args: [String]) {
    let size = Int(args[2])!
    let seed = Int(args[3])!
    let (elapsedUs, mem) = run_once(size: size, seed: seed)
    print("\(elapsedUs) \(mem)")
}

let args = CommandLine.arguments
if args.count > 1 && args[1] == "--run-once" {
    run_child(args)
} else {
    run_correctness_checks()

    print(
        "| \(padLeft("Size", 10)) | \(padLeft("Average time (s)", 16)) | \(padLeft("Maximum time (s)", 16)) | \(padLeft("Average memory (KiB)", 20)) | \(padLeft("Maximum memory (KiB)", 20)) |"
    )
    print("|----------:|----------------:|----------------:|--------------------:|--------------------:|")

    for power in MIN_POWER...MAX_POWER {
        let size = 1 << power

        var totalTime: UInt64 = 0
        var maxTime: UInt64 = 0

        var totalMem = 0
        var maxMem = 0

        for seed in 1...RUNS {
            let process = Process()
            process.executableURL = URL(fileURLWithPath: args[0])
            process.arguments = ["--run-once", "\(size)", "\(seed)"]
            let stdout = Pipe()
            let stderr = Pipe()
            process.standardOutput = stdout
            process.standardError = stderr

            do {
                try process.run()
            } catch {
                fatalError("failed to run benchmark child process: \(error)")
            }
            process.waitUntilExit()

            if process.terminationStatus != 0 {
                let err = String(data: stderr.fileHandleForReading.readDataToEndOfFile(), encoding: .utf8) ?? ""
                fatalError("benchmark child process failed: \(err)")
            }

            let data = stdout.fileHandleForReading.readDataToEndOfFile()
            let stdoutText = String(data: data, encoding: .utf8) ?? ""
            let fields = stdoutText.split(whereSeparator: \.isWhitespace)
            guard fields.count >= 2,
                  let elapsedUs = UInt64(fields[0]),
                  let auxMem = Int(fields[1]) else {
                fatalError("invalid child process output: \(stdoutText)")
            }

            totalTime += elapsedUs
            if elapsedUs > maxTime {
                maxTime = elapsedUs
            }

            totalMem += auxMem
            if auxMem > maxMem {
                maxMem = auxMem
            }
        }

        let avgTime = totalTime / UInt64(RUNS)
        // Memory is summed in bytes and converted to KiB once, after averaging.
        let avgMemKb = totalMem / RUNS / 1024
        let maxMemKb = maxMem / 1024

        print(
            "| \(padLeft(String(size), 10)) | \(padLeft(formatSeconds(avgTime), 16)) | \(padLeft(formatSeconds(maxTime), 16)) | \(padLeft(String(avgMemKb), 20)) | \(padLeft(String(maxMemKb), 20)) |"
        )
    }
}
SWIFT

RUN clang -O2 -fPIC -shared alloc_track.c -o liballoc_track.so -ldl

RUN swiftc -Ounchecked -whole-module-optimization \
    main.swift \
    -o swift-benchmark \
    -L. -lalloc_track \
    -Xlinker -rpath -Xlinker /app

ENV LD_PRELOAD=/app/liballoc_track.so
CMD ["./swift-benchmark"]
EOF

docker build -t swift-benchmark "$WORKDIR"
docker run --rm --init swift-benchmark