【从零开始学深度学习编译器】十九，MLIR的Pass机制实践

0x0. 前言

这个系列的前面几篇文章对MLIR的组件有了一些粗浅的认识，这篇文章不继续讲MLIR的架构。而是从实践的角度带读者来看一下，MLIR帮助我做了什么，这里仍然以OneFlow Dialect为例。在MLIR：摩尔定律终结的编译器基础结构 论文解读 这篇文章的评论部分已经简单介绍了OneFlow Dialect相关的组件是如何实现的。在实现了OneFlow Dialect的基础上，我继续来介绍一下MLIR的Pass机制是如何助力OneFlow模型训练和推理加速的。

从零开始学深度学习编译器系列的文章以及实验代码均整理在这个仓库：https://github.com/BBuf/tvm_mlir_learn，目前已收获300+ star 。感兴趣可以自行查看，如果能点个star就更好啦。

0x1. 背景

当前Transformer架构已经成为做AI的算法开发人员和工程师们不得不谈的基础架构。由Transformer基础架构派生出了一系列超大模型如Bert和GPT-2，在业界都有非常大的影响，并且也引领了大模型的潮流。然而大模型的高昂训练成本让很多人甚至很多公司望而却步，通常只能在预训练的大模型上做一些下游任务，因此如何加速大模型训练是十分重要的。在2019年，英伟达成功地构建并训练了最大的语言模型 GPT-2 8B，这一模型包含 83 亿参数量，是 BERT-Large 模型的 24 倍、GPT-2 的 5.6 倍。英伟达将这一模型称为「Megatron」（威震天），还开源了用来训练这一模型的 pytorch 代码：https://github.com/NVIDIA/Megatron-LM。

这篇论文中提到了很多加速大模型训练的手段，特别的如模型并行训练技术，但本人对分布式训练了解很少这里不做介绍。我这里唯一的关注点是在Megatron论文（https://arxiv.org/pdf/2104.04473.pdf）的4.2节中提到的编译优化加速模型训练：

这一节编译优化讲的主要是可以通过PyTorch JIT技术来做一些Op融合，比如将bias_add和gelu融合成一个算子，bias_add+dropout融合成一个算子。做了这些算子融合之后，不仅可以避免GPU重复读写数据减少显存占用，还可以减少cuda kernel launch次数对整个计算过程进行加速。

要实现论文中提到的编译优化，需要两个前置条件。一是框架提供了融合Op的实现，二是基于编译器实现一个优化Pass来自动寻找模型中可以融合的Pattern并将其重写为等价的融合Op，达到对计算图进行运行加速的目的。

0x2. BiasAdd Dropout以及融合算子简介

在OneFlow中为了对标Megatron的bias_add和dropout fuse，实现了一个fused_bias_add_mask_scale算子，做的事情就是将BiasAdd和Dropout融合成一个算子来加速。这个算子的实现过程这里不展开，重点是如何在模型中基于MLIR自动发现这种Pattern并自动将这种Pattern替换为fused_bias_add_mask_scale算子。

为了下一节更好的理解融合Pass的做法，这里对bias_add，dropout以及fused_bias_add_mask_scale这三种Op的参数列表进行简要介绍。

• bias_add 算子：

>>> import oneflow as flow
>>> x = flow.randn(2, 3)
>>> y = flow.randn(3)
>>> z = flow._C.bias_add(x, y, axis=1)

可以看到这个算子有3个参数，一个是输入Tensor，一个是bias Tensor，还有一个axis属性表示需要把bias Tensor附加到输入Tensor的哪个维度上。在Transformer结构中，带偏置的线性层（nn.Linear）就是通过一个矩阵乘法（matmul）算子和一个bias_add实现的。

• nn.Dropout 算子：Dropout算子相信大家非常熟悉，不需要多解释，可以参考下方OneFlow算子文档。

例子：

>>> import numpy as np
>>> import oneflow as flow

