CUDA Thrust max_element处理接近INT_MAX的大数组时触发非法内存访问错误

当使用Thrust的thrust::max_element查找最大值时，数组长度接近INT_MAX（2147483647）会触发以下错误：

what(): extrema failed to synchronize: cudaErrorIllegalAddress: an illegal memory access was encountered

测试表明，N=2140000000时程序正常运行，N=2147000000时触发错误。相关问题曾在Thrust内部问题追踪中被提及，多数同类问题已修复，但max_element疑似存在遗留的32位索引溢出问题。需求是基于Thrust的解决方案，避免自行编写CUDA Kernel导致性能下降。

基于Thrust的兼容解决方案 #

问题根源在于Thrust内部对大尺寸数组的中间计算使用32位整数，当数组长度接近INT_MAX时，索引计算会溢出，进而引发非法内存访问。通过分治拆分数组的方式可以规避该问题：将大数组拆分为多个长度小于INT_MAX的子数组，分别计算每个子数组的局部最大值，再对这些局部最大值求全局最大值。

实现步骤 #

• 将原数组划分为多个子段，每个子段长度不超过INT_MAX - 1（示例中取2140000000）
• 对每个子段调用thrust::max_element，记录局部最大值及其全局位置
• 收集所有局部最大值，再次调用thrust::max_element得到全局结果

修改后的测试代码 #

#include <iostream>
#include <cstdlib>
#include <limits>
#include "thrust/device_ptr.h"
#include "thrust/device_vector.h"
#include "thrust/fill.h"
#include "thrust/extrema.h"
#include "thrust/pair.h"

int main() {
    const size_t N = 2147000000; // 接近INT_MAX的测试长度
    const size_t SEGMENT_SIZE = 2140000000; // 安全子数组长度，避免索引溢出

    std::cout << "N  " << N << std::endl;

    thrust::device_vector<int> dArr(N);
    thrust::fill(dArr.begin(), dArr.end(), 1);
    dArr[N-2] = 2; // 设置全局最大值的位置

    // 存储每个子段的最大值（值+全局索引）
    thrust::device_vector<thrust::pair<int, size_t>> local_maxes;

    size_t offset = 0;
    while (offset < N) {
        size_t current_segment_len = std::min(SEGMENT_SIZE, N - offset);
        auto segment_start = dArr.begin() + offset;
        auto segment_end = segment_start + current_segment_len;

        // 查找当前子段的最大值
        auto max_iter = thrust::max_element(segment_start, segment_end);
        local_maxes.emplace_back(*max_iter, max_iter - dArr.begin());

        offset += current_segment_len;
    }

    // 从局部最大值中筛选全局最大值
    auto global_max_iter = thrust::max_element(local_maxes.begin(), local_maxes.end(),
        [](const thrust::pair<int, size_t>& a, const thrust::pair<int, size_t>& b) {
            return a.first < b.first;
        });

    std::cout << "max value in dArr  " << global_max_iter->first << std::endl;
    std::cout << " - max_index  " << global_max_iter->second << std::endl;

    return 0;
}

方案优势 #

• 完全基于Thrust实现，保留Thrust针对CUDA的性能优化
• 分治逻辑仅增加少量CPU端循环和局部结果收集，额外开销可忽略
• 兼容任意大尺寸数组，不受INT_MAX限制

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