MLIR一些背景知识

news/发布时间2024/5/17 6:43:46

MLIR一些背景知识

7.1.1背景

随着深度学习技术的发展，深度学习技术也逐渐从学术研究的方向转向了实践应用的方向，这不仅对深度模型的准确率有了较高的需求，也对深度模型的推理速度有了越来越高的需求。

目前深度模型的推理引擎按照实现方式大体分为两类：

1）解释型推理引擎：一般包含一个模型解析器与一个模型解释器，一些推理引擎可能还包含一个模型优化器。模型解析器负责读取与解析模型文件，并将其转换为适用于解释器处理的内存格式；模型优化器负责将原始模型变换为等价的、但具有更快的推理速度的模型；模型解释器分析内存格式的模型并接受模型的输入数据，然后根据模型的结构依次执行相应的模型内部的算子，最后产生模型的输出。

2）编译型推理引擎：一般包含一个模型解析器与一个模型编译器。模型解析器的作用与解释型推理引擎相同；模型编译器负责将模型编译为计算设备（CPU、GPU 等）可直接处理的机器码，并且可能在编译的过程中应用各种优化方法来提高生成的机器码的效率。由于机器码的模型可以直接被计算设备处理而无需特定的解释器的参与，其消除了解释器调度的开销。此外，相对于解释型推理引擎，由于生成机器码的过程更加靠底层，编译器有更多的优化机会以达到更高的执行效率。

由于现在业界对于推理引擎的执行速度有了更高的需求，编译型推理引擎也逐渐成为高速推理引擎的发展方向。目编译型推理引擎有 Apache TVM、oneDNN、PlaidML、Tensorflow XLA、Tensorflow Runtime 等。

为了便于优化，一般来说推理引擎会把模型转换为中间表示，然后对中间表示进行优化与变换，最终生成目标模型（对于解释型推理引擎）或目标机器码（对于编译型推理引擎）。此外，除了深度学习领域，在很早以前编程语言领域就引入了中间表示来做优化与变换。而新的编程语言层出不穷，因此就出现了各种各样的中间表示，如图6-34所示。

图6-34 各种推理引擎的中间表示

不同的推理引擎或者编译器都会有自己的中间表示与优化方案，而每种中间表示与优化方案可能都需要从头实现，最终可能会导致软件的碎片化与重复的开发工作。

7.1.2 MLIR 简介

MLIR（Multi-Level Intermediate Representation）是一种新型的用于构建可复用与可扩展的编译器的框架。MLIR 旨在解决软件碎片化、改善异构硬件的编译、降低构建领域特定编译器的成本，以及帮助将现有的编译器连接到一起。

MLIR 旨在成为一种在统一的基础架构中支持多种不同需求的混合中间表示，例如：

1）表示数据流图（例如在 Tensorflow 中）的能力，包括动态性状、用户可扩展的算子生态系统、Tensorflow 变量等。

2）在这些图中进行优化与变换（例如在 Grappler 中）。

1）适合优化的形式的机器学习算子内核的表示。

4）能够承载跨内核的高性能计算风格的循环优化（融合、循环交换、分块等），并能够变换数据的内存布局。

5）代码生成下降变换，例如 DMA 插入、显式缓存管理、内存分块，以及 1 维与 2 维寄存器架构的向量化。

6）表示目标特定操作的能力，例如加速器特定的高层操作。

7）在深度学习图中的做的量化与其他图变换。

MLIR 是一种支持硬件特定操作的通用中间表示。因此，对围绕 MLIR 的基础架构进行的任何投入（例如在编译器 pass 上的工作）都将产生良好的回报；许多目标都可以使用该基础架构，并从中受益。

尽管 MLIR 是一种强大的框架，它也有一些非目标。MLIR 不试图去支持底层机器码生成算法（如寄存器分配与指令调度）。这些更适合于底层优化器（例如 LLVM）。此外，MLIR 也不意图成为最终用户写算子内核的源语言（类似于 CUDA 与 C++）。另一方面，MLIR 提供了用于表示此类领域特定语言并将其集成到生态系统中的支柱。

MLIR 在构建时受益于从构建其他中间表示（LLVM IR、XLA HLO 与 Swift SIL）的过程中获得的经验。MLIR 框架鼓励现存的最佳实践，例如：编写与维护中间表示规范、构建中间表示验证器、提供将 MLIR 文件转储与解析为文本的功能、使用 FileCheck 工具编写详尽的单元测试、以及以一组可以以新的方式组合的模块化库的形式构建基础框架。

