MLIR设计与Dialect体系分析

特定领域 GPU MLIR 设计

Domain-Specific Multi-Level IR Rewriting for GPU

论文原文: https://arxiv.org/pdf/2005.13014.pdf

项目地址:  https://github.com/spcl/open-earth-compiler

多级（IRs）层次结构，逐级降低程序，执行代码转换。根据自有设计原则，开发原型编译器，设计stencil和 GPU dialect。在 LLVM MLIR 编译器基础架构上实现的两个域特定优化（500 行代码）。从本质上讲，多层次重写有望开启由域和特定目标dialect组成的专业编译器的时代。 

目标域特定方法正在彻底改变高性能设备的代码生成，激发了目标域特定语言 （DSL） 框架的发展，实现了通用编译器无法达到的性能。例如，Halide自动生成用于图像处理的高性能代码，XLA利用目标域特定的编译加速深度学习，Stella是将天气和气候模拟移植到 GPU 实现 2.9 × 加速。特定编译器的成功, 随着时间的推移也暴露了弱点：与通用编译器流程分离。Halide, XLA, Stella和其他专有方案，设计时没有考虑到重用性。可复用的编译器基础架构，如 ROSE 或 LLVM,  实现目标域特定功能，设计和维护成本很高。

因此, 如何设计一个目标域特定的编译器, 该编译器

(a) 与用户前端完全脱钩, 

(b) 便于实施域转换, 

(c) 明确分离通用组件?

 本方案设计和实现气象和气候建模编译器来构建多级IR结构。使用图像处理stencil计算模式, 但它需要完全不同的优化策略才能达到最优性能。地震成像是 3D 的, 实现了控制流的带宽限制的低阶stencil。相比之下, 图像处理流程在2D 数据结构上使用stencil, 地震成像stencil是高阶计算密集型。多维数组上的循环抽象, 算术优化 以及 生成 GPU 代码 这些方面大体相同, 但要重新实现。建议使用多层次 IR 重写来设计 DSL 编译器。

此方法是

(a) 基于SSA

(b)高级语义运算

(c) 渐进式IR的组合, 为可复用 高性能设备代码生成 提供了有效的框架。

基于 SSA IR , 可复用通用编译器的优化方法。高层运算主要编码目标领域属性, 使其可用，例如，SSA 数据流, 无需分析。渐进式层次下降保存目标领域信息、将转换表示为peephole优化 以及 引入可重用的低层级抽象。

a. 多级 IR 重写的编译器概述

