MLIR基础及开发初探分析

发布时间 2023-04-05 04:50:25作者: 吴建明wujianming

MLIR基础及开发初探分析

初识MLIR

0x1. 前言

最近开始做一些MLIR的工作,所以开始学习MLIR的知识。这篇笔记是对MLIR的初步印象,并不深入,适合想初步了解MLIR是什么的同学阅读,后面会继续分享MLIR的一些项目。这里要大力感谢中科院的法斯特豪斯(知乎ID)同学先前的一些分享,给了我入门MLIR的方向。

0x2. 什么是IR?

IR即中间表示(Intermediate Representation),可以看作是一种中介的数据格式,便于模型在框架间转换。在深度学习中可以表示计算图的数据结构就可以称作一种IR,例如ONNX,TorchScript,TVM Relay等等。这里举几个例子介绍一下:

首先,ONNX是微软和FaceBook提出的一种IR,他持有了一套标准化算子格式。无论你使用哪种深度学习框架(Pytorch,TensorFlow,OneFlow)都可以将计算图转换成ONNX进行存储。然后各个部署框架只需要支持ONNX模型格式就可以简单的部署各个框架训练的模型了,解决了各个框架之间模型互转的复杂问题。

但ONNX设计没有考虑到一个问题,那就是各个框架的算子功能和实现并不是统一的。ONNX要支持所有框架所有版本的算子实现是不现实的,目前ONNX的算子版本已经有10代以上,这让用户非常痛苦。IR可以类比为计算机架构的指令集,但我们是肯定不能接受指令集频繁改动的。另外ONNX有一些控制流的算子如If,但支持得也很有限。

其次,TorchScript是Pytorch推出的一种IR,它是用来解决动态图模式执行代码速度太慢的问题。因为动态图模式在每次执行计算时都需要重新构造计算图(define by run),使得性能和可移植性都比较差。为了解决这个问题,Pytorch引入了即时编译(JIT)技术即TorchScript来解决这一问题。Pytorch早在1.0版本就引入了JIT技术并开放了C++ API,用户之后就可以使用Python编写的动态图代码训练模型然后利用JIT将模型(nn.Module)转换为语言无关的模型(TorchScript),使得C++ API可以方便的调用。并且TorchScript也很好的支持了控制流,即用户在Python层写的控制流可以在TorchScript模型中保存下来,是Pytorch主推的IR。

最后,Relay IR是一个函数式、可微的、静态的、针对机器学习的领域定制编程语言。Relay IR解决了普通DL框架不支持control flow(或者要借用python 的control flow,典型的比如TorchScript)以及dynamic shape的特点,使用lambda calculus作为基准IR。

Relay IR可以看成一门编程语言,在灵活性上比ONNX更强。但Relay IR并不是一个独立的IR,它和TVM相互耦合,这使得用户想使用Relay IR就需要基于TVM进行开发,这对一些用户来说是不可接受的。

这几个例子就是想要说明,深度学习中的IR只是一个深度学习框架,公司甚至是一个人定义的一种中介数据格式,它可以表示深度学习中的模型(由算子和数据构成)那么这种格式就是IR

0x3. 为什么要引入MLIR?

目前深度学习领域的IR数量众多,很难有一个IR可以统一其它的IR,这种百花齐放的局面就造成了一些困境。我认为中科院的法斯特豪斯同学B站视频举的例子非常好,建议大家去看一下。这里说下我的理解,以TensorFlow Graph为例,它可以直接被转换到TensorRT的IR,nGraph IR,CoreML IR,TensorFlow Lite IR来直接进行部署。或者TensorFlow Graph可以被转为XLA HLO,然后用XLA编译器来对其进行Graph级别的优化得到优化后的XLA HLO,这个XLA HLO被喂给XLA编译器的后端进行硬件绑定式优化和Codegen。在这个过程中主要存在两个问题。

