GPU 编程实战——实用排序与搜索

GPU 加速的 Bitonic 排序

排序是数据处理、分析和模拟工作流中的基础操作。随着数据集规模不断增长，高效排序变得愈发重要。CPU 提供了多种通用排序算法，但它们并不总是适合映射到诸如 GPU 这样的大规模并行架构。

Bitonic 排序简介

Bitonic 排序是一种经典的基于比较的并行排序算法。与需要自适应分支和递归划分的快速排序或归并排序不同，Bitonic 排序使用固定顺序的比较-交换阶段，每个阶段都能并行执行，无需数据依赖分支。这种可预测性使其非常适合 GPU：成千上万的比较-交换操作可同时运行，带来高吞吐和最少同步开销。

• 每个阶段可并行运行大量比较-交换操作
• 固定、无条件分支的操作序列适合高效 GPU 内核
• 适用于对延迟敏感、需要实时排序的场景，如模拟管道、图形渲染或实时分析

算法核心是不断构建和合并 bitonic 序列（先单调增后单调减，或反之）的子序列，通过一系列固定的比较-交换步骤完成排序。阶段数与输入大小相关，可在主机端预先计算，并按序在 GPU 上逐步执行。

GPU 上的 Bitonic 排序实现

数据准备

import numpy as np
import pycuda.autoinit
import pycuda.driver as drv
import pycuda.gpuarray as gpuarray
from pycuda.compiler import SourceModule
import time

# 使用 2 的幂长度以获得最佳性能
N         = 2**15     # 32768 元素
arr_host  = np.random.rand(N).astype(np.float32)
arr_gpu   = gpuarray.to_gpu(arr_host)

Bitonic 排序单步内核

下面的内核实现了一个排序阶段：每个线程负责比较并在必要时交换一对元素，具体操作取决于当前的 j 和 k 值。

bitonic_kernel_code = """
__global__ void bitonic_step(float *data, int j, int k, int n)
{
    unsigned int idx = threadIdx.x + blockDim.x * blockIdx.x;
    if (idx >= n) return;
    unsigned int ixj = idx ^ j;
    if (ixj > idx) {
        if ((idx & k) == 0) {
            if (data[idx] > data[ixj]) {
                float tmp = data[idx];
                data[idx] = data[ixj];
                data[ixj] = tmp;
            }
        } else {
            if (data[idx] < data[ixj]) {
                float tmp = data[idx];
                data[idx] = data[ixj];
                data[ixj] = tmp;
            }
        }
    }
}
"""
mod           = SourceModule(bitonic_kernel_code)
bitonic_step  = mod.get_function("bitonic_step")

主机端控制与性能测量

在主机端，我们使用双重循环来执行所有排序阶段，依次传入 j 和 k：

def gpu_bitonic_sort(arr_gpu, n, threads_per_block=256):
    blocks_per_grid = (n + threads_per_block - 1) // threads_per_block
    for k in [2**i for i in range(1, int(np.log2(n)) + 1)]:
        for j in [2**i for i in range(int(np.log2(k)), -1, -1)]:
            bitonic_step(
                arr_gpu, np.int32(j), np.int32(k), np.int32(n),
                block=(threads_per_block, 1, 1),
                grid=(blocks_per_grid, 1)
            )

# GPU 排序计时
start_gpu = time.time()
gpu_bitonic_sort(arr_gpu, N)
drv.Context.synchronize()
end_gpu   = time.time()
sorted_gpu = arr_gpu.get()

# CPU 排序计时
arr_host_copy = arr_host.copy()
start_cpu     = time.time()
arr_host_copy.sort()
end_cpu       = time.time()

print(f"GPU bitonic 排序耗时：{end_gpu - start_gpu:.4f} 秒")
print(f"CPU numpy 排序耗时：{end_cpu - start_cpu:.4f} 秒")
print("结果一致吗？", np.allclose(arr_host_copy, sorted_gpu, atol=1e-6))

性能对比

在中等规模或批量处理场景下，GPU 加速的 Bitonic 排序常常多倍优于 CPU 排序。对于长度为 2 的幂的大数组，GPU 版本能显著发挥并行优势；在处理多个数组或排序是更大 GPU 工作流一部分的情况下，这一差距会更大，因为数据无需离开设备就可完成整个流程。

基于整数数组的基数排序

我们已经看到 Bitonic 排序因其规则的比较-交换结构和固定的并行步骤而能在 GPU 上高效运行。然而，对于大规模整数数组而言，基于比较的算法（如 Bitonic 排序）并不总是最优。基数排序则从根本上不同：它根据数字的二进制表示，一次处理若干位（digit），无需任意比较。

在 GPU 上，基数排序会以一系列快速、确定性的遍历来处理数据集：先按低位（或高位）对元素进行分发，然后再收集，利用每一步的并行性完成整个排序过程。

使用 CuPy 运行基数排序

首先准备一个大型整数数组：

import cupy as cp
import numpy as np
import time

N         = 2**20  # 1,048,576 个整数
arr_host  = np.random.randint(0, 10_000_000, size=N, dtype=np.int32)
arr_cupy  = cp.array(arr_host)