MLIR 编译器, 使我们能够获取 LLVM在过去 15 年中实践验证的 IR 设计原则, 实现生产级编译器 IR。Open Earch 编译器(图 1) 是一个端到端的编译流, 利用多层 IR 重写来实现高性能代码生成。其核心是一组 MLIR dialects, 即运算和变换的集合, 以及它们之间的转换。Open Earch编译器使用 peephole-style重写模式, 逐步将程序从高级别的域特定dialect转换为较低级别的平台特定dialect来优化代码。每个dialect都定义了一个抽象, 使相关分析成本较低, 便于转换。编译过程从 stencil dialect 开始, 该dialect 用于面向用户的DSL，以及用于转换的数据结构 (如 stencil inlining）。Stencils进一步向更低层次下沉, 使用一系列IR表示的操作 (Affine), 结构控制流(SCF) 和算术运算 (Standard), 这些操作在 MLIR中都有了, 以及循环转换和数值层级的转换（如循环展开或通用子表达消除）。这些 IR 使我们专注于结构化循环抽象, 而不是编译器后端常见的低级别的基于"goto"的 SSA IR。使用此结构设计通用GPU kernel dialect, 基于stencil级别保留的并行信息, 使用简单模式实现loop-to-kernel转换, 避免成本较高的 GPU mapping algorithms。这完整的流程可以将高级别代码转变为可高速执行的二进制文件。

虽然stencil dialect是通用的, 足以涵盖一系列应用 (如图像处理或地震成像）, 但本案例是要实现气候领域的高性能计算。基于Stencil operations的语义, 能够用通用指令转换（例如冗余消除）取代复杂的循环变换序列，仅需很少的分析以确保有效性。其他领域可以调整stencil dialect, 以适应需求变化, 或只重用中低层抽象。实践证明, 基于多级 IR 重写，开发具有可重用的组件的特定编译器非常简单。

多级 IR 重写设计模块化的编译器:

●  stencil language可表示为 MLIR dialect,将stencil program高级数据流编码为 SSA

●  多次转换, 用简单的peephole优化, 而非传统高级别层次上的循环转换, 来调优 stencil program

●  平台无关 GPU dialect编译流程, 转换为特定 GPU 代码

b. 多级 IR 重写

多级IR 重写, 定义一系列可重用的抽象, 实现转换来简化特定编译器的开发。目标是降低每级的复杂度, 直接在 IR 中编码, 保留转换的有效性, 降低分析成本。识别特定领域的抽象对于多层次重写至关重要。每一个新的抽象都会增加复杂度, 相反, 如果抽象不完整, 某些工作负载无法表达。使用 MLIR 开发 OpenEarth编译器, 对它进行抽象, 必要时引入新的抽象。目的是确保性能: 并行和数据本地化, 这两个是需要平衡的转换, 会产生复杂的优化问题。提出这些问题, 使用以下原则从具体问题中提取并行和数据本地化信息。

P1 转换驱动的语义

流程中的特定抽象，如 stencil 或 GPU kernels, 应该倾向于为转换准备而不需要程序员友好。主要是建立一堆中间表示, 使编译器能够对特定程序进行推理, 而无需多次重复分析, 如 循环提取 或 依赖分析。每个 IR 都侧重于处理一组特定的转换, 获取所有必要的信息。用户可用性可以直接放到 DSL 前端。

P2 渐进式转换

还要建立一个有效和简化的转换流程, 代码逐步从高层 IR 降低到低层目标 IR。应设计不同的 IR 抽象, 保持高级别语义信息, 避免对高层语义的恢复过于复杂。在较大的编译器中, 渐进式下降, 抽象应无缝组合, 共存于单模块中, 有选择地下降转换。

P3 显式分离

根据之前的原则的抽象组合性, 将各个抽象部分组合起来比解耦复杂的表示更容易。对IR分离的思路往往取决于特定的应用领域,这些区别要明确体现到IR中。特别是性能相关的抽象, 如并行度或内存分配, 应在 IR 中有所体现, 并彼此分离。同样, 编译和运行时抽象也应分开。从长远来看, 这种表示更适合现代搜索技术。渐进式降低的IR (P2）, 各级之间清晰的分离(P3), 如表I所示, 每级的转换易于实现(P1)。这种多层次的表示有助于将优化转换与层级降低分开。

 c. MLIR 基础架构

MLIR 编译器基础架构, 适合多层IR 重写, 它有dialect扩展以及对声明式重写模板的内置支持。因此,OpenEarth编译器可以作为一组 MLIR dialect实现, 并转换为重写模式。如果设计一些抽象, 重用 standard, Loop, Affine, 和 LLVM IR dialect, 以便它们能够组合在一起。

MLIR 核心概念包括operations,values, types, attributes, (basic)blocks 和 regions。operations是程序描述的原子单位。values表示运行时的数据, 始终与编译时已知的类型关联。operations使用values, 定义新 values。values只能定义一次, IR 需要遵守 SSA。types保留有关values编译时信息，而 attributes 将编译时信息添加到operations。block是一系列operations, 与其他block连在一起形成regions。region对应于含有语义operation。主要的控制流只处理 operation和对应的regions, 以及同一region 的 blocks 的出入口之间。特定 operation 定义了控制流的结构, 例如, block 的最后一个 operation 可以有条件或无条件地将控制流转移到另一个 block。

 图2将 64x64x64元素的定义值 %def 设置为负值 %use 的MLIR运算示例。

 图2 演示了stencil dialect 的一个 operation的语法。stencil.apply 操作用 %use使用值, 用 %def 定义值。types和 attributes 标记编译时运算。嵌套 region 由一个 block 组成, 用基本 block 参数 %arg 执行计算。这种分层组织分为 blocks 和 regions, 能够实现无限嵌套。没有一组固定的 operations, attributes, 或 types。相反, 每个 MLIR 用户可以定义自己的或重用其他人定义的。即使是 MLIR 内置功能很大程度依赖于这种扩展性。因此, stencil computation 可以通过提供自定义 types, attributes 和 operations, 构建stencil编译器。控制流操作, 比如循环或多维数组数据类型也是其构成部分。从架构角度来看, 自定义运算和类型与常见的数值运算和类型没有区别。dialect 是一组协同工作的operations,attributes 和 types。dialect的结构没有所谓正式或技术限制, 除非运算语义另有规定, 否则 region 可以包含不同dialect的运算, 运算可以引用不同 dialect 定义的 types 和 attributes。因此, 新的抽象可以作为新的dialect引入MLIR生态系统。

d. stencil dialect

OpenEarth编译器运行天气和气候模型。这些模型主要包括偏微分方程。包括数十个stencil 代码, 由多个 stencil operator 组成作用在不同的3D网格上。本案例将空间划分成三维网格, 每个网格与其他六个相邻。这样就可以通过索引处理单元格问题。在stencil 代码中, 不是优化单个stencil, 必须优化stencil chain 或整个程序链, 例如, 利用生产者-消费者融合来获得最优性能。设计stencil dialect来表示stencil 代码, 包括连接它们之间的stencil 运算, 以及数据流上的输入/输出数据结构, 可选控制流, 按照之前定义的多级 IR 重写原则实现。dialect分解为高级IR(P3), 构建运算和低层IR之间的数据流, 低层IR构建了单个运算的并行执行。其中我们模拟了单个操作员的并行执行。高层IR支持控制流路由转换, 它的stencil operator可被视为一个单位（与低级别IR算术指令集合相反）, 低层IR支持并行（P1）。层级分离也包含在渐进式下降中（P2）。 

dialect不是为用户设计的 DSL, 是支持变换（P1和P3）的编译器 IR, 通过:

(a) 保持 stencil 概念的高层语义, 以便它们可以变成一个单元进行处理, 

(b)不强加特定执行顺序以便平行, 

(c) 用值语义而不是分配存储对象, 避免缓存分析。

Dialect Overview

 图3：stencil 示例代码, 计算简单的stencil数组 %in, 将结果存储到数组 %out。

stencil dialect 主要考虑特有的概念, 用MLIR标准dialect来表达具体计算（P2）。图 3 stencil 代码, 处理 64x64x64的每一个元素, 添加输入array% 左右邻, 结果存储到 %out中。"stencil"前缀从dialect中识别 operations 和 types。Dialect 定义了两种 type: A !stencil.field 是一个多维数组, 存储网格中所有点。stencil 代码的输入和输出都是这个类型。A ！stencil.temp  是网格矩形子域上的多维元素集合。临时变量有值语义, 但不存储。这种类型的值要么指向输入数组的子域, 要么保留stencil operation的计算结果。这两种类型都存储1到3维数组上的单精度或双精度浮点元素（f32 或 f64）。stencil dialect 还定义了六种操作。stencil.assert操作判断数组的静态形状。stencil.load操作输入数组, 返回指向stencil的输入元素的临时变量。stencil.apply操作执行 stencil, 定义了保存计算结果的变量。循环迭代部分, 对某一点执行 stencil operator。stencil.access 操作以特定偏移读取输入临时变量, 而 stencil.return操作将输出设置为当前位置。在这两者之间, 使用标准 dialect 计算左右邻元素的和。stencil.store操作最终将计算结果存储到输出数组。那么, 其范围属性指定了stencil 代码编写的具体领域。编译器利用输出范围自动推导整个stencil 代码的访问范围和迭代域。

Shape & Domains

 图4：示例range表示为: （左）下界,上界,stencil accesses(右)

range对指定 stenciliteration 域和访问范围至关重要, 特别是考虑到stencil可能会访问其计算域之外的输入下标, 比如边界。图4显示二维range（左图）。原点表示计算域的下限, 所有坐标设置为零。包含下界和由冒号间隔的上界的绝对坐标, 以确定范围。GPU 中, 整型索引计算是一个重要的性能瓶颈。实际操作中, stencil 经常针对相同的问题反复执行stencil,stencil dialect 支持特定的JIT编译。内存大小和迭代域定义为数字类型编译时属性(P3）来实现这个功能。

Stencil Operators

stencil operator 对网格的所有元素进行元素计算, 除了某些常量宽度边界。相对于输出元素的坐标, 它以常量偏移访问输入数组的元素。stencil.apply操作实现了标量运算的stencil operator。标量运算在循环嵌套中作用到所有元素。stencil operator输入和输出, 与这个操作使用和定义的值相对应。Stencil.access 可以访问输入的各个元素, 以恒定偏移读取元素。然后, stencil dialect lowering 添加指针偏移到当前迭代。图4显示了二维stencil迭代的偏移计算（右图）。stenciloperator 的区域必须由某个stencil.return 终止, 输出元素值作为参数返回。Stencil.access 和 stencil.return一起指定了stencil operator的访存模式。两者都仅作为 stencil operator的一部分。实际的stencil gragrams通常会实现数十个stenciloperators。因此, stencil program  需要一些方法来编排它们。

Stencil Programs

stencil program 执行一系列相互依赖的 stencil operators。它从输入数组中加载数据, 添加stencil operators, 将结果存入输出数组。程序的 SSA 定义使用图指定了stencil operators（P2）之间的高层级数据流。将高级数据流和函数中的stencil operators, 进行代码转换, 消除复杂的程序分析（P1）。另外三个操作也是程序的一部分。stencil.assert指定输入或输出数组范围。需要定义所有输入和输出数组的索引范围。stencil.load 返回一个临时变量, 所有输入数组元素。stencil.store 存储输出数组。

 图5：是两个相互依赖stencil program。Stencil.load返回临时变量, 保存 %in 数组元素。第二个stenciloperator使用第一个的结果。最后, stencil.store 保存第二个stencil计算结果到数组 %out。所有stencilprogram 参数没有别名, 从unit加载或存储到unit。中间结果保存在 !stencil.temp, 不存储, 不命名。根据!stencil.temp的数值语义, def-use图对stencil operators之间的数据依赖关系进行编码（P1 和 P3）。

Control Flow

实践中的stencil application 不只是纯数据流语义。可以 eagerexecution 和 JIT编译处理大部分程序级别的控制流, 采用 MLIR 内置的结构控制流（SCF） dialect, 在 stenciloperators 内部实现动态控制流。

 图6：SCF dialect使 stencil operator内部控制流得以实现。

图 6显示了一个 stencil根据判断条件访问其中一个参数。scf.if执行"then"或"else"。与if-else相比, 操作返回 scf.yield 设置的结果。这个表示使数据流变得清晰, 每个stencil维护一个stencil.return。除了scf.if之外, select也可以基于条件选择计算值。要支持scf.if 不需要调整编译器（P2）。选择内置 MLIR SCF dialect可以渐进式逐层下降, 显式分离和可组合抽象, 使编译器组件在多级 IR 重写方案可以重用。

e. Stencil Transformations

stencil dialect 有3类转换: 

1）性能优化,  

2）准备降低,  

3）实际降低.

Optimizing Transformations

stencil dialect 的所有优化转换均在高层运行, 既不引入循环, 也不引入存储分配（P1）。

 图7: stencil program两种模式迭行生产者-消费者融合。def-use edges表示stenciloperations 之间的是数据流。

stencil inlining pass 在 def-use graph进行融合。反复使用生产者-消费者融合的stencil变体, 内联计算取代所有对生产者结果访存。如果消费者访问多个生产者, 需要对迭代域中的每一个点重计算。任意顺序inlining stencil可能会引入循环依赖。例如, 消费者的输入可能依赖于另一种stencil, 这种stencil会暂时依赖 producer stencil的输出。不需要开发一种算法将stencil按特定顺序融合, 而是匹配和重写小子图, 并使用 MLIR 一步步重写stencils。图7显示了inlining patterns。如果producer只有一个consumer, 则inlining pattern将生产者 P 和消费者 C 匹配。如果模式匹配成功, 删除producer stencil, 将计算内联到consumer。另外还要更新fused stencils的参数和结果列表。重定向的模式将生产者 P 及它的消费者 C1 与 CN 匹配。如果模式匹配成功, 将生产者的所有输出通过要执行的下一个消费者进行路由。红色箭头标记了重路由的数据依赖。前一种模式是真正的inlining, 后一种模式是准备一个inlining step。inlining实现引入重计算, 即使消费者多次访问相同的偏移, 始终 inline 整个producer, 即使只访问其中一个输出。死代码消除 dead code elimination 和子表达消除 common subexpression elimination 稍后会清理代码。这些转换依赖于stencil accesses 并无其他作用（stencil输入不可变, 不给输出起别名）。这个编译器目前不实现启发式融合和连续内联, 不管使用哪种模式。

 图8：沿着j维度展开stencil的两次迭代。

stencil unrolling pass 多次复制stencil operator, 一次更新多个网格。图 8显示了示例程序的unroll版本。Unrolling 是高层dialect实施循环转换的又一例证。这个实现不是改变循环, 而是标注高级stencil dialect, 直接降低到unrolled loops。只修改 stencil.apply 上的嵌套区域, 但不处理接口。刚开始的时候, 每次展开循环迭代只计算stencil一次, 调整访问偏移。使用stencil.return 返回所有展开循环迭代的结果, 用可选属性attribute标记展开因子和维度。unrolling pass支持所有展开维度和展开因子。但lowering目前仅限于域大小均匀划分的展开因子。

Inlining 和 unrolling 可提高stencil程序的性能。inlining 引入冗余计算, 减少了芯片外的数据移动。Unrolling 消除一部分冗余计算, 因为unrolled stencil operator 通常在同一偏移多次使用生产者。使用现有的通用子表达式消除pass,不是自动删除冗余计算。

Preparing the lowering

优化stencil程序后, 推断出所有访问范围和迭代域, 以准备lowering（P2）。

shape inference pass 获得输入数组和 stencil operators。stencil程序只定义了代码编写的输出范围, 需要做shape inference。pass从输出范围开始, 根据stencil程序依赖的use-def链, 扩展访问范围。算法反向遍历代码的所有运算, 用range attributes标记可计算范围。这些计算范围是最小边界框, 包含了这个当前运算需要用到的所有变量或值。如果消费者是stencil.load操作, 它的访问就和输出范围的属性是一样的。如果消费者是stencil.apply 操作, 它的访问范围就是所有要处理的变量或值的最小边界框的迭代区间。一旦stencil.load操作的访问范围确定了, 就可以知道输入数组是否足够大。

虽然这些访问区间分析看上去和渐进式lowering的思路相矛盾, 但它的目的不是恢复之前的高层信息, 它主要是自动化地对手动访问的地方做错误验证。形状推理使能lowering, 且不会有性能影响。

Lowering to Explicit Loops

 图9: 将stencil dialect 转换为MLIR SCF + Affine+ Standard dialect 进一步降低到GPU abstractions。

stencil lowering 使用转换模板将单一stencil 运算转换成对应的MLIR。这是编译流程中最后一个与特定领域相关的部分, 如图1中所示, 将high-level stencil程序向可执行代码的下降。即使是标准dialect层次上, MLIR 也提供了相当高层次的抽象。memref 是结构化的多维buffer抽象, 可以具有静态或动态大小, 如果layout不是行主格式, optional layout attribute 就可以定义索引计算。layoutattribute还可以将strided hyper-rectangular定义为memory buffer,例如, 每个维度都有偏移和非单位steps。另一个例子是scf.parallel, 可以并行多维循环。

图9演示了从stencil dialect 到MLIR standarddialect level 的lowering的过程。本例定义六种转换模板, 引入loops, index computations, memory accesses, temporarystorage 。lowering 之后, 通过分析检测 stencil operators或读取偏移。实现特定于域的转换变得更加困难。反过来, 通过引入循环loops 和临时存储 temporary storage, 实现了程序顺序执行, 但仍需为后续 GPU lowering 保留并行语义。

f. GPU Dialect

 图10: lowering of a kernel: 1)  inline form支持主机设备代码和其他转换, 2) function form 将设备代码隔离在一个可实现设备特定优化和单独主机/设备代码编译的独立模块中, 3)  二进制文件将kernel嵌入为常量数据。

GPU是实现高性能的首选平台, 可以构建了针对这些设备的多层级编译器。按照之前的原则进行设计, 实现 MLIR GPU dialect, 不依赖供应商库抽象 GPU 执行模型。泛化了 MLIR’s NVVM, ROCm 和 SPIR-V 表示, 将统一的平台独立设备映射（P1）与平台相关的代码生成（P3）code generation区分开来。GPU dialect 不是通用的 SIMT 执行模型（P3）, 也不是较低级别抽象（P2）的高层表示。

Dialect提供了一组GPU 的特定概念: 分层线程结构 (blocks, threads, warp）; synchronizationthrough barriers; 内存等级memory hierarchy (global, shared, private, constant memory);标准计算原语, 如 parallel reductions。单个模块（P3）中支持主机/设备单独编译。设备端编译由 MLIR 模块实现, 这些模块递归包含以不同方式处理的其他模块。

图10是 GPU lowering 期间, kernellaunch的两种形式。Inline form 使用gpu.launch 定义kernel inline。嵌套实现kernel, block 参数提供 thread 和 block 标识符。不需要处理参数, 嵌套之外定义的值可以处理。function form 使用 gpu.func 实现 kernel , 在单独的模块中实现, 通过 gpu.launch_func 激活 kernel。专有操作可访问线程和块标识符。所有非常量内核参数均可传递, 而常数则传递到内核函数中。inline 和 function form 均接受 GPU grid配置, 支持 GPU 内存等级的不同级别缓存的声明分配。Kernel code 按照 SIMT 模型进行单线程计算, 有专门的机制访问线程和块标识符。特定的GPU原语（barrier synchronization, shuffles, and ballots）只能存在于kernel内部。图10 演示了 GPU lowering 的主要步骤, 从内联形式（左）开始, 通过函数形式（中间）, 到编译好的二进制（右）。并行循环嵌套可以就地转换为内联形式, 使用循环边界作为 GPU grid 配置。转换之后, 使用通用子表达式和死码消除、规范化和  GPU kernel内常量传播, 最大限度地减少主机/设备内存流量。Kernel inline 无可见性限制, 基于 SSA 的转换可无缝地作用于主机/设备边界。Kernel 在专用 GPU 模块中实现独立的功能。内核调用的函数复制到模块中, 内核外定义的值作为函数参数传递进去。kernel位于单独的模块中, 实现独立的主机/设备优化和编译。Host code 程序内优化无法作用到kernel。GPU 模块最终通过专用dialect转换为平台特定表示（PTX）, 并使用供应商编译器（ptxas）将进一步编译为二进制文件。生成的二进制文件以全局常量嵌入到原有模块中。这种方法支持多种版本, 支持不同大小的工作负载的多个架构或特定 kernel。然后, 用二进制扩展的原始模块变成了host模块, 可以优化,编译,执行。Kernel 调用降低到对设备驱动程序库或运行时的调用。

机器学习编译器：MLIR Dialect体系

介绍编译器和中间表示 (IR) 演进趋势的整体理解， 也讨论了 LLVM IR, SPIR-V, 和 MLIR 所要解决的问题以及相应的设计着眼点。今天对 MLIR 做进一步展开，分析一下机器学习相关的 dialect 体系。值得注意的是，MLIR 是一个编译器基础设施，它可以用来编写各种领域专用编译器，并不限于机器学习。不过机器学习确实是 MLIR 得到最活跃开发和应用的领域，尤其是用来转换各种 ML 模型，以及支持各种异构硬件。

基础组件

编译器的一大优势是可组合性 (composability)。 如果功能甲、乙、丙分别得到了实现，那么它们的各种组合也自然而然会得到支持。 这种特性是编译器与算子库 (library) 的核心区别之一；在算子库中，不同的组合可能需要经由完全不同的手写代码来实现。 通过把指数级问题变成线性问题，编译器长期而言可以缩减大量的工程投入。

为实现这种可组合性，我们需要分解问题而后开发适宜的基础组件 (building block)。 在中间表示中， 我们一般把这些基础组件定义成各种操作 (operation)。 但对机器学习而言，仅用操作是很难组织出结构清晰优美的软件栈的， 因为输入模型和生成代码之间存在着巨大的语义鸿沟 (semantic gap)。除此之外，输入模型和目标硬件也种类繁多，有着各式各样的需求。 为此，MLIR 通过 dialect 机制实现了更高层次的基础组件。

一个 dialect 基本可以理解为一个命名空间。 在这个命名空间中，我们可以定义一系列互补协作的操作，和这些操作所需的类型 (type) 以及属性 (attribute) 等等。 特定的机器学习编译器只需要组合现有的 dialect，并加以自己的扩展或者定制。 这当中，MLIR dialect 有几个特性值得一提。

内嵌结构的操作（Operations carrying structures）

操作无论对表示（representation）还是转换（transformation）而言，都是编译器中的原子性组件。我们可以把操作放到基础块 (basic block) 中，然后把基础块放到函数 (function) 中。但这只是浅浅的两层结构；语义 (semantic) 其实还是依赖于每个单独的操作，模式匹配 (pattern matching) 依然发生在一个或者一组松散的操作之上。想要定制已有操作，或者把几个操作进行强结合以便给模式匹配设定清晰边界，依旧很困难。

 代表抽象层次的类型（Types signaling abstraction levels）

操作归根到底只是针对某种类型的值 (value) 所进行的某种计算 (computation)。类型才是抽象层次 (abstraction level) 的代表。举个栗子，张量 (tensor)、buffer、以及标量 (scalar) 都可以支持加减乘除等各种操作。这些操作在本质上并没有多少区别，但它们明显属于不同的抽象层次。张量存在于机器学习框架或者编程模型 (programming model) 这一高层抽象。Buffer 存在于执行系统 (system) 和内存体系 (memory hierarchy) 这一中层抽象。标量存在于执行芯片 (chip) 和寄存器 (register) 这一底层抽象。

一个 dialect 可以自由地定义各种类型。 MLIR 的核心基础设施会无差别地对待以及用统一的机制支持来自不同 dialect 的类型。 比如，type conversion [3]就是通用的转换类型的机制。 Dialect A 可以重用来自 dialect B 的类型，也可以对其进一步扩展和组合，例如将基础类型 (primitive type) 放入容器类型中 (container type)。 一个 dialect 也可以定义规则来实现自身类型和其他 dialect 类型的相互转换。 把这些规则加入到 type converter 中后，所有的规则会相互组合，由此 type conversion 机制会自行找出转换通路来实现转换。 不过总而言之，相较于操作的组合与转换，类型的组合以及转换通常有更多限制也更加复杂，毕竟类型的匹配奠定了操作可以衔接的基础。

不同建模粒度的Dialect（Dialects as modeling granularity）

通过定义和组织操作和类型，dialect 给编译器提供了粗粒度高层次的建模方式。如果两个 dialect 所涉及的类型相同，那么它们基本属于统一抽象层次。另一方面，对涉及不同类型的 dialect 进行转换本质上则是转换不同的抽象层次。

为简化实现，我们一般将高层次 (high-level) 抽象递降 (lower) 到低层次 (low-level) 抽象。递降的过程通常会进行某种形式的问题分解 (decomposition) 或者资源分配 (resource assignment) 来逐渐贴近底层硬件。前者的例子有 tiling, vectorization 等等；后者的例子有 bufferization, 寄存器分配 (register allocation) 等等。即便如此，递降依然不是一个简单的问题，因为不同的抽象层次有不同的目的以及对正确性和性能的理解。比如编程模型层考虑的是代码的表示能力以及简洁性，很少涉及具体硬件特性；而硬件层考虑的是资源的最佳使用，很少考虑易于编程。因此，在诸多 MLIR 机制中，dialect conversion [4]可能是最复杂的就并不奇怪了。

 Dialect体系

 以操作和类型的可组合性以及可扩展性为基础，dialect可以作为组合机器学习编译器的高层次基础组件。之前的讨论偏抽象，接下来我会具体地介绍具体现有的 dialect，并把它们放到统一的流程中。鉴于这里的目的是提供宏观的理解，讨论只会涉及主要的部分，而非对所有 dialect进行详细的分析。

• 在最上层，模型通常是基于某种框架用 Python 编写的。输入程序，或者说编程模型 (programming model)，通常是对高维张量 (high-D tensor) 进行操作。而在最底层，模型的主要计算部分通常是由某种具有向量或者 SIMD 单元的加速器执行的。底层硬件，或者说机器模型 (machine model)，只提供低维（通常是一或者二）向量 (low-D vector) 或者标量的指令。
• 现在有各种各样的框架可以用于编写机器学习模型，同样有许许多多的硬件可以执行它们。硬件可能会提供不同的计算和内存组织结构；在 CPU、GPU 以及各种加速器中，tile-based 架构是一种较常见的一种。整个模型的执行需要运行各种控制流以及同步机制，在这方面 GPU 或者一般的加速器通常都乏善可陈， 所以 CPU 依然处于中心进行调度协调。

真正的端到端的机器学习编译器需要将输入的模型同时转换成运行在加速器上的算子 (kernel) 和运行在CPU上的调度同步逻辑。 MLIR 生态中两部分都有其对应的 dialect 体系。本文的侧重是算子代码生成；调度同步相关的 dialect (例如 MLIR 中的 async dialect 和 IREE 中的 stream dialect) 与传统运行时系统 (runtime) 的功能相关， 值得另起文章来介绍。

 高层用于描述模型的 dialect自顶向下来看，原始模型是用某一框架来表示的。原始模型通过直接转换成这个框架相对应的 dialect（例如 TensorFlow 的 tf dialect, TFLite 的 tfl dialect, PyTorch 的 torch dialect） 来导入 (import) 到 MLIR 系统中。这些对应于具体框架的 dialect 的目的是准确地表示原模型的结构和语义，因为这种紧密的联系，它们通常存在于相应框架的代码库中。面对深度和广度的双重挑战，复杂度可控的编译器栈需要具有沙漏 (hourglass) 的结构。在模型导入之后，我们需要将各种框架表示转换成统一的用于表示模型的 dialect， 以便作为接下来的递降过程的输入。MLIR 在这一层的支持还在迅速演进中，将来希望能够看到一系列（存在于一个或者多个 dialect 中）协调的定义， 用于完整地表示来自于各种框架的各种模型，并且提供所需的兼容性支持。就目前而言，这一层有 mhlo dialect 和 tosa dialect。前者由 XLA [5]而生，是 TensorFlow 系与 MLIR 的桥梁；后者是 TOSA 规范的具体实现。TOSA 规范明确定义很多计算的数值要求，被越来越多的框架转换所采用。中间层用于递降的 dialect高层和低层的 dialect 通常处于 MLIR 系统的边界，所以需要准确地描述某一 MLIR 之外的标的。中间层的 dialect 则没有这样的限制，所以中间层的 dialect 具有更大的设计空间以及更高的设计灵活性。传统的中间表示，如 LLVM IR 或者 SPIR-V，通常都是完整 (complete) 的；它们包含所需的所有指令来表示整个 CPU 后者 GPU 程序。相较而言，中间层的 dialect 则可以认为是部分 (partial) 中间表示。这种组织结构有助于解耦 (decoupling) 和可组合性—我们可以通过混用这一层的不同 dialect 来表示原始模型，同时不同 dialect 可以独立发展演进。这些dialect 有的用来表示计算或者负载 (payload)，有的则表示控制流或者某种结构 (structure)。linalg dialect linalg dialect 是用以表示结构的重要 dialect 之一。linalg op 的本质是完美嵌套循环 (perfect loop nest)。linalg op 通过其 indexing map 来指定循环变量 (loop induction variable) 如何访问 (access) 操作数 (operand) 以及结果 (result)。linalg op region 内的负载操作则指定了循环内部所进行的计算。完美嵌套循环在 linalg op 中是隐性的，这一核心特性简化了很多的分析以及转换。例如，要融合 (fuse) 两个完美嵌套循环，传统上需要分析每个循环变量的范围 (range) 以及它们如何访问元素， 这是比较复杂的分析逻辑，之后的转换也同样比较复杂。用 linalg op 的 indexing map 来隐性表示嵌套循环则可以把上面的过程简化为 inverse(producerIndexMap).compose(consumerIndexMap) 这一步骤。除此之外，linalg dialect 的文档中还有许多其他不错的设计考虑值得一读。linalg dialect 中有很多结构化操作，它们分为两大类：“generic” op 以及 “named” op。前者只包括 linalg.generic op。这个是核心的结构化操作，也是所有结构化 linalg op 的原始未包装形式。Named op (例如 linalg.matmul op 以及各种 linalg.conv* op) 只是对具有特定 indexing map 以及负载操作的 linalg.generic op 进行包装的语法糖。我们可以很简单地把 named op 转换成其对应的 generic 形式。linalg op 背后统一的结构有助于简化转换的逻辑， 因为转换只需要针对 indexing map 以及 region 进行操作，而无需考虑这具体是哪一个 linalg op。linalg op 既可以操作张量也可以操作 buffer。这两者分别对应于 MLIR 中的 tensor and memref 类型。两者皆是高维的抽象，并都可以支持 dynamic shape。Tensor, tiling 以及 fusionmhlo 和 tosa dialect 都可以转换成 linalg dialect。这种转换会保持在张量这一抽象层级，所以其目的并非递降，而是为接下来的转换做准备。mhlo.dot_general op 和 tosa.matmul op 都可以表示 batch matmul，那么 linalg.batch_matmul op 的意义何在呢？因为隐性嵌套循环，tiling 和 fusion 这些对 tile-based 架构非常重要的转换在 linalg.batch_matmul op 上进行更加方便 — 我们只需要创建显式的嵌套循环，把之前的 linalg op 转移到其内并且缩小 linalg op 操作的范围到一个 slice 就可以了。比如下面的 tosa.conv2d op：

 转换成 linalg op 并进行 tiling 和 fusion 之后：

 在嵌套循环之内，我们依然维持着 linalg named op 的形态，以便于进一步的 tiling 和 fusion， 或者进行其他的模式匹配和转换。Buffer, distribution至此我们只讨论了张量。张量是不可变的 (immutable)、不可分割的 (integral) 个体，并且不具有 side effect。张量的 SSA def-use chain 可以用于数据流的分析。针对张量的转换因此比较简单。不过在整个流程的某一步，我们需要把张量转换成 buffer。这一转换在 MLIR 中称为 bufferization[6]。Buffer 是可变的 (mutable)、可 alias 的；buffer 上的转换通常需要比较复杂的依赖分析 (dependency analysis) 和别名分析 (alias analysis)。所以在 MLIR 中趋势是将 bufferization 尽量推后，放在 vectorization 之后。这样 bufferization 可能就会完全变成机械性 (mechanical) 的转换而无需特别复杂的逻辑。因为 vectorization 之后可以通过 read/write forwarding/cancelling 等手段来减少中间的张量结果，从而去掉对 buffer 的需求。Bufferization 是不同抽象层次之间的递降—抽象的值被安排到了具体的内存中的资源。如何进行这种分配既有技术方面也有策略 (policy) 方面的考量。我们即需要避免资源冲突 (hazard) 也要避免无谓的数据复制。这一部分在 MLIR 中还在快速演进中。Bufferization 之后，我们可以进行 distribution 来将 problem tile 分配到 hardware tile (比如 CPU thread, GPU workgroup, GPU workitem 等等)。沿用之前的例子，在 bufferization 以及 distribution 之后，代码会变成：

 tensor, memref, arith, math dialect上面的流程中也涉及了 tensor dialect, memref dialect, arith dialect 和 math dialect。tensor 和 memref dialect 包含了对张量和 buffer 的操作。比如在之前的流程中，tensor.*slice 和 memref.subview op 辅助表示 tiled IR structure。tensor 和 memref dialect 中还包含其他的用以产生 tensor/memref、对 shape 进行操作、 或者一些难以表示成结构性操作或者负载操作的操作。arith and math op 则用以表示整数或者浮点数计算。它们通常用以作为负载操作，并且可以支持各种抽象层级（包括张量、向量、标量等等）。所以在之前的流程图中，它们基本出现在每一层。vector dialectvector dialect 是除 linalg dialect 之外的另一个重要的用于结构化代码生成 (structured code generation) 的 dialect。如果说张量处于抽象的编程模型这一层级，buffer 属于具体的系统内存这一层级，那么向量则处于芯片寄存器这一层级。因为更加贴近硬件实现，向量有着更多的限制。在一个模型中，我们通常可以使用无限多个张量。Bufferization 是一次资源分配；它把张量分配到了内存中的 buffer。在这一过程中我们可以重用同一 buffer 来减小内存使用量或者避免不必要的数据复制。一般而言，内存比较灵活。比如，内存支持动态索引 (dynamic indexing)，并且容量 (capacity) 很大。向量则不同。一般向量要求静态索引 (static indexing)，并且通常数量稀少。如何有效利用寄存器和 vector/SIMD/SIMT 计算单元来高效存取和处理向量是另一次具有更多限制和折中的资源分配。在 MLIR 中，vector dialect 本身即是多层次的。除了支持机器原生 (machine-native) 的向量操作之外，它也支持高维的机器无关的 (machine-agnostic) 虚拟向量。核心想法是通过渐进递降 (progressive lowering) 来将高维的机器无关的虚拟向量分解成低维的机器原生的向量。

 Vectorization, unrolling, hoisting, canonicalization

linalg dialect 可以 tiling 来创建大小已知的 tile。之后我们可以通过 vectorization 来将每个 tensor/buffer tile 转换成同一 shape 的 vector tile。Vectorization 会生成 vector.transfer_read op 来读取 tensor/buffer 中的内容到虚拟 vector，然后生成 vector (例如 vector.contract op) 或者 arith op 来对虚拟 vector 进行计算， 最后会生成 vector.transfer_write op 来将结果写回 tensor/buffer。Vector transfer op 可以表示各种模式的内存读取与写入，也包括对 stride 和 padding 的支持。

这里所说的 vectorization 并非传统的 vectorization。传统的 vectorization 需要将抽象层级从标量递升 (raise) 到向量。 在 MLIR 中，vectorization 也同样转换抽象层级，不过是递降 (lower)，并且基本是机械性的 (mechanical)， 因为在 vectorization 过程中我们并不改变 shape。

Vectorization 之后，我们可以进一步使用 unrolling 和 decomposition 来将高维的向量递降到符合目标架构的低维向量。机器无关的 vector op 可以进一步转换成机器原生的，比如由 vector.contract op 到 vector.fma op。

Unrolling, decomposition 以及 lowering 之后，hoisting 以及各种 canonicalization 有助于进一步清理中间表示。 这其中各种 forward/cancel read/write 以及 insert/extract op 对的转换尤其重要。

在这些转换之后，之前的例子会变成：

 vector dialect 使用了 dialect 内部的转换来进行渐进式递降。因此，pattern 一般比较简单和机械化，组合起来却能真正显示它们的威力。不过要想恰当地排列并且使用它们却是比较难的。在这一方面 vector dialect 需要进一步改善。

 scf, cf dialect

linalg dialect 之后的各层使用 scf dialect 的架构化操作。scf dialect 包含有 structured control flow op, 比如用于分支的 scf.if op 以及用于循环的 the scf.for op。这些操作显示地表示循环变量的范围。不过它们依然使用 region 来标明边界，这有助于简化分析和转换。一旦我们生成了最终形态的控制流，trip-one distributed loop nest 可以被彻底消除， 剩下的循环可以进一步转换成基础块以及 cf op。

到此为止，我们已经接近整个转换流程的终点。下一步是进行整体 dialect conversion 来导出 (export) 到其他的系统。

底层用于描述目标的 dialect

在 MLIR 中目前有两个底层 dialect：llvm dialect 和 spv dialect。它们分别用来对 LLVM IR 和 SPIR-V 建模。转换成任何一个都是对导出到外部系统的准备。因为这两个 dialect 描述外部中间表示，它们在类型和指令方面受相应的限制。递降到 llvm 或者 spv dialect 需要进行整体的 dialect conversion；完成之后 IR 中不再有任何的非 llvm 或者 spv dialect 的操作。

一般而言，上层应该已经完成各种优化，在这个层次不会再有。这个层次的转换多是普适的 canonicalization 和清理，以及一些用以保障合法性的转换。

机器学习编译器面临深度和广度两方面的挑战。MLIR 提供了 dialect 这一更高层次的基础组件来应对这些挑战。理想情况下具体的机器学习编译器只需要组合现有的 dialect，并加以自己的扩展或者定制。相对于组合来自于各种层次的操作而言，这是更粗粒度的一种构建方式。每个 dialect 都包含一套协调的操作和类型，这种构建方式更加易于管理，会带来结构层次更加清晰的软件栈。当然，这一愿景需要一些时间来完全实现！文中的转换流程列出了与代码生成相关的主要 dialect 和转换。

参考资料

[1] “MLIR Dialect ”，https://mlir.llvm.org/docs/Dialects/ 

[2] “MLIR Region”，https://mlir.llvm.org/docs/LangRef/#regions 

[3] “Type Conversion”，https://mlir.llvm.org/docs/DialectConversion/#type-conversion 

[4] “Dialect Conversion”，https://mlir.llvm.org/docs/DialectConversion/

[5] "Tensorflow XLA”，https://www.tensorflow.org/xla

[6] "MLIR Bufferization”，https://mlir.llvm.org/docs/Bufferization/

[7] “Codegen Dialect Overview”，https://discourse.llvm.org/t/codegen-dialect-overview/2723

[8] “Composable and Modular Code Generation in MLIR: A Structured and Retargetable Approach to Tensor Compiler Construction”，https://arxiv.org/abs/2202.03293

参考文献链接

https://mp.weixin.qq.com/s/sgZCbWnCQEO8C8gXv7wV6Q

https://mp.weixin.qq.com/s/jdCdQaCoAIrdmvxPVlK-Ug

本栏目推荐文章
• 成绩管理系统软件体系结构设计说明书
• 01.接口测试价值与体系
• ShowMeBug X AfterShip | 构建高效招聘体系，助推企业人才战略发展
• 系统架构设计师-第1章计算机组成与体系结构
• 推荐系统之标签体系（转载）
• 书生浦语大模型全链开源体系介绍
• 一个ONNX-MLIR示例工程代码
• 【四大力学杂记】统计力学概要——叩开多体系统之门
• 体系化的进阶学习内容
• 书生·浦语大模型全链路开源体系——学习笔记

体系 Dialect MLIR体系dialect mlir lowering dialect mlir opt 源码dialect mlir dialect negotiated dialect support client mlir onnx-mlir 矩阵 算法lowering dialect 基础mlir 编译器operation dialect手册