MLIR

IR

LLVM1 通过引入 IR 的概念,减轻了传统编译器前后端之间的强耦合关系。与此同时也凸显出了模块化的概念,通过 IR 可以自由实现前后端的组合。

而 MLIR 和 LLVM IR 是同一类别的事物(都是 Intermediate Representation),差别在于前者更加通用可扩展,可以用于描述和表示程序的结构和语义信息。当提及 LLVM IR 时,容易产生混淆的是 LLVM DialectLLVM IR,LLVM Dialect 是在 MLIR 框架中对 LLVM IR 的一种扩展表示,其实从 dialect 这一名称就可以看出这一点,之所以称之为方言而非标准用语就是指这是在标准之外的内容,即对应 LLVM Dialect 是对 LLVM IR 进行扩展这一点。

如果想要利用 MLIR 实现特定领域的相关编译工作,一种可行的思路就是将 MLIR 转换为 LLVM IR(LLVM的中间表示形式),然后利用 LLVM 提供的优化器和代码生成器将 LLVM IR 编译成目标平台的机器代码,简单来说就是利用 LLVM 的后端,这样,MLIR 可以利用 LLVM 的强大优化和代码生成能力,为不同的编程模型和应用场景提供高效的编译器支持。

对于 MLIR 的理解

  1. 背景:在当前编译结构中,各种IR之间转换的效率和可迁移性不高

https://pic2.zhimg.com/80/v2-b6e62260b0faf1085f972d1eda6e4bb1_1440w.webp

  1. 引入 MLIR 所期望做到的:使用一种一致性强的方式,为各种DSL提供一种中间表达形式,将他们集成为一套生态系统,编译到特定硬件平台的汇编语言上

https://pic2.zhimg.com/80/v2-4b2fa235edc4387378a84b8a3587efd9_1440w.webp

MLIR 表达式组成:

https://pic4.zhimg.com/80/v2-6d75286d07a53555437f2c436f718083_1440w.webp

MLIR 并不是一个端到端(从计算图到最后可执行程序这个全流程)的框架,只是一个基础架构

多层次表达 是 MLIR 的优点,这有利于各种优化机制的实现(这并不难理解,中间过程越多,那么优化的入手点就越多),但是另一方面中间过程越多,出错的机会也同样越多,具体是指 compiler’s pass pipeline and toolchain are difficult to configure

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403110950934.png

中间表达的意义在于逐层解析,逐层降低抽象而偏向硬件。高级语言的 AST 到 IR 之间存在较大差距,通过在中间架设 IR 可以对高级抽象进行渐进式变换和递降,降低这种差距变化的梯度,降低阶段转换的难度,同时只要提供 IR 就相当于我们可以为任何领域的代码实现设计特定的编译器。

但是带来的一个问题就是,每当增加一种实现时都会出现一种全新的 IR,而实际上可能这些 IR 之间存在一些共性的东西,通过增设 MLIR 这一层类似标准化的层次,可以使得中间的转化流程变得更加规范

根据High Performance GPU Code Generation for Matrix-Matrix Multiplication using MLIR: Some Early Results2的说法,原有的IR基础设施并不能有效地解决自动生成特定领域库的问题。特别是,很难使用单个IR来表示和转换高,中,低级别的抽象。

MLIR 用于解决编程语言、编译器和硬件之间的交互问题,它的出现是为了应对日益复杂的编程语言和硬件架构

MLIR提供了一个统一的中间表示(IR),可以作为不同编程语言编译器和LLVM后端之间的桥梁。

通过自定义 dialect,可以实现语言扩充 和 实现特定领域的编译优化

如何评价MLIR项目中Linalg Dialect的设计思想? 各种回答中,有两点思想对于理解 MLIR 有一定帮助:

  1. MLIR 或许可以看为一个 IR Template
  2. MLIR 的核心工作在于各种 Dialect 之间的转换。对于目前的项目,假设存在一种机制可以从 C++ 源代码转化到某种 Dialect,之后在不同 Dialect 中来回转换。不过不是很理解的是编译层到底在整体架构中处于什么地位,如果说是以第三方库的形式为框架层提供服务,那么这种服务的接口又是什么,难道就像 MLIR 中原生的这些类型一样?我们需要做的就是编写td文件创建各种 Dialect、operation等?编译层的生成物应当是 host 端和 device 端的类汇编语言,通过具体硬件平台的工具链完成执行流程

降级转换机制

  1. Lowering(降级)的主要目标是将一个高级抽象表示转换为一个更低级别、更具体的表示。通常不会改变 IR 的属性,而只是将其转换为具有更具体语义的形式。例如,将高级的 vector dialect 转换为底层的 LLVM dialect,这只是对 IR 进行了降级,但它仍然是 MLIR IR,并且具有相同的属性和结构。

  2. Translation(翻译)是将一个表示转换为另一个表示的过程,可能会改变 IR 的属性,并将其转换为不同的语言或表示形式。例如,将 MLIR IR 翻译为 LLVM IR。

无论是降级还是翻译,本质上都可以算作一种转换,都需要依托 Pass 机制来实现,只是说两个阶段在转换时用到的 Pass 并不相同

如果仅仅是用 MLIR 原生的 Dialect,那么 lowering 过程是可以直接通过 mlir-opt 来实现的,translation 过程可以直接通过 mlir-translate 来实现

但是如果说需要用到 custom dialect,这个 lowering 和 translation 过程就需要通过代码实现了,通过 MLIR 提供的接口来进行接入


2024/04/22 对上述内容进行修正:

按照最新的理解,实际上 lowering 可能并不一定仅仅表示 Dialect 之间的转化。

lowering 更核心的含义是 表示从一种抽象级别到另一种更底层的抽象级别的转换,而从高层抽象 Dialect 转换为 低层抽象 Dialect,以及从一种形式的 IR 转换到另一种更底层的 IR,这两种情况显然都是符合这种从高层到低层抽象级别的转换概念的。

所以说不必过于纠结所谓的 lowering, transformation, conversation 这些概念之间到底有什么详细的区别

因为从本质上来说,MLIR 中最核心的工作就是转换,无论是 Dialect 转换还是 IR 转换,而这些转换工具都可以抽象说为是 pass,它们在不同的层次和不同的粒度上执行转换,不同的 pass 有不同的作用范围、优化目标和实现方式,但是它们都是用于转换和优化MLIR程序的工具,简单对 pass 进行分类,如下所示(可能不全)

  1. Dialect Conversion Passes:用于在不同的Dialect之间进行转换,例如将MLIR程序转换为LLVM Dialect。
  2. Lowering Passes:用于将高级抽象表示降低为更底层的表示,例如将高级Tensor表示降低为LLVM IR。
  3. Backend Passes:用于生成目标代码或与特定硬件/平台相关的表示。
  4. Transformation Passes:用于对程序进行结构上的转换,例如循环平铺、循环分块等。
  5. Optimization Passes:用于优化MLIR程序,例如常见的优化技术包括死代码消除、常量折叠、循环优化等。
  6. Analysis Passes:用于对程序进行静态分析,例如数据流分析、依赖分析等。
  7. Verification Passes:用于验证程序的正确性,例如检查程序是否符合特定的规范。
  8. Instrumentation Passes:用于在程序中插入额外的代码以收集性能数据或调试信息

MLIR 语法

The core concepts in MLIR include operations, attributes, regions, blocks, values, and types.

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202407101822361.png

MLIR 的源代码文档可见 mlir Namespace Reference,此文档是 mlir 的最外层的命名空间,因此可从此找到所有和 mlir 相关的代码

MLIR 采用类型后置表达,而且在需要用到的位置都需要注明常量或者变量的类型,例如:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
module {
  func.func @main() {
    %result = call @myfunc() : () -> index
    vector.print %result : index
    return
  }

  func.func @myfunc() -> index {
    %c0 = arith.constant 1 : index
    return %c0 : index
  }
}

注意看上述代码中的@myfunc() : () -> index(当然这里还用到了类似lambda中的后置返回值类型的表达方法),以及%result : index(虽然说目前还没有找到index代表常量的什么证据,但是可以姑且就这么认为),以及1:index,它们都是将类型放置在了常量或者变量的后面,即类型后置表达

需要明确的是在上述的各种语法中,点运算符后的实际是一个 operation,从偏向编程语言的角度来说,它们实际上就是一个个的 function,只是表现形式上有些类似变量声明的语法

1
2
3
4
5
// Integer constant
%1 = arith.constant 42 : i32

