XLA 的 MLIR 代码生成

XLA 在 HloInstruction 上进行运算并对此表示执行许多优化,同时在目标设备之间共享其中的许多优化。在某些时候,会计算线性调度,并将内存缓冲区静态地分配给每个值。特定于设备的代码生成在运算时会遍历此序列,并调用“发射器”以生成适合该设备的表示(例如,CPU 上每个 XLA 计算对应的单个 LLVM 函数,或封装 GPU 运算的“Thunk”序列,也可能是针对 GPU 时生成的 PTX)。

作为暂存步骤,我们目前会在 XLA 完成缓冲区分配阶段之后立即拦截该过程,转而在 lhlo 方言中发出 MLIR 模块。之后,我们会根据设备使用 MLIR 组件(主要是 Linalg、仿射和 GPU 方言)执行代码生成。

下面是通过使用 lhlo 作为代码生成输入来增量迁移 XLA/GPU 的记录计划。

任务

主机 设备
输入格式 HloInstruction*(任务 1) HloInstruction*(任务 1)
输出格式 xla::Thunk(任务 2) LLVM IR(任务 3)
  • 任务 1 将主机和设备的输入格式从 HloInstruction* 更改为 LHLO。
  • 任务 2 将主机的输出格式从 Thunk 更改为“主机的某种停机坪”(请参见下文)
  • 任务 3 将设备输出从 LLVM IR 迁移到某种形式的 MLIR。它在此项目中为可选项,有关详细信息,请参阅“迁移设备 LLVM IR”部分。

此项目尽可能优先考虑使用启用了 LHLO 发射器的端到端可运行模型。这意味着目标列表将按照优先级进行如下排序:

  • 使 XLA/GPU 可与 LHLO 发射器一起运行,而不修改现有的 Thunk 和发射器。
  • 在 LHLO 中根据具体情况消除对 HloInstruction* 的引用:
    • 将旧发射器切换为基于 MLIR 的发射器(例如 Linalg),或者
    • 以机械方式转换现有发射器以接受 MLIR 表示(使用 GPU 方言迁移到 Standard)。

迁移 Thunk(任务 2)

xla::gpu::Thunk 是一种具有如下特点的数据结构:

  • 可从主机调用 (xla::gpu::Thunk::ExecuteOnStream())
  • 在其子类中携带各种数据
  • 可与 BufferAllocation::Slice 和 StreamExecutor 进行交互
  • 可启动内核
  • 可调入所有运行时库

成本包括:

  • 表示特定于算子的配置数据(例如卷积配置)
  • 迁移算子形状和运算对象形状
  • 表示 Thunk 树(while、条件等)

迁移作业与 LHLO/发射器迁移无关。在资源有限的情况下,它的优先级位于 LHLO/发射器迁移之后。

关于如何从 LHLO 降级主机端部分,我们有几种选择:

  • TFRT
    • (利)可以使用出色的 CUDA 和 HIP 封装容器
    • (利)易于实现库调用(cuDNN、cuBLAS、cuFFT 等),因为 TFRT 算子由 C++ 代码解释
    • (弊)主机端正在开发且未经测试
  • 即时编译的 CPU 代码
    • (利)出色的降级能力。创建一些循环和条件即可完成
    • (弊)GPU 方言尚未对 chains/streams/asynchronicity/device 分配进行建模
    • (弊)仅对 CUDA/HIP 运行时提供最小支持(工具包路径、版本、动态加载等)
  • 现有的(解释)XLA 运行时

决策:采用 TFRT,但也支持在 TFRT 中即时编译 CPU 代码。

迁移设备 LLVM IR(任务 3)

元素发射器通过逐个填充元素来生成目标算子。每个输出元素都取决于来自运算对象的一组元素。所有元素都可以通过组合缓冲区和动态索引进行描述。这样足以描述几乎所有的“数学算子,但出于性能原因,(Cpu|Gpu)ElementalIrEmitter 中仅直接实现了“数学”算子的一个大子集。

