使用 Pallas 编写 TPU 内核#
本页面重点介绍在 Google TPU 上运行 Pallas 内核时需要注意的重要细节。首先，TPU 后端仍处于实验阶段，只有 JAX NumPy 的一个子集会被接受。此外，为 TPU 编写高性能代码可能需要仔细考虑硬件的原生能力。虽然许多对硬件来说不自然的模式会被接受，但它们可能最终需要软件模拟，并且会减慢计算速度。
警告
此功能仍应被视为实验性的，因为工作仍在进行中（特别是在改进错误消息方面）。
注意
虽然这里描述的所有功能都是实验性的，但我们对维护其正确性保持非常严肃的态度。因此，在尝试编写 TPU 内核时遇到"未实现"错误可能并不罕见。但是，如果一个内核被编译器接受了，它 必须 返回预期的结果。
如果你看到意外的输出，请与传入
interpret=True给pallas_call运行的内核进行比较。如果结果不一致，请提交 bug 报告。
什么是 TPU？#
TPU 是 Google 开发的硬件加速器。你可以把 TPU 想象成 GPU，但专门为机器学习工作负载而设计。因此，它们的架构有很大不同。然而，我们相信 Pallas 可以让编写 TPU 内核变得容易，即使你对底层硬件没有完全的了解。话虽如此，深入理解硬件肯定会让编写高性能内核变得更容易。
简而言之，TPU 和 GPU 之间的主要区别是 TPU 是具有非常宽向量寄存器的顺序执行机器（有点像 CPU！）。同时，它们允许软件在后台调度某些操作，使其相对于主指令流异步执行。这包括 HBM 内存访问（不能直接发出，而是必须由 DMA 子单元预取到内存层次结构的较低层级）、矩阵乘法（由 MXU 单元支持）或矩阵转置和置换（由 XLU 单元支持）。
如果你有兴趣详细了解 TPU 架构，我们推荐阅读多年来发表的一系列论文。虽然其中许多讨论的是特定的 TPU 代，但描述的许多思想也适用于后续代。
- A Domain-Specific Supercomputer for Training Deep Neural Networks
- The Design Process for Google’s Training Chips: TPUv2 and TPUv3
- Ten Lessons From Three Generations Shaped Google’s TPUv4i: Industrial Product
- TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings
值得注意的属性和限制#
BlockSpec 和 grid 迭代#
BlockSpec（参见 BlockSpec，又称如何切分输入）在 Pallas 中通常按预期行为——内核体的每次调用都获得输入的切片访问权，并负责初始化输出的一个切片。
注意
并非所有块形状都受支持。在 TPU 上，仅支持秩至少为 1 的块。此外，块形状的最后两个维度必须分别能被 8 和 128 整除，或者等于整体数组的对应维度。
Pallas TPU 内核的一个有趣方面是它们处理内存空间的方式：虽然 pallas_call 的输入通常驻留在 HBM（TPU 主存储器）中，但传递给内核体的引用指向内存层次结构较低级别的缓冲区（VMEM 或 SMEM）。这使得内核体能够以非常高的速度对它们进行读写，而所有与 HBM 的通信（具有非常高的延迟）由编译器处理并与计算重叠。
更重要的是，与 GPU 相比，TPU 实际上是高度顺序执行的机器。因此，grid 通常不是并行处理的，而是按字典序顺序处理的（但多核 TPU 配置部分有例外）。这解锁了一些有趣的能力：
当两个（字典序上）连续的 grid 索引使用输入的同一切片时，第二次迭代的 HBM 传输会被跳过，因为数据已经可用。
内核体的多次调用可以写入输出的同一切片，而没有任何竞争条件的风险。但是，我们要求写入特定切片的所有调用必须是连续的。
输出的"连续"限制通常意味着 grid 维度的某个前缀总是改变调用需要访问的输出切片，而输出窗口在剩余的后缀中保持不变。
例如，在实现矩阵乘法的 Pallas TPU 内核时，通常会使用一个 3 维 grid：前两个维度对应于沿左操作数的第一个轴和右操作数的第二个轴进行切片。第三个也是 最后一个 grid 轴会分块归约维度。对应于归约维度的 grid 轴必须是最后一个，因为输出窗口不沿该轴变化。输出引用然后可以用作部分结果的累加器。
注意
VMEM 对于如此低级的内存层次结构来说相当大（16MB+），使得可以使用较大的窗口大小。而且，通常窗口大小越大，最终的硬件利用率就越好。然而，有可能指定一个窗口大小（加上保存溢出的向量寄存器所需的空间）超过 VMEM 的大小。在这种情况下，你可能会看到一个低级编译器错误消息抱怨内存不足错误。
数组布局#
数组的维度顺序在 Pallas 中是有意义的。在 JAX 程序中，jax.jit 内部中间数组的顺序通常对性能没有影响，因为编译器可以自由地重新排列它们。然而，由于 Pallas 旨在暴露更底层的能力，维度顺序可能对生成代码的质量有很大影响。
TPU 在 2D 向量寄存器上执行大部分计算，对于 32 位值，这些寄存器通常大小为 8x128（截至 TPU v6）。当一个向量值从 VMEM 加载到寄存器中时（例如 x = x_ref[...]），数组的最后两个维度将被分块到寄存器中。Pallas 只会考虑将中间数组的最后两个维度映射到 8x128 向量寄存器维度（分别为 sublane 和 lane）。
以下是一个图形示例，展示如何使用 6 个 8x128 的 tile 对 12x320 数组进行分块：
分块布局对内核编写者有几个重要影响：
数组的最后两个轴与其他轴被区别对待。 例如，涉及最后两个轴的归约、reshape 和转置通常更昂贵。某些涉及最后两个维度的 reshape 不受支持并会导致编译器错误，但对于其他维度则是"免费的"并在编译时执行。
虽然有时不可避免，但在最后两个轴中有单例维度通常是浪费的，因为它们将在整个 tile 维度中只占用 1 个元素。消耗太多寄存器也可能导致寄存器溢出到 VMEM，从而降低内核性能。
与上述观点相关，所有向量计算都被填充到 tile 大小。 两个 1x1 数组相加的开销与两个 8x128 数组相加相同，而两个 8x128x1x1 数组相加将比两个 8x128 数组相加昂贵 1024 倍，因为 8x128x1x1 数组将被填充到 8x128x8x128。
多核 TPU 配置#
在较新的 TPU 代中，芯片上的两个核心通常被抽象为一个单一设备。为了利用多个核心，Pallas 必须打破顺序 grid 执行的保证，并需要在核心之间并行化某个 grid 轴。这是一个选择加入的过程。为此，pallas_call 需要一个名为 dimension_semantics 的额外参数：
pallas_call(
...,
compiler_params=pltpu.CompilerParams(
dimension_semantics=["parallel", "parallel", "arbitrary"]
),
)该参数是一个列表，条目数与 grid 中的轴数相同。只有 parallel 维度可以在核心之间分区。作为经验法则，维度是 parallel 的，除非输出窗口不变化。因此，dimension_semantics 总是若干个 parallel 轴后面跟着若干个 arbitrary 轴。
虽然在 2 核 TPU 设备上对内核进行分区通常会带来 2 倍加速，但实际上可能显著更小。如果内核体的不同实例开销差异很大，这一点尤其明显。如果所有昂贵的步骤都被映射到一个核心，而所有廉价的步骤被分配到另一个核心，第二个核心将一直空闲直到第一个核心完成任务。
Pallas TPU 通常倾向于分区大小为 TPU 核心数倍数的轴，并优先分区前导 grid 轴。
将操作数放置在 SMEM 中#
TPU 上的大部分计算将在向量单元上进行。尽管如此，在许多情况下执行一些标量操作是有用的，例如进行控制流。因此，TPU 配备了一个独立的标量单元和一个连接到它的独立标量内存（SMEM）。作为经验法则，任何用于做控制流决策的数据都应该放在 SMEM 中。
SMEM 是一个低延迟的内存，支持随机访问，但只允许你用单条指令读写 32 位值（与 VMEM 事务的 4KiB 粒度相比非常小，但由于没有对齐要求而灵活得多！）。
标量内存在实现不按规则模式访问输入 tile 的内核时也非常有用，例如编写块稀疏内核时。在 Pallas 中，这可以通过将 pallas_call 的 grid 参数替换为带有非零 num_scalar_prefetch 参数的 PrefetchScalarGridSpec 的 grid_spec 来实现。如果 num_scalar_prefetch 为 n，则 pallas_call 的前 n 个参数将被放置在 SMEM 中。不应为这些参数指定 BlockSpec。但是，所有后续参数的 BlockSpec 将不仅接收 grid 索引，还会接收前导操作数的 SMEM 引用。
参见 标量预取和块稀疏计算 了解使用此功能的示例。
支持的数据类型#
目前 Pallas TPU 支持以下数据类型：
jnp.float32jnp.bfloat16jnp.int*（所有精度，jnp.int4除外）jnp.uint*（所有精度）jnp.bool_
计算位置#
所有标量（即 0D）数组将存储在标量寄存器中，对它们的操作将在标量核心上执行。所有其他操作（即使是单元素但 1D+ 的数组）将在向量核心上执行。
支持的操作#
矩阵乘法#
矩阵乘法总是产生 float32 格式的结果。如果你的输入不是 float32，我们建议使用 lax.dot 并将 preferred_element_type 设置为 jnp.float32。
使用 lax.dot_general 时，可以将矩阵乘法操作数最后两个维度的转置融合到操作中，这可以提高整体内核性能。
精度控制#
Pallas TPU lowering 知道 jax.default_matmul_precision。为获得最佳性能（和最低精度），使用 bfloat16。如果你关心数值精度，你可能想将精度设置为 float32。
警告
即使你向矩阵乘法传入 32 位操作数，除非请求
float32精度，否则它们将被舍入为bfloat16。
转置#
如果值至少有 4 个维度，除最后两个轴外的所有轴的任意转置都是免费的。否则，仅实现了最后两个轴的转置。注意，最后两个维度的某些转置可以融合到矩阵乘法中。
访问内存#
引用的任意切片可以被读取或更新，受实现约束限制。目前，对 32 位宽的输入没有限制，但对于更窄的类型仅支持某些切片模式。在最后两个维度中对齐到 8 和 128 的倍数、且长度分别为 8 和 128 的倍数的读写总是受支持的。
对向量内存的读写通常以 (8, 128) 形状的 tile 进行。因此，当读写至少有两个维度的引用时，当内存访问的基础偏移量的索引可被 tile 大小整除，且读取区域的大小是 tile 大小的倍数时，可以获得最佳性能。
逐元素操作#
许多逐元素操作都受支持。值得注意的是，硬件通常只支持使用 32 位类型的逐元素计算。当加载使用较低精度类型的操作数时，通常应在应用逐元素操作之前将其向上转换为 32 位类型。
值得注意的是，它们的开销可能 显著 不同。因此，我们将支持的操作分为三类：便宜（🟢）、中等（🌕）和昂贵（🔴）。
| 操作 | 开销 |
|---|---|
jnp.add, + | 🟢 |
jnp.sub, - | 🟢 |
jnp.mul, * | 🟢 |
/, //, % | 🌕 |
jnp.max, jnp.min | 🟢 |
jnp.where（select） | 🟢 |
jnp.abs | 🟢 |
|, ^, &, ~ | 🟢 |
<<, >> | 🟢 |
比较运算（==, …） | 🟢 |
类型转换（.astype） | 🟢 |
jnp.exp | 🌕 |
jnp.tanh | 🌕 |
jnp.pow | 🌕 |
jnp.sin | 🔴 |
jnp.cos | 🔴 |
许多 JAX 函数是用其他 JAX 原语实现的，所以这个列表可能不全面。例如，jax.nn.relu 是用比较运算和 jnp.where 实现的，在 Pallas 内核中也能工作。
数组构造函数#
所有常量数组构造函数都受支持（jnp.ones、jnp.zeros、jnp.full）。
归约#
支持 sum、max、min（用于浮点值）归约，以及用于布尔值的 any 和 all。不支持整数归约。
沿最后一个数组维度的归约通常是最慢的。 沿倒数第二个维度的归约更快，但仍然比沿前导维度的归约慢。
广播#
广播的性能特征与归约非常相似。沿除最后两个维度以外的所有维度的广播总是受支持的且是免费的。沿倒数第二个维度的广播较慢，而沿最后一个维度的广播是最慢的。
Reshape#
一如既往，除最后两个维度以外的所有维度的 reshape 都受支持且是免费的。
当 reshape 可以修改数组最后两个维度时，仅支持两种情况：(1) 某些前导维度被展平到倒数第二个维度上，或 (2) 添加了一个刚刚被归约移除的维度。
随机数生成#
Pallas 支持 jax.random 模块中最常用的函数，如 uniform、normal 和 bernoulli。密钥应该是 threefry2x32 密钥，这是 JAX 中的默认设置。密钥可以直接传入内核，也可以在内核内部生成。
控制流#
TPU 后端目前对控制流的支持有限。当前支持的函数有 cond、fori_loop 和 for_loop。然而，循环原语目前在编译期间会被完全展开，所以请尽量保持循环次数在合理的小范围内。
过度使用控制流可能导致低级代码生成的严重退化，建议尽量将尽可能多的计算密集型操作挤入单个基本块中。