CUDA环境下基于Thrust库的ID对去重实现求助

核心思路 #

利用Thrust库的并行排序、转换、去重算法，根据same_id_src选项生成不同的去重键，完成设备端ID对的去重：

• 当same_id_src=false：直接以<id1, id2>有序对作为去重键，仅移除完全相同的ID对
• 当same_id_src=true：先将每个ID对转换为<min(id1,id2), max(id1,id2)的有序对作为键，此时<0,5>和<5,0>会生成相同键，去重时保留第一个出现的原始ID对

具体实现步骤 #

1. 包装设备指针为Thrust可操作的迭代器

你的SearchDataOnDevice结构中的id1和id2是设备端指针，先用thrust::device_ptr包装，让Thrust可以直接操作：

thrust::device_ptr<int> d_id1(search_data.id1);
thrust::device_ptr<int> d_id2(search_data.id2);
size_t original_npair = search_data.npair;

2. 分情况处理去重

情况A：same_id_src=false（仅完全相同的ID对去重）

直接对<id1, id2>有序对排序后去重：

// 将两个设备数组打包为pair迭代器，避免额外内存分配
auto zip_begin = thrust::make_zip_iterator(thrust::make_tuple(d_id1, d_id2));
auto zip_end = zip_begin + original_npair;

// 排序（去重前必须排序，Thrust的unique只处理连续重复元素）
thrust::sort(zip_begin, zip_end);

// 去重，返回去重后的末尾迭代器
auto new_zip_end = thrust::unique(zip_begin, zip_end);
size_t new_npair = new_zip_end - zip_begin;

// 提取去重后的ID对到新的设备内存
thrust::device_vector<int> new_id1(new_npair);
thrust::device_vector<int> new_id2(new_npair);
thrust::copy(thrust::get<0>(zip_begin), thrust::get<0>(new_zip_end), new_id1.begin());
thrust::copy(thrust::get<1>(zip_begin), thrust::get<1>(new_zip_end), new_id2.begin());

// 更新结果结构
SearchDataOnDevice result;
result.npair = new_npair;
result.id1 = thrust::raw_pointer_cast(new_id1.data());
result.id2 = thrust::raw_pointer_cast(new_id2.data());

情况B：same_id_src=true（无序ID对视为重复）

先生成<min, max>键，再按键排序原始数据，最后去重：

// 打包原始ID对为zip迭代器
auto zip_begin = thrust::make_zip_iterator(thrust::make_tuple(d_id1, d_id2));
auto zip_end = zip_begin + original_npair;

// 生成去重键：每个ID对转换为<较小值, 较大值>
thrust::device_vector<thrust::pair<int, int>> keys(original_npair);
thrust::transform(zip_begin, zip_end, keys.begin(),
    [] __device__ (const thrust::tuple<int, int>& t) {
        int a = thrust::get<0>(t);
        int b = thrust::get<1>(t);
        return thrust::make_pair(min(a, b), max(a, b));
    });

// 按键排序原始ID对，确保重复键的元素连续
thrust::sort_by_key(keys.begin(), keys.end(), zip_begin);

// 按键去重，保留第一个出现的原始ID对（比如<0,5>会保留，<5,0>被移除）
auto [new_key_end, new_zip_end] = thrust::unique_by_key(keys.begin(), keys.end(), zip_begin);
size_t new_npair = new_key_end - keys.begin();

// 提取去重后的ID对
thrust::device_vector<int> new_id1(new_npair);
thrust::device_vector<int> new_id2(new_npair);
thrust::copy(thrust::get<0>(zip_begin), thrust::get<0>(zip_begin) + new_npair, new_id1.begin());
thrust::copy(thrust::get<1>(zip_begin), thrust::get<1>(zip_begin) + new_npair, new_id2.begin());

// 更新结果结构
SearchDataOnDevice result;
result.npair = new_npair;
result.id1 = thrust::raw_pointer_cast(new_id1.data());
result.id2 = thrust::raw_pointer_cast(new_id2.data());

注意事项 #

• 内存管理：thrust::device_vector会自动管理设备内存，若需手动释放，可调用cudaFree(result.id1)和cudaFree(result.id2)（注意不要释放原始的search_data.id1/id2，除非你不再需要）
• 错误检查：每次Thrust操作后可添加cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { /* 处理错误 */ }排查问题
• 性能：Thrust的算法基于CUDA高度优化，适合大规模数据处理，无需手动编写CUDA核函数

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