其他的经验教训也已经整合到了设计中。例如，LLVM 有一个不明显的设计错误，其会阻止多线程编译器同时处理 LLVM 模块中的多个函数。MLIR 通过限制 SSA 作用域来减少使用-定义链，并用显式的符号引用代替跨函数引用来解决这些问题。

7.1.3 MLIR引论

MLIR 是一种新型的用于构建可复用、可扩展的编译器框架。MLIR 旨在解决上述说的软件碎片化、改善异构硬件的编译、降低构建特定领域编译器的成本、同时将现有的多种编译器链接到一起。

MLIR 最终想在统一的架构中支持多种不同需求的混合IR。

MLIR 是一种支持硬件特定操作的通用 IR。因此，对围绕 MLIR 的基础架构进行的任何操作（例如编译器的 pass）都将产生良好的效果；许多目标都可以使用该基础架构。

MLIR 虽然是一种强大的框架，但是不支持底层的机器码生成算法（如，寄存器分配与指令调度），这些通常由底层优化器 LLVM 负责。

MLIR 的设计中也被整合了其他的经验教训。例如，LLVM 有一个不明显的设计缺陷，会阻止多线程编译器同时处理 LLVM 模块中的多个函数。MLIR 通过限制 SSA 作用域来减少使用-定义链，并用显式的符号引用代替跨函数引用来解决这些问题。

MLIR 是一个重要应用领域是机器学习，这里用Tensorflow （使用数据流图作为数据结构）编译生态系统举例说明，如图6-36所示。

图6-36 TensorFlow编译生态系统示例

将 graph 转为 XLA 高级优化器（XLA HLO）表示，反之，这种表示亦可调用适合 CPU 或 GPU 的 LLVM 编译器，或者可以继续使用适合 TPU 的 XLA。

将 graph 转化为 tensorRT、nGraph 或另一种适合特定硬件指令集的编译器格式。

将 graph 转化为 Tensorflow Lite 格式，然后在 Tensorflow Lite 运行时内部执行此图，或者通过 Android 神经网络 API (NNAPI) 或相关技术将其进一步转化，以在 GPU 或 DSP 上运行。

整个编译流程先将 Tensorflow 的图转化为 XLA HLO，即一种类似高级语言的图的中间表达形式，可以基于此进行一些 High-Level 的优化。接着将 XLA HLO 翻译为 LLVM IR，使用 LLVM 编译成各种硬件的汇编语言，从而运行在硬件上进行数值计算。

除此之外，有时甚至会采用更复杂的途径，包括在每层中执行多轮优化。例如，Grappler 框架便能优化 Tensorflow 中的张量布局与运算。

下图的绿色阴影部分是基于 SSA（Static Single-Assignment，单静态赋值）的IR，然而这种编译方式的缺点是：构建这种编译系统的开销比较大，每一层的设计实现会有重复部分，同一个层次的IR彼此之间虽然相似，但是存在天生的隔离，升级优化缺乏迁移性（优化一个模块，并不能惠及到同层次的其他模块）。因此，目前的问题在于各种IR之间转换的效率与可迁移性不高。

SSA（Static Single-Assignment，单静态赋值）是一种高效的数据流分析技术，目前几乎所有的现代编译器，如GCC、Open64、LLVM都有将SSA技术的支持。在SSA中间表示中，可以保证每个被使用的变量都有唯一的定义，即SSA能带来精确的使用–定义关系，如图6-37所示。

图6-37 SSA单静态赋值的中间表示

对上述问题，MLIR 希望为各种 DSL（Domain-Specific Language，领域特定语言）提供一种中间表达形式，将他们集成为一套生态系统，使用一种一致性强的方式编译到特定硬件平台的汇编语言上。利用这样的形式，MLIR就可以利用它模块化、可扩展的特点来解决IR之间相互配合的问题，如图6-38所示。

图6-38 MLIR为各种DSL提供中间表示

7.1.4 MLIR中的方言：dialect

1.为什么要有方言

当前的编译结构的问题在于各种IR之间转换的效率与可迁移性不高。MLIR试图使用一种一致性强的方式，为各种DSL提供一种中间表达形式，将他们集成为一套生态系统，编译到特定硬件平台的汇编语言上。这样的目标是通过什么手段实现的呢？

从源程序到目标程序，要经过一系列的抽象以及分析，通过下译Pass来实现从一个IR到另一个IR的转换，这样的过程中会存在有些操作重复实现的情况，也就导致了转换效率低的问题。

