Codegen方言介绍

主要介绍CodeGen过程中使用的Dialect（方言），以及对设计演变的一些观察。

介绍

对MLIR基础架构中CodeGen功能的简化概述，特别是LLVM项目代码库中可用的部分(upstream或intree)。虽然偶尔会提到LLVM项目代码库之外的MLIR用户，但没有被深入分析，只是为了说明而引用。本文的目的是作为进入MLIR CodeGen基础架构的切入点，并讨论它在未来如何发展。

现状

分类

MLIR中与多个CodeGen相关的方言大致可以沿着两个维度组织: 张量/缓存和有效载荷/结构。

一种方言在张量/缓存维度上的位置，表明了数据的抽象是深度学习框架中的张量，还是传统的底层编译器所期望的内存数据缓冲区。张量被视为不可变的值，不一定与内存有关，也就是说，对张量的操作通常不会有副作用。这些操作之间的数据流可以用传统静态单一赋值(SSA)形式的使用定义链表示。这是使MLIR成为机器学习程序强大的转换工具的一个方面，它允许对张量操作进行简单的重写。另一方面，缓存是可变的，可能会被多个对象指定，例如，多个对象可能指向同一个底层内存。数据流只能通过额外的依赖关系和混叠分析（aliasing analyses）来提取。张量的抽象和缓存的抽象之间的转换是通过缓冲过程来完成的，缓冲过程逐步将张量与缓存关联起来，并最终替换它们。一些方言，如线性代数和标准，包含对张量和缓存的操作。一些线性代数操作甚至可以同时对两者进行操作。

一种方言在有效载荷/结构维度上的位置，表明了它是描述应该执行什么计算(payload)还是应该如何执行(structure)。例如，大多数标准方言中的数学运算指定了要执行的计算，例如arctangent，而没有进一步的细节。另一方面，SCF 方言定义了所包含的计算是如何执行的，例如，重复执行直到满足某个运行时条件，而不限制条件是什么和执行什么计算。类似地，异步方言表示适用于不同负载粒度级别的通用执行模型。

维度上的区分没有明确的界限，特别是在更高的抽象层次上。许多操作至少部分地指定了结构。例如，向量方言操作意味着SIMD执行模型。在编译过程中，如何执行的说明会变得更加详细和低级。同时，抽象堆栈的较低层次倾向于将结构（structure）操作与负载（payload）操作分离开来，以便只转换前者，同时只对负载保持抽象理解，例如，访问的数据或估计的成本。

兴趣方言

MLIR CodeGen 生成的管道要经过一系列中间步骤，这些步骤的特点是使用最近引入的方言。方言可以根据其特性的抽象级别粗略地组织成一个堆栈。将表示从更高层次的抽象转换为更低层次的抽象，即lowing，通常是直接的，而相反的过程可能不成立。

大多数管道通过线性代数方言进入 in-tree 方言的基础架构，线性代数方言的特点是对结构化数据进行结构化计算的通用表示。这种方言是专门为支持各种转换而设计的，只需进行最少的分析。这种方言中的操作既支持张量操作数，也支持缓冲区操作数，缓冲过程可以在不改变操作本身的情况下进行。此外，线性代数提供了具有特定负载的命名操作，如矩阵乘法和卷积，以及只定义结构的通用操作。这两种形式之间可以进行转换。线性代数方言操作的固有迭代结构允许它们被转换为向量操作，以及围绕向量或标量操作的(仿射方言)循环。

异步方言捕获了一个通用的异步编程模型，可能出现在不同的级别上:在较高的级别上，它用于跨设备和设备内部组织大型计算块，在较低的级别上，它可以包装原语指令序列。

向量方言(注意，向量方言类型属于内置方言，可以在向量方言之外使用)是SIMD(可能还有SIMT)执行模型的中级抽象。它利用MLIR的多维向量类型，通过专用的低层方言针对不同的平台。一项正在进行的工作研究了通过线程的显式表示将向量抽象用于目标GPU设备(SIMT)。

仿射方言是MLIR对多面体编译的一种尝试。它封装了相关编程模型的限制，并定义了相应的操作，即仿射循环、条件等控制流结构和仿射对应的内存操作。它的主要目标是实现多面体变换，如自动并行化、用于局部改进的循环融合和平铺，以及MLIR中的循环向量化。

SCF 方言(结构化控制流方言)包含了在比控制流图(CFG)分支更高的层次上表达的常见控制流概念，例如，(并行)for和while循环以及条件。这种方言用于表示(有时转换)计算的结构，而不影响有效负载。它是仿射和线性代数的一个常见的下译目标，也可以用作从较低层次表示(如C语言)到MLIR代码生成基础架构的入口点。

