指定结束位时cub::DeviceRadixSort排序函数崩溃问题求助

问题背景 #

我在使用CUB 1.16.0的GPU基数排序对10亿个32位无符号整数排序时，遇到了非确定性的非法内存访问崩溃（几乎必现）。这些整数仅使用最低有效位开始的k位（k=18、19、20），因此我调用cub::DeviceRadixSort::SortKeys时指定了位区间[0, k)。测试环境覆盖了Volta架构（CUDA 11.4）和Ampere架构（CUDA 11.2）的NVIDIA GPU。

最小可复现代码 #

编译命令：

nvcc -O3 -std=c++17 -Xcompiler -fopenmp cub_sort_test.cu -o cub_sort_test

代码：

#include <cub/cub.cuh>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/system/cuda/experimental/pinned_allocator.h>
#include <algorithm>
#include <chrono>
#include <iostream>
#include <parallel/algorithm>
#include <random>
#include <vector>
#include <iostream>
#define DEBUG
#ifdef DEBUG
#define CheckCudaError(instruction) \
{ AssertNoCudaError((instruction), __FILE__, __LINE__); }
#else
#define CheckCudaError(instruction) instruction
#endif
inline void AssertNoCudaError(cudaError_t error_code, const char* file, int line) {
    if (error_code != cudaSuccess) {
        std::cout << "Error: " << cudaGetErrorString(error_code) << " " << file << " " << line << "\n";
    }
}
template <typename T>
using PinnedHostVector = thrust::host_vector<T, thrust::system::cuda::experimental::pinned_allocator<T>>;
std::mt19937 SeedRandomGenerator(uint32_t distribution_seed) {
    const size_t seeds_bytes = sizeof(std::mt19937::result_type) * std::mt19937::state_size;
    const size_t seeds_length = seeds_bytes / sizeof(std::seed_seq::result_type);
    std::vector<std::seed_seq::result_type> seeds(seeds_length);
    std::generate(seeds.begin(), seeds.end(), [&]() {
        distribution_seed = (distribution_seed << 1) | (distribution_seed >> (-1 & 31));
        return distribution_seed;
    });
    std::seed_seq seed_sequence(seeds.begin(), seeds.end());
    return std::mt19937{seed_sequence};
}
int main(int argc, char* argv[]) {
    if (argc != 4) {
        std::cerr << "Usage: ./cub-sort-test <num_keys> <gpu_id> <bit_entropy>" << std::endl;
        return -1;
    }
    size_t num_keys = std::stoull(argv[1]);
    int gpu = std::stoi(argv[2]);
    size_t bit_entropy = std::stoi(argv[3]);
    cudaStream_t stream;
    CheckCudaError(cudaSetDevice(gpu));
    CheckCudaError(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
    PinnedHostVector<uint32_t> keys(num_keys);
#pragma omp parallel num_threads(64)
    {
        uint32_t max = (1 << bit_entropy) - 1;
        if (bit_entropy == sizeof(uint32_t) * 8) {
            max = std::numeric_limits<uint32_t>::max();
        } else if (bit_entropy == 1) {
            max = 2;
        }
        std::mt19937 random_generator = SeedRandomGenerator(2147483647 + static_cast<size_t>(omp_get_thread_num()));
        std::uniform_real_distribution<double> uniform_dist(0, max);
#pragma omp for schedule(static)
        for (size_t i = 0; i < num_keys; ++i) {
            keys[i] = static_cast<uint32_t>(uniform_dist(random_generator));
        }
    }
    thrust::device_vector<uint32_t> device_vector(num_keys);
    thrust::copy(keys.begin(), keys.end(), device_vector.begin());
    CheckCudaError(cudaDeviceSynchronize());
    size_t num_temporary_bytes = 0;
    cub::DeviceRadixSort::SortKeys(
        NULL, num_temporary_bytes,
        thrust::raw_pointer_cast(device_vector.data()),
        thrust::raw_pointer_cast(device_vector.data()),
        num_keys, 0, bit_entropy + 1, stream); // bit subrange is [begin_bit, end_bit), thus bit_entropy + 1
    uint8_t* temporary_storage = nullptr;
    CheckCudaError(cudaMalloc(reinterpret_cast<void**>(&temporary_storage), num_temporary_bytes));
    cub::DeviceRadixSort::SortKeys(
        (void*)temporary_storage, num_temporary_bytes,
        thrust::raw_pointer_cast(device_vector.data()),
        thrust::raw_pointer_cast(device_vector.data()),
        num_keys, 0, bit_entropy + 1, stream);
    CheckCudaError(cudaStreamSynchronize(stream));
    thrust::copy(device_vector.begin(), device_vector.end(), keys.begin());
    CheckCudaError(cudaFree(temporary_storage));
    if (std::is_sorted(keys.begin(), keys.end()) == false) {
        std::cout << "Error: Sorting failed." << std::endl;
    }
    return 0;
}

可能的原因及修复方法 #

1. 位区间参数错误（最可能的根源）

你的代码中对CUB位区间的理解存在偏差：cub::DeviceRadixSort::SortKeys的位区间参数[begin_bit, end_bit)表示包含begin_bit，不包含end_bit的位范围。如果你的整数仅使用低bit_entropy位（即0到bit_entropy-1位），那么正确的位区间应该是0, bit_entropy，而不是0, bit_entropy + 1。

错误地指定end_bit=bit_entropy+1会让CUB尝试处理第bit_entropy位（虽然该位全为0），这可能触发CUB内部针对非完整位区间的内存计算bug，尤其是在处理10亿级别的超大数据量时。

修正后的SortKeys调用：

cub::DeviceRadixSort::SortKeys(
    NULL, num_temporary_bytes,
    thrust::raw_pointer_cast(device_vector.data()),
    thrust::raw_pointer_cast(device_vector.data()),
    num_keys, 0, bit_entropy, stream); // 修正end_bit为bit_entropy，而非bit_entropy+1

2. 尝试改用Out-of-Place排序

CUB的In-place排序在极端数据量+有限位区间的组合下可能存在边缘case问题。你可以尝试使用独立的输出缓冲区，而非复用输入缓冲区：

// 新增输出device_vector
thrust::device_vector<uint32_t> device_output(num_keys);

// 调用SortKeys时使用独立的输入输出指针
cub::DeviceRadixSort::SortKeys(
    NULL, num_temporary_bytes,
    thrust::raw_pointer_cast(device_vector.data()),
    thrust::raw_pointer_cast(device_output.data()), // 输出到独立缓冲区
    num_keys, 0, bit_entropy, stream);

3. 升级CUB或CUDA版本

CUB 1.16.0可能存在针对大数量级有限位排序的已知bug，建议升级到CUB 1.17.0及以上版本。同时，CUDA 11.2/11.4也相对较旧，升级到CUDA 11.7+或最新版本，可能解决底层的兼容性问题。

4. 验证临时内存分配合理性

可以打印num_temporary_bytes的值，确认临时内存是否在GPU显存的可用范围内：

std::cout << "Temporary memory required: " << num_temporary_bytes / (1024*1024) << " MB" << std::endl;

如果临时内存超过GPU显存的剩余空间，可能需要分批排序，或者使用带有统一内存的设备。

总结 #

优先修正位区间参数，这是最可能导致崩溃的原因。如果问题仍然存在，尝试改用Out-of-Place排序，或者升级CUB/CUDA版本来解决潜在的框架bug。

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