Pallas Core 特定编程

Pallas 特定核心编程#

在本指南中，我们将探讨如何使用 pl.core_map 编写 Pallas 内核。与 pallas_call 相比，core_map 具有几个关键特性：

  • 单核心编程：你编写的是针对 TPU/GPU 单个核心的代码，而不是针对整个 JAX 设备。这让你能够完全控制在每个核心上运行的内容，以及核心之间如何通信和分配工作。

  • 集合通信 (Collectives)：core_map 显式建模了物理核心，因此可以安全地表达核心间的通信。

  • 平台通用：core_map 编程模型适用于 TPU（TensorCore 和 SparseCore）以及 GPU，只需极少的样板代码更改。

本指南侧重于 TPU。关于如何在 GPU 上使用 core_map 以获得更高的线程灵活性，请查看我们的 Pallas GPU core_map 教程。

环境设置#

现代加速器通常在单个设备下拥有多个核心。对于较新的 TPU 芯片（v4, v5p），每个 JAX 设备可能包含 2 个 TensorCore（即 Megacore）。一些 TPU（v5p, v6e, 7x）还包含 SparseCores，每个 SparseCore 由许多子核心组成。

本指南基于 v5p 芯片编写，该芯片包含 4 个设备（每个设备有 2 个 TensorCore）和 4 个 SparseCore，每个 SparseCore 有 16 个子核心。

from functools import partial

import jax
from jax.sharding import NamedSharding
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu
from jax.experimental.pallas import tpu_sc as plsc
import jax.numpy as jnp
import numpy as np


num_devices = jax.local_device_count()
assert num_devices > 1, "Please run this notebook with more than one device."

tpu_info = pltpu.get_tpu_info()  # This notebook only runs on TPU.
print(f"Running on {num_devices} TPU {tpu_info.chip_version} devices.")

Running on 4 TPU v5p devices.

除了典型的 TPU 设备网格之外，你需要创建一个核心网格。可以将其视为一个名为 core 的附加维度（长度为 2），叠加在你使用的 4 设备网格之上。总计即 8 个核心。 

# Mesh of devices
mesh = jax.make_mesh((jax.device_count(),), ('device',))
print(mesh)

# Mesh of cores, within a JAX device
tc_mesh = pltpu.create_tensorcore_mesh('core')
print(tc_mesh)

num_devices = mesh.size
num_cores = len(tc_mesh.devices)
print(f"There are {num_devices} devices, and {num_cores} cores each.")

Mesh('device': 4, axis_types=(Explicit,))
TensorCoreMesh(devices=array([TensorCore(id=0), TensorCore(id=1)], dtype=object), axis_names=('core',))
There are 4 devices, and 2 cores each.

简单的单核心内核#

pl.core_map 允许你编写单核心本地代码，就像 jax.shard_map 允许你编写单设备代码一样。

在下面的内核示例中，每个核心拥有自己的 VMEM 和信号量分配。与常规内核一样，你可以使用 pltpu.async_copy 发起 HBM 和 VMEM 引用之间的拷贝。

核心间通信

在核心间通信之前，建议执行一个屏障操作（使用 pl.semaphore_signal），以确保资源已分配，且两个核心在程序运行期间处于同步点。

一旦核心同步，使用 pltpu.make_async_remote_copy 在它们之间发送数据。device_id 关键字参数通用地允许发送到任何设备上的任何核心，但如果你只传入 {'core': other_core_id}，它将执行设备内的核心间拷贝（其他轴名称保持不变）。 