从SCF方言中可以获得各种编程模型，即GPU/SIMT、异步、OpenMP和OpenACC。每个模型都由相应的方言表示，其中的操作很少受到进一步优化转换的影响。然而，这些表示是实现特定于编程模型的转换的机会，例如，目前正在探索的异步方言的转换。

SCF也可以转换为标准CFG表示，用块之间的分支替换结构化控制流。分支操作目前包含在标准方言中，以及各种抽象级别上的许多其他操作。例如，标准方言还包含对张量和向量的点操作，缓冲区和张量之间的转换，标量的三角运算，等等。因此，标准方言正在被分裂成多个明确定义的方言。

最终，标准方言的部分(标量和向量的操作，以及分支)被转换为特定于目标的方言，这些方言主要用作MLIR代码生成基础架构的出口点。这些包括LLVM、NVVM、ROCDL、AVX 、Neon、SVE和SPIR-V 方言，所有这些方言都对应于外部格式、IR或指令集。除了规范化之外，这些方言不需要进行转换。

最后，Shape 方言用于描述独立于负载或(主要的)结构的数据形状。它出现在代码生成管道的入口层，通常被下译到地址算术或规范化。

PDL (模式描述语言)和PDLInterp 方言被用作下一代MLIR的特定转换模式与重写的基础架构。因此，它们从不出现在CodeGen管道中，但在描述其操作时可能是必要的。

管道一些现有管道

张量流内核生成器

张量流内核生成器项目，从张量flow (TF) 方言开始，最近已经转向从MHLO (Meta HLO，由于像隐式广播移除等特性，更适合编译，并支持动态形状，其中HLO是高级优化器表示，源自XLA)而不是LMHLO(后期MHLO，与MHLO相同，但在缓冲区而不是张量上)，并在线性代数上调用缓冲之前，在该级别上执行融合。进一步的循环转换(如tiling)发生在SCF级别，然后转换为特定于目标的GPU方言，而有效负载操作则转换为以标准为中介的LLVM方言。现在已经退役的原型已经尝试使用LMHLO方言针对线性代数-on-buffer，并在SCF上执行所有转换，其中SCF可能比张量抽象更复杂。

在生成多个内核时，与张量flow相关的流预计将使用异步方言来编排计算。

IREE编译器（LLVM目标）

IREE(中级表示执行环境)有它自己的高级表示，它有一组方言，从代码生成的目的来说，这些方言目前正在向面向张量上的语言进化。特定于iree的方言主要用于组织计算有效载荷，可以(目前)表示为MHLO、TOSA、线性代数 -on-tensor等。大多数转换都发生在线性代数中，要么是张量级，要么是缓冲级。可执行文件的首选路径通过向量方言，在这里可以进行额外的转换。当从线性代数下译时，SCF可用于围绕向量操作的控制流，但对这些操作不执行任何转换。针对SCF本质上意味着不再进行进一步的结构优化。向量方言可以逐步降低为它所包含的不那么复杂的抽象，直到最终针对LLVM方言。

IREE编译器 (SPIR-V目标)

SPIR-V (标准可移植中间表示，Khronos组标准)是IREE编译器的主要目标。顶层流程类似于针对LLVM IR的流程，大多数转换发生在张量和向量层的线性代数上。从这里开始，较低的转换倾向于直接转到具有丰富操作集的SPIR-V，该操作集跨越多个抽象级别:高级操作、结构化控制流和类指令的原语。该流程通过GPU方言进行设备操作，如工作项标识符提取，并依赖IREE的运行时来管理GPU内核。

最近的工作允许IREE从向量方言转换为GPU方言，将GPU线程暴露为向量通道(在warp或block级别)。类似地，有些转换可以直接从线性代数和向量方言转换到SPIR-V，绕过中间阶段，但可能会被更渐进的下译方法取代。

多面体编译器

从HLO开始并绕过线性代数的多面体编译流可以通过转换来自LMHLO的仿射方言或任何其他缓冲形式的操作的方式来实现。在这种情况下，大多数转换发生在仿射方言上，这是目前由多面体转换支持的主要抽象。然后代码被下译到SCF控制流和标准内存操作，然后被进一步转换为平台特定的抽象，如OpenMP或GPU。

多面体模型还支持早期的矢量化，从仿射控制流构造到向量方言，而不是后来发现循环矢量化的机会。

对于多面体编译器来说，能够接受以较低级抽象(如C编程语言)表示的代码一直是非常重要的。目前可以通过逐步增加来实现。

分析

Crossing方言

