网格和 BlockSpecs#

grid,也称为循环中的内核#

当使用 jax.experimental.pallas.pallas_call() 时,内核函数在不同输入上执行多次,如 pallas_callgrid 参数所指定。概念上

pl.pallas_call(some_kernel, grid=(n,))(...)

映射到

for i in range(n):
  some_kernel(...)

网格可以推广到多维,对应于嵌套循环。例如,

pl.pallas_call(some_kernel, grid=(n, m))(...)

等效于

for i in range(n):
  for j in range(m):
    some_kernel(...)

这可以推广到任何整数元组(长度为 d 的网格对应于 d 个嵌套循环)。内核执行的次数与 prod(grid) 相同。默认网格值为 (),导致一次内核调用。每个调用被称为“程序”。要访问内核当前正在执行的程序(即网格的哪个元素),我们使用 jax.experimental.pallas.program_id()。例如,对于调用 (1, 2)program_id(axis=0) 返回 1,而 program_id(axis=1) 返回 2。您还可以使用 jax.experimental.pallas.num_programs() 获取给定轴的网格大小。

以下是一个使用 gridprogram_id 的内核示例。

>>> import jax
>>> from jax.experimental import pallas as pl

>>> def iota_kernel(o_ref):
...   i = pl.program_id(0)
...   o_ref[i] = i

现在我们使用 pallas_call 执行它,并添加 grid 参数。

