如何基于谓词对 warp 中的数据进行分区，以便所有保留项都是连续的

Question

我有一个充满数据的扭曲，其中一些我想保留，一些我想丢弃。
我想将保留项目存储在连续的内存中。

例如，假设我只想保留素数

input      0 1 2 3 4 5 6 7 8 9 10 11
keep prime     2 3   5   7        11
output     2 3 5 7 11

对于单个扭曲（32 项）数据执行此操作的最佳（最快/最简单）方法是什么？

Answer 1

我使用

lanesBefore

函数，它使用整数内在函数来生成不同场景的目的地：

__device__ uint32_t lanemask_lt() { //the other lanemask functions work in a similar way.
    uint32_t result;
    asm ("mov.u32 %0, %%lanemask_lt;" : "=r"(result));
    return result;
}

enum where { lt, le, eq, neq, gt, ge, biggest, smallest, nle, nlt, next };
template <where Where = lt>
__device__ int LanesBefore(const int Mask) {
    static_assert((Where >= lt) && (Where <= next));
    switch (Where) {
        case lt:       return __popc(Mask    & lanemask_lt());
        case le:       return __popc(Mask    & lanemask_le());
        case eq:       return bool(Mask    & lanemask_eq());
        case neq:      return bool((~Mask) & lanemask_eq());
        case gt:       return __popc(Mask    & lanemask_gt());
        case ge:       return __popc(Mask    & lanemask_ge());
        case nle:      return __popc((~Mask) & lanemask_le());
        case nlt:      return __popc((~Mask) & lanemask_lt());
        case biggest:  return 31 - __clz(Mask);
        case smallest: return __ffs(Mask) - 1;
        case next:     return __ffs(Mask & lanemask_ge()) -1;
        default: assert(false); return 0;
    }
}

lanesBefore

期望一个位掩码，其中包含用于保留项目的

1

和用于丢弃项目的

0

，并将从该掩码和

lanemask_lt()

生成目的地，其中laneid - 1个设置。例如：lanemask_lt(lane 7) = 0b1111111（七个 1）。 对于我们的示例数据：

Lane 7: data = 7, keepmask = 100010101100, lanemask_lt = 00001111111

keepmask:     100010101100
lanemask_lt   000001111111
and           ------------
popcount(     000000101100) destination = 3

pos:           0 1 2 3
desired output 2 3 5 7 <<-- yes `7` needs to be in position 3.

完整代码首先读入原始数据，然后使用谓词和

__ballot_sync

生成掩码。
该掩码通过

lanesBefore

函数运行，该函数计算我们之前保留项目的数量以生成目标索引。

const auto laneid = threadIdx.x % 32;
const auto data = numbers[laneid];
numbers[laneid] = 0; //optional: clear out non primes
const auto keep = isPrime(data);
const auto keepMask = __ballot_sync(Everyone, keep); 
const auto dest = lanesBefore<lt>(keepMask);
numbers[dest] = data; //put back primes
__syncwarp();
const auto validCount = popc(keepMask); 

因此，扭曲数据被压缩在几条指令中。