事后看来，似乎在这个分类中跨维度边界的方言(GPU、线性代数和向量方言)在被接受为核心生态系统的一部分之前需要进行最多的讨论和迭代。即使是现在，MLIR基础架构的用户报告说，理解一些方言的定位仍然具有挑战性。例如，IREE使用与设备上执行相关的部分GPU方言，而不是与管理数据和来自主机的内核相关的部分，它们与结构的关系更密切，而不是与负载的关系。类似地，关于将张量和memref抽象与相应操作连接起来的讨论需要很大的努力才能收敛。

这表明，如果将新的方言或较小的IR概念与其他方言和设计空间明确定位，那么它们可以更容易地进行讨论并达成共识。当有必要跨越抽象之间的鸿沟时，最好单独讨论它，并以在不同方言之间泛化为目标(例如，缓冲过程)。

中心线性代数

线性代数方言是MLIR代码生成管道的主要入口点之一。它最近的发展使它同时作用于张量和缓冲区，使缓冲成为方言内部的转换。它有足够的关于操作的高级信息来执行转换，而不需要昂贵的分析，特别是在将张量作为值进行操作时。一些转换，如融合元素操作和tiling，并结合它们来生成不完全嵌套的计算，可以捕获针对广泛架构所需的足够数量的转换。此外，命名操作的概念支持构建在线性代数计算模式上的有效负载操作。

当生产用户开始依赖线性代数时，在几乎所有的编译管道中使用它会增加线性代数的维护压力和稳定性要求。事实上，它跨缓冲区和张量工作，并且可能同时捕获有效载荷和计算结构，这使得理解线性代数在某种程度上成为理解任何MLIR代码生成管道的必要条件。虽然有一个良好定义和维护的入口点是有好处的，但必须特别注意确保线性代数仍然可以与其他方言组合，并且转换算法不会为它过度设计。

管道互补转换的差异

生态系统中出现了几个并行编译管道。以GPU为目标的案例尤其具有说明意义:优化转换、并行检测和设备映射决策可以在各种不同的方言中发生(线性代数、仿射、SCF)，这可能最终会部分地重新实现彼此的功能;设备映射可以使用SCF或向量方言从显式循环或向量派生SIMT线程。最后，GPU库支持更高级的操作，如收缩和卷积，这些操作可以从代码生成管道的入口点直接针对它们。

现在比以往任何时候都更重要的是确保表示和转换组合并相互补充，以实现MLIR统一编译基础架构的承诺，而不是构建独立的并行流。这并不一定意味着立即重用所有组件，但要尽可能避免可组合性差的模式。特定于领域和目标的编译器的效用是不可否认的，但是项目可能需要投资于横切表示，使用属性和接口的通用机制。以GPU为例，设备映射策略表示为可以附加到不同操作(例如线性代数泛型或并行SCF循环)的属性，可以通过接口进行转换，而不需要知道特定操作的细节。

构建小型可重用摘要

可以观察到在代码生成管道的更高级别上执行大多数转换的一个不足为奇的趋势:在这些级别上很容易获得或容易提取必要的有效性信息，而不需要进行复杂的分析。然而，更高层次的抽象通常对可表示的内容有更严格的限制。在追求这类抽象的好处时，重要的是要记住表现力和至少在较低级别执行一些转换的能力，以快速增加表现力，而无需重新实现顶级抽象(HLO中的动态形状就是一个很好的例子)。

需要更多结构

另一个正在出现的趋势是更大量的方言和逐步下译执行不离开方言的边界。一些例子包括对张量进行线性代数运算(转换为使用缓冲区)，GPU方言块缩减(分解为shuffle)，以及标准中的数学运算(可以扩展为更小的运算或使用其他标准运算近似)。正在到达这样一个点:一种方言包含几个松散连接的操作子集。这导致了将标准方言拆分为多个组成部分的提议。然而，分裂可能并不总是可取的，甚至是可行的:相同的线性代数操作接受张量和缓冲区，所以在没有大量重复的情况下分离T线性代数和B线性代数是具有挑战性的。这就需要以编程上可以理解的方式在方言中构造操作。

主机/设备或程序/内核作为附加轴

最后，在更大的、主要与ml相关的流程中，MLIR的使用促使与计算的整体组织相关的方面(例如，在分布式系统上反映模型映射的操作，以及与嵌入MLIR的框架的交互，通常在主机上)和单个计算的组织(例如，与内核或另一个大型计算单元的内部对应的操作，可能卸载到一个设备)。对于宿主部分，代码生成可能也是必要的，并且通常需要不同于内核部分的转换。

这种分离的具体例子包括GPU方言，它包含从主机控制执行的操作和在GPU上执行的操作，这可能会受益于分离成独立的方言，以及异步方言，它在两个层次上使用:组织独立的内核的执行和通过针对LLVM协程在内核内并行执行。

参考文献链接

https://discourse.llvm.org/t/codegen-dialect-overview/2723