// Equivalent generic form
%1 = "arith.constant"() {value = 42 : i32} : () -> i32

在官方文档的 arith Dialect 中有这么一段代码,它们两个是等价的,考虑下面那一种形式

1
2
3
arith.constant() {
	value = 42 : i32
} : () -> i32

考虑 MLIR 采用的后置类型表达,这实际就等价于编程语言中的

1
2
3
4
i32 arith::constant() {
	i32 value = 42;
	return value;
}

不同点就在于类型以及作用域的表达方式发生了改变

MLIR 内置的数据类型:Builtin Dialect

对于 mlir-opt 命令携带的选项中那些 convert,目前的理解是代码中使用到了哪些 Dialect,如果想要利用 mlir-cpu-runner 执行,那就需要添加一个从这个 Dialect low 到 LLVM Dialect 的选项,例如:

如果代码中使用到了 vector Dialect 和 math Dialect,那么就需要添加两个选项:

  • -convert-vector-to-llvm
  • -convert-math-to-llvm

从 MLIR 提供的原生 Dialect convert (官方说法是 standard dialect) 到 LLVM Dialect 的相关介绍可以参考 “Conversion to the LLVM Dialect

Dialect 理解

https://img2024.cnblogs.com/blog/1898659/202404/1898659-20240414150832414-966035897.png

Codegen Dialect Overview

MLIR Interface

背景问题

简单来说MLIR多层级的设计为其带来了便利,但是也同样带来了问题。

在不同的Dialect层次进行Operation转换或者做变换(Pass)的时候我们需要明确每个Dialect下的每个Operation的具体语意,否则就可能会转换或变换失败

Interface并不是Operation的核心,而是一些通用变换的核心

MLIR中的Interfaces

MLIR概述

MemRef

In-memory representation of a tensor

The only way to access the element of a memref is through load and store operation.

使用 MLIR 实现矩阵乘法

所需的基本元素:

  1. 多维数组定义及取数操作
1
2
%v = arith.constant dense<[[1, 2], [2, 3]]> : vector<2x2xi32>
vector.print %v : vector<2x2xi32>
  1. 循环语句

  2. 加减乘除数值运算

MLIR Trait

MLIR 的 Trait 是用于描述和约束 Operation 行为的一种机制。Traits 提供了一种方式来共享操作之间的行为和属性,而无需在每个操作定义中重复相同的代码。

Trait 的几种类型:

  • Interface Traits: 定义一组接口方法,操作可以实现这些方法来提供特定的行为。(需要重写 Interface 对应的方法)
  • Property Traits: 提供操作的一些静态属性或特性,比如操作是否是无副作用的(NoSideEffect)。
  • Utility Traits: 提供一些通用的功能,比如打印操作、验证操作等。

使用示例:

MLIR 编译

  1. 如果需要用到Python Bindings,在MLIR Getting Started给出的编译命令之外是需要添加额外的编译选项的,完整示例如下所示:

同时需要注意文档中提到的对于部分python包的依赖,这部分需要在编译之前进行安装

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
cmake -G Ninja ../llvm \
    -DLLVM_ENABLE_PROJECTS=mlir \
    -DLLVM_BUILD_EXAMPLES=ON \
    -DLLVM_TARGETS_TO_BUILD="Native" \
	-DCMAKE_BUILD_TYPE=Debug \
	-DLLVM_USE_SPLIT_DWARF=ON \
	-DLLVM_ENABLE_ASSERTIONS=ON \
	-DCMAKE_C_COMPILER=clang \
	-DCMAKE_CXX_COMPILER=clang++ \
	-DLLVM_ENABLE_LLD=ON \
	-DMLIR_ENABLE_BINDINGS_PYTHON=ON \
	-DPython3_EXECUTABLE="/opt/homebrew/bin/python3"
  1. compiler这个项目所包含的CMakeLists.txt中写死了一些路径,这些都需要进行修改

  2. compiler项目在CMakeLists.txt中使用到了一些find_package(),涉及到numpy和pybind11,需要通过CMAKE_PREFIX_PATH手动指定路径

MLIR Tutorials

在 llvm 的项目源码中,同MLIR Tutorials相关的需要关注的是两个路径下的文件:

  • /Users/gaohongyu/llvm/mlir/test/Examples/Toy:存放的是 toy 语言的源程序
  • /Users/gaohongyu/llvm/mlir/examples/toy:存放的是实现 toy 语言获取 AST,MLIR 表达式等工具的源代码 (产生的工具可执行文件在build/bin目录中)

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403191134543.png

Chapter 1

Toy Language and AST

这一部分主要内容就是创造了一种新的语言,叫做Toy,然后实现了一个简易的语法分析器,仅仅能够获得toy源程序对应的抽象语法树。所以说这一节的重点就是抽象语法树

需要使用到的命令:~/llvm/build/bin/toyc-ch1 /Users/gaohongyu/llvm/mlir/test/Examples/Toy/Ch1/ast.toy --emit=ast

意思是通过toyc-ch1程序生成ast.toy程序的抽象语法树

实际上,这里采用 Toy 语言进行分析可能会给初学者带来一点困惑,因为这很容易让初学者陷入这个语言本身,而忽略了 MLIR 到底想做一件什么事情。

实际上,MLIR 采用一种自创的 Toy 语言,这是为了完整的展示从最开始的源程序,首先转换为 AST,然后进一步翻译为 MLIR 表达式,然后通过 lowering,逐步转换到可执行程序的全流程。而采用 Toy 并不是说 MLIR 只是应用于编译一门新的语言(当然它有这个功能),将 Toy 改为 python 或者 C++ 等其他语言也是一样的,关键并不在于 MLIR 所面对的是什么语言所形成的应用,而是 MLIR 面对的是一个应用这件事情本身,它要做的就是把这个应用通过 MLIR 的各个阶段,逐步翻译到我们所期望得到的东西,而 Chapter 1 所展现的从 源代码 转换为 AST 只是这个过程中的一环,并且这一环和 MLIR 并没有什么关系,因为 源代码 -> AST 往往是语言本身所支持的。MLIR 关注的是从 AST 到 MLIR表达式的转换,以及 MLIR 表达式后续的变化。

上面提到,从 源程序 到 AST 的转换过程实际并不需要 MLIR 的介入,对于常见的编程语言,例如 python,官方是有提供 ast 包来实现这一点的。只不过由于 toy 是一个自创的语言,所以 Chapter 1 的代码主要做的就是解析抽象语法树。

Chapter 2

Emitting Basic MLIR

在第一节中已经生成了 Toy 语言源程序的 AST,这一节就是要根据 AST 结合 Dialect 来生成 MLIR表达式

https://mmbiz.qpic.cn/mmbiz_png/SdQCib1UzF3tN9fRfXZhWRgL2OLr400ESibMbgibPJfUrSLDicq855g64h5cz6CHn4lstoRPJ2KjGbG2q43ANqSPmg/640?wx_fmt=png&wxfrom=5&wx_lazy=1&wx_co=1

Ch2/toyc.cpp 相较于 Ch1/toy.cpp相比,有一个区别就是程序接收的参数多了一个:

这是Ch1/toy.cpp中和程序参数相关的内容:

1
2
3
static cl::opt<enum Action>
    emitAction("emit", cl::desc("Select the kind of output desired"),
               cl::values(clEnumValN(DumpAST, "ast", "output the AST dump")));

这是Ch2/toy.cpp中和程序参数相关的内容:

1
2
3
4
static cl::opt<enum Action> emitAction(
    "emit", cl::desc("Select the kind of output desired"),
    cl::values(clEnumValN(DumpAST, "ast", "output the AST dump")),
    cl::values(clEnumValN(DumpMLIR, "mlir", "output the MLIR dump")));

区别就在于编译生成的toyc-ch2不仅能够生成toy源程序的抽象语法树,还能够生成对应的mlir表达式

关于 Ops.td 是如何同 Dialect模块(Dialect.h,Dialect.cpp)相结合的?**

通过Ops.td可以生成以下这些.inc文件:

  1. Dialect.h.inc: Dialect Declarations
  2. Dialect.cpp.inc: Dialect Definitions
  3. Ops.h.inc: Op Declarations
  4. Ops.cpp.inc: Op Definitions

而 Dialect 模块则使用到了这些文件:

  1. Dialect.h 负责引入 .h.inc,即 Declarations, 包括 Dialect.h.inc 和 Ops.h.inc
  2. Dialect.cpp 负责引入 .cpp.inc,即 Definitions,包括 Dialect.cpp.inc 和 Ops.cpp.inc

