c++ - 有效评估CUDA中的索引映射功能

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

映射如下所示:



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

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

密钥的数量不限于3(这只是一个示例!)，并且仅在运行时知道，每个密钥的相应“宽度”也是如此。

如何有效地在CUDA内核中找到相应的密钥？
我想到了以下两种选择:

在内核内部使用二进制搜索(内存高效)

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

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

最佳答案

还有另一种算法可以使您在内存和运行时效率方面达到中位数:
假设线程总数为N。让我们以接近M的数字sqrt(N)并将每个线程按M线程分别划分为组(最后一个不完整)。现在，仅针对每个组中的第一个线程预先计算密钥(它们的idx将为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]进行搜索。