能否通过CUB::DeviceSelect同时返回数组选中与未选中部分？

我想通过CUB::DeviceSelect同时返回数组的选中部分和未选中部分，目前想到的办法是用相反的筛选逻辑再调用一次CUB::DeviceSelect，但想知道有没有更高效的实现方式。

示例代码（原需求）：

struct Positive {
    __host__ __device__ __forceinline__
    bool operator()(const int a) const {
        return (a > 0);
    }
};

int len = 6;
int *d_in = {0, 2, -2, -1, 1, 3};
int *d_out;
int *d_num_select;
void *temp = NULL;
size_t temp_size = 0;

// 期望结果：d_out = {2, 1, 3, 0, -2, -1} 且 *d_num_select = 3
// 但当前调用仅得到 d_out = {2, 1, 3}
cub::DeviceSelect::If(temp, temp_size, d_in, d_out, d_num_select, len, Positive());
cudaMalloc(&temp, temp_size);
cub::DeviceSelect::If(temp, temp_size, d_in, d_out, d_num_select, len, Positive());

更高效的实现方案 #

1. 使用cub::DevicePartition替代两次DeviceSelect #

CUB提供的DevicePartition接口可以一次性将输入数组划分为满足条件和不满足条件的连续两段，完全匹配你的需求。相比两次调用DeviceSelect，它只需要遍历输入数组一次，省去了额外的内核调用和内存读取开销，效率更高。

优化后的示例代码：

#include <cub/cub.cuh>

struct Positive {
    __host__ __device__ __forceinline__
    bool operator()(const int a) const {
        return (a > 0);
    }
};

int main() {
    const int len = 6;
    int h_in[] = {0, 2, -2, -1, 1, 3};
    int *d_in, *d_out, *d_num_selected;
    void *temp = nullptr;
    size_t temp_size = 0;

    // 分配设备内存
    cudaMalloc(&d_in, len * sizeof(int));
    cudaMalloc(&d_out, len * sizeof(int));
    cudaMalloc(&d_num_selected, sizeof(int));
    cudaMemcpy(d_in, h_in, len * sizeof(int), cudaMemcpyHostToDevice);

    // 1. 查询所需临时内存大小
    cub::DevicePartition::If(temp, temp_size, d_in, d_out, d_num_selected, len, Positive());
    // 2. 分配临时内存
    cudaMalloc(&temp, temp_size);
    // 3. 执行分区操作
    cub::DevicePartition::If(temp, temp_size, d_in, d_out, d_num_selected, len, Positive());

    // 验证结果
    int h_num_selected;
    int h_out[len];
    cudaMemcpy(&h_num_selected, d_num_selected, sizeof(int), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_out, d_out, len * sizeof(int), cudaMemcpyDeviceToHost);

    // 此时 h_num_selected = 3，h_out = {2, 1, 3, 0, -2, -1}，符合预期

    // 释放资源
    cudaFree(d_in);
    cudaFree(d_out);
    cudaFree(d_num_selected);
    cudaFree(temp);

    return 0;
}

2. 效率优势说明 #

• 单次内核调用：DevicePartition仅启动一个内核，遍历一次输入数组即可完成筛选和分区，避免了两次DeviceSelect带来的两次内核启动与内存遍历。
• 更低的内存开销：两次DeviceSelect需要分别申请临时内存，而DevicePartition只需要一份临时内存，减少了内存分配的资源占用和开销。

注意事项 #

• 确保使用的CUB版本在1.8.0及以上（该版本开始稳定支持DevicePartition）。
• 输出数组d_out的容量需等于输入数组长度，以容纳所有选中和未选中元素。

内容的提问来源于stack exchange，提问作者Nicolás Tsu