所以说,关于 dialect 本身 和 其 操作 等内容,都是在.td文件中说明,然后通过 tablegen 工具结合选项生成所需要的内容

Dialect.h 负责引入头文件这件事情并不难理解,那么如何理解 Dialect.cpp 和 Ops.cpp.inc 都具备的对于 Op Definition 功能这件事情?

TransposeOp::build 为例:

Ops.cpp.inc 中的相关内容为:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
void TransposeOp::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type resultType0, ::mlir::Value input) {
  odsState.addOperands(input);
  odsState.addTypes(resultType0);
}

void TransposeOp::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value input) {
  odsState.addOperands(input);
  assert(resultTypes.size() == 1u && "mismatched number of results");
  odsState.addTypes(resultTypes);
}

void TransposeOp::build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes) {
  assert(operands.size() == 1u && "mismatched number of parameters");
  odsState.addOperands(operands);
  odsState.addAttributes(attributes);
  assert(resultTypes.size() == 1u && "mismatched number of return types");
  odsState.addTypes(resultTypes);
}

Dialect.cpp 中的相关内容为:

1
2
3
4
5
void TransposeOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                        mlir::Value value) {
  state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
  state.addOperands(value);
}

目前还未搞懂这两者之间有何关联,不过 Ops.cpp.inc 还有额外的工作,就是在 Dialect.cpp 最开始的 dialect 初始化时为 dialect 指定操作:

1
2
3
4
5
6
void ToyDialect::initialize() {
  addOperations<
#define GET_OP_LIST
#include "toy/Ops.cpp.inc"
      >();
}

关于创建一个 Op 所需要用到的类

以 Transpose Op 为例,根据 Ops.cpp.inc 的内容,和 transpose op 相关的有以下内容:

  1. TransposeOpGenericAdaptorBase
  2. TransposeOpGenericAdaptor
  3. TransposeOpAdaptor
  4. TransposeOp

随想疑问

一直在思考,toy tutorial 给出的这个示例到底对于实际的应用场景有何启发作用?

实际上toy是一门新创造出的语言,是为了展示MLIR的各项特点而出现的,获取抽象语法树,获取MLIR表达式的工具都是项目自行实现的。

如果考虑真实的场景,例如用C++实现的代码,哪怕是框架,通过cmake都可以完成编译,引入MLIR难道就是为了替换掉整个编译流程中的前端和中端,用人工实现MLIR的方式逐层降级?首先编译层的上层是框架层,在我的理解中,我们可以把这个框架和Tensorflow或者Pytorch等价来看,但是在我看来,无论是不是新的框架,只要不是新的语言,理论上都可以用现有的编译工具完成编译,非要做这个编译层,非要用MLIR逐层降级只是为了提高性能,优化性能,使得每一层的优化都能够做到极致,而不是仅仅依靠现有编译器采用的那些优化。

那么有一个问题是:我们要站在哪个起点来分析这件事。toy是一门新的语言,所以说最基础的解析器都需要全新实现,也就是说面对toy这个场景就应当是一无所有的。但是现有项目是用C++实现的,MLIRGen这个工具对于现在的场景是存在的,肯定需要做的是写Dialect,那结合Dialect生成MLIR表达式这个过程是谁负责的?

突然感觉我们的重点就是写Dialect,结合编程框架中的一些数据结构,在Dialect中把需要完成的内容完成,考虑一层一层是怎么降级的。我觉得像:

1.是怎么从源程序解析得到的AST?(这个问题不考虑,暂时不重要) 2.MLIRGen是怎么结合Dialect生成的一层表达(每一次Dialect的加入都会形成一层中间层),主要是类和类之间是怎么关联起来的 3. .td文件通过tablegen能生成很多东西,这些东西都怎么用? 4. .td结合tablegen生成的一大堆东西 和 Dialect.cpp又是什么关系,或者说Dialect.cpp到底负责干什么的?

  1. 我现在比较纠结的就是各个部分的联动关系,我不知道要写的内容到底要怎么发挥出它的作用
  2. 不太理解如果想要完成一层降级都需要做些什么,只知道要写dialect,但是这dialect是怎么对降级起到作用的?通过生成MLIR表达式或许是,所以说就是让dialect为生成MLIR表达式做出贡献

如果对toy ch2这个内容来说,是否可以假设现在有MLIRGen的工具,但是还没有相关的dialect,任务就是编写dialect实现MLIR表达式的生成。(因为后续教程的内容是例如表达式优化和降级的内容,如果让dialect在生成MLIR表达式这个过程中发挥作用我觉得是整个任务的基础)

如何理解 Dialect?

Dialects provide a grouping mechanism for abstraction under a unique namespace

MLIR is designed to allow all IR elements, such as attributes, operations, and types,

自定义 Dialect 同 MLIR 结合的方式:

C++实现Dialect:

1
2
3
4
5
6
7
8
class ToyDialect : public mlir::Dialect {
public:
  explicit ToyDialect(mlir::MLIRContext *ctx);

  static llvm::StringRef getDialectNamespace() { return "toy"; }

  void initialize();
};

借助tablegen工具实现Dialect:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
def Toy_Dialect : Dialect {
  let name = "toy";

  let summary = "A high-level dialect for analyzing and optimizing the "
                "Toy language";

  let description = [{
  }];

  let cppNamespace = "toy";
}

因为这一节的内容的主要目的就是探索从 抽象语法树AST 转为 MLIR表达式 的过程,所以说 Ch2 相较于 Ch1 增加的从模块上来将,就将增加 生成MLIR表达式所需的3个模块

具体模块对应的文件是:

  • MLIRGen模块:mlil/MLIRGen.cppinclude/toy/MLIRGen.h
  • Dialect模块:mlir/Dialect.cppinclude/toy/Dialect.h
  • TableGen模块:include/toy/Ops.td
  1. 从抽象语法树 生成 MLIR表达式
1
2
3
4
5
# AST 中 toy源程序中 transpose(a) 语句对应的内容

Call 'transpose' [ @codegen.toy:5:10
  var: a @codegen.toy:5:20
]

经过 MLIRGen 以下代码的加工处理,将 AST 转变为了 MLIR 表达式,其中的关键在于if (callee == "transpose"),即当程序扫描到transpose关键词时,就会返回构造出的 transpose 的 MLIR表达式中的节点,即第一步的内容

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

不过存在一个问题是 在构造MLIR表达式节点时,利用到了一个 TransposeOp类,它应当表示的是源程序和MLIR表达式中的 transpose操作,这个类从何而来?

我目前的理解是,所有数据类型都是在td文件中以一种声明式语法来说明的,后续需要用到哪些类型的文件则通过mlir-tblgen工具来生成。那么如何理解Dialect模块,我的理解是TableGen模块提供的只是一些基本原料,但是这些原料究竟怎么用到项目中,是由Dialect模块来决定的或者说设置的。所以说我们需要通过td文件对于源程序中涉及到的结构进行抽象。

实际上,从 mlir-tblgen 的 help 信息可以看出,tb文件可以生成很多格式的信息文件

 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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
--gen-attr-interface-decls                        - Generate attribute interface declarations
--gen-attr-interface-defs                         - Generate attribute interface definitions
--gen-attr-interface-docs                         - Generate attribute interface documentation

--gen-attrdef-decls                               - Generate AttrDef declarations
--gen-attrdef-defs                                - Generate AttrDef definitions
--gen-attrdef-doc                                 - Generate dialect attribute documentation

--gen-avail-interface-decls                       - Generate availability interface declarations
--gen-avail-interface-defs                        - Generate op interface definitions

--gen-bytecode                                    - Generate dialect bytecode readers/writers

--gen-convertible-llvmir-intrinsics               - Generate list of convertible LLVM IR intrinsics

--gen-dialect-decls                               - Generate dialect declarations
--gen-dialect-defs                                - Generate dialect definitions
--gen-dialect-doc                                 - Generate dialect documentation

--gen-directive-decl                              - Generate declarations for directives (OpenMP/OpenACC etc.)

--gen-enum-decls                                  - Generate enum utility declarations
--gen-enum-defs                                   - Generate enum utility definitions
--gen-enum-from-llvmir-conversions                - Generate conversions of EnumAttrs from LLVM IR
--gen-enum-to-llvmir-conversions                  - Generate conversions of EnumAttrs to LLVM IR

