从一维数组表示计算三维索引的高效方法

Question

我有一份存储在1D数组中的3D数据。我通过以下方式计算1D索引：

index = i + j * WIDTH + k * WIDTH * HEIGHT

然后我需要从 index 中找回原始的 i,j,k 索引。显而易见的方法就是像这样：

k = index / (WIDTH * HEIGHT) 
j = (index % (WIDTH * HEIGHT)) / WIDTH
i = index - j * WIDTH - k * WIDTH * HEIGHT

但是我想知道，有没有更有效的方法来做到这一点？至少不用取模...

这个问题的背景是，我在CUDA中有一个内核，我访问数据并计算i、j、k索引（索引对应于唯一的线程ID）。也许有一些特定于CUDA的方法可以解决这个问题？我猜这是一个相当普遍的问题，但我找不到更好的方法来解决它...

感谢您的想法！

- Jaa-c

---

我认为没有办法将这些值取回，你自己提出的建议肯定是错误的。例如：i=2 j=3 k=4 WIDTH=100 HEIGHT=100。这将使索引=2+300+40000=40302。你说k=index/(WIDTH*HIEGHT)=40302/10000=4,0302。正如你所看到的，这与原始的k不同。 - Kevin

如果您期望i，j，k是整数，则它将始终有效。 - Jaa-c

你现在的代码已经很好了；如果你想避免使用模运算（因为在GPU上非常昂贵），你可以像对待i一样对待j：j = (index - (k*WIDTH*HEIGHT))/WIDTH。如果你想让它看起来更清晰，并且不需要原始索引，你可以这样做：k = index/(WIDTH*HEIGHT); index -= k*WIDTH*HEIGHT; j = index/WIDTH; index -= j*WIDTH; i = index。 - Jonathan Dursi

我会忽略关于将数组大小舍入为2的幂次方的建议。在这种情况下，虽然影响不算太大，但会增加3D数组的内存需求60％。而且由于GPU上的内存空间通常已经很紧张，将数组大小设置为2的幂次方可能会导致访问模式产生银行冲突问题。 - Jonathan Dursi

@JonathanDursi：这是一个很好的回答。我认为你应该将它添加为一个答案。在GPU上，由于硬件直接支持的三个维度通常不足，所以经常出现这种情况，需要将更多的维度打包到现有的维度中。每当我不得不使用模数时，我都感到非常难过，但我并没有想到有更好的方法。 - Roger Dahl

Answer 1

您现有的代码已经很好；如果您想避免使用模数（因为在GPU上非常耗费资源），您可以像处理i变量那样处理j变量：

j

j = (index - (k*WIDTH*HEIGHT))/WIDTH

如果你希望逻辑更加清晰，并且不需要原始的index，你可以这样做：

k = index/(WIDTH*HEIGHT); 
index -= k*WIDTH*HEIGHT; 

j = index/WIDTH; 
index -= j*WIDTH; 

i = index/1;

这个方法可以轻松地扩展到任意维度。你可以尝试通过预先计算 WIDTH * HEIGHT 来调整上述方法，但我建议只需增加优化程度，并相信编译器会为您完成。

关于将数字舍入至2的幂次方的建议是正确的，因为这会加快索引计算的速度，但代价很高。在这种情况下（WIDTH=HEIGHT=100），它会将您的 3D 数组的内存需求增加 60％ （WIDTH=HEIGHT=128），而 GPU 上的内存通常已经很紧张了；并且让您的数组大小成为2的幂次方可能会根据您的访问模式导致银行冲突问题。

- Jonathan Dursi

Answer 2

尝试将您的尺寸向上舍入到下一个二次幂。然后您可以使用位移和掩码来代替乘法、除法和模运算。

index = i | (j | k << HEIGHT_BITS) << WIDTH_BITS;

k = index >> (WIDTH_BITS + HEIGHT_BITS);
j = (index >> WIDTH_BITS) & ((1 << HEIGHT_BITS) - 1);
i = index & ((1 << WIDTH_BITS) - 1);

- MvG

Answer 3

仅适用于维度为2的幂次方的情况。使用位掩码。例如，如果第一个索引的最大值为4，则应在索引中使用前两个位。

- Leonid Volnitsky