第一,IR的数量太多,开源要维护这么多套IR,每种IR都有自己的图优化Pass,这些Pass可能实现的功能是一样的,但无法在两种不同的IR中直接迁移。假设深度学习模型对应的DAG一共有10种图层优化Pass,要是为每种IR都实现10种图层优化Pass,那工作量是巨大的。
第二,如果出现了一种新的IR,开发者想把另外一种IR的图层优化Pass迁移过来,但由于这两种IR语法表示完全不同,除了借鉴优化Pass的思路之外,就丝毫不能从另外一种IR的Pass实现受益了,即互相迁移的难度比较大。此外,如果你想为一个IR添加一个Pass,难度也是不小的。举个例子你可以尝试为onnx添加一个图优化Pass,会发现这并不是一件简单的事,甚至需要我们去较为完整的学习ONNX源码。
第三,在上面的例子中优化后的XLA HLO直接被喂给XLA编译器后端产生LLVM IR然后Codegen,这个跨度是非常大的。这里怎么理解呢?我想到了一个例子。以优化GEMM来看,我们第一天学会三重for循环写一个naive的矩阵乘程序,然后第二天你就要求我用汇编做一个优化程度比较高的矩阵乘法程序?那我肯定是一脸懵逼的,只能git clone了,当然是学不会的。但如果你缓和一些,让我第二天去了解并行,第三天去了解分块,再给几天学习一下SIMD,再给几个月学习下汇编,没准一年下来我就可以真正的用汇编优化一个矩阵乘法了。所以跨度太大最大的问题在于,我们这种新手玩家很难参与。我之前分享过TVM的Codegen流程,虽然看起来理清了Codegen的调用链,但让我现在自己去实现一个完整的Codegen流程,那我是很难做到的。

针对上面的问题,MLIR(Multi-Level Intermediate Representation)被提出。MLIR是由LLVM团队开发和维护的一套编译器基础设施,它强调工具链的可重用性和可扩展性。下面我们具体分析一下:

针对第一个问题和第二个问题,造成这些深度学习领域IR的优化Pass不能统一的原因就是因为它们没有一个统一的表示,互转的难度高。因此MLIR提出了Dialect,我们可以将其理解为各种IR需要学习的语言,一旦某种IR学会这种语言,就可以基于这种语言将其重写为MLIR。Dialect将所有IR都放在了同一个命名空间里面,分别对每个IR定义对应的产生式以及绑定对应的操作,从而生成MLIR模型。关于Dialect我们后面会细讲,这篇文章先提一下,它是MLIR的核心组件之一。

针对第三个问题,怎么解决IR跨度大的问题?MLIR通过Dialect抽象出了多种不同级别的MLIR,下面展示官方提供的一些MLIR IR抽象,我们可以看到Dialect是对某一类IR或者一些数据结构相关操作进行抽象,比如llvm dialect就是对LLVM IR的抽象,tensor dialect就是对Tensor这种数据结构和操作进行抽象:

 

 官网提供的MLIR Dialect

除了这些,各种深度学习框架都在接入MLIR,比如TensorFlow,Pytorch,OneFlow以及ONNX等等,大家都能在github找到对应工程。

抽象了多个级别的IR好处是什么呢?这就要结合MLIR的编译流程来看,MLIR的编译流程大致如下:

 

 图源法斯特豪斯,侵删

对于一个源程序,首先经过语法树分析,然后通过Dialect将其下降为MLIR表达式,再经MLIR分析器得到目标程序。注意这个目标程序不一定是可运行的程序。比如假设第一次的目标程序是C语言程序,那么它可以作为下一次编译流程的源程序,通过Dialect下降为LLVM MLIR。这个LLVM MLIR即可以被MLIR中的JIT执行,也可以通过Dialect继续下降,下降到三地址码IR对应的MLIR,再被MLIR分析器解析获得可执行的机器码。

因此MLIR这个多级别的下降过程就类似于我们刚才介绍的可以渐进式学习,解决了IR到之间跨度太大的问题。比如我们不熟悉LLVM IR之后的层次,没有关系,我们交给LLVM编译器,我们去完成前面那部分的Dialect实现就可以了。

MLIR Toy Tutorials学习笔记

Toy语言和AST

MLIR提供了一种Toy语言来说明MLIR的定义和执行的流程。Toy语言是一种基于张量的语言,我们可以使用它来定义函数,执行一些数学计算以及输出结果。下面要介绍的例子中限制Tensor的维度是<=2的,并且Toy语言中唯一的数据类型是64位浮点类型,对应C语言中的"double"。另外Values是不可以重写的,即每个操作都会返回一个新分配的值,并自动管理释放。直接看下面这个例子:

def main() {

  # Define a variable `a` with shape <2, 3>, initialized with the literal value.

  # The shape is inferred from the supplied literal.

  var a = [[1, 2, 3], [4, 5, 6]];

 

  # b is identical to a, the literal tensor is implicitly reshaped: defining new

  # variables is the way to reshape tensors (element count must match).

  var b<2, 3> = [1, 2, 3, 4, 5, 6];

 

  # transpose() and print() are the only builtin, the following will transpose

  # a and b and perform an element-wise multiplication before printing the result.

  print(transpose(a) * transpose(b));

}