# This runs on every core
def swap_cores_kernel(in_hbm, out_hbm,
                      in_vmem, scratch_vmem, out_vmem,
                      sem, send_sem, recv_sem):
  core_index = jax.lax.axis_index('core')
  num_cores = jax.lax.axis_size('core')
  slc_size = in_hbm.shape[-1] // num_cores
  slc = pl.ds(core_index * slc_size, slc_size)

  # Copy in a core-dependent slice of the input
  pltpu.async_copy(in_hbm.at[:, slc], in_vmem, sem).wait()

  # A barrier to make sure all cores have entered run_scoped.
  # You won't need this if not doing inter-core communications.
  dst_core = (core_index + 1) % num_cores
  sem0 = pltpu.get_barrier_semaphore()
  pl.semaphore_signal(sem0, 1, device_id={'core': dst_core})
  pl.semaphore_wait(sem0, 1)

  # Swap data between core 0 and core 1
  the_copy = pltpu.make_async_remote_copy(
      in_vmem, scratch_vmem, send_sem, recv_sem, device_id={'core': dst_core},
  )
  the_copy.start()
  the_copy.wait()

  # Core-local compute
  out_vmem[...] = scratch_vmem[...] * 2

  # Copy out the output
  pltpu.async_copy(out_vmem, out_hbm.at[:, slc], sem).wait()

获得本地内核后

  • 从 HBM 引用开始你的顶层 JAX 代码，并在需要时分配输出引用。

  • 使用 pl.core_map（接收 TensorCore 网格）来启动单核心编程。

    • 你需要 collective_id 来设置屏障信号量。

  • 在 pl.core_map 内部，调用 pl.run_scoped 来分配单核心暂存空间（VMEM 和信号量）并运行本地内核。

input_shape = (32, 256)
local_vmem_shape = (32 // num_devices, 256 // num_cores)
in_spec = jax.P('device', None)
sharding = NamedSharding(mesh, in_spec)

@jax.jit
@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec,
         check_vma=False)
def swap_cores(x):
  # Get buffers out of the input and output
  x_hbm_ref = jax.new_ref(x)
  o_hbm_ref = jax.new_ref(jax.lax.empty(x.shape, x.dtype))

  @pl.core_map(tc_mesh, compiler_params=pltpu.CompilerParams(collective_id=0))
  def _():
    pl.run_scoped(
        partial(swap_cores_kernel, x_hbm_ref, o_hbm_ref),
        *([pltpu.VMEM(local_vmem_shape, x.dtype)] * 3),  # VMEM allocations
        *([pltpu.SemaphoreType.DMA] * 3),                # semaphores
    )
  return o_hbm_ref[...]


x = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)
x = jax.device_put(x, sharding)
y = swap_cores(x)

np.testing.assert_array_equal(y[:, 128:], x[:, :128] * 2)
np.testing.assert_array_equal(y[:, :128], x[:, 128:] * 2)

保存样板代码#

你可以使用 pl.kernel 装饰器来封装诸如 core_map、run_scoped 和输出缓冲区分配等样板代码。

注意，这段代码应该在顶层的任何 jax.shard_map 内部运行。 

@jax.jit
@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)
def swap_cores(x):
  scratch_shapes = [pltpu.VMEM(local_vmem_shape, x.dtype)] * 3 + [pltpu.SemaphoreType.DMA] * 3
  return pl.kernel(swap_cores_kernel, out_shape=x, mesh=tc_mesh,
                   scratch_shapes=scratch_shapes,
                   compiler_params=pltpu.CompilerParams(collective_id=0))(x)

y = swap_cores(x)
np.testing.assert_array_equal(y[:, 128:], x[:, :128] * 2)
np.testing.assert_array_equal(y[:, :128], x[:, 128:] * 2)

使用 core_map 进行流水线处理#

注意，上述内核仅执行简单的拷贝和计算，没有通过 Pallas grid 和 BlockSpec 进行自动流水线处理。要在 core_map 内部进行流水线处理，请在核心本地内核中使用 pltpu.emit_pipeline。

自动在核心间并行化工作

最简单的方法是将块轴标注为 pltpu.PARALLEL，Pallas 将自动沿该轴并行化工作。pl.pallas_call 和 pltpu.emit_pipeline 均支持通过 core_axis 和 dimension_semantics 参数实现此功能。pallas_call 的示例位于另一篇指南中，emit_pipeline 的情况如下所示。

