【从零开始学深度学习编译器】十三，如何在MLIR里面写Pass？

最新推荐文章于 2025-10-13 14:05:52 发布

原创

最新推荐文章于 2025-10-13 14:05:52 发布 · 2k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能 #计算机视觉

【GiantPandaCV导语】这篇文章是学习了比较久然后按照自己的理解步骤重新总结了下来，主要是MLIR Toy Tutorials第3，4篇文章的内容。这里主要讲解了如何在MLIR中自定义Pass，这里主要以消除连续的Transpose操作和Reshape操作，内联优化Pass，形状推导Pass 4个例子来介绍了在MLIR中定义Pass的各种技巧，实际上也并不难理解。但要入门MLIR掌握这些Pass实现的技巧是有必要的。我在从零开始学习深度学习编译器的过程中维护了一个project：https://github.com/BBuf/tvm_mlir_learn ，主要是记录学习笔记以及一些实验性代码，目前已经获得了150+ star，对深度学习编译器感兴趣的小伙伴可以看一下，能点个star就更受宠若惊了。

前言

在【从零开始学深度学习编译器】十一，初识MLIR 和【从零开始学深度学习编译器】十二，MLIR Toy Tutorials学习笔记一这两篇文章中，我们已经初步了解了MLIR为何物，并且讲到了Toy语言从源文件生成MLIR的具体过程，以及在这个过程中MLIR中的MLIRGen，Dialect，Operation以及TableGen这几个MLIR的核心组成部分以及它们是如何相互作用的。

这篇笔记将基于Toy Tutorials总结MLIR中的表达式变形是如何实现的。

Chapter3: MLIR中的表达式变形(如何写Pass)

在Chapter2中我们已经生成了初级的合法MLIR表达式，但MLIR表达式一般还可以被进一步处理和简化，可以类比于TVM的Pass对Relay IR的优化。这里我们来看看要对初级的MLIR表达式进行变形是如何做的？在MLIR中是基于表达式匹配和重写来完成MLIR表达式变形的。这个教程中分别介绍使用C++模板匹配和重写以及基于DRR框架（https://mlir.llvm.org/docs/DeclarativeRewrites/）来定义表达式重写规则，然后使用ODS框架来自动生成代码。

使用C++模式匹配和重写的方法优化转置（Transpose）操作

这里的目标是要消除两个具有相互抵消效果的转置序列：transpose(transpose(X)) -> X，即对同一个输入进行连续的Transpose操作肯定存在冗余的操作。该操作对应的源码如下（在mlir/test/Examples/Toy/Ch3/transpose_transpose.toy中）：

def transpose_transpose(x) {
   
   
  return transpose(transpose(x));
}

如果不使用任何优化Pass，我们看下这个Toy源程序生成的MLIR表达式是什么样子的，使用下面的命令产生MLIR：./toyc-ch3 ../../mlir/test/Examples/Toy/Ch3/transpose_transpose.toy -emit=mlir。

func @transpose_transpose(%arg0: tensor<*xf64>) -> tensor<*xf64> {
   
   
    %0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64>
    %1 = toy.transpose(%0 : tensor<*xf64>) to tensor<*xf64>
    toy.return %1 : tensor<*xf64>
  }

可以看到生成的MLIR表达式中对x进行了两次真正的transpose操作，并且返回了两次transpose之后的Tensor。但实际上这两次transpose是不必要的，因为输出的结果其实就是传入的x。所以为了优化这种情况，我们先使用C++方式来写出表达式匹配和重写的代码（在mlir/examples/toy/Ch3/mlir/ToyCombine.cpp中）：

/// This is an example of a c++ rewrite pattern for the TransposeOp. It
/// optimizes the following scenario: transpose(transpose(x)) -> x
struct SimplifyRedundantTranspose : public mlir::OpRewritePattern<TransposeOp> {
   
   
  /// We register this pattern to match every toy.transpose in the IR.
  /// The "benefit" is used by the framework to order the patterns and process
  /// them in order of profitability.
  SimplifyRedundantTranspose(mlir::MLIRContext *context)
      : OpRewritePattern<TransposeOp>(context, /*benefit=*/1) {
   
   }