ElementalIrEmitter 的独特之处在于:

  • XLA/GPU 和 CPU 之间会共享大部分代码
  • 它代表了在模型中看到的大部分算子,包括所有元素算子
  • 大多数融合完全取决于 ElementalIrEmitter
  • 结构简单,因为它描述的是算子元素和运算对象元素之间的数据依赖 DAG
  • 它通常可移植且高级(例如,不同于 GPU kReduce 和 GPU kCopy)
  • 对动态形状的支持至少对于元素级算子来说很容易

现在,对于所有算子(无论是否以元素为单位进行发射),每个 XLA 算子的结束状态都有几种形式:

  1. 设备代码保持为 LLVM IR
  2. 将旧发射器重构为类似 LHLO -> MLIR LLVM 方言的形式:
    • (成本)如果我们最终要迁移到 Standard,那将是一劳永逸的工作
    • (收益)简单而机械,可在短期内完成
    • (收益)与 (1) 相比并没有更多收益
  3. 将旧发射器重构为类似 LHLO -> MLIR GPU + Standard + 循环的形式:
    • (成本)将现有发射器提升为 Standard 会带来一些挑战。指针和 GEP 需要转换为 MemRef 和 SubView。确保 AMDGPU 的完整性也是一个问题。
    • (成本)XLA/GPU 严重依赖 LLVM 元数据:
      • 用于块/线程索引的 range
      • 用于加载/存储的 aligndereferenceableinvariant.loadalias.scopenoalias
      • 用于顺序循环的 llvm.loop.unroll.disablellvm.loop.unroll.fullllvm.loop.vectorize.enable
    • (收益)可为长期。更好的可移植性。
  4. 将旧发射器重构为 LHLO -> Linalg,并编写新的 Linalg 发射器
    • (成本)这视情况而定。与之前的选项相比,匹配 XLA 性能的新实现需要通过基准测试 <-> 优化工作流,这对于某些算子而言是一项巨大的开销。
    • (收益)统一堆栈;社区支持;可移植性;更大的优化潜力。

结论:

  • 不要选择 (2)。(1) 或 (3) 都比 (2) 好。(2) 的成本比 (1) 高,因为它需要大量的机械重构。如果选择 (1),我们仍然可以实现使 XLA 拾取 MLIR 发射器这一目标。具体方法为:LHLO -> LLVM IR -> 运行旧的设备发射器。
  • ElementalIrEmitter 算子适合选择 (4),但无法递增。因为以元素发出的所有算子都连接到同一个计算图中,所以无法逐算子执行。此项工作还可以作为几个正在发展的方案(xla/service/mlir_gpu、内核生成器,以及 Linalg)的统一点。
  • 所有其他算子应选择 (1)。作为延伸目标,它们可能会迁移到 (3) 或 (4)。

优先级

虽然上述所有三个任务都可并行化,但在有限的资源下,它们必须被序列化。优先级重点是完成每个任务的可见结果。

优先级为:任务 1(旧发射器的 LHLO)> 任务 2(Thunk)> 任务 3(MLIR 发射器)。

在任务 1 结束时,XLA 的用户可以生成 LHLO(例如内核生成器)并执行它们。编译格式不是可序列化的 MLIR。

在任务 2 结束时,LHLO 降级为合适的可序列化的 MLIR。这样可以进行离线编译。

在任务 3 结束时,所有 XLA 发射器的实现均基于 MLIR。

详细设计

第 1 步:(任务 1)完成 LHLO 并让旧发射器接受 LHLO

此步骤使所有现有的 XLA/GPU 发射器与 MLIR 算子 进行交互。此步骤是纯重构和 NFC。

