jax.experimental.pallas.mosaic_gpu.planar_snake

jax.experimental.pallas.mosaic_gpu.planar_snake(lin_idx, shape, minor_dim, tile_width)

将线性索引转换为 shape 中的索引，尝试优化局部性。

此函数计算的“空间填充曲线”将次要维度（minor dimension）分割成长度为 tile_width 的块。每隔一个块，其主要维度（major dimension）都会反转，这样当从一个块移动到另一个块时，迭代顺序会“蛇形”移动。

对于 shape 为 (8, 8)，minor_dim=0 和 tile_width=2，迭代顺序如下：

 0   2   4   6   8  10  12  14
 1   3   5   7   9  11  13  15
30  28  26  24  22  20  18  16
31  29  27  25  23  21  19  17
32  34  36  38  40  42  44  46
33  35  37  39  41  43  45  47
62  60  58  56  54  52  50  48
63  61  59  57  55  53  51  49

请注意，每两行构成一个块（minor_dim=0，tile_width=2），并且在从一个块移动到另一个块时，其中一个块的列索引会增加，而另一个块的列索引会减少。

参数:

• lin_idx (Array)

• shape (tuple[int, int])

• minor_dim (int)

• tile_width (int)