Dialect及Operation详解

参考资料：

[MLIR] Dialect及Operation详解 - 知乎 (zhihu.com)

2.  Dialect 及Operation

2.1Dialect

2.1.1Dialect 是什么？

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

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

2.1.2 dialect 是怎么工作的？

MLIR中的Dialect可以看作是一种扩展机制，它可以用于表示不同的编程语言和领域特性。在MLIR中，每个Dialect都会定义自己的语法规则和操作，这些规则和操作可以用于构建对应的IR表示。为了避免命名冲突，MLIR将不同的Dialect的IR定义在不同的命名空间中。

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

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

2.1.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 语言的结构进行建模，并为高级分析和转换提供方便的途径。

3.1.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();
};

3.1.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>

3.3.1使用 C++ 语言手动编写

operation 类是继承于 CRTP 类，有一些可选的 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>();
}

3.3.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为例，总结添加流程如下：

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

① (在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"