此页面由 Cloud Translation API 翻译。
Switch to English

XLA的MLIR CodeGen

XLA在HloInstruction运行, HloInstruction这种表示形式进行了许多优化,在目标设备之间共享了许多优化。从某种意义上说,会计算线性计划,并且将内存缓冲区静态分配给每个值。特定于设备的代码生成器通过遍历此序列并调用“发射器”来生成适合于该设备的表示(例如,CPU上每个XLA计算使用单个LLVM函数,或者封装GPU操作的一系列“ thunk”,可能在生成时生成PTX)定位GPU)。

作为暂存步骤,我们目前正在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 Dialect迁移到Standard)。

迁移暴徒(任务2)

xla :: gpu :: Thunk是以下数据结构:

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

费用包括:

  • 表示特定于操作的配置数据(例如卷积配置)。
  • 迁移op形状和操作数形状。
  • 代表一堆笨拙的树(while,condition等)。

迁移工作独立于LHLO /发射器迁移。在资源有限的情况下,它优先于LHLO /发射器迁移。

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

  • TFRT
    • (专业版)很棒的CUDA和HIP包装器。
    • (Pro)易于实现的库调用(cuDNN,cuBLAS,cuFFT等),因为TFRT操作由C ++代码解释。
    • (Con)主机端正在开发中,未经测试。
  • 中央处理器代码
    • (亲)极大的降低能力。创建一些循环和条件,即可完成。
    • (缺点)GPUDialect尚未对链/流/异步/设备分配建模。
    • (缺点)CUDA / HIP运行时支持最少(工具包路径,版本,动态加载等)。
  • 现有(解释)XLA运行时

决策:采用TFRT,但也支持在TFRT中压缩CPU代码。

迁移设备LLVM IR(任务3)

元素发射器通过逐个元素填充来生成目标op。每个输出元素都取决于操作数中的一组元素。通过组合缓冲区和动态索引来描述所有元素。几乎描述所有“数学”操作就足够了,但是出于性能原因,直接在(Cpu | Gpu)ElementalIrEmitter中直接实现了很大一部分“数学”操作。

ElementalIrEmitter的独特之处在于:

  • XLA / GPU和CPU之间共享大部分代码。
  • 它代表了模型中看到的大部分操作,包括所有基于元素的操作。
  • 大多数融合仅取决于ElementalIrEmitter。
  • 它在结构上很简单,因为它描述了op元素和操作数元素之间的数据依赖DAG。
  • 它主要是便携式的和高级的(例如,不同于GPU kReduce和GPU kCopy)。
  • 动态形状支持至少对于元素级操作很容易。

现在,对于所有操作(无论是否元素发射),每个XLA操作的结束状态都有几种形式:

  1. 设备代码保持为LLVM IR。
  2. 将旧的发射器重构为LHLO-> MLIR LLVM方言:
    • (成本)如果我们最终要迁移到标准版,那将是一项繁琐的工作。
    • (优点)这很容易并且机械。可以在短期内完成。
    • (优点)与(1)相比,它没有更多的好处。
  3. 重构旧的发射器,使其像LHLO-> MLIR GPU + Standard + Loops:
    • (成本)将现有的排放者提升为标准排放带来了一些挑战。指针和GEP需要转换为MemRefs和SubViews。确保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 ops适用于(4),但不能递增。无法通过op进行操作,因为所有元素发射的op都连接到同一图形中。这项工作还可以用作多个持续性力量的统一点(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之间存在以下差异:

  • 每个Hlo指令都可以直接访问其操作数(数据流DAG)。相反,每个LHLO op只能访问其操作数缓冲区(ops和缓冲区之间的二分之一)。 LHLO操作必须通过use-def链来访问其操作数操作。
  • 根据经验,未嵌套的旧发射器几乎永远不会访问其操作数。唯一的例外是kReduce。
  • 未嵌套的旧发射器仅用于获取切片而访问BufferAssignment,而不用于访问诸如dataflow_analysis()或alias_analysis()之类的辅助数据结构。 llvm_ir基于切片信息构建自己的alias_analysis()。

结论是LHLO应该在没有重大麻烦的情况下就可以正确地入内。

步骤2 :(可选)分析支持

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

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

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

支持MLIR中的概要分析的方法大致有三种:

  • 端到端运行探查器
  • 使用注入的探查器为LHLO中的每个op添加一个探查op。

“端到端”方法对MLIR是透明的,但存在使XLA最初不使用它的问题:探查器(nvprof / ...)收集的库调用无法轻松地与HLO相关联行动。例如,cuDNN为每个HLO启动多个内核,很难区分哪个内核对应于哪个HLO。

“注入的探查器”方法要求:

  • LHLO以探查器作为参数。
  • 在每个操作之前和之后插入profile.start / profile.end。
  • 从较低的配置文件开始。{start,end}到C ++实现。

对于MLIR生成的操作,无法轻松进行准确的性能分析,因为:

  • MLIR没有计时器,也不取决于TFRT / StreamExecutor。
  • MLIR不会轻易调用具有复杂参数的C函数。

步骤3 :(任务2)迁移暴徒

请注意,大约有三种类型的重击:

  • KernelThunk,它启动一个内核。
  • 控制流重击,具有主机控制流逻辑(有条件的,同时的,顺序的)和启动主体内核。
  • 库重击:cuDNN,cuBLAS,cuFFT,NCCL等

该计划是:

  • 使Thunks(反序列化)。
  • 帮助将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已经进行了大量工作,因此需要完成的实际工作量尚待观察。