各种IR组成一个统一的流水线，怎么才能让源语言变成汇编语言，然后在机器上运行呢？统一IR的第一步就是要统一语言，各个IR原来配合不默契，谁也理解不了谁，就是因为语言不通，没法用统一的语言调度流水线工作。MLIR自主创建一个Dialects语言！让各个IR学习这个Dialects语言，这样一来，不光能调度流水线高效工作了，还能随意扩展更改分工，从此各种IR就可以完美地分工协作。

为区分不同的硬件与软件受众，MLIR 提供方言，其中包括：

1）Tensorflow IR，代表 Tensorflow 图中可能存在的一切；

2）XLA HLO IR，旨在利用 XLA 的编译功能（输出到 TPU 等）；

3）实验性仿射方言，侧重于多面表示与优化；

4）LLVM IR，与 LLVM表示之间存在 1:1 映射，可使 MLIR 通过 LLVM 发出 GPU 与 CPU 代码；

5）Tensorflow Lite，将会转换以便在移动平台上运行代码。

每种方言均由一组存在不变性的已定义操作组成，如：这是一个二进制运算符，输入与输出拥有相同类型。

2.将方言添加至 MLIR

MLIR 没有众所周知的固定或内置的操作列表（无内联函数）。方言可完全定义自定义类型，即 MLIR 如何对 LLVM IR 类型系统（拥有一流汇总）、域抽象（对量化类型等经机器学习 (ML) 优化的加速器有着重要意义），乃至未来的 Swift 或 Clang 类型系统（围绕 Swift 或 Clang 声明节点而构建）进行建模。

如果想要连接新的低级编译器，则需要创建新方言，以及 Tensorflow 图方言与方言之间的降阶。如此一来，硬件及编译器制造商便可畅通无阻。甚至可以在同一个模型中定位不同级别的方言；高级优化器将保留 IR 中不熟悉的部分，并等待较低级别的优化器来处理此类部分。

对于编译器研究者与框架制造者，则可以借助 MLIR 在每个级别进行转换，甚至是在 IR 中定义自己的操作与抽象，从而针对试图解决的问题领域构建最佳模型。由此看来，MLIR 比 LLVM 更像是纯编译器基础设施。

虽然 MLIR 充当 ML 的编译器，但也看到，MLIR 同样支持在编译器内部使用机器学习技术！这一点尤为重要，因为在进行扩展时，开发数字库的工程师无法跟上 ML 模型或硬件的多样化速度。MLIR 的扩展性有助于探索代码降阶策略，并在抽象之间执行逐步降阶。

3.dialect是如何工作的

dialects是将所有的IR放在了同一个命名空间中，分别对每个IR定义对应的产生式以及绑定相应的操作，从而生成一个MLIR的模型。整个的编译过程，从源语言生成AST（Abstract Syntax Tree，抽象语法树），借助Dialects遍历AST，产生MLIR的表达式，此处可为多层IR通过下译 Pass依次进行分析，最后经过MLIR分析器，生成目标语言，如图6-39所示。

图6-39 MLIR为各种DSL提供中间表示

7.1.5 MLIR方言

MLIR方言中的方言在部分文献中被翻译为方言。

MLIR 通过方言来定义不同层次的中间表示，每一个方言都有自己唯一的名字空间。开发者可以创建自定义方言，并在方言内部定义操作、类型、属性以及它们的语义。

MLIR 推荐使用方言对 MLIR 进行扩展。有一个统一的中间表示框架降低了开发新编译器的成本。

除了可以使用 C++ 语言对方言进行定义之外，MLIR 也提供了一种声明式的方式来定义方言，即用户通过编写 TableGen 格式的文件来定义方言，然后使用 TableGen 工具生成对应的 C++ 头文件、源文件以及对应的文档。

MLIR 也提供了一个框架用于在方言之间或者方言内部进行转换。

MLIR 使用操作来描述不同层次的抽象与计算，操作是可扩展的，用户可以创建自定义的操作并规定其语义。MLIR 也支持用户通过声明式的方式（TableGen）来创建自定义操作。

MLIR 中的每个值都有对应的类型，MLIR 内置了一些原始类型（比如整数）与聚合类型（张量与内存缓存）。MLIR 的类型系统也允许用户对其进行扩展，创建自定义的类型以及规定其语义。

MLIR 中，用户可以通过指定操作的属性控制操作的行为。操作可以定义自身的属性，比如卷积操作的 stride 属性等。

1.MLIR方言转换

MLIR 中定义操作时可以定义其规范化的行为，以便后续的优化过程更为方便地进行。MLIR 以一种贪婪策略不断地应用规范化变换，直到 IR 收敛为止。比如将 x+2 与 2+x 统一规范化为 x+2。