此步骤主要是机械步骤,但值得注意的是,未嵌套的 HloComputation 和 LHLO 之间存在以下差异:

  • 每个 HloInstruction 都可以直接访问它的运算对象(数据流 DAG)。相反,每个 LHLO 算子只能访问它的运算对象缓冲区(算子和缓冲区之间的二分体)。LHLO 算子必须通过使用-定义链访问它们的运算对象算子。
  • 从经验上看,未嵌套的旧发射器几乎从不访问其运算对象。唯一的例外是 kReduce。
  • 未嵌套的旧发射器访问 BufferAssignment 只是为了获取切片,而非访问辅助数据结构(如 dataflow_analysis() 或 alias_analysis())。llvm_ir 会基于切片信息构建自己的 alias_analysis()。

结论是 LHLO 应该能够轻而易举地满足要求。

第 2 步:(可选)性能分析支持

仅当我们开始舍弃某些 XLA Thunk 逻辑(请参见下一个步骤)时才需要执行此步骤。

在实际打开任何基于 MLIR 的发射器之前,我们需要对基于 MLIR 的发射器进行性能分析。

目前,XLA 会通过调入 StreamExecutor 的计时器来执行自己的性能分析。后台计时器会在内核启动前后插入两个事件,并测量这两个事件之间的同步时间。

要在 MLIR 中支持性能分析,大致有三种方式:

  • 端到端地运行性能分析器
  • 使用注入的性能分析器为 LHLO 中的每个算子添加一个性能分析算子

“端到端”方式对 MLIR 透明,但存在一个问题,XLA 一开始不使用该方式也是因为这个问题:由性能分析器 (nvprof/...) 收集的库调用无法轻松关联到 HLO 算子。例如,cuDNN 会为每个 HLO 启动多个内核,很难区分哪个内核对应于哪个 HLO。

“注入的性能分析器”方式要求:

  • LHLO 将性能分析器作为参数
  • 在每个算子之前和之后插入 profile.start/profile.end
  • 将 profile.{start,end} 降级到 C++ 实现的传递

无法对 MLIR 生成的算子轻松进行精确的性能分析,因为:

  • MLIR 没有计时器,也不依赖 TFRT/StreamExecutor
  • MLIR 无法轻松调入具有复杂参数的 C 函数

第 3 步:(任务 2)迁移 Thunk

请注意,Thunk 大致分为三种类型:

  • KernelThunk,启动内核
  • 控制流 Thunk,具有主机控制流逻辑(conditional、while、for 和 sequence)并启动主体内核
  • 库 Thunk:cuDNN、cuBLAS、cuFFT、NCCL 等

计划是:

  • 使 Thunk 可(反)序列化
  • 帮助将 TFRT 改进到能够支持这些语义的状态
  • 随着状态的改进,以增量方式迁移各个 Thunk

这些操作项仅进行了部分排序。实际执行顺序/工程并行性将即时进行评估。

第 4 步:(任务 3)迁移后的 ElementalIrEmitter

性能分析完成后,我们就可以完成并调整 MLIR 中所有基于 ElementalIrEmitter 的发射器。然后,我们默认将它们打开,假设所有这些基于 MLIR 的发射器都使用一个流。

请注意,迁移 XLA/CPU 的 ElementalIrEmitter 也是有益的,因为它们会共享大部分代码。

完成所有基准测试和性能搜寻(TODO:定义性能奇偶校验)后,我们打开新的基于 MLIR 的元素发射器,然后删除旧的 ElementalIrEmitter。

此步骤还为以后的迁移提供了简单的融合过渡(嵌套算子)。

第 5 步:多流支持或删除

我们不能删除某些发射器,除非我们在 MLIR 中对其提供支持,或者删除该功能。在这种情况下,MLIR 的工作量相对较大,而 XLA 的收益却相对较小。我们应该调查多流 XLA/GPU 的当前用户,并在合理的情况下尝试删除此功能。

第 6 步:(任务 3)迁移后的设备算子

此步骤会迁移所有未嵌套的算子,之后我们可以删除所有未嵌套的发射器。

这要求重写/重构 kCopy 和 kReduce。kReduce 已经进行了大量工作,因此需要完成的实际工作量还有待观察。