如何使用Thrust基于另一向量的谓词条件实现选择性元素复制

核心实现逻辑 #

你描述的按掩码向量筛选元素的场景，Thrust 已经内置了原生支持的thrust::copy_if接口，不需要手动实现遍历逻辑，两种常用实现方式如下：

实现方式1：自动分配结果空间（写法最简单） #

不需要提前计算结果向量C的长度，用thrust::back_inserter自动扩容，适合小数据量快速开发场景。

示例代码 #

#include <thrust/device_vector.h>
#include <thrust/copy.h>
#include <thrust/functional.h> // 用于thrust::identity

int main() {
    // 初始化输入向量
    thrust::device_vector<int> A = {2, 3, 6, 0, 11};
    thrust::device_vector<int> B = {1, 0, 1, 1, 0};
    thrust::device_vector<int> C;

    // 执行选择性复制：B对应位置为1时复制A的元素到C
    thrust::copy_if(
        A.begin(), A.end(), // 待筛选的A向量迭代区间
        B.begin(), // 掩码向量B的起始迭代器
        thrust::back_inserter(C), // 结果写入迭代器，自动扩容
        thrust::identity<int>() // 谓词规则：B的元素非0即判定为真，执行复制
    );

    // 此时C的内容为 [2,6,0]，符合要求
    return 0;
}

实现方式2：预分配结果空间（性能更优） #

大数据量场景下可以先对掩码向量B做归约计算得到C的长度，提前分配空间，避免自动扩容的性能开销：

示例代码 #

#include <thrust/device_vector.h>
#include <thrust/copy.h>
#include <thrust/reduce.h>
#include <thrust/functional.h>

int main() {
    thrust::device_vector<int> A = {2, 3, 6, 0, 11};
    thrust::device_vector<int> B = {1, 0, 1, 1, 0};

    // 预计算C的长度：统计B中1的总数
    int c_size = thrust::reduce(B.begin(), B.end(), 0);
    thrust::device_vector<int> C(c_size);

    // 直接写入预分配的C空间
    thrust::copy_if(
        A.begin(), A.end(),
        B.begin(),
        C.begin(),
        thrust::identity<int>()
    );

    return 0;
}

自定义谓词拓展 #

如果你的筛选规则不是B等于1，比如需要B大于5才复制，只需要替换最后一个谓词参数即可，示例：

thrust::copy_if(A.begin(), A.end(), B.begin(), thrust::back_inserter(C), 
    [] __device__ (int b_val) {
        return b_val > 5; // 自定义筛选规则
    }
);

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