CUDA Thrust：如何结合copy_if与transform且不物化数据？

现有两个输入：一个thrust设备端数组和一个可在内核访问的位图：

thrust::device_vector<point_t> points;
Bitset bits; // 内核中可访问该对象

需要把bits.is_active(i)返回true的points[i]复制到输出向量：

thrust::device_vector<point_t> output;

逻辑等价于串行循环：

for(size_t i = 0; i < points.size(); i++) {
    if (bits.is_active(i)) {
        output.push_back(points[i]);
    }
}

尝试用copy_if、transform_iterator、zip_iterator组合实现，但遇到问题：copy_if的谓词需要访问原索引，而当前代码的转换迭代器已经把索引和点的元组转换成了单独的点，导致谓词无法获取索引。以下是尝试的错误代码：

auto get_point =
        [] __device__(const thrust::tuple<size_t, point_t>& t) {
          return thrust::get<1>(t);
        };

    auto it1 = thrust::make_transform_iterator(
        thrust::make_zip_iterator(thrust::make_tuple(
            thrust::counting_iterator<size_t>(0), points.begin())),
        get_point);
    auto it2 = thrust::make_transform_iterator(
        thrust::make_zip_iterator(thrust::make_tuple(
            thrust::counting_iterator<size_t>(points.size()), points.end())),
        get_point);

    thrust::copy_if(
        thrust::cuda::par.on(stream), it1, it2,
        output.begin(),
        [=] __device__(const thrust::tuple<size_t, point_t>& t) {
          auto index = thrust::get<0>(t);

          return bits.is_active(index);
        });

你的代码核心问题是：提前用transform_iterator把索引和点的元组转换成了单独的point_t，导致copy_if的谓词只能拿到点，拿不到原索引。

正确的思路是：先带着索引和点的元组做筛选，筛选完成后再把点提取出来写入输出向量。

方法一：使用带转换的输出迭代器 #

直接构造包含索引和点的zip_iterator，用transform_output_iterator在输出阶段提取点，谓词直接用元组里的索引判断：

// 构造包含索引与点的zip迭代器
auto zip_begin = thrust::make_zip_iterator(
    thrust::make_tuple(
        thrust::counting_iterator<size_t>(0),
        points.begin()
    )
);
auto zip_end = zip_begin + points.size();

// 先计算有效元素数量，避免output越界
size_t valid_count = thrust::count_if(
    thrust::cuda::par.on(stream),
    thrust::counting_iterator<size_t>(0),
    thrust::counting_iterator<size_t>(points.size()),
    [=] __device__(size_t idx) {
        return bits.is_active(idx);
    }
);
output.resize(valid_count);

// 执行筛选并输出
thrust::copy_if(
    thrust::cuda::par.on(stream),
    zip_begin, zip_end,
    thrust::make_transform_output_iterator(
        output.begin(),
        [] __device__(const thrust::tuple<size_t, point_t>& t) {
            return thrust::get<1>(t); // 从元组中提取点
        }
    ),
    [=] __device__(const thrust::tuple<size_t, point_t>& t) {
        return bits.is_active(thrust::get<0>(t)); // 用索引判断有效性
    }
);

方法二：先筛选元组再转换 #

如果需要分开步骤，也可以先把符合条件的元组复制到临时容器，再提取点，但效率略低：

thrust::device_vector<thrust::tuple<size_t, point_t>> temp;
temp.reserve(points.size());

// 筛选符合条件的元组
thrust::copy_if(
    thrust::cuda::par.on(stream),
    zip_begin, zip_end,
    std::back_inserter(temp),
    [=] __device__(const thrust::tuple<size_t, point_t>& t) {
        return bits.is_active(thrust::get<0>(t));
    }
);

// 从元组提取点到output
output.resize(temp.size());
thrust::transform(
    thrust::cuda::par.on(stream),
    temp.begin(), temp.end(),
    output.begin(),
    [] __device__(const thrust::tuple<size_t, point_t>& t) {
        return thrust::get<1>(t);
    }
);

关键注意事项 #

• 必须提前计算有效元素数量并resizeoutput，否则会因为输出迭代器越界导致未定义行为；
• 确保Bitset对象在设备端可访问（比如是设备端分配的，或者通过thrust的设备可访问容器实现）；
• 使用thrust::cuda::par.on(stream)可以指定CUDA流，提升异步执行的灵活性。

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