# 在 GPU 上执行 CuPy 的排序（对足够大的整数数组，底层使用基数排序）
start_gpu     = time.time()
sorted_cupy   = cp.sort(arr_cupy)
cp.cuda.Stream.null.synchronize()
end_gpu       = time.time()

# 拷回主机以供后续比较
sorted_host_radix = cp.asnumpy(sorted_cupy)

使用 PyCUDA 运行 Bitonic 排序

为了与 Bitonic 排序进行公平的技术比较，我们先将整数数组转换为浮点数，然后调用之前实现的 gpu_bitonic_sort：

import pycuda.gpuarray as gpuarray
import pycuda.autoinit
import time

# 转为 float32 并拷贝到 GPU
arr_gpu_bitonic = gpuarray.to_gpu(arr_host.astype(np.float32))

# 启动 Bitonic 排序
start_bitonic = time.time()
gpu_bitonic_sort(arr_gpu_bitonic, N)  # 前文定义的函数
pycuda.autoinit.context.synchronize()
end_bitonic   = time.time()

# 拷回结果
sorted_host_bitonic = arr_gpu_bitonic.get()

性能与正确性对比

# 打印耗时
print("CuPy 基数排序耗时：     {:.4f} 秒".format(end_gpu - start_gpu))
print("PyCUDA Bitonic 排序耗时：{:.4f} 秒".format(end_bitonic - start_bitonic))

# 与 CPU 排序结果比较
reference       = np.sort(arr_host)
match_radix     = np.array_equal(reference, sorted_host_radix)
match_bitonic   = np.allclose(reference, sorted_host_bitonic, atol=1)

print("基数排序结果正确吗？", match_radix)
print("Bitonic 排序结果正确吗？", match_bitonic)

在实际测试中，CuPy 内部对大规模整数数组采用的基数排序往往显著超越 Bitonic 排序，成为 GPU 上高吞吐整数排序的首选算法，尤其在需要对多组数组或流水线式 GPU 工作流中排序时更是如此。

并行线性搜索内核

为什么要并行线性搜索？

线性搜索是数据库查询、模式匹配和数据分析中的经典操作。通常在线性搜索中，CPU 会逐个元素地检查数组，直到找到目标。对小数据集而言，这种方式简单有效；但对大规模数据集而言，由于每个 CPU 周期只能检查一个元素，开销显著变大。

GPU 拥有数千个并行线程，非常适合此类任务。只需为每个线程分配一个或多个数组元素，即可同时检查大量元素，一旦任一线程发现目标即可结束搜索。该模式适用于目标位置未知的大数据集，以及需要同时进行多次搜索的场景。虽然对于已排序数组，二分搜索或扫描搜索可能更快，但对于无序数据或大批量模式查询，并行线性搜索相比传统 CPU 方法能显著提速。

编写并行线性搜索 Kernel

数据准备

import numpy as np
import pycuda.autoinit
import pycuda.driver as drv
import pycuda.gpuarray as gpuarray
from pycuda.compiler import SourceModule

N            = 2_000_000
target_value = 123456

# 随机生成数组并确保目标存在
arr_host   = np.random.randint(0, 1_000_000, N).astype(np.int32)
rand_index = np.random.randint(0, N)
arr_host[rand_index] = target_value

arr_gpu = gpuarray.to_gpu(arr_host)

并行搜索内核

每个线程检查自己对应的元素，若匹配则使用原子操作将其索引写入输出缓冲，确保记录最小索引并避免竞态：

kernel_code = """
__global__ void parallel_search(const int *data, int *found_idx, int target, int n)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n && data[idx] == target) {
        atomicMin(found_idx, idx);  // 记录最小匹配索引
    }
}
"""
mod             = SourceModule(kernel_code)
parallel_search = mod.get_function("parallel_search")

准备输出缓冲

# 用 N 初始化，表示尚未找到
found_idx = gpuarray.to_gpu(np.array([N], dtype=np.int32))

启动内核

threads_per_block = 256
blocks_per_grid   = (N + threads_per_block - 1) // threads_per_block

parallel_search(
    arr_gpu, found_idx, np.int32(target_value), np.int32(N),
    block=(threads_per_block, 1, 1),
    grid =(blocks_per_grid,   1   )
)
drv.Context.synchronize()

result = found_idx.get()[0]
if result < N:
    print(f"在索引 {result} 处找到了目标")
else:
    print("未找到目标")

与串行搜索对比

import time

# CPU 串行搜索
start_cpu = time.time()
for i, val in enumerate(arr_host):
    if val == target_value:
        cpu_idx = i
        break
end_cpu   = time.time()

print(f"CPU 搜索索引：{cpu_idx}，耗时：{end_cpu - start_cpu:.4f} 秒")

对于大数组，并行内核将搜索时间大幅缩短，因为所有元素同时进行检查。实际场景中，整体延迟主要取决于最慢线程的完成时间及内核启动/同步开销，通常相比 CPU 串行搜索要快数个量级。