--gen-intr-from-llvmir-conversions                - Generate conversions of intrinsics from LLVM IR

--gen-llvmir-conversions                          - Generate LLVM IR conversions
--gen-llvmir-intrinsics                           - Generate LLVM IR intrinsics

--gen-op-decls                                    - Generate op declarations
--gen-op-defs                                     - Generate op definitions
--gen-op-doc                                      - Generate dialect documentation
--gen-op-from-llvmir-conversions                  - Generate conversions of operations from LLVM IR
--gen-op-interface-decls                          - Generate op interface declarations
--gen-op-interface-defs                           - Generate op interface definitions
--gen-op-interface-docs                           - Generate op interface documentation

--gen-pass-capi-header                            - Generate pass C API header
--gen-pass-capi-impl                              - Generate pass C API implementation
--gen-pass-decls                                  - Generate pass declarations
--gen-pass-doc                                    - Generate pass documentation

--gen-python-enum-bindings                        - Generate Python bindings for enum attributes
--gen-python-op-bindings                          - Generate Python bindings for MLIR Ops

--gen-rewriters                                   - Generate pattern rewriters

--gen-spirv-attr-utils                            - Generate SPIR-V attribute utility definitions
--gen-spirv-avail-impls                           - Generate SPIR-V operation utility definitions
--gen-spirv-capability-implication                - Generate utility function to return implied capabilities for a given capability
--gen-spirv-enum-avail-decls                      - Generate SPIR-V enum availability declarations
--gen-spirv-enum-avail-defs                       - Generate SPIR-V enum availability definitions
--gen-spirv-serialization                         - Generate SPIR-V (de)serialization utilities and functions

--gen-type-interface-decls                        - Generate type interface declarations
--gen-type-interface-defs                         - Generate type interface definitions
--gen-type-interface-docs                         - Generate type interface documentation

--gen-typedef-decls                               - Generate TypeDef declarations
--gen-typedef-defs                                - Generate TypeDef definitions
--gen-typedef-doc                                 - Generate dialect type documentation

怎么理解这个方言?我目前的看法就是适配器,方言也同时说明了不同的人有不同的说法,同样的对于transpose来说,细节可能也不同,所以才被称之为方言

但是说通过多层IR逐步降级,这个降级是怎么表现的还是不太明白

根据This is the C++ definition of a dialect, but MLIR also supports defining dialects declaratively via tablegen.的说法,td文件写的其实就是dialect,这个dialect可以直接通过C++代码进行定义,同时也可以使用tablegen结合td文件生成

Chapter 2 展现出的问题

以下是通过 toyc-ch2 codegen.toy --emit=mlir生成的 mlir表达式

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
module {
  toy.func @multiply_transpose(%arg0: tensor<*xf64> loc("codegen.toy":4:1), %arg1: tensor<*xf64> loc("codegen.toy":4:1)) -> tensor<*xf64> {
    %0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64> loc("codegen.toy":5:10)
    %1 = toy.transpose(%arg1 : tensor<*xf64>) to tensor<*xf64> loc("codegen.toy":5:25)
    %2 = toy.mul %0, %1 : tensor<*xf64> loc("codegen.toy":5:25)
    toy.return %2 : tensor<*xf64> loc("codegen.toy":5:3)
  } loc("codegen.toy":4:1)
  toy.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("codegen.toy":9:17)
    %1 = toy.reshape(%0 : tensor<2x3xf64>) to tensor<2x3xf64> loc("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("codegen.toy":10:17)
    %3 = toy.reshape(%2 : tensor<6xf64>) to tensor<2x3xf64> loc("codegen.toy":10:3)
    %4 = toy.generic_call @multiply_transpose(%1, %3) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("codegen.toy":11:11)
    %5 = toy.generic_call @multiply_transpose(%3, %1) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("codegen.toy":12:11)
    toy.print %5 : tensor<*xf64> loc("codegen.toy":13:3)
    toy.return loc("codegen.toy":8:1)
  } loc("codegen.toy":8:1)
} loc(unknown)

可以看出这个层次下的 mlir表达式 是严格按照代码逐句翻译后的结果,例如下面两句是 源代码 和 对应的mlir表达式,实际上从 tensor<*xf64> transpose到 tensor<*xf64> 显然是没有意义的,这是一步冗余的操作,所以 Chapter 3 就着手通过 表达式优化 来解决这个问题

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

%0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64> loc("codegen.toy":5:10)

Chapter 3

High-level Language-Specific Analysis and Transformation

关于表达式变形消除冗余在整个流程中的位置:

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403191137970.png

Chapter 3 的展开逻辑本质上是遵循了 Chapter 2 遗留下来的问题,不过其重新提出了一个 transpose(transpose(X)) -> X 的问题,讲述如何解决这个问题。需要注意的是本节一共举了为两个操作添加优化方法的示例,它们采用了不同的方法, transpose操作 采用了 C++ 实现(对应ToyCombine.cpp),reshape操作 采用了 DRR模块(对应ToyCombine.td)

这一节中,共涉及到整体架构的三部分; https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403201132203.png

这一节对应到的是 MLIR 的 Pattern Rewrite机制,用于对IR做一些通用变换优化,还负责Op的规范化以及Dialect间以及Dialect内部的Op转换,我感觉就是架构图中提到的 transformation 和 canonicalization

C++ 实现 canonicalization

表达式的变形优化 在 MLIR 的场景下被称为 rewrite

想要实现 rewrite 需要进行以下几点:

  1. 实现一个 RewritePattern,即一个继承了mlir::OpRewritePattern<>的类
1
struct SimplifyRedundantTranspose : public mlir::OpRewritePattern<TransposeOp>

当继承这个OpRewritePattern类后,需要重写matchAndRewri方法

  1. rewrite能够被应用,需要通过 canonicalization pass。需要做的工作一方面是让 规范化器知道有这个 rewrite 的存在,另一方面就是让其运用上这一 rewrite。所以这一步就是实现注册工作
1
2
3
4
void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
  results.add<SimplifyRedundantTranspose>(context);
}

不过有一点疑问是,这里只是类方法的定义,并没有找到对此方法的调用,那么是什么时候真正执行的注册?

关于这个问题,官方文档中有提到一个 canonicalization framework,虽然在第2点中说注册是为了让 规范化器 知道有这个rewrite的存在,但是本质上可能是为了让这个 优化框架 了解到这个rewrite的存在。有一种可能就是只要我们为一个操作定义了getCanonicalizationPatterns方法,那么 MLIR 在面对这个操作类的时候,就一定能检测到它是具备这个优化方法的,那么在下一步就能够真正实现这种优化。

真实的实现机制需要去看 MLIR 实现的源码是怎么处理的,我们暂且可以认为这样做就是满足 MLIR对于实现一个方法的优化 所做出的规定,也就是说 MLIR 提供了一种标准化的模版,我们只需要按照模版的要求去做,那么就可以为这个操作实现这种优化(实际上 Chapter 4 给出了支持这种看法的证据,详见Chapter 4)

  1. 规范化器 的 代码原型 就是 mlir::PassManager,即我们需要利用它来真正运用上第一步实现的 rewrite
1
2
mlir::PassManager pm(module.get()->getName());
pm.addNestedPass<mlir::toy::FuncOp>(mlir::createCanonicalizerPass());

使用 DRR模块 实现 transformation

相较于C++实现,使用 DRR模块 的优点似乎就在于省去了后两步的注册和应用到规范化器的两项操作,只需要实现一个继承关系即可

1
2
3
4
5
6
7
class Pattern<
    dag sourcePattern, list<dag> resultPatterns,
    list<dag> additionalConstraints = [],
    dag benefitsAdded = (addBenefit 0)>;

def ReshapeReshapeOptPattern : Pat<(ReshapeOp(ReshapeOp $arg)),
                                   (ReshapeOp $arg)>;

不过教程中还额外给出了另外两种实现方式,那两种只是说在基础之上,额外添加了其他信息:

  1. 增加了参数限制:只需要实现一个Constraint类的子类,并且将该子类在创建模式时添加到模版参数之中
  2. 当限制条件比较复杂时,可以通过实现一个NativeCodeCall的子类添加原生的C++语句实现,同样在创建模式时将其添加到模版参数之中

对 Chapter 3 的理解和感悟

这样来说,是不是可以理解为 rewrite 就是一种 pass

