MLIR初识 —— Dialect及Operation详解

MLIR初识 —— Dialect及Operation详解

论文链接： MLIR: A Compiler Infrastructure for the End of Moore’s Law

1. MLIR简介

MLIR 全称是 Multi-Level Intermediate Representation (多级中间表示)，是一种全新的编译器框架。

1.1 IR是什么

IR即 Intermediate Representation，可以看作是一种数据格式，作为从端到端转换中的中间表示。例如深度学习模型一般表示为计算图，能够表示计算图的数据结果就可以称为一种IR，例如ONNX、TorchScript、TVM Relay等等。

• ONNX(Open Neural Network Exchange) : ONNX 协议首先由微软和Meta提出，它定义了一组和环境、平台均无关的标准格式（如算子功能）。在训练完成后可以将支持框架(Pytorch、Tensorflow等)的模型转化为 ONNX 文件进行存储，ONNX 文件不仅存储了神经网络模型的权重，也存储了模型的结构信息以及网络中每一层的输入输出等信息。
• TorchScrpit : PyTorch 最大的卖点是它对动态网络的支持，比其他需要构建静态网络的框架拥有更低的学习成本。但动态图模式在每次执行计算时都要重新构造计算图，非固定的网络结构给网络结构分析并进行优化带来了困难。TorchScript 就是为了解决这个问题而诞生的工具，包括代码的追踪及解析、中间表示的生成、模型优化、序列化等各种功能。
• Relay IR : 与 TVM 框架绑定，是一个函数式、可微的、静态的、针对机器学习的领域定制编程语言，解决了普通DL框架不支持 control flow 以及 dynamic shape 的特点，使用 lambda calculus 作为基准IR。

1.2 常见的IR表示系统

在上图表示中：

(1) 由于 C、C++ 源码直接转成 AST 时，并不会进行语言特定的优化，程序的优化主要集中于 LLVM IR 阶段。但 LLVM IR 表示层级较低，会丢失源码中的部分信息(如报错信息)，会导致优化不充分。

(2) 类似于Tensorflow、Keras等框架，会先转化为计算图Computation Graph形式，然后会基于图做一定的优化。但图阶段缺少硬件部署的相关信息，所以后续会转化为某个后端的内部表示，根据不同的硬件(TPU、Phone)，进行算子融合等等优化。

…

可见，当前IR表示系统主要的问题

• 可复用性差：针对不同种类IR开发的Pass(优化)可能重复，但不同IR的同类Pass可能并不兼容。
• 不透明：前层IR所作的Pass优化在后层中不可见，可能导致优化重复。
• 变换开销大：转换过程中存在多种IR，这些不同类型的IR转换时开销很大。

1.3 MLIR的提出

Tensorflow 团队较早时采用了多种IR的部署，这样导致软件碎片化较为严重。

因此 Tensorflow 团队就提出了 MLIR，主要是为了统一各类IR格式，协调各类IR的转换，带来更高的优化效率。

2. Dialect及Operation详解

参考：

Chapter 2: Emitting Basic MLIR - MLIR (llvm.org)

MLIR Toy Tutorials 第二章 生成初步的 MLIR 代码 - 知乎 (zhihu.com)

MLIR & python binding简介_哔哩哔哩

2.1 Dialect

1. Dialect 是什么？

从源程序到目标程序，要经过一系列的抽象以及分析，通过 Lowering Pass 来实现从一个IR到另一个IR的转换。但IR之间的转换需要统一格式，统一IR的第一步就是要统一“语言”，各个IR原来配合不默契，谁也理解不了谁，就是因为“语言”不通。

因此 MLIR 提出了Dialect，各种IR可以转换为对应的 mlir Dialect，不仅方便了转换，而且还能随意扩展。不妨将dialect看成各种具有IR表达能力的黑盒子，之后的编译流程就是在各种dialect之间转化。

2. dialect 是怎么工作的？

dialect 将所有的IR放在了同一个命名空间中，分别对每个IR定义对应的产生式并绑定相应的操作，从而生成一个MLIR的模型。

每种语言的 dialect(如tensorflow dialect、HLO dialect、LLVM IR dialect)都是继承自 mlir::Dialect，并注册了属性、操作和数据类型，也可以使用虚函数来改变一些通用性行为。

整个的编译过程：从源语言生成 AST（Abstract Syntax Tree，抽象语法树），借助 dialect 遍历 AST，产生 MLIR 表达式（此处可为多层IR通过 Lowering Pass 依次进行分析），最后经过 MLIR 分析器，生成目标硬件程序。