>>> def iota(size: int):
...   return pl.pallas_call(iota_kernel,
...                         out_shape=jax.ShapeDtypeStruct((size,), jnp.int32),
...                         grid=(size,), interpret=True)()
>>> iota(8)
Array([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32)

在 GPU 上,每个程序在独立的线程块上并行执行。因此,我们需要考虑对 HBM 写入的竞争条件。一个合理的方法是编写我们的内核,使不同的程序写入 HBM 中的不同位置,以避免这些并行写入。

在 TPU 上,程序以并行和顺序的组合执行(取决于体系结构),因此有一些略微不同的考虑因素。参见 Pallas TPU 文档

BlockSpec,也称为如何将输入分成块 #

结合 grid 参数,我们需要向 Pallas 提供有关如何切片每个调用的输入的信息。具体来说,我们需要提供一个映射,将循环的迭代映射到要操作的输入和输出的块。这是通过 jax.experimental.pallas.BlockSpec 对象提供的。

在深入研究 BlockSpec 的细节之前,您可能需要重新访问 Pallas 快速入门 BlockSpecs 示例

BlockSpec 通过 in_specsout_specs 提供给 pallas_call,每个输入和输出各一个。

首先,我们讨论当 indexing_mode == pl.Blocked()BlockSpec 的语义。

非正式地说,BlockSpecindex_map 以调用索引(与 grid 元组的长度相同)作为参数,并返回块索引(每个整体数组轴一个块索引)。每个块索引然后乘以来自 block_shape 的相应轴大小,以获得对应数组轴上的实际元素索引。

注意

并非所有块形状都受支持。

  • 在 TPU 上,仅支持等级至少为 1 的块。此外,块形状的最后两个维度必须等于整体数组的相应维度,或者分别被 8 和 128 整除。对于等级为 1 的块,块维度必须等于数组维度,或者被 128 * (32 / bitwidth(dtype)) 整除。

  • 在 GPU 上,块本身的大小不受限制,但每个操作必须对大小为 2 的幂的数组进行操作。

如果块形状不能均匀地划分整体形状,则每个轴上的最后一次迭代仍将接收对 block_shape 块的引用,但超出范围的元素将在输入上填充并在输出上丢弃。填充值的含义未定义,您应该假设它们是垃圾。在 interpret=True 模式下,我们使用 NaN 填充浮点值,让用户有机会发现访问超出范围的元素,但这项行为不应依赖于此。请注意,每个块中至少有一个元素必须在范围内。

更准确地说,形状为 x_shape 的输入 x 的每个轴的切片计算如下面的 slice_for_invocation 函数所示

>>> def slices_for_invocation(x_shape: tuple[int, ...],
...                           x_spec: pl.BlockSpec,
...                           grid: tuple[int, ...],
...                           invocation_indices: tuple[int, ...]) -> tuple[slice, ...]:
...   assert len(invocation_indices) == len(grid)
...   assert all(0 <= i < grid_size for i, grid_size in zip(invocation_indices, grid))
...   block_indices = x_spec.index_map(*invocation_indices)
...   assert len(x_shape) == len(x_spec.block_shape) == len(block_indices)
...   elem_indices = []
...   for x_size, block_size, block_idx in zip(x_shape, x_spec.block_shape, block_indices):
...     start_idx = block_idx * block_size
...     # At least one element of the block must be within bounds
...     assert start_idx < x_size
...     elem_indices.append(slice(start_idx, start_idx + block_size))
...   return elem_indices

例如

>>> slices_for_invocation(x_shape=(100, 100),
...                       x_spec = pl.BlockSpec((10, 20), lambda i, j: (i, j)),
...                       grid = (10, 5),
...                       invocation_indices = (2, 4))
[slice(20, 30, None), slice(80, 100, None)]

>>> # Same shape of the array and blocks, but we iterate over each block 4 times
>>> slices_for_invocation(x_shape=(100, 100),
...                       x_spec = pl.BlockSpec((10, 20), lambda i, j, k: (i, j)),
...                       grid = (10, 5, 4),
...                       invocation_indices = (2, 4, 0))
[slice(20, 30, None), slice(80, 100, None)]

>>> # An example when the block is partially out-of-bounds in the 2nd axis.
>>> slices_for_invocation(x_shape=(100, 90),
...                       x_spec = pl.BlockSpec((10, 20), lambda i, j: (i, j)),
...                       grid = (10, 5),
...                       invocation_indices = (2, 4))
[slice(20, 30, None), slice(80, 100, None)]

以下定义的 show_program_ids 函数使用 Pallas 显示调用索引。 iota_2D_kernel 将用十进制数字填充每个输出块,其中第一位数字表示第一个轴上的调用索引,第二位数字表示第二个轴上的调用索引

>>> def show_program_ids(x_shape, block_shape, grid,
...                      index_map=lambda i, j: (i, j),
...                      indexing_mode=pl.Blocked()):
...   def program_ids_kernel(o_ref):  # Fill the output block with 10*program_id(1) + program_id(0)
...     axes = 0
...     for axis in range(len(grid)):
...       axes += pl.program_id(axis) * 10**(len(grid) - 1 - axis)
...     o_ref[...] = jnp.full(o_ref.shape, axes)
...   res = pl.pallas_call(program_ids_kernel,
...                        out_shape=jax.ShapeDtypeStruct(x_shape, dtype=np.int32),
...                        grid=grid,
...                        in_specs=[],
...                        out_specs=pl.BlockSpec(block_shape, index_map, indexing_mode=indexing_mode),
...                        interpret=True)()
...   print(res)

例如

>>> show_program_ids(x_shape=(8, 6), block_shape=(2, 3), grid=(4, 2),
...                  index_map=lambda i, j: (i, j))
[[ 0  0  0  1  1  1]
 [ 0  0  0  1  1  1]
 [10 10 10 11 11 11]
 [10 10 10 11 11 11]
 [20 20 20 21 21 21]
 [20 20 20 21 21 21]
 [30 30 30 31 31 31]
 [30 30 30 31 31 31]]

>>> # An example with out-of-bounds accesses
>>> show_program_ids(x_shape=(7, 5), block_shape=(2, 3), grid=(4, 2),
...                  index_map=lambda i, j: (i, j))
[[ 0  0  0  1  1]
 [ 0  0  0  1  1]
 [10 10 10 11 11]
 [10 10 10 11 11]
 [20 20 20 21 21]
 [20 20 20 21 21]
 [30 30 30 31 31]]

>>> # It is allowed for the shape to be smaller than block_shape
>>> show_program_ids(x_shape=(1, 2), block_shape=(2, 3), grid=(1, 1),
...                  index_map=lambda i, j: (i, j))
[[0 0]]

当多个调用写入输出数组的相同元素时,结果将依赖于平台。

在以下示例中,我们有一个 3D 网格,最后一个网格维度未使用在块选择中 (index_map=lambda i, j, k: (i, j))。因此,我们对同一个输出块迭代 10 次。以下显示的输出是在 CPU 上使用 interpret=True 模式生成的,该模式目前按顺序执行调用。在 TPU 上,程序以并行和顺序的组合执行,并且此函数生成显示的输出。参见 Pallas TPU 文档

>>> show_program_ids(x_shape=(8, 6), block_shape=(2, 3), grid=(4, 2, 10),
...                  index_map=lambda i, j, k: (i, j))
[[  9   9   9  19  19  19]
 [  9   9   9  19  19  19]
 [109 109 109 119 119 119]
 [109 109 109 119 119 119]
 [209 209 209 219 219 219]
 [209 209 209 219 219 219]
 [309 309 309 319 319 319]
 [309 309 309 319 319 319]]

block_shape 中作为维度值出现的 None 值的行为类似于值 1,除了相应的块轴被压缩。在以下示例中,观察到 o_ref 的形状为 (2,),而块形状被指定为 (None, 2)(前导维度被压缩)。

>>> def kernel(o_ref):
...   assert o_ref.shape == (2,)
...   o_ref[...] = jnp.full((2,), 10 * pl.program_id(1) + pl.program_id(0))
>>> pl.pallas_call(kernel,
...                jax.ShapeDtypeStruct((3, 4), dtype=np.int32),
...                out_specs=pl.BlockSpec((None, 2), lambda i, j: (i, j)),
...                grid=(3, 2), interpret=True)()
Array([[ 0,  0, 10, 10],
       [ 1,  1, 11, 11],
       [ 2,  2, 12, 12]], dtype=int32)

当我们构造一个 BlockSpec 时,我们可以对 block_shape 参数使用值 None,在这种情况下,整体数组的形状将用作 block_shape。如果对 index_map 参数使用值 None,则使用一个返回零元组的默认索引映射函数: index_map=lambda *invocation_indices: (0,) * len(block_shape)

>>> show_program_ids(x_shape=(4, 4), block_shape=None, grid=(2, 3),
...                  index_map=None)
[[12 12 12 12]
 [12 12 12 12]
 [12 12 12 12]
 [12 12 12 12]]

>>> show_program_ids(x_shape=(4, 4), block_shape=(4, 4), grid=(2, 3),
...                  index_map=None)
[[12 12 12 12]
 [12 12 12 12]
 [12 12 12 12]
 [12 12 12 12]]

“未阻塞”索引模式 #

上面记录的行为适用于 indexing_mode=pl.Blocked()。当使用 pl.Unblocked 索引模式时,索引映射函数返回的值将直接用作数组索引,而不会先按块大小缩放它们。当使用未阻塞模式时,您可以指定数组的虚拟填充,作为每个维度的低-高填充元组:行为就像整体数组在输入时填充一样。对于未阻塞模式的填充值,不保证其值,类似于当块形状不能划分整体数组形状时,阻塞索引模式的填充值。

未阻塞模式目前仅在 TPU 上受支持。

>>> # unblocked without padding
>>> show_program_ids(x_shape=(8, 6), block_shape=(2, 3), grid=(4, 2),
...                  index_map=lambda i, j: (2*i, 3*j),
...                  indexing_mode=pl.Unblocked())
    [[ 0  0  0  1  1  1]
     [ 0  0  0  1  1  1]
     [10 10 10 11 11 11]
     [10 10 10 11 11 11]
     [20 20 20 21 21 21]
     [20 20 20 21 21 21]
     [30 30 30 31 31 31]
     [30 30 30 31 31 31]]

>>> # unblocked, first pad the array with 1 row and 2 columns.
>>> show_program_ids(x_shape=(7, 7), block_shape=(2, 3), grid=(4, 3),
...                  index_map=lambda i, j: (2*i, 3*j),
...                  indexing_mode=pl.Unblocked(((1, 0), (2, 0))))
    [[ 0  1  1  1  2  2  2]
     [10 11 11 11 12 12 12]
     [10 11 11 11 12 12 12]
     [20 21 21 21 22 22 22]
     [20 21 21 21 22 22 22]
     [30 31 31 31 32 32 32]
     [30 31 31 31 32 32 32]]