MLIR 中进行方言内部或之间的转换时：

1）用户首先要定义一个转换目标。转换目标规定了生成的目标中可以出现哪些操作。

2）用户需要指定一组重写模式。重写模式定义了操作之间的转换关系。

3）框架根据用户指定的转换目标与重写模式执行转换。

这个转换过程会自动检测转换方式，例如如果指定了 A → B 与 B → C 的重写模式，框架会自动完成 A → C 的转换过程。

MLIR 也支持用户通过声明式的方式（TableGen）来创建自定义的重写。

MLIR用户

ONNX MLIR：将 ONNX 格式的深度学习网络模型转换为能在不同二进制执行的二进制格式。

PlaidML：一个开源的张量编译器，允许在不同的硬件平台上运行深度学习模型。

Tensorflow：Tensorflow 项目的 XLA 与 Tensorflow Lite 模型转换器用到了 MLIR。

Tensorflow Runtime (TFRT)：全新的高性能底层进行时。旨在提供一个统一、可扩展的基础架构层，在各种领域特定硬件上实现一流性能。高效利用多线程主机的 CPU，支持完全异步的编程模型，同时专注于底层效率。

Verona：一种新的研究型的编程语言，用于探索并发所有权。其提供了一个可以与所有权无缝集成新的并发模型。

2.MLIR 方言（Dialect）

MLIR 通过方言来定义不同层次的中间表示，每一个方言都有自己唯一的名字空间。开发者可以创建自定义方言，并在方言内部定义操作、类型与属性，以及它们的语义。MLIR 推荐使用方言来对 MLIR 进行扩展。有这样一个统一的中间表示框架降低了开发新的编译器的成本。除了可以使用 C++ 语言对方言进行定义之外，MLIR 也提供了一种声明式的方式来定义方言，即用户通过编写 TableGen 格式的文件来定义方言，然后使用 TableGen 工具来生成对应的 C++ 头文件与源文件，以及对应的文档。MLIR 也推荐使用这种声明式的的方式来定义方言。此外，MLIR 也提供了一个框架用于在方言之间或者方言内部进行转换。

为了方便开发，MLIR 也内置了一些方言可供直接使用，见表6-14。

表6-14 MLIR内置了一些方言

acc	affine	async	avx512	gpu	linalg
llvm	nnvm	omp	pdl	pdl_interp	quant
rocdl	scf	shape	spv	std	vector

MLIR 使用操作来描述不同层次的抽象与计算。MLIR 中的操作也是可扩展的，用户可以创建自定义的操作并规定其语义。例如目标无关操作、仿射操作与目标特定操作。MLIR 也支持用户通过声明式的方式（TableGen）来创建自定义操作。

MLIR 中的每个值都有其对应的类型，MLIR 内置了一些原始类型（比如整数）与聚合类型（张量与内存缓存）。MLIR 的类型系统也允许用户对其进行扩展，创建自定义的类型以及规定其语义。

此外在 MLIR 中，用户可以通过指定操作的属性的值来控制操作的行为。操作可以定义自身的属性，比如卷积操作的stride属性等。

方言的变换

在 MLIR 中定义操作时可以定义其规范化的行为，比如将x + 2与2 + x统一规范化为x + 2，以便后续的优化过程更为方便地进行。MLIR 以一种贪婪地策略，不断地应用规范化变换，直到中间表示收敛为止。

在 MLIR 中进行方言内部或方言之间的转换时，用户首先要定义一个转换目标。转换目标规定了生成的目标中可以出现哪些操作。然后用户需要指定一组重写模式，这些重写模式定义了操作之间的转换关系。最后框架根据用户指定的转换目标与重写模式执行转换。这个转换过程会自动检测转换方式，例如如果指定了A → B与B → C的重写模式，框架会自动完成A → C的转换过程。MLIR 也支持用户通过声明式的方式（TableGen）来创建自定义的重写模式。当转换的方言之间有着不同的类型系统，用户可以使用类型转换器来完成类型之间的转换。

7.1.6 MLIR实践

1.MLIR应用工具

1）ONNX MLIR：将 ONNX 格式的深度学习网络模型转换为能在不同二进制执行的二进制格式。

1） PlaidML：一个开源的张量编译器，允许在不同的硬件平台上运行深度学习模型。

2） Tensorflow：Tensorflow 项目的 XLA 与 Tensorflow Lite 模型转换器用到了 MLIR。

3） Tensorflow Runtime：一种新的 Tensorflow 运行时。