  /// This method attempts to match a pattern and rewrite it. The rewriter
  /// argument is the orchestrator of the sequence of rewrites. The pattern is
  /// expected to interact with it to perform any changes to the IR from here.
  mlir::LogicalResult
  matchAndRewrite(TransposeOp op,
                  mlir::PatternRewriter &rewriter) const override {
   
   
    // Look through the input of the current transpose.
    mlir::Value transposeInput = op.getOperand();
    TransposeOp transposeInputOp = transposeInput.getDefiningOp<TransposeOp>();

    // Input defined by another transpose? If not, no match.
    if (!transposeInputOp)
      return failure();

    // Otherwise, we have a redundant transpose. Use the rewriter.
    rewriter.replaceOp(op, {
   
   transposeInputOp.getOperand()});
    return success();
  }
};

可以看到在matchAndRewrite函数中，首先获取当前操作的操作数，然后判断当前位置的操作数对应的操作是否为转置，如果是就将表达式重写为内层转置操作的操作数，不然就不需要进行优化，保持现状。

接下来，需要在归范化框架（Canonicalization Framework）中注册刚刚创建的匹配重写模式，使得框架可以调用它。对于Canonicalization 更多的介绍请看https://mlir.llvm.org/docs/Canonicalization/，注册的代码如下（代码仍在：mlir/examples/toy/Ch3/mlir/ToyCombine.cpp）：

/// Register our patterns as "canonicalization" patterns on the TransposeOp so
/// that they can be picked up by the Canonicalization framework.
void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   
   
  results.add<SimplifyRedundantTranspose>(context);
}

在我们将表达式重写规则添加到了规范化框架后，我们还需要修改一下定义Operator的td文件，启用规范化框架，同时在定义Operator添加一个“无副作用的”(NoSideEffect)新特征，现在Transpose操作的定义如下：

def TransposeOp : Toy_Op<"transpose", [NoSideEffect]> {
   
   
  let summary = "transpose operation";

  let arguments = (ins F64Tensor:$input);
  let results = (outs F64Tensor);

  let assemblyFormat = [{
   
   
    `(` $input `:` type($input) `)` attr-dict `to` type(results)
  }];

  // Enable registering canonicalization patterns with this operation.
  let hasCanonicalizer = 1;

  // Allow building a TransposeOp with from the input operand.
  let builders = [
    OpBuilder<(ins "Value":$input)>
  ];

  // Invoke a static verify method to verify this transpose operation.
  let verifier = [{
   
    return ::verify(*this); }];
}

最后，我们需要在主程序中将基于规范化框架的优化添加到运行流程里，这部分代码在mlir/examples/toy/Ch3/toyc.cpp中的dumpMLIR函数里面。如下图的红框部分：

下降MLIR的时候启用优化Pass

至此，我们就完成了基于C++的MLIR表达式匹配和重写，我们可以通过下面的命令来看下经过上面transpose表达式的重写后产生的MLIR表达式是否已经去掉了transpose。命令为：./toyc-ch3 ../../mlir/test/Examples/Toy/Ch3/transpose_transpose.toy -emit=mlir -opt。结果为：

func @transpose_transpose(%arg0: tensor<*xf64>) -> tensor<*xf64> {
   
   
    toy.return %arg0 : tensor<*xf64>
  }

可以看到优化后的MLIR表达式已经去掉了transpose操作了，达到了优化效果。

使用 DRR 优化张量变形（Reshape）操作

MLIR还提供了一种表达式重写的方法，是基于DDR规则的方式来自动生成表达式匹配和重写函数，代码生成的部分仍然基于ODS框架实现。DRR（Declarative, Rule-based Pattern-match and Rewrite）：声明性、基于规则的模式匹配和重写方法。它是一种基于 DAG 的声明性重写器，提供基于表格的模式匹配和重写规则的句法。

这里以消除MLIR表达式中冗余的张量reshape操作为例，对应的Toy源文件如下（在mlir/test/Examples/Toy/Ch3/trivial_reshape.toy中）：

def main() {
   
   
  var a<2,1> = [1, 2];
  var b<2,1> = a;
  var c<2,1> = b;
  print(c);
}

使用下面的命令先产生对应的MLIR表达式看看：./toyc-ch3 ../../mlir/test/Examples/Toy/Ch3/trivial_reshape.toy -emit=mlir

module  {
   
   
  func @main() {
   
   
    %0 = toy.constant dense<[1.000000e+00, 2.000000e+00]> : tensor<2</

最低0.47元/天解锁文章