3. dialect 内部构成

dialect主要是由自定义的 Type、Attribute、Interface 以及 operation 构成。operation 细分为Attribute、Type、Constraint、Interface、Trait（属性、类型、限制、接口、特征）。同时存在 ODS 和 DRR 两个重要的模块，这两个模块都是基于 tableGen 模块，ODS 模块用于定义 operation ，DRR 模块用于实现两个 dialect 之间的 conversion。

2.2 Operation

Operation 是 Dialect 的重要组成部分，是抽象和计算的核心单元，可以看成是方言语义的基本元素。

下面的例子可理解为：

生成的结果是 %t_tensor，toy dialect，执行的是 transpose 操作，输入数据是 %tensor，能够将 tensor<2x3xf64> 的数据转换成tensor<3x2xf64> 的数据，该 transpose 的位置在 “example/file/path”，第12行，第1个字符：

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

结构拆分解析：

（1）%t_tensor：定义结果名称，SSA值，由%和<t_tensor>构成，一般<t_tensor>是一个整数型数字。

IR 是 LLVM 的设计核心，它采用 SSA（Single-Static Assignments，静态单赋值）的形式，并具备两个重要特性：

• 代码被组织成三地址指令
• 有无限的寄存器

（2）"toy.transpose"：操作的名称，应该是唯一的字符串，方言空间以.开头；指明为 Toy Dialect 的transpose 操作；.之前的内容是 Dialect 命名空间的名字，.后面是操作的名称。

（3）(%tensor)：输入的操作数的列表，多个操作数之间用逗号隔开。

（4）{inplace = true}：属性字典，定义一个名为inplace的布尔类型，其常量值为true。

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

（6）loc("example/file/path":12:1)：此操作的源代码中的位置。每个操作都有与之关联的强制性源位置，在 MLIR 中是核心要求，并且 API 依赖并操纵他。例如：如果一个转换将操作替换成另一个操作，必须在新的操作中附加一个位置，可以追踪该操作的来源。所以，在使用工具链 mlir-opt 中默认没有这个位置信息，添加 -mlir-print-debuginfo 标志指定要包含位置。

更一般的格式可见下图：

3. 创建新的dialect(添加新的operation)

本节创建新的dialect包括 手动编写C++创建 以及 利用ODS框架生成

ODS 全称 Operation Definition Specification，操作者只需要根据 operation 框架定义的规范，在一个.td文件中填写相应的内容，使用 mlir 的 tableGen 工具就可以自动生成上面的 C++ 代码。

本节完全参考官方文档 ：Chapter 2: Emitting Basic MLIR - MLIR (llvm.org)

本节将以Toy语言为例，演示构造 Toy Dialect并添加相应的Operation的流程。

Toy语言是为了验证及演示MLIR系统的整个流程而开发的一种基于Tensor的语言。
Toy 语言具有以下特性：

• Mix of scalar and array computations, as well as I/O
• Array shape Inference
• Generic functions
• Very limiter set of operators and features

3.1 定义 Toy Dialect

Dialect 将对 Toy 语言的结构进行建模，并为高级分析和转换提供方便的途径。

1. 使用 C++ 语言手动编写

// 下面是官方给出的Toy Dialect定义
// 默认位置为 ../mlir/examples/toy/Ch2/include/toy/Dialect.h
class ToyDialect : public mlir::Dialect {
public:
  explicit ToyDialect(mlir::MLIRContext *ctx);

  /// Provide a utility accessor to the dialect namespace.
  static llvm::StringRef getDialectNamespace() { return "toy"; }

  /// An initializer called from the constructor of ToyDialect that is used to
  /// register attributes, operations, types, and more within the Toy dialect.
  void initialize();
};

2. 使用 ODS 框架自动生成

在使用 ODS 定义操作的这些代码，都在Ops.td中，默认位置为 …/mlir/examples/toy/Ch2/include/toy/Ops.td

下面的代码块定义一个名字为 Toy 的 Dialect 在 ODS 框架中，使用let <...> = "..."/[{...}];方式依次明确 name、summary、description 和 cppNamespace（对应 Dialect 类所在的 C++ 命名空间）各个字段的定义。

