【从零开始学深度学习编译器】二十，MLIR的Pattern重写机制

0x0. 前言

这篇文章对MLIR的Pattern重写机制进行梳理和汇总，会结合实际例子把MLIR的两篇文档转化成容易看懂的形式。这两篇文档分别是https://mlir.llvm.org/docs/PatternRewriter/ 和 https://mlir.llvm.org/docs/Rationale/RationaleGenericDAGRewriter/ 。做这件事的动机是因为在我的开发过程中已经大量使用了MLIR的这个Pattern Rewrite机制，也经常回看这两篇文档所以翻译+梳理+总结一下。

0x1. Generic DAG Rewriter Infrastructure Rationale

题目可以翻译为通用的Dag重写架构的基本原理。对应 https://mlir.llvm.org/docs/Rationale/RationaleGenericDAGRewriter/ 这篇文档的内容。这里主要介绍了用于MLIR的通用Dag-to-Dag重写架构背后的基本原理。

0x1.1 介绍和动机

编译器IR目标是在各种抽象级别上表示代码 ，这在表示能力和易于变换方面提出了不同的折衷。 但是，表示代码的能力本身并不是很有用——您还需要能够实现这些变换。

编译器的变换有很多，这里主要介绍的是一种对MLIR目标非常重要且反复出现的变换：匹配一系列Op组成的Dag，然后将其替换为另外一个Dag。这是很多学习编译器不可或缺的一部分，对于诸如“消除identity（直连）节点”或者使用"x"替换"x+0"这种优化，通用规范化框架（比如LLVM的指令组合(Instruction Combiner)），以及为编译器在多个中间IR上实现优化算法提供了一个有用的抽象。

MLIR 的一个特殊优势（以及与 LLVM、GCC、XLA、TensorFlow 等其他编译器基础架构的主要区别）是它使用单个编译器 IR 来表示多个抽象级别的代码：MLIR 操作可以是“TensorFlow operation”、“XLA HLO”、仿射循环嵌套、LLVM IR 指令（可传递地包括 X86、Lanai、PTX 和其他目标特定指令）或 MLIR 算子系统可以合理表达的任何其它内容。 鉴于 MLIR 跨越了如此广泛的不同问题范围，用于执行图到图重写的单一基础架构可以帮助解决许多不同的领域挑战。

像 MLIR 这样的基于静态单赋值 (SSA) 的IR可以轻松访问Op的操作数和“users”。 因此，这些图到图重写的自然抽象是 DAG 模式匹配的抽象：客户端定义 DAG tile模式（其中tile是定义 DAG 子图的一系列Op），并且每个模式都包含一个产生的结果 DAG 和产生它的成本（或者相反，叫作进行替换的好处(benifit)）。 一个通用的基础设施可以有效地找到并执行重写。

虽然上面提到的概念很简单，但细节很微妙。 这篇文档里定义并探索了可以解决范围广泛的不同问题的一组抽象，并预计可以应用于 MLIR 随着时间的推移将面临的许多不同类型的问题。

常量折叠（Constant Folding）

DAG 到 DAG 模式匹配的一个退化但常见的情况是常量折叠：操作数包含常量的Op通常可以折叠为结果常量值。

MLIR 的Op可能会覆盖fold来实现，与一般的 DAG 到 DAG 模式匹配器相比，它暴露了一个更简单的 API，并适用于通用的匹配器不适用的情况。 例如，DAG 重写可以删除当前函数中的任意节点，这可能会使迭代器无效。 作为 API 的常量折叠则不会删除任何节点，它只是提供一个常量值（列表）并允许客户端根据需要更新其数据结构。

关于常量折叠请看一下后面的0X3节的示例讲解，是这篇https://mlir.llvm.org/docs/Canonicalization 文档的翻译。

相关工作

考虑到几乎每个现有的编译器都必须多次解决这个问题，因此需要考虑大量相关工作。 一个统一的问题是，所有这些系统都旨在解决一个特定的、通常是狭窄的问题：另一方面，MLIR 希望在单个基础设施中解决许多这些问题。 以下是一些相关的Pattern Rewriter系统，以及它们工作的优缺点（与 MLIR 中存在的基础设施最相似的设计是 LLVM DAG-to-DAG 指令选择算法）。

• AST 级模式匹配器：文本中存在大量的source-to-source的翻译器用来做等价变换以提升性能（比如把x*0变成0）。一个较大的例子是GCC fold函数，它对AST进行了很多优化。Clang具有应用于表达式的简单常量折叠的类似例子（C++的要求），但并不会对AST执行常见的优化。
  AST 优化器的主要缺点是我们无法看到具有多种用途的Op。 众所周知，DAG 模式匹配比树模式匹配更强大，但另一方面，DAG 模式匹配会导致重复计算。
• 第二种就不介绍了，感兴趣可以看官方文档。
• LLVM’s DAG-to-DAG Instruction Selection Infrastructure：LLVM 中的指令选择子系统是多年迭代和研究的结果，这是由于 LLVM 需要支持大量的目标代码生成、现代指令集（例如 X86）的代码生成器的复杂性以及狂热的追求跨目标重用代码。 Eli Bendersky 写了一篇关于它如何工作的简短概述，LLVM 文档更深入地描述了它，包括它的优点和局限性。 它允许编写这样的模式。

def : Pat<(or GR64:$src, (not (add GR64:$src, 1))),
          (BLCI64rr GR64:$src)>;

此示例为 X86 目标描述中的“blci”指令定义了一个匹配器，该文件中还有许多其他指令（查找 Pat<> 模式，因为它们没有纠缠于编译器的细节，如汇编器/反汇编器生成逻辑）。

下面说了一些LLVM的这个DAG-to-DAG 指令选择机制的好处和坏处，截图放在下方。

小结

MLIR 面临着广泛的模式匹配和图重写问题，在多个级别上使用通用代码表示的主要优势之一是它允许投资并高度利用单一基础设施来完成此类工作。

这里后续还介绍了一些Dag重写机制的目标，包括它解决了哪些问题以及使用的匹配策略，以及良好的报错信息等等。

0x2. Pattern Rewriting : Generic DAG-to-DAG Rewriting</