CUDA thrust：如何实现支持“stencil”的“partition”？

Question

假设有一个整数数组：

A[]={2, 2, 9, 8, 5, 7, 0, 6}

还有一个模板:

B[]={1, 0, 0, 1, 1, 1, 0, 1}

我的问题是怎么样才能根据 B[] 重新排列 A[]，使得如果 B[i]==1, B[j]==0，那么在新的数组中可以保证 A[i] 在 A[j] 之前，新数组应该长成这个样子：

C[]={2, 8, 5, 7, 6, 2, 9, 0}

PS：我发现“partition”函数几乎是答案，只是它只支持谓词。有什么解决方法吗？

非常感谢任何提示！

- biubiuty

---

这不就是按键进行稳定排序吗？ - talonmies

一开始我并没有意识到那点。非常感谢@talonmies! - biubiuty

似乎partition-with-stencil是API中已知的漏洞：https://github.com/thrust/thrust/issues/49 - Jared Hoberock

Answer 1

现在，使用stencil实现了thrust::partition和thrust::stable_partition（可能需要从official Thrust repository获取源代码），可以通过以下方式实现：

#include <thrust/partition.h>

struct is_one
{
  __host__ __device__
  bool operator()(const int &x)
  {
    return x == 1;
  }
};

// Partition values on device thanks to stencil
thrust::stable_partition(d_A.begin(),
                         d_A.end(),
                         d_B.begin(),
                         is_one());

这导致：

    A =   0  1  2  3  4  5  6  7  8  9
    B =   0  1  1  0  0  1  0  0  1  0
    C =   1  2  5  8  0  3  4  6  7  9

此实现更有效率，因为我们没有对两个分区中的值进行排序。类似且更复杂的示例可在此处（答案中提供了更多细节）找到。

- BenC

Answer 2

可以使用 thrust::stable_sort_by_key() 来实现。

- harrism