类型检查是通过类型推断静态执行的。Toy语言只需在必要时指定Tensor形状的类型声明。下面定义了一个multiply_transpose函数,注意这个函数里面参数a和b的形状我们预先都是不知道的,只有调用这个函数时我们才知道,可以关注一下下面例子中的shape变化。

# User defined generic function that operates on unknown shaped arguments.

def multiply_transpose(a, b) {

  return transpose(a) * transpose(b);

}

 

def main() {

  # Define a variable `a` with shape <2, 3>, initialized with the literal value.

  var a = [[1, 2, 3], [4, 5, 6]];

  var b<2, 3> = [1, 2, 3, 4, 5, 6];

 

  # This call will specialize `multiply_transpose` with <2, 3> for both

  # arguments and deduce a return type of <3, 2> in initialization of `c`.

  var c = multiply_transpose(a, b);

 

  # A second call to `multiply_transpose` with <2, 3> for both arguments will

  # reuse the previously specialized and inferred version and return <3, 2>.

  var d = multiply_transpose(b, a);

 

  # A new call with <3, 2> (instead of <2, 3>) for both dimensions will

  # trigger another specialization of `multiply_transpose`.

  var e = multiply_transpose(b, c);

 

  # Finally, calling into `multiply_transpose` with incompatible shape will

  # trigger a shape inference error.

  var f = multiply_transpose(transpose(a), c);

}

然后我们可以使用下面的命令来产生这个Toy语言程序的AST:

 

cd llvm-project/build/bin

./toyc-ch1 ../../mlir/test/Examples/Toy/Ch1/ast.toy --emit=ast

前提是要构建好llvm-project工程,构建过程按照https://mlir.llvm.org/getting_started/ 这里的方法操作即可,这里再列一下完整过程:

 

$ git clone https://github.com/llvm/llvm-project.git

$ mkdir llvm-project/build

$ cd llvm-project/build

$ cmake -G "Unix Makefiles" ../llvm \

     -DLLVM_ENABLE_PROJECTS=mlir \

     -DLLVM_BUILD_EXAMPLES=ON \

     -DLLVM_TARGETS_TO_BUILD="host" \

     -DCMAKE_BUILD_TYPE=Release \

     -DLLVM_ENABLE_ASSERTIONS=ON

$ cmake --build . --target check-mlir

上面Toy程序产生的AST长下面这样:

 

Module:

    Function

      Proto 'multiply_transpose' @../../mlir/test/Examples/Toy/Ch1/ast.toy:4:1

      Params: [a, b]

      Block {

        Return

          BinOp: * @../../mlir/test/Examples/Toy/Ch1/ast.toy:5:25

            Call 'transpose' [ @../../mlir/test/Examples/Toy/Ch1/ast.toy:5:10

              var: a @../../mlir/test/Examples/Toy/Ch1/ast.toy:5:20

            ]

            Call 'transpose' [ @../../mlir/test/Examples/Toy/Ch1/ast.toy:5:25

              var: b @../../mlir/test/Examples/Toy/Ch1/ast.toy:5:35

            ]

      } // Block

    Function

      Proto 'main' @../../mlir/test/Examples/Toy/Ch1/ast.toy:8:1

      Params: []

      Block {

        VarDecl a<> @../../mlir/test/Examples/Toy/Ch1/ast.toy:11:3

          Literal: <2, 3>[ <3>[ 1.000000e+00, 2.000000e+00, 3.000000e+00], <3>[ 4.000000e+00, 5.000000e+00, 6.000000e+00]] @../../mlir/test/Examples/Toy/Ch1/ast.toy:11:11

        VarDecl b<2, 3> @../../mlir/test/Examples/Toy/Ch1/ast.toy:15:3

          Literal: <6>[ 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00] @../../mlir/test/Examples/Toy/Ch1/ast.toy:15:17

        VarDecl c<> @../../mlir/test/Examples/Toy/Ch1/ast.toy:19:3

          Call 'multiply_transpose' [ @../../mlir/test/Examples/Toy/Ch1/ast.toy:19:11

            var: a @../../mlir/test/Examples/Toy/Ch1/ast.toy:19:30

            var: b @../../mlir/test/Examples/Toy/Ch1/ast.toy:19:33

          ]

        VarDecl d<> @../../mlir/test/Examples/Toy/Ch1/ast.toy:22:3

          Call 'multiply_transpose' [ @../../mlir/test/Examples/Toy/Ch1/ast.toy:22:11

            var: b @../../mlir/test/Examples/Toy/Ch1/ast.toy:22:30

            var: a @../../mlir/test/Examples/Toy/Ch1/ast.toy:22:33

          ]

        VarDecl e<> @../../mlir/test/Examples/Toy/Ch1/ast.toy:25:3

          Call 'multiply_transpose' [ @../../mlir/test/Examples/Toy/Ch1/ast.toy:25:11

            var: b @../../mlir/test/Examples/Toy/Ch1/ast.toy:25:30

            var: c @../../mlir/test/Examples/Toy/Ch1/ast.toy:25:33

          ]

        VarDecl f<> @../../mlir/test/Examples/Toy/Ch1/ast.toy:28:3

          Call 'multiply_transpose' [ @../../mlir/test/Examples/Toy/Ch1/ast.toy:28:11

            Call 'transpose' [ @../../mlir/test/Examples/Toy/Ch1/ast.toy:28:30

              var: a @../../mlir/test/Examples/Toy/Ch1/ast.toy:28:40

            ]

            var: c @../../mlir/test/Examples/Toy/Ch1/ast.toy:28:44

          ]

      } // Block