def Toy_Dialect : Dialect {
  // The namespace of our dialect, this corresponds 1-1 with the string we
  // provided in `ToyDialect::getDialectNamespace`.
  let name = "toy";

  // A short one-line summary of our dialect.
  let summary = "A high-level dialect for analyzing and optimizing the "
                "Toy language";

  // A much longer description of our dialect.
  let description = [{
    The Toy language is a tensor-based language that allows you to define
    functions, perform some math computation, and print results. This dialect
    provides a representation of the language that is amenable to analysis and
    optimization.
  }];

  // The C++ namespace that the dialect class definition resides in.
  let cppNamespace = "toy";
}

然后在编译阶段，由框架自动生成相应的 C++ 代码。当然也可以运行下面的命令 直接得到生成的 C++ 代码。

${build_root}/bin/mlir-tblgen -gen-dialect-decls ${mlir_src_root}/examples/toy/Ch2/include/toy/Ops.td -I ${mlir_src_root}/include/

下图中右侧是 ODS 中的定义，左侧是自动生成的 C++ 代码。

3.2 加载到 MLIRContext 中

定义好 Dialect 之后，需要将其加载到 MLIRContext 中。默认情况下，MLIRContext 只加载内置的 Dialect，若要添加自定义的 Dialect，需要加载到 MLIRContext。

// 此处的代码与官方文档中的稍有不同，但实际意义相同。
// 在代码文件 toyc.cpp 中，默认位置为 ../mlir/examples/toy/Ch2/toyc.cpp。
int dumpMLIR() {
...
  // Load our Dialect in this MLIR Context.
  context.getOrLoadDialect<mlir::toy::ToyDialect>();
...
}

3.3 定义operation

有了上述的 Toy Dialect，便可以定义操作(operation)。官方文档围绕 Toy toy.ConstantOp 的定义介绍如何使用 C++ 的方式直接定义 operation。

# 此操作没有输入，返回一个常量。
%4 = "toy.constant"() {value = dense<1.0> : tensor<2x3xf64>} : () -> tensor<2x3xf64>

1. 使用 C++ 语言手动编写

operation 类是继承于 [CRTP](Curiously recurring template pattern - Wikipedia) 类，有一些可选的 traits 来定义行为。下面是 ConstantOp 的官方定义：

// `mlir::Op` is a CRTP class
class ConstantOp : public mlir::Op<
                     ConstantOp, // The ConstantOp 
                     mlir::OpTrait::ZeroOperands, // takes zero input operands
                     mlir::OpTrait::OneResult, // returns a single result.
                     mlir::OpTraits::OneTypedResult<TensorType>::Impl> {
 public:
  // Op inherit the constructors from the base Op class.
  using Op::Op; 
  // Return a unique name of the operation
  static llvm::StringRef getOperationName() { return "toy.constant"; }
  // Return a value by fetching it from the attribute
  mlir::DenseElementsAttr getValue(); 
  // Operations may provide additional verification beyond what the attached traits provide.
  LogicalResult verifyInvariants(); 

  // Provide an interface to build this operation from a set of input values.
  // mlir::OpBuilder::create<ConstantOp>(...)
  // Build a constant with the given return type and `value` attribute.
  static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
                    mlir::Type result, mlir::DenseElementsAttr value);
  // Build a constant and reuse the type from the given 'value'.
  static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
                    mlir::DenseElementsAttr value);
  // Build a constant by broadcasting the given 'value'.
  static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
                    double value);
};

定义好 operation 的行为后，我们可以在 Toy Dialect 的 initialize 函数中注册(register)，之后才可以正常在 Toy Dialect 中使用 ConstantOp。

// 位于../mlir/examples/toy/Ch2/mlir/Dialect.cpp
void ToyDialect::initialize() {
  addOperations<ConstantOp>();
}

2. 使用 ODS 框架自动生成

首先在 ODS 中定义一个继承自 Op 类的基类 Toy_Op。

Operation 和 Op的区别

Operation：用于对所有操作的建模，并提供通用接口给操作的实例。

Op：每种特定的操作都是由 Op 类继承来的。同时它还是 Operation * 的 wrapper，这就意味着，当我们定义一个 Dialect 的 Operation 的时候，我们实际上是在提供一个 Operation 类的接口。

Op 类的定义在 OpBased.td 文件中，默认位置为 …/mlir/include/mlir/IR/OpBased.td。

下面的代码都在Ops.td中，默认位置为 …/mlir/examples/toy/Ch2/include/toy/Ops.td

class Toy_Op<string mnemonic, list<OpTrait> traits = []> :
    Op<Toy_Dialect, mnemonic, traits>;
