本文档概述了 XLA 高级优化器 (HLO) 模块从初始状态到最终可执行文件的整个过程。有时，我们会省略“模块”，而只将其称为“HLO”。
优化前的 HLO
我们从预优化 HLO 模块开始。优化前的 HLO 不包含被视为 XLA 内部操作（操作）的运算，例如 fusion 或 bitcast。操作在此阶段没有布局，或者即使有，也会被忽略。预优化 HLO 通常由 TensorFlow 和 JAX 等更高级别的框架生成。使用 XLA 标志 -xla_dump_to 时，预优化 HLO 会转储到文件名为“before_optimizations.txt”的文件中。
优化 HLO 模块
XLA:GPU 流水线通过运行一系列遍数将优化前的 HLO 转换为优化后的 HLO。这些传递可以按语义分组，并按以下顺序运行：
分片相关卡券
这包括 Shardy Partitioner 等传递，或用于 SPMD 分片的传递。
优化次数
这可以包括合法化传递和简化传递。
集体优化传递
与优化遍数类似，但侧重于集合运算。
布局分配传递
每个 HLO 操作都会分配一个布局，该布局是指令形状的一部分。布局控制张量在内存中的物理布局方式。
包含布局的形状示例：
f32[10,20,30]{2,0,1}
在元素类型之后，是形状的逻辑维度，然后是按从小到大的顺序排列的布局置换。在此示例中，最小维度为 30，次小维度为 10，主要维度为 20。
布局分配的目标是使用贪婪策略最大限度地减少所需的物理转置次数。它从某些布局限制（例如，cuDNN/cuBLAS 库需要连续的维度）开始，然后将布局“向下”传播到 HLO 图，再“向上”传播。在布局传播结束时，某些指令可能具有冲突的布局，一个来自操作数，一个来自用户。为了解决此冲突，系统会插入一条 copy HLO 指令，将布局从操作数布局更改为指令布局。
布局归一化传递
鉴于确定实体形状有些困难，布局归一化会尝试重写形状，使其使用默认布局 {rank-1, rank-2, …, 0}。在上面的示例中，归一化后的形状为 f32[20,10,30]{2,1,0}。更改布局的复制操作会被重写为 transpose 和 bitcast 的组合。鉴于目前我们无法对所有操作进行归一化处理，因此仍有一些操作可能具有非默认布局，最值得注意的是 gather 和 dot。在归一化操作和非归一化操作之间的边界处，将有 bitcast 个表示转置的操作，即分配了布局的转置，使其在物理上成为空操作。
布局归一化还会使一些隐式转置变为显式转置，这非常重要，因为代码生成可以利用专用发射器处理显式转置。例如，从技术上讲，reshape 允许操作数和结果之间存在不同的物理布局（例如，由于不同的秩）。作为布局规范化传递的一部分运行的 ReshapeDecomposer 传递会将 reshape 转换为 transpose、reshape bitcast 和 transpose 的序列。
布局分配后优化传递
这里最重要的 pass 是 Triton 融合（GEMM 融合 + Softmax/Layernorm 融合）或重写为库调用。自动调优也会在此步骤中运行，其中 XLA 会在不同的发射器之间进行选择，为卷积或点积挑选最佳算法，为由 Triton 发射器处理的融合找到最佳平铺等。
Fusion 卡券
两个主要传递是 PriorityFusion 和 Multi-Output 融合。
在 PriorityFusion 中，我们根据费用模型形成融合。在融合时，如果操作可以融合到所有用户中，我们会允许复制具有多个用户的操作。如果可能，我们还会允许扩展现有的 Triton Softmax 融合。
Multi-Output 融合是一个单独的传递，允许融合共享运算对象的运算/融合。它还可以通过添加额外的输出将操作数/操作数融合融合到用户中，而不会重复，因此要融合的操作的其他用户可以重定向到这些输出。此传递需要注意不要在 HLO 图中引入环路。
在多输出融合之后，系统会运行通用子表达式消除（HloCSE 传递），如果之前重复的运算最终位于同一融合中，则可能会将它们重新合并在一起。
多次融合后传递
与集合相关的多次传递（例如将集合转换为异步，或强制执行集合的特定相对顺序）。
最后，我们运行 CopyInsertion，添加副本以确保就地操作不会覆盖其他地方仍需要的数据。
在优化结束时，如果使用标志 -xla_dump_to，则会将优化后的 HLO 转储到文件名为“after_optimizations.txt”的文件中。如果您想在实际更改 HloModule 的中间传递之后转储 HLO，可以使用标志 -xla_dump_hlo_pass_re=.*（或使用特定的正则表达式将其限制为某些传递）。
时间安排
没有时间表的 HLO 模块在处理操作的顺序方面仍然具有一定的自由度。任何尊重操作数/结果关系和控制依赖关系的拓扑排序都是有效的。调度决定了要使用的具体顺序。此阶段的主要考虑因素是取决于张量生命周期的最大内存消耗。在初始步骤中，我们会尝试不同的调度器算法，并选择应能最大限度减少峰值内存消耗的调度。请注意，此时我们尚未使用物理缓冲区（这将在“缓冲区分配”中实现），而是模拟内存使用情况。
然后，LatencyHidingScheduler 传递运行并尝试最大限度地提高计算-通信重叠。但这样可能会再次增加内存用量。
最后，如果内存消耗峰值高于可用内存量，我们会运行 HloRematerialization。此传递会尝试以牺牲性能为代价来减少内存使用量，例如，可能会拆分某些融合，并可能会复制某些操作以缩短缓冲区生命周期。如果发生重新实物化，则最好调查一下如何减少模型端的内存需求（例如，使用较小的批次大小）。
缓冲区分配
在降级到 LLVM IR 之前，我们会立即运行缓冲区分配传递，将缓冲区切片分配给 HLO 图中的每个指令。缓冲区分配分多个步骤进行：
HloDataflowAnalysis会将HloValues（本质上是逻辑缓冲区）分配给指令。对于就地操作，可以重复使用操作数的HloValue。一个 op 可以定义多个HloValue（例如，使用元组结果形状）。HloAliasAnalysis尝试合并用于别名操作的缓冲区，并计算从HloValue到HloBuffer的映射。BufferAssignment计算HloBuffers到大缓冲区内缓冲区切片的映射，以确保同一缓冲区切片不会用于具有重叠生命周期的不同HloBuffers。对于可能存在别名的操作，可以存在轻微的重叠（一个HloBuffer的结束时间可能与另一个HloBuffer的开始时间重合）。使用标志-xla_dump_to时，有关缓冲区分配的一些信息会转储到名称后缀为“after_optimizations-buffer-assignment.txt”的文件中。
Thunk
在 HLO 图经过优化和调度后，它会被降为特定后端（CPU 或 GPU）的线性 thunk 序列。
在 XLA 中，Thunk 是运行时执行的独立工作单元的抽象。它可能是编译后的内核启动、特定操作、库调用、控制流构造、集体通信等。Thunk 序列表示特定后端的整个可执行文件。
Thunk 排放
将已调度的 HLO 计算转换为 thunk 序列的过程称为“thunk 发射”。这由每个后端中的专用发射器类处理。
对于 GPU 后端，这由 IrEmitterUnnested 处理。
EmitHloComputation 会遍历计算中已调度的 HLO 指令列表，并调度到专门的 Emit... 方法（例如，EmitFusion、EmitConvolutionThunk、EmitWhile）。这些方法中的每一种都会构建相应的 Thunk 对象，并将其附加到 Thunk 序列中。
对于 CPU 后端，ThunkEmitter 扮演着类似的角色，并且以类似的方式进行组织。最终的 ThunkSequence 嵌入在 CpuExecutable 中。
请注意，HLO 模块的条目计算中的每条指令可能对应于最终 thunk 序列中的零个（kTuple、kConstant、...）、一个或多个（例如排序指令）thunk。
命令缓冲区：优化 GPU 上的执行
借助现代 GPU 硬件，您可以记录一次 GPU 操作序列（内核启动、内存复制等），然后以极低的 CPU 开销多次重放该序列。这是一项关键的性能优化，尤其适用于包含许多快速启动的小型内核的工作负载。XLA 使用命令缓冲区作为 CUDA 图或 HIP 图的抽象。核心接口在 GpuCommandBuffer 中定义。
命令缓冲区在 thunk 序列中由 CommandBufferThunk 表示。
发射器不会直接从 HLO 指令生成此 thunk。相反，这是通过在 ThunkSequence 本身上运行的 CommandBufferConversionPass 来完成的。
该 pass 会识别兼容 chunk 的连续子序列（例如，一系列 KernelThunk 和 GemmThunk）。然后，它会将找到的子序列替换为单个 CommandBufferThunk。新 thunk 将原始 thunk 的逻辑封装为轻量级 CommandBufferCmd 对象的列表。当 CommandBufferThunk 首次在给定的 GPU 流上执行时，它会将其命令序列“记录”到硬件命令缓冲区中。在所有后续执行中，它只需向 GPU 发出一条命令，即可“重放”记录的序列。这样可避免启动每个单独的内核所产生的 CPU 开销。
可执行文件
XLA 编译流水线的最终产品是一个独立的平台专用可执行文件。此对象封装了在目标设备（例如 CPU 或 GPU）上运行已编译程序所需的所有信息。它是编译器和运行时之间的桥梁。PJRT 等现代运行时使用稍高层次的抽象（请参阅 PjRtExecutable），但这些抽象最终会封装特定于后端的执行程序。
Executable 包含在编译期间生成的几条关键信息。虽然具体内容因后端而异，但通常包括：
已编译的代码：这是将在设备上运行的低级机器代码。 对于 CPU，这通常是一个或多个对象文件。对于 GPU，这是 PTX 或 HSACO 格式的已编译设备代码，会在运行时加载到 GPU 上。
执行计划 (ThunkSequence)：运行时逻辑的核心。这是一个 Thunk 对象的线性序列。每个 thunk 都表示一个工作单元，例如启动内核、调用库函数（例如 cuBLAS）或处理控制流。运行时通过迭代此序列来执行程序。
内存布局 (BufferAssignment)：这是由 BufferAssigner 生成的关键元数据，用于描述计算的完整内存布局。它指定了每个缓冲区的大小，以及如何为参数、输出和临时值分配和重用内存。运行时使用此信息来分配设备内存，并将正确的指针传递给每个 thunk。
（可选）HLO 模块：为了进行调试和分析，可执行文件通常会保留对其编译自的最终优化 HloModule 的引用。
最终可执行文件的创建由编译器针对每个特定后端进行编排。编译器实现的 RunBackend 方法是编译流程的最后一步，它会将所有已编译的制品打包到可执行对象中。GpuCompiler 和 CpuCompiler 分别以 GPU 和 CPU 为目标平台。
当用户对可执行文件调用 Execute... 时，运行时会使用 BufferAssignment 分配内存，然后调用 ThunkSequence 以使用编译后的代码在设备上启动操作。