推力：流压缩仅复制前 N 个有效元素

Question

我有一个 const 推力向量，我想从中提取 最多 N 个传递谓词（以任何顺序）的元素，其中推力向量size 和 N 在编译时是已知的。在我的具体情况下，我的向量是 500k 元素，而 N 是 100k。

我最初的想法是使用thrust::copy_if 来获取通过谓词的所有 个元素，然后只使用前N 个元素进行后续计算。但是，在这种情况下，我必须分配两个包含 500k 元素的向量（一个用于初始向量，一个用于copy_if 的输出），并且我必须处理每个元素。

由于这是我必须多次跨多个 CUDA 流执行的操作，我想知道是否有一种方法可以获得 N 个输出元素，同时最大限度地减少所需的内存占用，理想情况下，尽量减少需要处理的元素数量（即，一旦找到 N 个有效元素就中断处理）。

Answer 1

One possible method 执行流压缩操作是执行谓词前缀和，然后执行条件索引副本。通过将“单片”操作分解为这两个部分，在输出大小上插入所需的限制行为变得相当容易。

前缀和是一个相当复杂的操作。我们将为此使用推力。条件索引副本相当简单，因此我们将为此编写我们自己的 CUDA 内核，而不是尝试与thrust::copy_if 操作搏斗来获得恰到好处的复制逻辑。这个内核是我们将在输出大小上插入限制行为的地方。

这是一个有效的例子：

$ cat t34.cu
#include <thrust/scan.h>
#include <thrust/copy.h>
#include <thrust/device_vector.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/iterator/counting_iterator.h>
#include <iostream>

using namespace thrust::placeholders;

typedef int mt;

__global__ void my_copy(mt *d, int *i, mt *r, int limit, int size){
  int idx = threadIdx.x+blockDim.x*blockIdx.x;
  if (idx < size){
    if ((idx == 0) && (*i == 1) && (limit > 0))
      *r = *d;
    else if ((idx > 0) && (i[idx] > i[idx-1]) && (i[idx] <= limit)){
      r[i[idx]-1] = d[idx];}
  }
}

int main(){
  int rs = 3;
  mt d[] = {0, 1, 0, 2, 0, 3, 0, 4, 0, 5};
  int ds = sizeof(d)/sizeof(d[0]);
  thrust::device_vector<mt> data(d, d+ds);
  thrust::device_vector<int> idx(ds);
  thrust::device_vector<mt> result(rs);
  auto my_cmp = thrust::make_transform_iterator(data.begin(), 0+(_1>0));
  thrust::inclusive_scan(my_cmp, my_cmp+ds, idx.begin());
  my_copy<<<(ds+255)/256, 256>>>(thrust::raw_pointer_cast(data.data()), thrust::raw_pointer_cast(idx.data()), thrust::raw_pointer_cast(result.data()), rs, ds);
  thrust::host_vector<mt> h_result = result;
  thrust::copy_n(h_result.begin(), rs, std::ostream_iterator<mt>(std::cout, ","));
  std::cout << std::endl;
}
$ nvcc -std=c++14 -o t34 t34.cu -arch=sm_52
$ ./t34
1,2,3,
$

（CUDA 11.0、Fedora 29、GTX 960）

请注意，此代码仅用于演示目的。您不应假定它没有缺陷或适用于任何特定目的。使用它需要您自担风险。

使用分析器进行一些研究会表明，thrust::inclusive_scan 操作确实会“在后台”执行cudaMalloc 和cudaFree 操作。因此，即使我们在这里“公开”了大部分分配，推力显然仍需要执行单个临时分配（大小未知）以支持扫描操作。

回答以下 cmets 中的问题。要理解这一点：0+(_1>0)，有两点需要注意：

1. 一般语法使用thrust::placeholders。这种推力允许我们内联编写简单的一元或二元函数，避免使用 lambda 或编写单独的函子。

2. 0+ 的原因如下。如果我们简单地使用(_1>0)，那么thrust 将使用它作为它的一元函数，对通过取消引用迭代器返回的项目进行布尔测试，比较为零。比较的结果是一个布尔值，如果我们保持这种状态，前缀和最终将使用布尔算术计算，这是我们不想要的。我们希望将布尔大于测试的结果（即真/假）转换为整数，以便使用整数算术执行后续前缀和。在 (_1>0) 布尔测试前加上 0+ 可以实现这一点。