方言 Dialect

MLIR 是一个设计完全可扩展的基础设施,它的可扩展性体现在 IR 的各个元素,包括 operationstypesattributes 等都是可以进行扩展的。这种可扩展性是通过方言(dialect)来实现的,方言为 IR 提供了一种通过 namespace 分组的机制,可以赋予操作(operation)新的语义,实现自定义的行为。

在 IR 中,如果想要为操作赋予新的语义以实现自定义行为,通常需要添加新的属性或者重新定义输入输出。这时,IR 的可扩展性就变得非常重要。

举例来说,考虑一个 matmul 算子,如果想要为它添加一个属性以进行后端图优化,就需要查看该算子的 attributes 字段是否可扩展,以及 op builder 是否允许传入该属性。只有两者都支持,才能满足需求。否则,通常需要创建一个自定义的 matmul 算子,并将原有的 matmul 算子替换为自定义算子。

而在 MLIR 中,你可以定义一个特定的方言,比如 toy dialect,然后在该方言中定义一个 matmul 算子,这样就可以获得 toy.matmul 算子。在图优化过程中,可以根据方言的不同对算子进行分层优化,使用不同方言的转换来实现多层次的优化。

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 分析器,生成目标硬件程序。

v2 84ca774231e37e4c20d5fba20c9e06a9 r

dialect组件

一个 Dialect 主要由以下几个核心组件构成:

  • Type:类型是表示数据的基本方式。每个 Dialect 都可以定义自己特定的类型。例如,标准 Dialect 中的整数类型、浮点类型,或者自定义 Dialect 中的矩阵类型、张量类型等。
  • Attribute:属性是附加到操作上的不可变数据。属性通常用于表示静态数据,如常量值、尺寸、维度等信息。
  • Operation:操作表示具体的计算行为或转换。每个操作都可能有不同的属性、输入和输出类型,并且可以具有限制(Constraints),接口(Interface),特征(Trait)等。操作可以是算术运算、控制流、内存操作等。
  • Interface:接口用于定义操作间的交互方式。它定义了一个操作能够暴露出来的行为,并确保在实现这些操作时满足一定的协议。
  • Trait:特征是附加到操作的元数据或行为描述。它允许我们为操作添加功能性特性,比如支持某种优化策略或特定硬件特性等。

Operation

与操作定义格式相关的两种描述方式如下:

  1. ODS描述格式:用于在 MLIR 中定义操作的完整语法和行为。
  2. 操作签名(Operation Signature):是一个更加简化的描述,主要侧重于操作的输入输出以及特性,而不涉及操作的完整定义细节。

ODS描述格式

ODS(Operation Definition Schemes)是用于在 MLIR 中定义操作的一种格式,它允许通过特定的语法来描述操作的输入输出、属性、约束、特征等。通过 ODS,操作的细节定义会生成自动化代码。

示例 ODS 格式:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def AddOp : Arith_Op {
let summary = "Add two integers";
let description = "This operation adds two integer values and returns the result.";

let arguments = (
Arg<lhs, i32>,
Arg<rhs, i32>
);

let results = (
Result<result, i32>
);

let attributes = (
Attr<some_attribute, i32>
);

let constraints = [
lhs.type == rhs.type
];

let traits = [
"HasSideEffect"
];

let interfaces = [
"MemoryInterface"
];
}

关键元素:

  • summary 和 description:简要描述操作的功能。
  • arguments 和 results:定义输入和输出的操作数。
  • attributes:定义与操作相关的附加信息。
  • constraints:定义操作数之间的关系和限制条件。
  • traits:标识操作的附加特性,例如副作用等。
  • interfaces:定义操作间的交互方式,例如内存接口。

操作签名

与 ODS 定义不同,操作签名通常是对操作的简化描述,主要关注于操作的输入输出类型,而不涉及具体的操作细节。这种方式在一些语境中可能用于描述操作的接口或者操作的高层结构。

示例操作签名:

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

关键元素:

  • **{inplace = true}**:操作的特性(例如原地修改输入张量)。

  • **(tensor<2x3xf64>) -> tensor<3x2xf64>**:定义输入输出张量的类型(输入是 2x3 形状的 f64 张量,输出是 3x2 形状的 f64 张量)。

  • **loc(“example/file/path”:12:1)**:指定该操作在源代码中的位置(文件路径和行列号)

ODS和ODDR

同时存在 ODS 和 DRR 两个重要的模块,这两个模块都是基于 tableGen 模块,ODS 模块用于定义 operation ,DRR 模块用于实现两个 dialect 之间的 conversion

ODS

ODS 模块通过描述操作的属性输入输出类型约束等信息,自动生成代码,以简化操作的定义和管理。通过这种方式,可以高效地为特定的 Dialect 定义复杂的操作,而不需要手动编写大量的 boilerplate 代码。

  • ODS 生成代码:ODS 基于定义好的描述生成操作的定义、类型检查、构造函数、序列化、反序列化等代码。

DDR

ODS 模块通过描述操作的属性输入输出类型约束等信息,自动生成代码,以简化操作的定义和管理。通过这种方式,可以高效地为特定的 Dialect 定义复杂的操作,而不需要手动编写大量的 boilerplate 代码。

  • ODS 生成代码:ODS 基于定义好的描述生成操作的定义、类型检查、构造函数、序列化、反序列化等代码。

定义方言

方言可以被理解为一组操作(op),因此定义方言就意味着定义操作。操作的定义通常包括以下内容:

  • 定义操作的静态信息,包括输入输出、属性和类型等,这些信息描述了操作的语义。
  • 创建操作的构造方法,用于创建一个操作并将其添加到IR中。
  • 创建操作的实现方法,用于在IR执行时调用。

MLIR提供了一种领域特定语言(DSL),使得我们可以通过声明的方式描述和定义操作的输入输出、属性、类型和行为。利用MLIR提供的 tablegen 工具(mlir-tblgen),我们可以基于 Operation Definition Specification(ODS)框架自动生成操作类的声明和实现代码。这种声明式的定义风格使得操作的定义更加清晰易懂,我们只需要关注操作的语义定义即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
llvm-project/
└── mlir/
├── include/
│ └── mlir/
│ └── Dialect/
│ └── MyDialect/
│ ├── IR/
│ │ ├── MyDialect.h # 方言定义
│ │ ├── MyOp.h # 操作定义
│ │ ├── MyDialect.td # 方言的 TableGen 定义
│ ├── Transforms/ # 优化定义
│ │ ├── MyTransform.h # 优化声明
│ │ └── CMakeLists.txt # 优化相关 CMake 文件
│ └── CMakeLists.txt # 方言定义相关 CMake 文件
├── lib/
│ └── mlir/
│ └── Dialect/
│ └── MyDialect/
│ ├── IR/
│ │ ├── MyDialect.cpp # 方言实现
│ │ ├── MyOp.cpp # 操作实现
│ └── Transforms/
│ ├── MyTransform.cpp # 优化实现
│ └── CMakeLists.txt # 优化实现 CMake 文件
└── CMakeLists.txt