通过这一部分内容,更加加深了对于 MLIR 似乎就和 SpringBoot 那种东西一样,我们似乎完全可以把他们当作普通的框架,为操作添加优化就是必须要为操作类定义getCanonicalizationPatterns方法,这就好像在其他那些框架中所要求做的类似

compiler pattern-match transformations 分类

  1. local

有两种方法用于实现transformation:

  • Imperative, C++ pattern-match and rewrite
  • Declarative, rule-based pattern-match and rewrite(Declarative Rewrite Rules (DRR)模块)(此时要求之前的op定义是通过ods模块实现的,否则只能采用C++代码手动实现transformation了)

到这里,架构图中涉及到的两个模块就都出现了

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403191146250.png

  1. global

关于 Pattern Rewrite机制 的疑问

在这个机制下,存在一些名称极为类似暂时还不清楚用途的类,例如

  • RewriterBase
  • PatternRewriter

以上两者关系见

https://mlir.llvm.org/doxygen/classmlir_1_1RewriterBase__inherit__graph.png

  • RewritePattern
  • OpRewritePattern

上面这 4 个类是不同的,区别在于前者核心是 rewriter,后者核心是 pattern

这两个 pattern 之间的继承关系如下:

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403211147117.png

Chapter 4

Enabling Generic Transformation with Interfaces

在 Chapter 3 的 C++ 实现 transformation 一节中留下了一个疑问,即在给一个 operation 应用一种 模式patter,或者说希望重写它时,我们是定义了TransposeOp::getCanonicalizationPatterns这一方法,当时我们疑问在于代码中只是定义了但是并没有在任何位置看到调用它的代码,那么其是如何发挥作用的?

在 Chapter 4 开头给出,getCanonicalizationPatterns 实际是一个钩子函数,作用于operations之上,通过钩子函数实现为 operation 添加额外功能的方法并不难理解,这也是软件开发中常见的一种方法

不过钩子函数从何而来,是需要我们自行创建(如果是这样的话,又需要在哪个环节进行创建),还是 MLIR 本身对于 operation 就声明了一些固定的钩子函数?

关于这个问题,我们进行了如下的探究:

  1. 关于 TransposeOp,源头应当位于 Ops.td,所以我们回到这一文件

文件中主要包含 2 部分内容,一是对于 Dialect 的说明,二是对于 operation 的说明

对于 Dialect 的说明:

1
2
3
4
def Toy_Dialect : Dialect {
  let name = "toy";
  let cppNamespace = "::mlir::toy";
}

对于 operation 的说明:

1
2
3
4
5
6
class Toy_Op<string mnemonic, list<Trait> traits = []> :
    Op<Toy_Dialect, mnemonic, traits>;

def TransposeOp : Toy_Op<"transpose">{

}

由此可见,关于 operation 的继承关系是 TransposeOp -> Toy_Op -> Op

  1. 由于 td 只是声明性语言,并不代码最终的代码实现,所以我们对 cpp 文件下的实际的 TransposeOp类 进行了追踪,发现其实际声明位于 Ops.h.inc,同时声明内容为
1
2
3
class TransposeOp : public ::mlir::Op<> {
  static void getCanonicalizationPatterns(::mlir::RewritePatternSet &results, ::mlir::MLIRContext *context);
}

由于 Ops.h.inc 是有 Ops.td 生成的,但是在 Ops.td 中并没有显示给出 getCanonicalizationPatterns 的声明,说明这一方法应当是通过继承关系得到的(在llvm源码中,确实有找到 和 getCanonicalizationPatterns 相关的内容,不过还未完全分析透彻逻辑关系),说明只要是继承了 mlir::Op 的子类,都应当是能够使用这个钩子函数的

官方定义的毕竟是有限且不够灵活的,所以 MLIR 就提出了 interface 的概念,用户可以自定义实现各种钩子函数(感觉这就是一个组件化工具正常的发展流程,官方提供一些,然后官方开放接口,第三方可自行实现接入)

本 Chapter 提及的也是 Dialect 的一个子部分

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403201134960.png

引入的情景是 tensor 的形状推导

目前感觉 Chapter 4 所介绍的 Interface 只是对于 Chapter 3 所介绍的用于表达式优化的 pattern 在功能性上的一种扩充,所以说并不需要把其当作一种新的技术点

MLIR Python Bindings

MLIR 表达式的构建在此之前都是通过 C++ 代码实现的,

法斯特豪斯博客里面给的也是一个创建 operation 的流程,应该是将 c++源程序 解析为 MLIR表达式这个过程,而不是对 MLIR 表达式进行降级这个过程

按照我目前的理解,python binding 是用在 frontend 的,用在将 python源程序 解析为 MLIR表达式这个过程,就像将 C++源程序 解析为 MLIR表达式

将源程序解析为 MLIR 表达式是 frontend 的工作,而对 MLIR 表达式进行降级是 middleend 的工作

python binding 关注的是 frontend 的工作,或者说其就是用于实现 MLIRGen 模块的

只要理解了 frontend + backend 这个结构,python binding存在的意义就不难理解了,后续的关键问题就是为了实现将 python 解析为 MLIR 表达式,需要写哪些内容

而 frontend 的重点就是 MLIRGen,这包含两部分工作:一是 源代码 生成 AST,这一步和 MLIR 并没有关联;二是 AST 生成 MLIR 表达式,在这一步中 MLIR 才开始介入

因为 MLIR 本身是 C++ 实现的,tablegen 工具最终生成的也同样为 C++ 代码,所以通过 C++ 实现从一种 语言源程序 转换为 MLIR 表达式,这个过程是很自然的(C++ 代码调用 C++ 代码)。而使用 python 实现 MLIR 表达式生成,需要利用 tablegen 生成的 C++ 代码内容,但是又无法直接使用,所以就需要通过 python binding,将 C++ 内容生成 python 接口,以供 python 实现的 MLIRGen 模块使用,因此我们首先了解一下 C++ 代码实现的 MLIRGen 模块所需要进行的工作,以便更好迁移到 python 实现上。

生成 MLIR 表达式 所需的模块:3

  1. TableGen模块(生产线的零件): 通过.td文件定义了各种操作的类(这部分也叫做Operation Definition Specification (ODS)框架)(我理解这部分也可以通过手动编写C++代码实现,只是说可能写起来比较繁琐,同时在不同的场景下可能存在类似的需求,如果总是手动编写会带来很大的重复工作量,所以说一般通过td文件结合TableGen工具来生成)
  2. Dialect模块(生产线的机械臂): 负责定义各种操作和分析,为操作添加相应的类型和操作数的值
  3. MLIRGen模块(生产线履带): 遍历抽象语法树(AST),构造 MLIR 节点

TableGen模块在编译时向Dialect模块提供支持

我理解着上述这 3 个模块之间的关系,或者说作用流程,大概是 TableGen 模块生成的是一种类似模版的东西,然后 Dialect 像是具体的适配器,为模版添加不同的属性,这样两阶段分离的设计可以做到 模版 和 数据 解耦合,同时做到 模版 的复用,最后再通过 MLIRGen 将添加了具体属性的模版 生成为 MLIR 节点

https://pic4.zhimg.com/80/v2-95f6bf7b8482ab5a4d8541d16ba6cf7b_1440w.webp

关注图中TransposeOp的指向

从项目的编译结果来看,build/python 下的文件都是通过工具生成的,那么直观来理解,我们只需要编写好 tablegen 的内容以及 pybind 所涉及到的接口,通过 pybind 提供的 cmake 即可根据 tablegen 生成的 c++ 内容生成对应的 python 内容,然后直接在 MLIRGen 模块中使用即可。

mlir “hello world”

假设现在有下面这样一段 mlir,我们的目标是让其能够在机器上跑起来

