c++ - 在 CUDA 中高效评估索引映射函数

Question

在 CUDA 内核中，我需要找到一个映射到 threadIdx 的键。

映射可以是这样的:

key -> threadIdx

0 -> {0,1,2,3,4}

1 -> {5,6,7}

2 -> {8,9,10}

...

每个键 k_i 都映射到 n_i (可变的，任意 n_i 和 n_i>0 )线程。该键将用于检索全局数组中的相应值。然后，该值将用于该内核中的后续计算。

映射可以绘制为分段常数函数:

键的数量不限于 3 个(这只是一个例子!)并且只有在运行时才知道，以及每个键各自的“宽度”。

我怎样才能高效找到CUDA内核中对应的key？我想到了以下两种选择:

1. 在内核中使用二进制搜索(内存高效)

2. 为每个 threadIDx 预先计算映射，然后启动内核(运行时高效)

   0 0 0 0 0 1 1 1 2 2 2 ...

有没有更好的方法来实现这一目标？

Answer 1

还有另一种算法可以让您在内存和运行时效率方面处于中等水平:假设线程总数为 N。让我们取接近 sqrt(N) 的编号 M 并将所有线程按 M 线程分组(最后一个将是不完整的) ).现在，仅为每个组中的第一个线程预先计算 key (它们的 idxes 将为 0、M、2M 等等)。它为我们提供了 O(sqrt(N)) 内存渐近。现在，在内核中，我们可以轻松找到当前组 (groupIdx = threadIdx/M) 和下一组 (groupIdx + 1) 的索引。对于它们中的每一个，我们知道预先计算的键 key[groupIdx] 和 key[groupIdx + 1]。现在你可以做 BS，但是拿 [key[groupIdx]; key[groupIdx + 1]] 段用于搜索而不是 [1; MAX_KEY_VALUES]。