这一原理同样可扩展到多目标搜索、计数出现次数，甚至更复杂的并行匹配操作，具备极好可扩展性。

在主机端合并结果

为什么要合并结果？

到目前为止，我们已经学会了如何使用 PyCUDA 和 CuPy 在 GPU 上加速排序和搜索任务。每个内核都会并行处理大数组，并在许多情况下生成部分结果或每块输出，而要完成整个工作流，必须在主机（CPU）端收集、合并或最终处理这些结果。无论是对大数据集分段排序，还是并行搜索查询，主机端合并都是保证正确性、可扩展性并与后续任务集成的关键模式。

常见场景包括：

• 分段排序后合并：数组被拆分成多个块并行排序后，需要在主机端将这些已排序片段合并成一个完整的有序数组。
• 并行搜索结果汇总：多个线程或线程块可能同时找到目标的位置，要在主机端收集所有匹配索引（或仅最小索引），并进行汇总。
• 边界元素处理：当数据分散到不同块或多次 kernel 调用时，接近块边界的元素可能会在多个块中出现，需要在合并时避免丢失或重复。

从设备收集结果

在 GPU 上完成排序或搜索后，首先用 .get()（PyCUDA）或 cp.asnumpy()（CuPy）将结果数组拷回主机：

# 排序结果
sorted_gpu  = arr_gpu.get()              # PyCUDA
sorted_cupy = cp.asnumpy(arr_cupy_sorted)  # CuPy

# 搜索结果（部分索引）
partial_indices_gpu = found_idx_gpu.get()

合并已排序的片段

假设我们对数组做了分段排序，每个片段单独调用内核。主机端只需按归并排序的方式，合并这些有序片段。例如，合并两个已排序的数组：

import numpy as np

segment1 = np.sort(np.random.randint(0, 100, size=5000))
segment2 = np.sort(np.random.randint(0, 100, size=5000))
merged   = np.empty(len(segment1) + len(segment2), dtype=segment1.dtype)

i = j = k = 0
while i < len(segment1) and j < len(segment2):
    if segment1[i] < segment2[j]:
        merged[k] = segment1[i]; i += 1
    else:
        merged[k] = segment2[j]; j += 1
    k += 1

# 复制剩余元素
merged[k:] = segment1[i:] if i < len(segment1) else segment2[j:]

若要合并多于两个的片段，可使用 heapq.merge() 或者 SciPy/Python 库中的多路归并函数。

合并搜索结果

如果并行搜索为每个线程块写入了一个匹配索引（未找到则写入 -1），主机端需要收集所有有效索引并选出最小者：

partial_indices = np.array([12, -1, -1, 78, -1])  # GPU 输出示例
hits            = partial_indices[partial_indices != -1]

if len(hits) > 0:
    print("目标出现于索引：", hits)
    print("第一次出现位置：", hits.min())
else:
    print("未找到目标")

合并时要注意：

• 去重：防止边界重叠导致重复索引。
• 保持顺序：合并排序片段时，确保输出数组整体有序。
• 索引修正：搜索结果须为全局索引，而非各自片段内的局部索引。

合并后校验

合并完成后，最好与 CPU 完整排序或搜索结果进行对比，以确保正确性：

# 排序结果校验
reference_sorted = np.sort(complete_input_array)
if np.array_equal(merged, reference_sorted):
    print("合并后的排序结果正确！")
else:
    print("合并后的排序结果有误。")

# 搜索结果校验
if all(complete_input_array[idx] == target_value for idx in hits):
    print("所有搜索索引均有效。")
else:
    print("存在无效的搜索索引。")

通过在主机端合并 GPU 各部分结果，我们能将 PyCUDA 的低级控制与 CuPy 的高层便利结合起来，构建既高效又可扩展的排序、搜索和大规模数据分析解决方案。

总结

总的来说，我们探索了 GPU 加速在排序与搜索上的应用，并发现在处理大规模数据时效果非常显著。

1. Bitonic 排序

   • 了解了 Bitonic 排序的分阶段、固定并行结构，适合 GPU 大规模同时执行成千上万次比较-交换操作。
   • 使用 PyCUDA 实现并计时，与 CPU 排序对比，体会到其预测性强、面向并行硬件的设计优势。

2. 基数排序

   • 针对整数数组，利用 CuPy 内部的基数排序（基于数字位的划分汇聚），无需比较即可高效排序。
   • 性能对比表明，在大规模整数集合上，基数排序明显优于 Bitonic 排序。

3. 并行线性搜索

   • 设计并实现了并行线性搜索 kernel，让 GPU 上的每个线程同时检查数组元素，一旦命中即可原子记录索引。
   • 与传统的串行搜索相比，在海量数组上展现出数倍的速度优势。

4. 主机端结果合并

   • 演示了如何在主机端合并 GPU 端分段排序输出或并行搜索索引，确保正确性并与 CPU 参考结果对齐。

通过这些实战练习，我们掌握了构建高效、可扩展的 GPU 加速排序与搜索管道的方法，能够胜任真实世界中对大数据处理的挑战。