1
2
3
4
5
func.func @main() {
    %0 = "hello.constant"() {value = dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64>
    "hello.print"(%0) : (tensor<2x3xf64>) -> ()
    return
}

定义 Dialect

这里面涉及到一个 hello Dialect,其至少需要 constant 和 print 两个 operation

从项目文件结构上来说,和目前需要的文件包含两部分:HelloDialect.[h/td] 和 HelloOps.[h/td],前者对应 Dialect,后者对应 operation

主要需要做的是两件事:1. 创建新的 Dialect 2. 创建 operation 基类

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
def Hello_Dialect : Dialect {
    // 定义名字空间 namespace,对应 C++ 的 getDialectNamespace 方法返回值
    let name = "hello";
    ...
    let cppNamespace = "::hello";
    // 该设置用于激活 materializeConstant 方法,这使得可以例如 Canonicalize 优化
}

class Hello_Op<string mnemonic, list<Trait> traits = []> :
        Op<Hello_Dialect, mnemonic, traits>;

定义 operation

使用 mlir-tblgen 工具对 HelloOps.td 进行转换可以得到对于 operation 的声明,其中包含的类的继承关系为

  • constant operation

ConstantOpAdaptor -> ConstantOpGenericAdaptor -> ConstantOpGenericAdaptorBase (前者继承后者)

核心的 ConstantOp 继承的是 ::mlir::Op,不过其会利用到上述 3 种类

  • print operation

PrintOpAdaptor -> PrintOpGenericAdaptor -> PrintOpGenericAdaptorBase

新的 operation 一般至少包含 4 个元素(实际上并不知道这几种字段都负责什么)

  • summary
  • builders
  • arguments
  • results
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
def ConstantOp : Hello_Op<"constant", [Pure]> {
  // 一行关于这个 Op 的介绍
  let summary = "constant";
  ...
  let builders = [
    OpBuilder<(ins "mlir::DenseElementsAttr":$value), [{
      build($_builder, $_state, value.getType(), value);
    }]>,
    OpBuilder<(ins "double":$value)>
  ];
  ...
  let arguments = (ins F64ElementsAttr:$value);
  let results = (outs F64Tensor);
}

创建 pass

如果只看 mlir-hello,那么 pass 实际上只用于了 lowering 的过程,但是实际上 lowering 和 transformation 两个过程使用到的都属于 pass,后续的描述由于是在看 mlir-hello 时所写,所有有可能偏向于 lowering,但是你需要知道的是 pass 也可同样应用于 transformation。简单点来理解,只要涉及到形式上的变化,都可以使用 pass 来完成。

同时 pass 的实现可以通过手写 C++ 实现,也可以通过利用 DRR(Declarative Rewrite Rule)框架来实现

降级到底在降什么?

HelloDialect 是我们自定义出的 dialect,而我们最终的目标是对接到 LLVM backend,这些自定义 dialect 显然是做不到的,但是 MLIR 本身提供的一些 dialect 可以做到,可能 LLVM 为它们提供了对接方式。

所以说我们所需要做的,一方面是在自定义 dialect 之间实现转化,另一方面是需要把 自定义dialect 转换到 MLIR提供的 dialect(可以理解为标准库)上

这部分工作除了需要用到 dialect 和 operation 作为基本原料,还需要 HelloPasses.h 来注册pass,即从一种 dialect 转换到另一种 dialect。仅注册还不行,还需要 LowerToAffine.cpp 的具体实现


以上提到了两个文件 HelloPasses.h 和 LowerToAffine.cpp,不过在实际的项目中并不一定就只有这两个文件,可能会进行很多拆分,不过我们只需要掌握这两个文件所要完成的核心工作即可适应各种项目

方法 1 - 全人工实现

此部分遵循 mlir-hello 项目的实现方式

  1. pass 注册
1
std::unique_ptr<mlir::Pass> createLowerToAffinePass();

所谓注册,其实就是创建了一个 pass 的函数钩子

  1. pass实现

2.1 编写操作的 lowering

1
2
3
4
5
6
7
class ConstantOpLowering : public mlir::OpRewritePattern<hello::ConstantOp> {

  mlir::LogicalResult
  matchAndRewrite(hello::ConstantOp op, mlir::PatternRewriter &rewriter) const final {

  }
};

首先不考虑具体的实现细节,观察这部分代码的核心有以下几点:

  • 定义为 class XxxOpLowering
  • 继承自 mlir::OpRewritePatternxxx::XxxOp
  • 重载 matchAndRewrite 函数,做具体实现
  • XxxOpLowering 最终将作为模板参数传入新 pass 的 mlir::RewritePatternSet(这一点并不在此函数中实现,而实在下一部分)

2.2 实现 dialect 到 dialect 的转换(将 lowering 加入 pass)

实际项目中,这一步会划分为 声明(.h 文件中) 和 实现(.cpp 文件中) 两部分,一般 runOnOperation() 的实现会置于 cpp 文件中,这里为了描述上的简洁就合并在一起了

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class HelloToAffineLowerPass
    : public mlir::PassWrapper<HelloToAffineLowerPass,
                               mlir::OperationPass<mlir::ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(HelloToAffineLowerPass)

  void getDependentDialects(mlir::DialectRegistry &registry) const override {
    registry.insert<mlir::affine::AffineDialect, mlir::func::FuncDialect,
                    mlir::memref::MemRefDialect>();
  }

  void runOnOperation() final {
    mlir::RewritePatternSet patterns(&getContext());
    patterns.add<ConstantOpLowering, PrintOpLowering>(&getContext());
  };
};

这里的重点在于实现了 runOnOperation() 方法, 在此方法中完成上一步所提及的 lowering 作为模版参数传入 pass

2.3 完成函数钩子

1
2
3
std::unique_ptr<mlir::Pass> hello::createLowerToAffinePass() {
  return std::make_unique<HelloToAffineLowerPass>();
}

如何理解这一 function hook?

official website - Pass Infrastructure - Declarative Pass Specification 中可以找到在使用 td 声明 pass 时,指定 constructor 的 作用为

A constructor must be provided to specify how to create a default instance of MyPass.

而为 constructor 指定的内容就是我们创建的 function hook,因此这一 function hook 的作用即为 指出如何创建一个 pass 的实例


所以综合以上所有内容来看,创建一个 pass,其实就是 create 了一个 unique_str,这个智能指针指向了一个类,这个类实现了 runOnOperation()方法,并且在这个方法中有指定操作的匹配重写规则

方法 2 - 使用 tablegen

在人工实现中,创建 pass 时,我们是手动继承的 mlir::PassWrapper 类。

而使用 tablegen 时,首先会在 td 文件中给出一个对 pass 的继承,如以下代码所示。

1
2
3
4
5
6
7
def ConvertGraphToMatrix : Pass<"convert-graph-to-matrix", "ModuleOp"> {
  let summary = "Convert Graph Ops to Matrix Ops";
  let constructor = "compiler::createConvertGraphToMatrix()";
  let dependentDialects = ["compiler::graph::GraphDialect", "compiler::matrix::MatrixDialect"];
  let options = [
  ];
}

而这部分代码经过 tablegen 转换后所形成的类,以及相关的继承关系为:ConvertGraphToMatrixBase -> mlir::OperationPass

这里出现了两种不同的基类(我个人理解全人工实现下也一样可以采用 OperationPass),PassWrapperOperationPass,它们二者之间也有着一定继承关系

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202405231613671.png

之后的操作其实就与全人工实现的方式类似了,不同之处在于创建包含 runOnOperation() 方法的类时,不再去继承 mlir::PassWrapper 类,而是直接继承 ConvertGraphToMatrixBase

然后其他需要做的事情并没有改变。

那么不禁让人产生一种疑问,当我们采用 tablegen 时,除了在创建包含runOnOperation() 方法的类时,多增加了编写 td 文件的一步,其他流程并没有减少,相当于使用 td 反而让流程变得更加复杂了。

实际上,如果我们在 td 中给定了 constructor,我们看 td 文件生成的内容时,就会发现其中包含着 pass 的 function hook,但是目前简单的测试下,还无法正常直接使用这些 function hook,利用的仍然是人工给出的内容,后续可以再进行其他测试,按照理解,当使用 tablegen 并且指定 constructor 时,应当就无需人工声明和定义 function hook 了,即只需要重写 matchAndRewrite()runOnOperation(), 从而通过 tablegen 简化代码编写。

同时需要注意的是,mlir-hello 这里为了降低理解上的难度,并没有采用 td 文件来实现 pass 的声明,而是直接利用了 C++ 代码

方法 3 - 借助 Canonicalization Framework

canonicalization function 同 pass 的关系为,canonicalization function 可以认为是一种 pass,但是 pass 所执行的并不一定是 canonicalization 功能。

以 transpose 操作为例,如果为此 operation 指定了一个消除重复转置的规范化器 canonicalization,那么在发生重复转置后可以对此操作进行规范化,以进行表达式优化,然而这种效果是通过自定义 pass 同样可以实现的。

但是 pass 的概念则更加宽泛一些,不同 Dialect 之间的转换以及 IR 之间的转换都可以称之为 pass,但是这从含义上来说就不属于规范化。