AST的解析具体实现在mlir/examples/toy/Ch1/include/toy/Parser.h和mlir/examples/toy/Ch1/include/toy/Lexer.h中,感兴趣的读者可以看一下。我对这一块并不熟悉,就暂时不深入下去了,但这个AST看起来还是比较直观的,首先有两个Function对应了Toy程序里面的multiply_transpose和main,Params表示函数的输入参数,Proto表示这个函数在ast.toy文件中的行数和列数,BinOp表示transpose(a) * transpose(b)中的*是二元Op,并列出了左值和右值。其它的以此类推也比较好理解。

 

简单介绍了一下Toy语言的几个特点以及Toy示例程序产生的AST长什么样子,如果对AST的解析感兴趣可以去查看代码实现。

生成初级MLIR

MLIR 被设计成完全可扩展的基础框架,没有封闭的属性集、操作和类型。MLIR 通过Dialect(https://mlir.llvm.org/docs/LangRef/#dialects)的概念来支持这种可扩展性。Dialect在一个特定的namespace下为抽象提供了分组机制。

在MLIR里面,Operation是抽象和计算的核心单元,在许多方面与 LLVM 指定类似。具有特定于应用程序的语义,并且可以用于表示 LLVM 中的所有核心的 IR 结构:指令、globals(类似function)和模块。下面展示一个Toy语言产生的的transpose Operation。

%t_tensor = "toy.transpose"(%tensor) {inplace = true} : (tensor<2x3xf64>) -> tensor<3x2xf64> loc("example/file/path":12:1)

结构拆分解释:

%t_tensor:这个Operation定义的结果的名字,前面的%是避免冲突,见https://mlir.llvm.org/docs/LangRef/#identifiers-and-keywords 。一个Operation可以定义0或者多个结果(在Toy语言中,只有单结果的Operation),它们是SSA值。该名称在解析期间使用,但不是持久的(例如,它不会在 SSA 值的内存表示中进行跟踪)。

"toy.transpose" :Operation的名字。它应该是一个唯一的字符串,Dialect 的命名空间前缀为“.”。这可以理解为Toy Dialect 中的transpose Operation。

(%tensor):零个或多个输入操作数(或参数)的列表,它们是由其它操作定义或引用块参数的 SSA 值。

{ inplace = true }:零个或多个属性的字典,这些属性是始终为常量的特殊操作数。在这里,我们定义了一个名为“inplace”的布尔属性,它的常量值为 true。

(tensor<2x3xf64>) -> tensor<3x2xf64>:函数形式表示的操作类型,前者是输入,后者是输出。<2x3xf64>号中间的内容描述了张量的尺寸2x3和张量中存储的数据类型f64,中间使用x连接。

loc("example/file/path":12:1):此操作的源代码中的位置。

了解了MLIR指令的基本结构后,我们把目光放到Chapter2要做什么事情上?即生成初级MLIR。我们执行下面的命令为Chapter2测试例子中的codegen.toy产生MLIR。

 

./toyc-ch2 ../../mlir/test/Examples/Toy/Ch2/codegen.toy -emit=mlir -mlir-print-debuginfo

其中codegen.toy的内容为:

 

def multiply_transpose(a, b) {

  return transpose(a) * transpose(b);

}

 

def main() {

  var a<2, 3> = [[1, 2, 3], [4, 5, 6]];

  var b<2, 3> = [1, 2, 3, 4, 5, 6];

  var c = multiply_transpose(a, b);

  var d = multiply_transpose(b, a);

  print(d);

}

产生的MLIR为:

 

module  {

  func @multiply_transpose(%arg0: tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":4:1), %arg1: tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":4:1)) -> tensor<*xf64> {

    %0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":5:10)

    %1 = toy.transpose(%arg1 : tensor<*xf64>) to tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":5:25)

    %2 = toy.mul %0, %1 : tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":5:25)

    toy.return %2 : tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":5:3)

  } loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":4:1)

  func @main() {

    %0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":9:17)

    %1 = toy.reshape(%0 : tensor<2x3xf64>) to tensor<2x3xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":9:3)

    %2 = toy.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":10:17)

    %3 = toy.reshape(%2 : tensor<6xf64>) to tensor<2x3xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":10:3)

    %4 = toy.generic_call @multiply_transpose(%1, %3) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":11:11)

    %5 = toy.generic_call @multiply_transpose(%3, %1) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":12:11)

    toy.print %5 : tensor<*xf64> loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":13:3)

    toy.return loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":8:1)

  } loc("../../mlir/test/Examples/Toy/Ch2/codegen.toy":8:1)

} loc(unknown)