当提供 PARALLEL 标注时，相应的网格维度将被逻辑拆分并在独立的核心上执行。（保证了哪些网格维度在哪个核心上执行的具体语义）。

暂存空间分配

注意，在下面的示例中，顶层的 pl.run_scoped（封装在 kernel 中）没有分配任何 VMEM 暂存缓冲区。相反，pltpu.emit_pipeline 会在 VMEM 中分配自己的暂存缓冲区，并将其用于多重缓冲。 

def add_one_body(in_vmem, out_vmem):
  out_vmem[...] = in_vmem[...] + 1

input_shape = (1024, 1024)
in_spec = jax.P('device', None)

def add_one_kernel(x_hbm_ref, o_hbm_ref):
  in_shape = x_hbm_ref.shape
  pltpu.emit_pipeline(
      add_one_body,
      grid=(in_shape[0] // 8, in_shape[1] // 128),
      in_specs=[pl.BlockSpec(
          block_shape=(8, 128), index_map=lambda i, j: (i, j),
      )],
      out_specs=[pl.BlockSpec(
          block_shape=(8, 128), index_map=lambda i, j: (i, j),
      )],
      core_axis_name='core',
      dimension_semantics=(pltpu.PARALLEL, pltpu.ARBITRARY),
  )(x_hbm_ref, o_hbm_ref)


@jax.jit
@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)
def add_one(x):
  return pl.kernel(add_one_kernel, out_shape=x, mesh=tc_mesh, scratch_shapes=[])(x)


x = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)
x = jax.device_put(x, NamedSharding(mesh, in_spec))
y = add_one(x)

np.testing.assert_array_equal(y, x + 1)

标量预取#

下面的代码扩展了上述内核，但使用了标量预取和动态块索引来选择输入的特定子切片。

这涉及预先分配一个 SMEM 缓冲区（通过 kernel 内部的 pl.run_scoped 调用），并在流水线启动前使用 sync_copy 填充缓冲区。在 index_map 内部闭合动态索引值以使用它。

手动在核心间委派工作

下面的代码示例还展示了 core_map 如何让你自定义工作在核心间的分拆方式，而无需依赖上述的自动 API。

为实现这一点，请自定义你的 index_map，使用核心索引在不同核心上操作不同的切片。 

input_shape = (1024, 1024)
in_spec = jax.P('device', None)
output_shape = (1024, 512)