// Toy_Dialect : 父类 Dialect 操作
// mnemonic : 注记符号，一般是一个字符串型的单词，代表了该操作的含义
// traits : 该操作的一些特征，放在一个列表中

其次以声明的方式定义相应操作：

def ConstantOp : Toy_Op<"constant", [NoSideEffect]> {
  // "constant"就是注记符号，[NoSideEffect]说明了该操作的一个特点
  // Provide a summary and description for this operation. 
  let summary = "constant";
  let description = [{
    Constant operation turns a literal into an SSA value. The data is attached
    to the operation as an attribute. For example:
    ```mlir
      %0 = toy.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]>
                        : tensor<2x3xf64>
    ```
  }];

  /*
  arguments和results：定义参数和结果,参数可以是SSA操作数的属性或类型。
  通过为参数或结果提供名称，ODS将自动的生成匹配的访问器。
  arguments一般模板(results同理): 
  let arguments = (ins <data_type><data_attribute>:$<variable_name>);
  - ins: 输入 (results中该参数为 outs)
  - <data_type>: 数据类型
  - <data_structure>: 数据属性
  - ElementsAttr: 稠元(dense element)
  - <variable_name>: 变量名
  */
  // The constant operation takes an attribute as the only input.
  // `F64ElementsAttr` corresponds to a 64-bit floating-point ElementsAttr.
  let arguments = (ins F64ElementsAttr:$value);
  // The constant operation returns a single value of TensorType.
  let results = (outs F64Tensor);

  // Divert the printer and parser to `parse` and `print` methods on our operation.
  let hasCustomAssemblyFormat = 1;
  /*
  // 自定义程序的组装格式，使最终输出的 IR 格式更精简、易读
  let parser = [{ return ::parseConstantOp(parser, result); }];
  let printer = [{ return ::print(p, *this); }];
  */
    
  // ODS 可以自动生成一些简单的构建方法，用户也可自定义添加一些构造方法
  let builders = [
    // Build a constant with a given constant tensor value.
    OpBuilderDAG<(ins "DenseElementsAttr":$value), [{
      build($_builder, $_state, value.getType(), value);
    }]>,
    // Build a constant with a given constant floating-point value.
    OpBuilderDAG<(ins "double":$value)>
  ];

  // Add additional verification logic to the constant operation.
  // will generate a `::mlir::LogicalResult verify()`
  let hasVerifier = 1;
}

然后在编译阶段，由框架自动生成相应的 C++ 代码。当然也可以运行下面的命令 直接得到生成的 C++ 代码。

${build_root}/bin/mlir-tblgen -gen-op-defs ${mlir_src_root}/examples/toy/Ch2/include/toy/Ops.td -I ${mlir_src_root}/include/

下图中右侧是 ODS 中的定义，左侧是自动生成的 C++ 代码。

官方的文档在这时候没提及需要在 Toy Dialect 的 initialize 函数中注册生成的Op

3.4 创建流程总结(使用ODS)

整个 tableGen 模块是基于 ODS (Operation Definition Specification)框架进行编写以及发挥作用。tableGen 模块促进了自动化生成，减少了 operation 的手动开发，并且避免了冗余开发。

我们以添加 Toy Dialect为例，总结添加流程如下：

Ops.td文件默认位置为 …/mlir/examples/toy/Ch2/include/toy/Ops.td

① (在Ops.td中) 定义一个和 Toy Dialect 的链接

def Toy_Dialect : Dialect {
    let name = "toy";
    ...
    let cppNamespace = "toy";
}

② (在Ops.td中) 创建 Toy Dialect Operation 基类

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

③ (在Ops.td中) 创建 Toy Dialect 中各种 Operation

def ConstantOp : Toy_Op<"constant", [NoSideEffect]> {
  let summary = "constant";
  let arguments = (ins F64ElementsAttr:$value);
  let results = (outs F64Tensor);
  let builders = [
    OpBulider<"Builder *b, OperationState &state, Value input">
  ];
  let verifier = [{ return ::verify(*this); }];
}

④ 通过 mlir-tblgen 工具生成 C++ 文件

使用 mlir-tblgen -gen-dialect-decls 命令生成对应的 Dialect.h.inc 文件。

使用 mlir-tblgen -gen-op-defs 命令生成对应的 Ops.h.inc 文件。

使用 #include 直接引用生成文件

#include "toy/Dialect.h.inc"
#include "toy/Ops.h.inc"

下一篇：Toy接入MLIR

MLIR深入 —— 转换流程详解(以Toy接入为例)