5）Verona：一种新的研究型的编程语言，用于探索并发所有权。其提供了一个可以与所有权无缝集成新的并发模型。

MLIR是一种新型的编译器框架，其设计从已有的编译器的实现中吸取了经验与教训，包括了中间表示的定义、转换以及优化等功能，极大地方便了新的编译器的开发与调试工作。同时，MLIR 也包含了很多现成的工具可直接使用（batteries included）。MLIR 包揽了编译器设计中的通用部分，使得编译器的开发人员可以专注于核心的语义分析、中间表示的设计与变换，以此降低开发成本，提高开发效率与提高成品质量。

2.外部链接

1）MLIR 主页：https://MLIR.llvm.org/。

2）MLIR 语言参考：https://MLIR.llvm.org/docs/LangRef/。

MLIR 是 LLVM 项目的子项目，要编译 MLIR，首先获取 LLVM 的源代码。

LLVM 的源码可从 GitHub 获取：

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

附录：编译与安装 MLIR

1）下载 MLIR

用户也可以直接下载源码包：https://github.com/llvm/llvm-project/releases。

假定 LLVM 的源码目录为 $LLVM_SRC。

2）编译 MLIR

首先用户需要指定一个路径用于存放编译中间产物，假定其路径为 $LLVM_BUILD。然后使用下列命令对 LLVM 进行配置：

cmake -S "$LLVM_SRC" -B "$LLVM_BUILD" -DLLVM_ENABLE_PROJECTS=MLIR -DCMAKE_BUILD_TYPE=Release

默认情况下，LLVM 禁用了异常处理与运行时类型信息。如果应用程序需要依赖这些功能，可指定在配置时指定 LLVM_ENABLE_EH 与 LLVM_ENABLE_RTTI CMake 变量的值为 ON：

cmake -S "$LLVM_SRC" -B "$LLVM_BUILD" -DLLVM_ENABLE_PROJECTS=MLIR -DLLVM_ENABLE_EH=ON -DLLVM_ENABLE_RTTI=ON -DCMAKE_BUILD_TYPE=Release

更多的 LLVM 配置参数参见 https://llvm.org/docs/CMake.html。代码如下：

//第1章/sum.c

执行完配置过程后使用下列命令执行编译：

cmake --build "$LLVM_BUILD"

安装 MLIR

使用如下命令将 LLVM 安装到 /usr/local 目录：

cmake --install "$LLVM_BUILD"

如果想指定另外一个安装目录，例如 $INSTALL_DIR，可以使用 --prefix 命令行参数来指定：

cmake --install "$LLVM_BUILD" --prefix "$INSTALL_DIR"

在 CMake 项目中使用 MLIR

用户可以在 CMake 项目文件中使用下列语句添加查找 MLIR 依赖：

find_package(MLIR REQUIRED CONFIG)

如果 MLIR 被安装到了系统目录（比如 /、/usr、/usr/local 等），CMake 无需特定的配置就能找到 MLIR；如果 MLIR 被安装到了非系统目录，可以在 CMake 的配置过程通过 CMake 的 MLIR_DIR 变量来指定 MLIR 的安装位置：

cmake "$MY_PROJECT_DIR" -DMLIR_DIR="$INSTALL_DIR"

成功之后用户可以直接使用 MLIR 的库作为编译目标的依赖：

add_executable(my-executable main.cpp)

target_include_directories(my-executable SYSTEM PRIVATE ${MLIR_INCLUDE_DIRS})

target_link_libraries(my-executable PRIVATE MLIRIR)

其中 MLIR_INCLUDE_DIRS 是自动生成的变量，其指向 MLIR 的包含目录。

在使用 CMake 定义可执行文件目标时，如果 LLVM 禁用了运行时类型信息，那么依赖于 LLVM 的可执行文件目标，也需要禁用运行时类型信息，否则可能会编译失败。LLVM 提供了一个 CMake 帮助函数 llvm_update_compile_flags 可以自动完成这个配置。这个函数定义在 LLVM 提供的 AddLLVM.cmake 文件中。用户可以使用下列语句导入 AddLLVM.cmake 文件：

list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")

include(AddLLVM)

导入 AddLLVM.cmake 文件后就可以对编译目标进行配置了：

llvm_update_compile_flags(my-executable)

完整的 CMake 项目文件示例如下：

cmake_minimum_required(VERSION 3.15)

project(my-executable)

find_package(MLIR REQUIRED CONFIG)

list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")

include(AddLLVM)

add_executable(my-executable main.cpp)

target_include_directories(my-executable SYSTEM PRIVATE ${MLIR_INCLUDE_DIRS})