>>> m = flow.nn.Dropout(p=0)
>>> arr = np.array(
...    [
...        [-0.7797, 0.2264, 0.2458, 0.4163],
...        [0.4299, 0.3626, -0.4892, 0.4141],
...        [-1.4115, 1.2183, -0.5503, 0.6520],
...    ]
... )
>>> x = flow.Tensor(arr)
>>> y = m(x)
>>> y 
tensor([[-0.7797,  0.2264,  0.2458,  0.4163],
        [ 0.4299,  0.3626, -0.4892,  0.4141],
        [-1.4115,  1.2183, -0.5503,  0.6520]], dtype=oneflow.float32)

• fused_bias_add_mask_scale：fused_bias_add_mask_scale算子需要bias_add算子的输入a和b（bias），然后还需要一个由输入a调用random_mask_like Op产生的掩码Tensor mask作为它的第三个输入，最后还需要bias_add算子的axis属性和Dropout的p属性。

这里需要解释一下为什么需要mask。其实Dropout算子在实现的时候也会产生两个输出，一个是输出Tensor，一个是mask。这是因为Dropout会根据p和我们输入的随机数种子产生一个mask来决定哪些位置的神经元应该保留，哪些位置的神经元置0，为了正确的反向传播的需要我们必须保留这个mask来求取输入Tensor对应的梯度。因此在fused_bias_add_mask_scale Op中，需要将mask显示的传给这个Op，因为这个Op的输出只有一个，不会再输出一个额外的mask了。而这个mask的生成是利用oneflow内部的random_mask_like Op来生成的，这个Op接受一个输入Tensor和p以及一个随机数种子来产生一个具有一定概率分布的掩码Tensor mask。

0x3. Pattern匹配和重写

在了解了这些Op的操作数，属性以及输出之后，我们就可以基于MLIR来做针对BiasAdd和Dropout的Patten自动匹配和重写了。这个功能实现在：https://github.com/Oneflow-Inc/oneflow/pull/7709。

首先，我们需要在oneflow/ir/include/OneFlow/OneFlowPatterns.td这个文件中基于MLIR的DRR框架写出自动匹配和重写的模板，实现如下：

def GetDefaultSeed :
  NativeCodeCall<"mlir::oneflow::GetDefaultSeed($_builder)">;

def FusedBiasAddMaskScale :
  NativeCodeCall<"mlir::oneflow::CreateFusedBiasAddMaskScale($_builder, $0, $1, $2)">;

def IsAddToOutputNone: Constraint<CPred<"mlir::oneflow::IsAddToOutputNone($0)">, "">;

def FusedBiasAddDropoutPattern : Pattern<
  (
    OneFlow_DropoutOp: $dropout_res
    (
      OneFlow_BiasAddOp: $bias_add_res
        $a,
        $b,
        $bias_add_op_name,
        $bias_add_device_tag,
        $bias_add_device_name,
        $bias_add_scope_symbol_id,
        $bias_add_hierarchy,
        $bias_add_op_axis
    ),
    $_add_to_output,
    $dropout_op_name,
    $dropout_device_tag,
    $dropout_device_name,
    $dropout_scope_symbol_id,
    $dropout_hierarchy,
    $dropout_op_rate
  ),
  [
    (
      FusedBiasAddMaskScale
      $dropout_res__0,
      $bias_add_res,
      (
        OneFlow_RandomMaskLikeOp : $mask
          $a,
          $bias_add_op_name,
          $dropout_device_tag,
          $dropout_device_name,
          $dropout_scope_symbol_id,
          $dropout_hierarchy,
          $dropout_op_rate,
          (GetDefaultSeed)
      )
    ),
    (replaceWithValue $mask)
  ],
  [(IsAddToOutputNone $_add_to_output)]
>;

NativeCodeCall是一个占位代码，我们可以通过NativeCodeCall调用我们在Dialect下手写的C++函数。比如：

def GetDefaultSeed :
  NativeCodeCall<"mlir::oneflow::GetDefaultSeed($_builder)">;

这里就调用了我们在OneFlow Dialect下手写的GetDefaultSeed函数，它返回一个OneFlow的DefaultAutoGenerator类生成的随机种子，这个随机种子在Pattern里面作为RandomMaskLikeOp的一个属性被使用：

mlir::IntegerAttr GetDefaultSeed(::mlir::PatternRewriter