我们需要弄清楚codegen.toy是如何产生的MLIR文件。也即下图的AST到MLIR表达式那部分(包含Dialect)。

 

 图源知乎法斯特豪斯,侵删

生成MLIR的流程

 

 

从AST到MLIR由是和Dialect相关的这部分

这里首先有一个MLIRGen函数负责遍历AST。具体在mlir/examples/toy/Ch2/mlir/MLIRGen.cpp文件中实现,里面有一个mlirGen函数,实现如下:

 

/// Dispatch codegen for the right expression subclass using RTTI.

  mlir::Value mlirGen(ExprAST &expr) {

    switch (expr.getKind()) {

    case toy::ExprAST::Expr_BinOp:

      return mlirGen(cast<BinaryExprAST>(expr));

    case toy::ExprAST::Expr_Var:

      return mlirGen(cast<VariableExprAST>(expr));

    case toy::ExprAST::Expr_Literal:

      return mlirGen(cast<LiteralExprAST>(expr));

    case toy::ExprAST::Expr_Call:

      return mlirGen(cast<CallExprAST>(expr));

    case toy::ExprAST::Expr_Num:

      return mlirGen(cast<NumberExprAST>(expr));

    default:

      emitError(loc(expr.loc()))

          << "MLIR codegen encountered an unhandled expr kind '"

          << Twine(expr.getKind()) << "'";

      return nullptr;

    }

  }

这个函数会根据AST中的节点类型递归调用其它的mlirGen子函数,并在各个子函数完成真正的转换MLIR表达式的操作。以上面codege.toy的transpose(a)操作为例,对应的mlirGen子函数为:

 

/// Emit a call expression. It emits specific operations for the `transpose`

  /// builtin. Other identifiers are assumed to be user-defined functions.

  mlir::Value mlirGen(CallExprAST &call) {

    llvm::StringRef callee = call.getCallee();

    auto location = loc(call.loc());

 

    // Codegen the operands first.

    SmallVector<mlir::Value, 4> operands;

    for (auto &expr : call.getArgs()) {

      auto arg = mlirGen(*expr);

      if (!arg)

        return nullptr;

      operands.push_back(arg);

    }

 

    // Builtin calls have their custom operation, meaning this is a

    // straightforward emission.

    if (callee == "transpose") {

      if (call.getArgs().size() != 1) {

        emitError(location, "MLIR codegen encountered an error: toy.transpose "

                            "does not accept multiple arguments");

        return nullptr;

      }

      return builder.create<TransposeOp>(location, operands[0]);

    }

 

    // Otherwise this is a call to a user-defined function. Calls to

    // user-defined functions are mapped to a custom call that takes the callee

    // name as an attribute.

    return builder.create<GenericCallOp>(location, callee, operands);

  }

我们可以看到if (callee == "transpose")这里是对函数签名进行判断,如果是transpose 那么就需要新建一个TransposeOp类型的MLIR节点,即builder.create<TransposeOp>(location, operands[0])。这行代码涉及到MLIR的Dialect和TableGen,我们详细解释一下。

 

