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

Question

我有一个元素的 const 推力向量，我想从中提取最多 N 个传递谓词的元素（在任何顺序），其中推力矢量大小和 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 操作。因此，即使我们在这里“公开”了大部分分配，thrust 显然仍然需要执行单个临时分配（大小未知）来支持扫描操作。

正在回答下面评论中的问题。要理解这个：0+(_1>0)，有两点需要注意：

一般语法是使用thrust::placeholders。 Thrust 的这种能力允许我们编写简单的一元或二元内联函数，避免使用 lambda 或编写单独的仿函数。
0+的原因如下。如果我们简单地使用 (_1>0)，那么 thrust 将使用一个布尔测试作为其一元函数，该测试通过取消引用迭代器返回的项目与零进行比较。该比较的结果是一个布尔值，如果我们保留它，前缀和最终将使用我们不想要的布尔运算来计算。我们希望将布尔大于测试（即 true/false）的结果转换为整数，以便使用整数算术执行后续前缀和。在 (_1>0) 布尔测试前加上 0+ 可以实现这一点。

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

Thrust: Stream compaction copying only first N valid elements

cuda

thrust