因此在实际应用中,是否采用 Canonicalization Framework 只需要考虑是否进行的一种对于操作的规范化,如果是则采用规范化器,反之如果仅仅是转换,则应当采用 pass。

据说只需要实现 matchAndRewrite 方法即可,暂时没有实验过

2024/04/24,在【从零开始学深度学习编译器】十三,如何在 MLIR 里面写 Pass? 这篇文章中,看到了另一种使用 pass 的方式,即借助于 归范化框架(Canonicalization Framework),与上述的流程相比,相同之处在于仍然需要给出实现了 matchAndRewrite 方法的类,但是 function hook 和 实现了 runOnOperation 方法的类都无需给出。只需要在创建 operation 时添加一个选项 let hasCanonicalizer = 1;,并且给出 xxxOp::getCanonicalizationPatterns() 的实现,最后在 opt 工具中注册另一种 pass 即可。

这里给出一个示例,展示在使用 Canonicalization Framework 时所需要进行的所有操作

  1. operation tablegen 文件中添加 let hasCanonicalizer = 1;
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
def TransposeOp : Toy_Op<"transpose", [Pure]> {
  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)>
  ];

  // Indicate that additional verification for this operation is necessary.
  let hasVerifier = 1;
}
  1. 重写 matchAndRewrite 方法
 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
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LogicalResult.h"
#include "toy/Dialect.h"
using namespace mlir;
using namespace toy;

struct SimplifyRedundantTranspose : public mlir::OpRewritePattern<TransposeOp> {
  SimplifyRedundantTranspose(mlir::MLIRContext *context)
      : OpRewritePattern<TransposeOp>(context, /*benefit=*/1) {}

  mlir::LogicalResult
  matchAndRewrite(TransposeOp op,
                  mlir::PatternRewriter &rewriter) const override {
    mlir::Value transposeInput = op.getOperand();
    TransposeOp transposeInputOp = transposeInput.getDefiningOp<TransposeOp>();

    if (!transposeInputOp)
      return failure();

    rewriter.replaceOp(op, {transposeInputOp.getOperand()});
    return success();
  }
};
  1. 定义 getCanonicalizationPatterns 方法

在第一步中指定 let hasCanonicalizer = 1; 就是在 tablegen 的生层文件中生成了 getCanonicalizationPatterns 方法的声明,需要由人工给出定义,并且在其中添加第二步中重写了 matchAndRewrite 方法的类

1
2
3
4
void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
  results.add<SimplifyRedundantTranspose>(context);
}

易错点及总结

无论利用什么辅助工具,在最终编译出的代码中都一定包含了方法1-全人工实现所涉及到的全部组件,辅助工具只是帮助我们在背后实现了它们

pass 这里很容易产生一个思维误区,即实现一个 pass 必须包含全部流程。从graph_compiler backend 2.执行 pass pipeline的函数调用栈中,我们可以看出 pass 的执行流程,function hook 紧跟着的是 runOnOperation(),然后是 matchAndRewrite(),实际上这里调用 matchAndRewrite() 并不是必须的流程,graph_compiler 的 matrix 部分涉及到的两个 pass:RegisterAllocPassCodeGenPass 就并没有使用到 matchAndRewrite()

创建 opt

分析到这里会给人一种迷惑,似乎我们已经实现了这个降级的过程。

但是仔细考虑一下,截止到目前我们都创建了什么?dialect, operation,operation重写规则 和 pass

而我们的目的是什么?将 mlir 转变为可以对接到 LLVM backend 的内容。

可以发现,我们目前创建的这些东西就像是原材料一样,我们需要一个媒介,将它们和 mlir 进行链接,实现内容转换,而 opt 起到的就是这种作用。简单来说,opt 就是利用已有的工作接口完成一个 lowering 的流程描述

如何理解 hello-opt 这类工具?

mlir 片段确实是摆在这里了,所需要的 Dialect 和 对应的 operation 也都写好了,但是问题是他们之间如何建立起联系来?

hello-opt 充当的就是这个桥梁作用,简单理解它可以将 mlir片段 作为输入,利用声明好的 Dialect 和 operation 对 mlir 进行一系列 lowering 和 transformation 操作,从而最终转化为 LLVM IR 或者其他什么格式,从而接入到 LLVM backend 中


opt 类工具采用的一般的操作流程:

  1. 向 contex 中 加载 dialect
1
context.loadDialect<mlir::func::FuncDialect>();
  1. 使用 pass 完成 lowering 和 transformation
1
2
3
mlir::PassManager pm(&context);;
pm.addPass(compiler::createConvertGraphToMatrix());
pm.run(module_.get());

对于 mlir-hello 来说,lowering 和 transformation 是两个界限清晰的过程。它使用 PassManager 仅仅完成了 lowering 的过程,transformation 的过程是通过其他 llvm 接口来实现的。不过对于 graph_compiler 来说,pass 不仅完成了 lowering,也完成了 transformation 的过程

CMakeLists.txt解读

– CMAKE_BINARY_DIR: 项目的 build 目录: /home/xxx/project/build

– LLVM_INCLUDE_DIRS: 采用的 llvm 下的两个和 llvm 相关的 include 路径: /home/xxx/llvm/llvm/include;/home/xxx/llvm/build/include

– MLIR_INCLUDE_DIRS: 采用的 llvm 下的两个和 mlir 相关的 include 路径: /home/xxx/llvm/mlir/include;/home/xxx/llvm/build/tools/mlir/include

– PROJECT_SOURCE_DIR: 项目根路径: /home/xxx/project

– PROJECT_BINARY_DIR: 项目 build 路径: /home/xxx/project/build

疑惑

在 mlir-hello 的 LowerToAffine.cpp 中,存在一个如下的方法调用

1
rewriter.modifyOpInPlace(op, [&] { op->setOperands(adaptor.getOperands()); });

但是编译老是报错

error: ‘class mlir::ConversionPatternRewriter’ has no member named ‘modifyOpInPlace’

但是通过 clangd 又可以定位到此方法的位置,在 llvm 的 install 目录下,同时发现在源码中确实没有这一方法,这就非常神奇了

buddy-mlir

introduction

  1. 同 MLIR 的关联:reuse the MLIR infrastructure and LLVM backend tools

  2. MLIR 之外的特点:

  • optimization tool
  • auto-config mechanism

Buddy-MLIR 包含一些基于 MLIR infrastructure 实现的算法,需要利用到 MLIR Dialect 和 Op

MLIR 作为一种基础设施,实际上是不去负责端到端的这个编译流程的,那么说上层应用(例如pytorch、tensorflow等)如何对接到下层形式各样的硬件上就是目前的难题,尤其是目前形式各样的异构芯片的出现更加剧了这种问题。

关于不负责端到端的编译流程的具像化理解,从目前的使用和理解来看,MLIR 并没有提供前端,其只提供了和 MLIR 表达式相关的一些工具和一些原生的 Dialect,只用这些是不足以完成从上层应用到下层硬件之间的对接的,所以说其不负责端到端的编译流程。更多的感觉是提供了一种框架,给出了一种约束,开发者可以利用这种模式来实现各种各样的编译流程。

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403281051700.png

但是 Buddy-MLIR 对于解决这种问题提供了什么样的解决方案呢?从上图来看,如果说 MLIR 是一些积木的话,那么 Buddy-MLIR 的意图似乎是利用这些积木,搭建起一个城堡,通过为前中后端提供 Dialect,帮助上层的例如 pytorch 应用能够部署到底层具体的硬件上

但是我没太明白在 Buddy Compiler 的项目中插入如此多的 Example 是什么意思,感觉这些 Example 更多的应该是 MLIR 所提供的,用于展示其 Dialect 的用法,不知道放在这里和其本身的设计思想之间有什么关系。

不过如果你看了代码,或者说Buddy Compiler As A Service (Buddy-CAAS),你会发现 Buddy-MLIR 项目中给出的 IR Level Examples 其实就是 MLIR 原生提供的 Dialect 的使用示例,在项目中,对应着 examples/MLIR*这些目录。这些示例是通过 mlir-opt, mlir-translatemlir-cpu-runner 可以直接 lowering、translation 以及 run 的。所以我理解就是首先通过这些示例让学习者可以更好地理解 MLIR 中各种概念的存在形式,所以入门 MLIR 时不妨直接从了解 MLIR 的原生 Dialect 入手,逐渐搞懂各种概念,然后再去了解如何自行创建 Dialect 或者 operation 等。官方教程从 Toy 的示例开始,直接就自定义 Dialect,确实让人很难理解各种模块存在的内在机理