def indexed_add_one_kernel(in_refs, out_refs, i_smem_ref):
  (x_hbm_ref, i_hbm_ref), o_hbm_ref = in_refs, out_refs
  in_shape = x_hbm_ref.shape
  pltpu.sync_copy(i_hbm_ref, i_smem_ref)

  core_idx = jax.lax.axis_index('core')
  core_slc_size = in_shape[0] // num_cores
  i_map = lambda i: core_idx * core_slc_size // 8 + i  # split work among cores
  j_map = lambda j: i_smem_ref[0] // 128 + j           # use the prefetched offset

  pltpu.emit_pipeline(
      add_one_body,
      grid=(core_slc_size // 8, output_shape[1] // 128),
      in_specs=[pl.BlockSpec(
          block_shape=(8, 128), index_map=lambda i, j: (i_map(i), j_map(j)),
      )],
      out_specs=[pl.BlockSpec(
          block_shape=(8, 128), index_map=lambda i, j: (i_map(i), j),
      )]
  )(x_hbm_ref, o_hbm_ref)


@jax.jit
@partial(jax.shard_map, mesh=mesh,
         in_specs=(in_spec, jax.P()), out_specs=in_spec, check_vma=False)
def indexed_add_one(x, index):
  out_shape = jax.ShapeDtypeStruct((x.shape[0], x.shape[1] // 2), x.dtype)
  return pl.kernel(indexed_add_one_kernel,
                   out_shape=out_shape, mesh=tc_mesh,
                   scratch_shapes=[pltpu.SMEM((1,), jnp.int32)])((x, index))


xs = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)
xs = jax.device_put(xs, NamedSharding(mesh, in_spec))
idx = 256
y = indexed_add_one(xs, jnp.array([idx]))

np.testing.assert_array_equal(y, xs[:, idx:(idx+512)] + 1)

在 SparseCore 上进行映射#

TPU v5p 包含 4 个 SparseCore，它们专门用于稀疏内存访问和操作。本指南不会深入探讨 SparseCore 的全部功能，而是展示如何在 SparseCore 上运行程序，其语义与 TensorCore 代码基本一致，只需极少改动。

首先了解芯片的基本 SparseCore 规格，并为向量操作创建一个 VectorSubcoreMesh。注意，TPU v5p 上的每个 SparseCore 有 16 个（或其他数量）子核心，core_map 将在每一个子核心上以 SPMD 方式运行你的代码。 

sc_info = pltpu.get_tpu_info().sparse_core
assert sc_info is not None
print(sc_info)

sc_mesh = plsc.VectorSubcoreMesh(
    core_axis_name="core", subcore_axis_name="subcore",
    num_cores=sc_info.num_cores
)
sc_num_cores = sc_info.num_cores
sc_num_subcores = sc_info.num_subcores

SparseCoreInfo(num_cores=4, num_subcores=16, num_lanes=8)

下面的代码与我们之前编写的 add_one_kernel 非常相似，只有几点不同：

  1. 你需要拆分所有子核心之间的工作，因此需要几行代码来计算每个子核心的特定切片。

  2. SparseCore 寄存器计算允许更小的切片（int32 下最大为 4x16），因此你需要在计算阶段使用嵌套循环来迭代切片。

input_shape = (4096, 128)
SC_REG_OP_SHAPE = (4, 16)

def sc_add_one_body(in_vmem, out_vmem):
  @pl.loop(0, in_vmem.shape[0], step=SC_REG_OP_SHAPE[0])
  def _reg_loop_0(c0):
    @pl.loop(0, in_vmem.shape[1], step=SC_REG_OP_SHAPE[1])
    def _reg_loop_1(c1):
      slc = (pl.ds(c0, SC_REG_OP_SHAPE[0]), pl.ds(c1, SC_REG_OP_SHAPE[1]))
      out_vmem[slc] = in_vmem[slc] + 1


def sc_add_one_kernel(x_hbm_ref, o_hbm_ref):
  in_shape = x_hbm_ref.shape
  core_idx = jax.lax.axis_index('core')
  subcore_idx = jax.lax.axis_index("subcore")
  cm_idx = core_idx * sc_num_subcores + subcore_idx  # index on the core_map
  slc_size = in_shape[0] // (sc_num_subcores * sc_num_cores)
  index_map = lambda i, j: (
      pl.ds(pl.multiple_of(cm_idx * slc_size + i * 8, 8), 8), j)

  pltpu.emit_pipeline(
      sc_add_one_body,
      grid=(slc_size // 8, in_shape[1] // 128),
      in_specs=[pl.BlockSpec(
          block_shape=(pl.BoundedSlice(8), 128), index_map=index_map,
      )],
      out_specs=[pl.BlockSpec(
          block_shape=(pl.BoundedSlice(8), 128), index_map=index_map,
      )]
  )(x_hbm_ref, o_hbm_ref)


@jax.jit
@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)
def sc_add_one(x):
  return pl.kernel(sc_add_one_kernel, out_shape=x, mesh=sc_mesh, scratch_shapes=[])(x)


x = jax.random.randint(jax.random.key(0), input_shape, 0, 64, jnp.int32)
x = jax.device_put(x, NamedSharding(mesh, in_spec))
y = sc_add_one(x)

np.testing.assert_array_equal(y, x + 1)