在【从零开始学深度学习编译器】十一,初识MLIR 中已经说过,MLIR是通过Dialect来统一各种不同级别的IR,即负责定义各种Operation和解析,同时还具有可扩展性。在Toy语言中我们也定义了Dialect,定义这个Dialect的时候是通过TableGen规范来定义到mlir/examples/toy/Ch2/include/toy/Ops.td中的。

 

// Provide a definition of the 'toy' dialect in the ODS framework so that we

// can define our operations.

def Toy_Dialect : Dialect {

  let name = "toy";

  let cppNamespace = "::mlir::toy";

}

在MLIR中,Dialect和Operation(也可以说算子)的定义是框架是基于TableGen(一种声明性编程语言)规范构造的,在源码中它以.td的格式存在,在编译时会自动生成对应的C++文件,生成定义好的Dialect。使用TableGen的好处不仅是因为它是声明性的语言让新增Dialect和Operation变得简单,而且容易修改和维护。可能我解释得不是很直观,但我们可以直接结合Chapter2的代码mlir/examples/toy/Ch2/include/toy/Ops.td 来理解。后面我们会看到在Toy语言的示例中,.td文件的组成以及TableGen是如何自动解析.td生成C++代码的。

 

这里首先在td中定义一下Toy Dialect,并建立和Dialect的链接,它负责将后续在Toy Dialect空间下定义的所有Operation联系起来。即:

 

// Provide a definition of the 'toy' dialect in the ODS framework so that we

// can define our operations.

def Toy_Dialect : Dialect {

  let name = "toy";

  let cppNamespace = "::mlir::toy";

}

然后构造一个Toy_Op类代表Toy Dialect下所有Operation的基类,后面新增Operation都需要继承这个类。

 

// Base class for toy dialect operations. This operation inherits from the base

// `Op` class in OpBase.td, and provides:

//   * The parent dialect of the operation.

//   * The mnemonic for the operation, or the name without the dialect prefix.

//   * A list of traits for the operation.

class Toy_Op<string mnemonic, list<OpTrait> traits = []> :

    Op<Toy_Dialect, mnemonic, traits>;

下面给出transpose Operation的定义感受一下:

 

def TransposeOp : Toy_Op<"transpose"> {

  let summary = "transpose operation";

 

  let arguments = (ins F64Tensor:$input);

  let results = (outs F64Tensor);

 

  let assemblyFormat = [{

    `(` $input `:` type($input) `)` attr-dict `to` type(results)

  }];

 

  // 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); }];

}

在继承Toy_Op的基础上,还使用TableGen语法定义了描述信息,参数,值,builder,verfier这些元素。

编写完td文件之后,就可以使用mlir-tblgen工具生成C++代码,先使用下面的命令生成Dialect的C++代码:./mlir-tblgen -gen-dialect-decls llvm-project/mlir/examples/toy/Ch2/include/toy/Ops.td -I ../../mlir/include/

 

 

 自动生成的Toy Dialect C++代码

把上面的命令换成./mlir-tblgen -gen-op-defs llvm-project/mlir/examples/toy/Ch2/include/toy/Ops.td -I ../../mlir/include/ 就可以生成Operation的C++代码。感兴趣的读者可自行查看。

与工具链 toyc-ch2 的联系,查看CMakeList.txt 文件(默认位置为 llvm-project/mlir/examples/toy/Ch2/include/toy):

set(LLVM_TARGET_DEFINITIONS Ops.td)

mlir_tablegen(Ops.h.inc -gen-op-decls)

mlir_tablegen(Ops.cpp.inc -gen-op-defs)

mlir_tablegen(Dialect.h.inc -gen-dialect-decls)

mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs)

add_public_tablegen_target(ToyCh2OpsIncGen)

使用mlir-tblgen搭配 -gen-op-decls 和 -gen-op-defs 选项,生成 Ops.h.inc 声明代码和 Ops.cpp.inc 定义代码,将两者作为构建工具链 toyc-ch2 的代码依赖。

 总结一下,Chapter2主要介绍了MLIR中的MLIRGen,Dialect,Operation以及TableGen这几个MLIR的核心组成部分以及它们是如何相互作用的。它们的关系可以借用中科院Zhang Hongbin同学的PPT来更好的描述:

 

 

 图源知乎法斯特豪斯,为了方便理解借用到这里,侵删

 

参考文献链接

https://mp.weixin.qq.com/s/4pD00N9HnPiIYUOGSnSuIw

https://mp.weixin.qq.com/s/jMHesvKmAUU5dYH0WznulA