llvm-as 可以将 LLVM IR 转变为 LLVM IR bitcode

llc 可以将 LLVM IR bitcode 转变为 assembler source text, ASCII text

llc 也可以直接将 LLVM IR 转变为 assembler source text, ASCII text

但是并没有真正运行起来,不过就是说即便不考虑后续如何执行,翻译到 LLVM IR 这一层级也就 OK 了

Build

在第一步构建过程中遇到几个坑点:

  1. The target building platform of MLIR is uncompleted,because MLIR Getting Started asks that we can use the build option -DLLVM_TARGETS_TO_BUILD="Native;NVPTX;AMDGPU", but the buddy-mlir needs the RISCV platform data, so we need to recompile the MLIR

  2. I learn about the process of building buddy-compiler needs the mlir compiling data from the option -DMLIR_DIR=$PWD/../llvm/build/lib/cmake/mlir and -DLLVM_DIR=$PWD/../llvm/build/lib/cmake/llvm, so I think a idea that create a soft link of llvm project for the buddy-compiler. But when I build the buddy-compiler, I get the error message that cmake cannot find some files about RISCV, until I get some tips from the slack

https://img2024.cnblogs.com/blog/1898659/202403/1898659-20240324100837179-389572672.png

执行之后就能够正常build了,但是里面不确定的因素有二,其一是 我手动创建软链接确实也起到了增加submodule的功能,因为执行 git submodule update –init 之后并没有重复 git clone llvm- project,但是不太清楚执行之后到底产生了什么其他的额外影响(主要怀疑会不会执行后修改了什么变量),使得 cmake 就能够找到相关文件了;其二是 submodule 并不仅仅只有 llvm 这一个,还存在另外一个,不确定是不是因为之前缺少这个子模块从而导致build失败(不过这一点是可以验证的,只需要注释掉这部分,只 git submodule llvm-project 那部分,查看是否可以完成 build 就可以,按理说应该不太行)

Structure

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403241200921.png

目前的困惑在于 buddy-mlir 是否可以看为是对 mlir 的一种封装,如果是的话,封装了啥,如果不是,那它相较于 mlir 又有何区别或者说设计的意义在哪里

source code structure

https://cdn.jsdelivr.net/gh/gaohongy/cloudImages@master/202403251627454.png

understanding

Bud Dialect 展示了自定义 Dialect,并降级到 MLIR提供的基础设施上(从代码来看就是一些 llvm 项目下的 Dialect)

include目录下包含了.td,lib下的 conversion 中实现了降级的代码

pass 机制,pattern rewrite机制,interface

目前粗浅的理解就是 pattern rewrite 就是所谓表达式优化和降级的过程(即无论是表达式优化还是降级,都是通过rewrite实现的),这个过程中使用到的东西从概念上来讲就是各种 pass,而各种 pass 在实际命名上就是各种 dialect。而operation的概念在整个流程中是一直存在的。每一个 dialect 或许就可以看为是一个层次,每个层次下都有不同抽象级别的,与当前层次对应的 operation 表述

(这篇文章如何在MLIR里面写Pass

从 examples 下的 IR Level Examples 了解到了通过 MLIR tool chain 可以直接执行 MLIR 表达式

才领悟到我们似乎可以将 MLIR 当作一门普通的语言,其同样具备相关的语法,我们可以借助它的语法实现一些算法,但是带来的疑问是如果仅仅通过 MLIR 就能够完成所有操作,那我们为什么不直接写 MLIR 了,而是还需要去写 C++ 代码,那些 operation 又要起到什么作用。

关于这个疑问,我们大概可以从 mlir-opt --help 命令的输出结果中得到一点启发,其中 --convert-vector-to-llvm 选项的内容是

–convert-vector-to-llvm - Lower the operations from the vector dialect into the LLVM dialect

注意上面提到两个 dialect:vector dialect 和 LLVM dialect

同时在命令输出中还包含了以下内容:

Available Dialects: acc, affine, amdgpu, amx, arith, arm_neon, arm_sme, arm_sve, async, bufferization, builtin, cf, complex, dlti, emitc, func, gpu, index, irdl, linalg, llvm, math, memref, mesh, ml_program, mpi, nvgpu, nvvm, omp, pdl, pdl_interp, quant, rocdl, scf, shape, sparse_tensor, spirv, tensor, test, test_dyn, tosa, transform, ub, vector, x86vector, xegpu

那么也就是说上面这些以及mlir-opt命令能够携带的参数中使用到的dialect,实际都是 MLIR 项目中自带的一些 dialect,或者说是原生的

但是,在实际的应用场景下,我们所需要用到的 dialect 并不仅仅是这些,我们往往需要用到其他类型的,结合实际场景的,所以这就是我们为什么需要自己编写 dialect 的原因

所以首先通过直接借用 mlir 的相关工具和原生的 dialect,直接编写 MLIR 代码,能够帮助我们了解

Buddy-MLIR 提供的 IR Level Examples,主要包含三部分:

  • low
  • translate
  • run

在降级环节(low)中,始终都是 MLIR,即便是 mlir-opt 应用了例如 --convert-func-to-llvm 这类选项,但是也只是说用 llvm dialect 进行了降级,但是其属于 MLIR 的本质并没有改变,证据可见

–convert-func-to-llvm - Convert from the Func dialect to the LLVM dialect

然而在翻译环节(translate),注意是translate而非transformation,才是真正完成从 MLIR 到 其他 IR 的转变,证据可见 mlir-translate 采用的选项 --mlir-to-llvmir

–mlir-to-llvmir - Translate MLIR to LLVMIR

而在运行环节,由于 mlir-cpu-runner JIT 实际运行的是 MLIR,所以其只能处理

这里给出的这些 example,利用到的 Dialect 都是 LLVM 项目,准确地说是 MLIR 中预置的,所以可以通过各种 mlir 工具来实现降级和翻译。如果想要实现自定义的Dialect,那就和创建一个例如 compiler 的项目相同了。

Buddy-MLIR 项目框架重点

  1. include下的td文件负责给出 Dialect 和 Operation 的说明,后续在编译过程中通过 llvm-tblgen 生成所需的各种文件

  2. lib 下的 Conversion 负责给出 lowering pipeline,将 custom dialect’s operation lowering 到 MLIR’s standard dialect,从而将 custom dialect 接入 MLIR ecosystem

  3. example 下给出的实际就是如何使用 buddy-mlir 的示例。Buddy-MLIR 为使用者在frontend中提供了上层应用和 buddy-mlir 进行交互的接口数据结构,然后Buddy-MLIR 应用了MLIR 提供的一个 C/C++ 的前端接口功能,完成了端到端的应用构建

不过有一个疑问是,buddy-mlir 向上层提供接口的方式是提供了一个 Memref 数据结构,难道通过这一个数据结构就能够完成

简单来说就是 buddy-mlir 向上层提供了C函数接口,上层C应用只需要直接调用函数即可(不过在 ConvOpt/edge-detection.cpp 的示例中使用到了_mlir_ciface_conv_2d函数,但是此函数并没有找到相关定义,不过它利用到了buddy-opt以及llvm的相关工具,可能和它们有关,具体的实现机制还有待商榷)。所以说项目需要做的就是提供好相关 dialect、operation,为上层应用提供好调用接口

从 Buddy-MLIR 提供的 BudDialect example 就可以看出,在 examples/BudDialect 提供的是一些 mlir 文件,其中用到的一些 mlir 语法实际上都是属于 BudDialect 的。但是通过 buddy-opt 就可以降级到 MLIR,这里涉及到两个问题:

  1. buddy-opt 哪里来?
  2. 为什么 buddy-opt 可以把自定义的 BudDialect 降级到 MLIR?

buddy-opt 是在 tools 中实现的一个工具,我们可以先简单认为它利用了在 midend 中给出的和 BudDialect 相关的内容,这其中就包含不同 dialect 转换所需要用到的内容。也就是说我们只需要实现了 dialect 体系中的相关内容,通过利用 MLIR 提供的接口就可以简单创建出相关的工具,从而实现 lowering 或者 transformation 的过程

对于 BudDialect 的构建过程,理论上在不添加 buddy-opt 的情况下,虽然什么都做不了,但是也应当可以正常通过编译,只是说没有 buddy-opt 确实没办法做出什么效果。然后再加入 buddy-opt 完成编译。

Examples Introductoin 4

Reference

0%