TPU-MLIR中的Lowering过程

Lowering 的作用#

从高级语言过渡到低级的机器码,在 MLIR 中是通过 Lowering(降低)来实现的。

在上期视频中我们提到 MLIR 的 dialect conversion 是指 IR 在不同 dialect 间进行转换的过程,而 lowering 其实就是从 high-level dialect 到 low-level dialect 的 dialect conversion。

以 TPU-MLIR 为例,lowering 在其中扮演的作用就是将处理器无关层的 mlir 模型转换到处理器相关层

模型的转换其实本质上就是算子的转换,Top 层的算子可以分为 F32 与 INT8 两种,大多数网络中算子都是 F32 的,部分如 tflite 等量化过的网络存在 INT8 算子。

Top 层的 F32 算子可以被直接转换为 F32,F16 或 BF16 的 tpu 层算子,如果要转换为 int8 算子,则需要经过校准量化,而 int8 的 top 层算子则只能转换为 int8 的 tpu 层算子

另外,在进行混合精度运算,也就是前后算子的类型不一致时,为了保持运算精度相同,TPU-MLIR 会在算子间插入 CastOp 对算子的输出类型进行转换。

如何实现 Top 层到 Tpu 层的转换过程?#

  1. 首先我们就需要对目标 dialect 以及转换 pass 进行定义, 这些定义我们都可以在 td 文件中完成,其中 ConvertTopToTpu pass 定义中的 mode 指的是对总体算子的量化模式 F32/INT8/F16/BF16。而 qtable 则用来明确模型中某一层算子的量化模式,使用 qtable 就可能会产生我们之前提到的前后算子类型不一致的情况,这时候就需要插入 CastOp 来帮助混合精度运算的实现。processor 则是指定需要 lower 到的特定处理器,当 mode 为 int8 时,又存在对称与非对称 int8 量化两种情况,所以即使都在同一层 dialect 里,转换过程与最终生成的 IR 也会根据 option 里的内容不同而有所不同。

  1. 之后我们由 OpRewritePattern 派生一个 TopLowering 类来定义算子 rewrite 的逻辑,类中定义的一系列虚函数用于将算子 lowering 到某个特定处理器上时被实现,例如我们可以在 LoweringBM1684X 头文件中看到,下降到 BM1684x 处理器的中所需实现的量化模式。

之后针对每个 BM1684X 所支持的算子我们需要实现相应的 pattern,而每个 pattern 中就包含了不同量化模式的具体 lowering 实现

  1. 接下来就是实现 pass 中的功能,大致的流程也与我们这期视频中提到的相同。

  • 首先就是定义合法的 dialect 与 op,在转换时我们会保留这些合法的 op
  • 然后根据目标处理器与算子的量化类型添加相应的 rewrite pattern
  • 最后通过转换接口迭代地遍历 Top 层模型并将相应的 patterns 应用到转换过程中。另外,这个转换过程并不会只进行一次,我们会在整个流程中多次添加不同的 patterns,进行转换,然后清空 patternSet,为下一次转换做准备。
  • 在转换结束后,我们还会对转换过后的模型进行类型验证,并在相应的位置添加 CastOp 以应对混精度的情况。
  • 至此,整个 Top 层到 Tpu 层的 lowering 流程结束