欢迎来到尧图网

客户服务 关于我们

您的位置:首页 > 科技 > IT业 > 全文 - MLIR: A Compiler Infrastructure for the End of Moore’s Law

全文 - MLIR: A Compiler Infrastructure for the End of Moore’s Law

2025/3/22 10:09:57 来源:https://blog.csdn.net/eloudy/article/details/146422045  浏览:    关键词:全文 - MLIR: A Compiler Infrastructure for the End of Moore’s Law

MLIR: 为摩尔定律终结而设计的编译器基础设施

摘要

        本文介绍了MLIR,一种构建可重用和可扩展编译器基础设施的新方法。MLIR旨在解决软件碎片化问题,改善异构硬件的编译,显著降低构建领域特定编译器的成本,并帮助连接现有编译器。MLIR促进了在不同抽象层次以及跨应用领域、硬件目标和执行环境的代码生成器、翻译器和优化器的设计和实现。本文的贡献包括:

(1) 讨论MLIR作为一个研究工件,旨在扩展和演化,并识别其在设计、语义、优化规范、系统和工程方面的挑战和机遇。

(2) 评估MLIR作为一个通用基础设施,如何降低构建编译器的成本——描述多样的用例,以展示未来编程语言、编译器、执行环境和计算机架构的研究和教育机会。本文还介绍了MLIR的基本原理、最初的设计原则、结构和语义。

1 引言

编译器设计是一个成熟的领域,拥有广泛的知名算法,应用于代码生成、静态分析、程序转换等。该领域还开发了许多成熟的技术平台,这些平台在编译器社区中实现了大规模的重用,包括LLVM编译器基础设施、Java虚拟机(JVM)等。这些流行系统的一个共同特征是其“通用”方法——一个单一的抽象层次与系统接口:LLVM中间表示(IR)大致是“带有向量的C”,而JVM提供了“带有垃圾收集器的面向对象类型系统”抽象。这种“通用”方法非常有价值——实际上,从普遍存在的源语言(分别是C/C++和Java)到这些领域的映射是直接的。


同时,许多问题在更高或更低的抽象层次上建模更好,例如,在LLVM IR上进行C++代码的源级分析非常困难。我们观察到,许多语言(包括Swift、Rust、Julia、Fortran等)开发了自己的IR,以解决领域特定的问题,如语言/库特定的优化、流敏感类型检查(例如线性类型),以及改进下降过程的实现。同样,机器学习系统通常使用 “ML 图” 作为领域特定的抽象。

虽然领域特定IR的开发是一门研究透彻的艺术,但其工程和实现成本仍然很高。对于这些系统的实现者来说,基础设施的质量并不总是首要任务(或容易证明其合理性)。因此,这可能导致质量较低的编译器系统,包括用户可见的问题,如编译时间慢、实现有缺陷、诊断质量不佳、优化代码的调试体验差等。

        MLIR项目旨在直接解决这些编程语言设计和实现挑战——通过定义和引入新的抽象层次,使成本变得非常低,并提供 “开箱即用” 的基础设施来解决常见的编译器工程问题。MLIR 通过

(1) 标准化基于静态单赋值(SSA)的IR数据结构,

(2) 提供一个声明性系统来定义 IR 方言,

(3) 提供广泛的通用基础设施(包括文档、解析和打印逻辑、位置跟踪、多线程编译支持、Pass管理等)

来实现这一目标。

本文探讨了MLIR系统的各种设计点,介绍了我们将其应用于多个不同问题的经验,并讨论了这项工作可能对语言设计和教育产生的影响。

1.1 贡献

虽然 MLIR 系统的大部分是由众所周知的编译器算法构建的,但其设计点足够新颖,提供了有趣的研究机会。本文的贡献包括:

描述了一种具有重要工业和研究应用的新型编译器基础设施;
构建可扩展和模块化编译器系统的新方法;
探索 MLIR 在不同领域的选定应用,展示系统的通用性;
分享了基于 MLIR 基础设施构建系统的经验。

1.2 MLIR的起源

        MLIR 的工作始于认识到现代机器学习框架由许多不同的编译器、图技术和运行时系统组成(见图1)——它们没有共享一个共同的基础设施或设计点,并且并非所有这些都遵循编译器设计的最佳实践。这在多个用户可见的方面表现出来,包括错误信息不佳、边缘情况失败、性能不可预测,以及难以将堆栈泛化以支持新硬件。

        我们很快意识到,整个编译器行业也面临类似的问题:现有系统如 LLVM 在统一和集成不同语言实现方面非常成功,但现代高级语言往往最终构建自己的高级 IR,并为更高层次的抽象重新发明许多相同类型的技术(见图2)。同时,LLVM 社区经常面临如何最好地表示并行构造、如何共享常见前端下降基础设施(例如C调用约定,或跨语言特性如OpenMP)的实现的问题,而没有可用的令人满意的解决方案。

        面对这一挑战和视角,我们感觉无法承担实现 N 个改进的编译器实例的工程努力,我们需要构建一个更通用的解决方案。我们推断,这将使我们能够投资于一套高质量的基础设施,这将使多个领域受益,使我们能够逐步升级现有系统,使我们更容易解决异构编译等紧迫问题,并为定义和探索提供有趣的研究机会。
         现在,我们在构建和部署基于MLIR的系统方面积累了大量经验,我们能够回顾基础设施的基本原理和设计,并讨论为什么选择了这个方向。

2 设计原则

现在让我们探讨指导 MLIR 设计的需求。

极少内置,完全可定制

该系统基于少量的基本概念,允许大部分中间表示(IR)完全可定制。少数抽象——类型、操作和属性,这些是 IR 中最常见的——应该用于表达其他所有内容,从而允许更少且更一致的抽象,这些抽象易于理解、扩展和采用。总体而言,可定制性确保系统能够适应不断变化的需求,并更有可能适用于未来的问题。在这个意义上,我们应该将 IR 构建为一个丰富的基础设施,具有可重用的组件和编程抽象,支持其中间语言的语法和语义。
定制化的成功标准是能够表达多样化的抽象集,包括机器学习图、抽象语法树(AST)、数学抽象(如多面体)、控制流图(CFG)和指令级 IR(如 LLVM IR),而无需将这些抽象的概念硬编码到系统中。
当然,可定制性会因不兼容的抽象而带来内部碎片化的风险。虽然生态系统碎片化问题不太可能有纯技术解决方案,但系统应鼓励设计可重用的抽象,并假设它们将在其初始范围之外使用。

SSA 和区域

        静态单赋值(SSA)形式是编译器 IR 中广泛使用的表示。它提供了许多优势,包括使数据流分析简单且稀疏,因其与续延传递风格的关系而被编译器社区广泛理解,并在主要框架中确立。虽然许多现有 IR 使用平坦的线性化 CFG,但表示更高级别的抽象推动引入嵌套区域作为 IR 中的一级概念。这超越了传统的区域形成,以提升更高级别的抽象(例如,循环树),加速编译过程或提取指令或 SIMD 并行性。为了支持异构编译,系统必须支持表达结构化控制流、并发构造、源语言中的闭包等多种用途。一个具体的挑战是使基于 CFG 的分析和转换能够在嵌套区域上组合。


        在这样做时,我们旨在牺牲 LLVM 的规范化,有时是规范化属性。能够将各种数据和控制结构降低为较小的规范化表示集合是保持编译器复杂性可控的关键。具有前导、头部、锁存、主体的规范循环结构是前端语言中各种循环构造的线性化控制流表示的典型案例。我们旨在为用户提供选择:根据感兴趣的编译算法,在编译流中的 pass,嵌套循环可以被捕获为嵌套区域,或作为线性化控制流。通过提供这样的选择,我们偏离了仅规范化的 LLVM 方向,同时保留了在需要时处理更高级别抽象的能力。反过来,利用这样的选择引发了关于如何控制抽象规范化的问题,这是下一段的目的。

渐进式降低

        系统应支持渐进式降低,即从更高级别的表示逐步降低到最低级别,降低过程在多个抽象层次上进行。多层次抽象的需求源于通用编译器基础设施必须支持的平台和编程模型的多样性。

         以前的编译器在其管道中引入了多个固定的抽象层次——例如,Open64 WHIRL 表示有五个层次,Clang 编译器从 AST 降低到 LLVM IR,再到 SelectionDAG,再到 MachineInstr,最后到 MCInst。这些方法以刚性的方式完成,需要更灵活的设计来支持可扩展性。
        这对转换的阶段顺序有深远的影响。随着编译器专家开始实现越来越多的转换 pass,这些 pass 之间的复杂交互开始出现。早期就显示出结合优化 pass 使编译器能够发现更多关于程序的事实。结合 pass 的好处之一的最早例证是混合常量传播、值编号和不可达代码消除。更一般地说,编译器 pass 可以大致分为四个角色:(1) 优化转换,(2) 启用转换,(3) 降低和 (4) 清理。系统应允许在单个操作的粒度上混合和匹配这些角色,而不是在整个编译单元上顺序执行 pass。

保持高级语义

        系统需要保留计算的高级语义和结构,这对于分析或优化性能是必需的。尝试在降低后提升语义是脆弱的,并且将此信息强行插入低级别通常是侵入性的(例如,在使用调试信息记录结构的情况下,所有 pass 需要验证/重新审视)。相反,系统应保持计算结构并逐步降低到硬件抽象。结构的丧失是有意识的,并且仅在不再需要结构以匹配底层执行模型时发生。例如,系统应在相关转换过程中保留结构化控制流(如循环结构);移除此结构,即转向基于 CFG 的控制流,实质上意味着在此级别上不会再执行进一步的转换。当前在生产编译器中建模并行计算构造的最新技术突显了这项任务在一般情况下可能有多困难。
        作为推论,在同一 IR 中混合不同层次的抽象和不同的概念是系统的一个关键属性,允许表示的一部分保持在更高级别的抽象中,而另一部分被降低。这将使得,例如,一个自定义加速器的编译器能够重用系统定义的一些高级结构和抽象,同时与加速器特定的原始标量/向量指令一起使用。

IR 验证

        生态系统的开放性要求广泛的验证机制。虽然验证和测试对于检测编译器错误和捕获 IR 不变量是有用的,但在可扩展系统中对强大验证方法和验证工具的需求被放大。该机制应旨在使其易于定义并尽实际可能地声明式的,提供单一的真实来源。
         一个长期目标是重现翻译验证的成功和编译器测试上现代方法的成功。在可扩展编译器生态系统的背景下,这两个问题目前都是开放问题。

声明性重写模式

         定义表示修改器应该像新抽象一样简单;编译器基础设施的好坏取决于它支持的转换。常见的转换应该可以实现为以声明性方式表达的重写规则,以机器可分析的格式推理重写的属性,如复杂性和完成性。重写系统因其健全性和效率而被广泛研究,并应用于从类型系统到指令选择的众多编译问题。由于我们旨在实现前所未有的可扩展性和增量降低能力,这为将程序转换建模为重写系统开辟了许多途径。这也引发了关于如何表示重写规则和策略,以及如何构建能够通过多个抽象层次引导重写策略的机器描述的有趣问题。系统需要在保持可扩展性的同时解决这些问题,并强制执行健全、单调和可重复的行为。

源位置跟踪和可追溯性

        操作的来源——包括其原始位置和应用的转换——应在系统中易于追溯。这旨在解决复杂编译系统中常见的透明度缺失问题,在这些系统中,几乎不可能理解最终表示是如何从原始表示构建的。这在编译安全关键和敏感应用时尤其成问题,在这些应用中,跟踪降低和优化步骤是软件认证程序的一个基本组成部分。当操作安全代码(如加密协议或处理隐私敏感数据的算法)时,编译器经常面临看似冗余或繁琐的计算,这些计算嵌入了源程序功能语义未完全捕获的安全或隐私属性:这段代码可能会防止侧信道的暴露或增强代码以抵御网络或故障攻击。优化可能会改变或完全无效化这些保护;这种透明度缺失在安全编译中被称为 WYSINWYX。准确地将高级信息传播到低级别的一个间接目标是帮助支持安全和可追溯的编译。

3 IR 设计细节

本节描述了 MLIR 中 IR 的设计,遵循前一节的原则。

操作

        在 MLIR 中,语义的基本单位是“操作”,简称为 Op。从“指令”到“函数”再到“模块”都在这个系统中被建模为 Ops。MLIR 没有固定的 Ops 集合,而是允许(并鼓励)用户定义扩展——编译器 pass 保守地处理未知的 Ops,MLIR 通过特性、特权操作钩子和优化接口(如第 6.1 节所述)为 pass 描述 Ops 的语义提供了丰富的支持。

        Ops(见图 3)具有唯一的操作码,文本上是一个标识其方言和操作的字符串。Ops 接受和产生零个或多个值,分别称为操作数和结果,这些值以 SSA 形式维护。所有值都有一个类型,类似于 LLVM IR。除了操作码、操作数和结果外,Ops 还可能具有属性、区域、块参数和位置信息。图 4 展示了值和 Ops,%-标识符是命名值的(包),如果超过一个则用“:”指定包中的数量,用“#”指定特定值。在通用文本表示中,操作名称是带引号的字符串字面量,后跟括号中的操作数。

属性

         MLIR 属性是结构化的编译时静态信息,例如整数常量值、字符串数据或常量浮点值列表。属性是有类型的,每个 Op 实例都有一个从字符串名称到属性值的开放键值字典。在通用语法中,属性位于 Op 操作数和其类型之间,作为大括号括起来的逗号分隔的键值对列表。例如,图 4 使用属性定义了已知为常量仿射形式的循环界限:{lower_bound = () -> (0), step = 1 : index, upper_bound = #map3},其中 lower_bound、upper_bound 和 step 是属性名称。() -> (0) 表示用于内联仿射形式,在这种情况下生成一个仿射函数,产生一个常量 0 值。#map3 表示用于属性别名,允许将属性值与标签预先关联,并在需要属性值的任何地方使用该标签。
        与操作码一样,没有固定的属性集合。属性的意义要么来自 Op 语义,要么来自与之关联的方言(第 3 节)。属性也是可扩展的,允许直接引用外部数据结构,这对于与现有系统集成非常有用。例如,一个属性可以引用 ML 系统中(在编译时已知的)数据存储的内容。

位置信息

         MLIR 提供了紧凑的位置信息表示,并鼓励在整个系统中处理和传播此信息。它可以用于保留生成 Op 的源程序堆栈跟踪,以生成调试信息。它标准化了从编译器发出诊断信息的方式,并被广泛的测试工具使用。
        位置信息也是可扩展的,允许编译器引用现有的位置信息跟踪系统、高级 AST 节点、LLVM 风格的文件-行-列地址、DWARF 调试信息或其他高质量实现所需的内容。

区域和块

        Op 的一个实例可能有一个附加区域列表。区域为 MLIR 中的嵌套结构提供了机制:一个区域包含一个块列表,一个块包含一个操作列表(可能包含区域)。与属性一样,区域的语义由其附加的操作定义,然而区域内的块(如果多于一个)形成一个控制流图(CFG)。例如,图 4 中的 affine.for 操作是一个循环,单块体附加为一个区域,位于 ({ 和 }) 分隔符之间。Op 指定了跨区域的控制流。在此示例中,主体重复执行,直到达到上限。


        每个区域的主体是一个块列表,每个块以一个终止操作结束,该操作可能有后继块,控制流可以转移到这些块。每个终止符(例如“switch”、“条件分支”或“unwind”)定义其自己的语义。它可以选择将控制流转移到同一区域的另一个块,或将其返回到包含该区域的 Op。后继者的图定义了一个 CFG,允许在区域内进行标准的基于 SSA 的控制流。


        MLIR 使用 SSA 的函数形式而不是 φ 节点,其中终止符将值传递给由后继块定义的块参数。每个块都有一个(可能为空的)类型化块参数列表,这些参数是常规值并遵循 SSA。终止符 Ops 的语义定义了控制转移后块的参数将采用的值。对于区域的第一个(入口)块,值由包含 Op 的语义定义。例如,affine.for 使用入口块参数 %arg4 作为循环归纳变量。

值的支配和可见性

        Ops 只能使用在作用域内的值,即根据 SSA 支配、嵌套和包含操作施加的语义限制可见的值。如果它们遵循标准的 SSA 支配关系,控制保证在到达使用之前通过定义,则值在 CFG 中是可见的。

        基于区域的可见性是基于区域的简单嵌套定义的:如果 Op 的操作数在当前区域之外,则它必须在使用区域的上方和外部定义。这允许 affine.for 操作中的 Ops 使用在外部作用域中定义的值。

        MLIR 还允许将操作定义为与上方隔离,指示该操作是一个作用域屏障——例如,“std.func” Op 定义了一个函数,函数内的操作不能引用在函数外定义的值。除了提供有用的语义检查外,包含与上方隔离的 Ops 的模块可以由 MLIR 编译器并行处理,因为没有使用-定义链可以跨越隔离屏障。这对于利用多核机器进行编译非常重要。

符号和符号表
        Ops 可以附加一个符号表。该表是将名称(表示为字符串)与 IR 对象(称为符号)关联的标准化方式。IR 不规定符号的用途,留给 Op 定义。符号对于不需要遵循 SSA 的命名实体最有用:它们不能在同一表中重新定义,但可以在定义之前使用。例如,全局变量、函数或命名模块可以表示为符号。没有这种机制,就不可能定义,例如,递归函数在其定义中引用自身。如果附加了符号表的 Op 具有包含类似 Ops 的关联区域,则符号表可以嵌套。MLIR 提供了一种从 Op 引用符号的机制,包括嵌套符号。

方言

        MLIR 使用方言管理可扩展性,方言在唯一命名空间下提供 Ops、属性和类型的逻辑分组。方言本身不引入任何新语义,而是作为逻辑分组机制,并可用于提供方言通用的 Op 支持(例如,方言中所有 ops 的常量折叠行为)。方言命名空间在操作码中显示为点分隔的前缀,例如,图 4 使用仿射和 std 方言。
        将 Ops、类型和属性分为方言是概念性的,类似于设计一组模块化库。例如,一个方言可以包含用于操作硬件向量的 Ops 和类型(例如,shuffle、插入/提取元素、掩码),另一个方言可以包含用于操作代数向量的 Ops 和类型(例如,绝对值、点积等)。这两个方言是否使用相同的向量类型以及该类型属于何处是留给 MLIR 用户的设计决策。
        虽然可以将所有 Ops、类型和属性放在一个方言中,但由于同时存在的大量概念和名称冲突等问题,这将很快变得难以管理。尽管每个 Op、类型和属性仅属于一个方言,MLIR 明确支持方言的混合以实现渐进式降低。不同方言的 Ops 可以在 IR 的任何级别共存,它们可以使用在不同方言中定义的类型等。方言的混合允许更大的重用、可扩展性,并提供灵活性,否则这将需要开发人员诉诸于各种不可组合的变通方法。

类型系统

        MLIR 中的每个值都有一个类型,该类型在生成值的 Op 中或在将值定义为参数的块中指定。类型为 IR 提供编译时语义。MLIR 中的类型系统是用户可扩展的,并且可以引用现有的外部类型系统(例如,llvm::Type 或 clang::Type)。MLIR 强制执行严格的类型相等性检查,不提供类型转换规则。Ops 使用尾随函数样式语法列出其输入和结果类型。在图 4 中,std.load 从内存引用和索引类型映射到其加载的值的类型。
        从类型理论的角度来看,MLIR 仅支持非依赖类型,包括简单、参数化、函数、和积类型。虽然可以通过结合 Ops 与符号和用户定义类型以 Curry-Howard 同构的字面解释实现依赖类型系统,但这些类型对 IR 是不透明的。

标准类型

        此外,MLIR 提供了一组标准化的常用类型,包括任意精度整数、标准浮点类型和简单的常见容器——元组、多维向量和张量。这些类型仅仅是对方言作者有用的便利,但不要求使用它们。

函数和模块

        与传统 IR 类似,MLIR 通常被结构化为函数和模块。然而,这些在 MLIR 中不是新的或独立的概念:它们作为内置方言中的 Ops 实现。
        模块是一个具有单个区域的 Op,该区域包含一个块,并以不转移控制流的虚拟 Op 终止。模块定义一个符号,可以被引用。像任何块一样,其主体包含一个 Ops 列表,这些 Ops 可以是函数、全局变量、编译器元数据或其他顶级构造。
        函数是一个具有单个区域的 Op,其参数对应于函数参数。它定义一个符号,可以通过名称引用。使用函数调用 Op 将控制流转移到函数中。一旦进入,控制流遵循区域中块的 CFG。“返回”终止符没有后继者,而是终止区域执行,将控制流转移回函数的调用点。“返回”终止符 Op 的任何操作数都是函数的返回值。


4 IR 基础设施

        除了 IR 本身,MLIR 还提供了定义 IR 元素(如方言、Ops、模式重写、验证和可重用 passes)的基础设施。MLIR 的基础设施对于在定义新抽象和使用 MLIR 作为优化工具包时提供可扩展性和易用性至关重要。

4.1 操作描述

        MLIR 使用基于 TableGen 的规范来进行操作描述(ODS),以声明的方式定义 Op 的结构及其验证器的组件。TableGen 是一种数据建模工具,旨在帮助定义和维护领域特定信息的记录,在 LLVM 中被广泛使用。我们选择它来建模 Ops 和重写模式,以利用其在业界的接受度。ODS 可以看作是嵌入到 TableGen 输入语言中的 MLIR Op 定义的 DSL,因此 ODS 语法由 TableGen 强加,但 MLIR 特定的语义由 ODS 提供。ODS 定义最终被翻译成 C++ 代码(包括具有命名访问器、验证等的 Op 类),与系统的其余部分互操作。

        Ops 在 ODS 中使用 TableGen Op 类建模。图 5 显示了一个 Op ODS 定义的示例。每个定义的 Op 都有一个名称,这是一个唯一标识符,一个描述 Op 属性的特性列表,一个指定 Op 操作数和属性的参数列表,以及一个 Op 的结果列表。参数和结果具有名称和类型约束(例如,固定形状的浮点或 int32 张量)。Op 定义还可以为文档指定人类可读的 Op 描述。以及为其生成自定义打印机/解析器的(有限)自定义文本形式。当 Op 定义需要比 ODS 提供的更细粒度的控制时,可以通过构建器、打印机、解析器、验证器子句注入额外的 C++ 代码。Op 特性可以是通用的,例如“无副作用”,也可以是方言或 ODS 特定的,例如“具有自定义导出器”。ODS 中的特性可以由定义特性行为的 C++ 类支持。没有固定的特性集,但一些特性为 ODS(例如,“形状结果和操作数类型”表示一个完全捕获给定输入类型的输出类型的约束)或优化器所知(例如,“无副作用”,见第 6.1 节)。

        类型约束检查参数/结果类型的属性,并且是用户/方言可扩展的。MLIR 基础设施还提供了许多预定义的类型约束,例如“任何类型”、“元素满足给定约束的张量”、“给定秩的向量”等。ODS 还有限支持使用特性引发的约束自动推导操作数结果的返回类型,更多信息见第 4.2 节。

4.2 声明性重写

        许多 MLIR 转换涉及 Op 操作,而一些转换需要对 IR 进行复杂修改,许多其他转换可以表示为基于 SSA 使用-定义关系定义的 DAG 上的简单重写。MLIR 提供了一个图重写框架,辅以声明性重写规则(DRR)系统,使得表达模式变得简单。
        与 ODS 类似,DRR 是嵌入到 TableGen 语言中的 DSL。DRR 表达源和目标 DAG 模式以及约束(包括动态约束)和模式优先级的好处。模式可以捕获和重用 Op 的参数。从概念上讲,DRR 表达在特定约束下 DAG 的等价性。图 6 给出了一个 DRR 模式的示例,该模式将图 5 中定义的 Op 转换为由比较和选择组成的常见低级实现。
        DRR 被转换为 C++ 代码,可以与使用通用图重写框架直接在 C++ 中定义的更复杂的模式混合。这种能力允许 MLIR 保持常见情况简单而不限制框架的通用性。

4.3 Pass 管理器

        MLIR pass 管理器组织和处理一系列在各种粒度上操作的 IR passes 的高效执行。现有系统中的 pass 管理通常在固定粒度上定义(例如,模块、函数或循环 pass 管理器),而在 MLIR 中,模块和函数并不特殊——它们只是具有区域的 Ops,并且可以有多个变体。因此,MLIR pass 管理器也不专门针对固定的 ops 集,而是可以在任意嵌套级别上处理任意 Ops。

并行编译

        MLIR 的一个重要要求是需要利用多核机器以加快编译速度。pass 管理器支持中间表示的并发遍历和修改,这通过操作的“与上方隔离”属性提供的不变量成为可能,因为 SSA 使用-定义链不能跨越这些 ops 的区域边界。具有此行为的操作(例如“std.func”操作)因此定义了一个可以并行处理的区域树。
        这个要求是为什么(与例如 LLVM 相比)MLIR 不具有整个模块使用-定义链的原因。全局对象通过符号表条目引用,常量实现为具有相关属性的操作。

4.4 可往返的文本 IR 形式

        MLIR 中的 IR 和 Ops 具有完全反映内存表示的文本表示,这对于调试、理解转换过程中的 IR 以及编写测试用例至关重要。图 4 中显示的原始 IR 形式冗长且难以理解。因此,MLIR 允许为 Ops 定义自定义打印和解析格式,这允许示例如图 8 所示进行打印和解析,这更易于使用。
        这两种形式都是完全可往返的,每个编译器 pass 可以单独测试,使用文本形式作为输入和输出。由于没有隐藏状态,运行单个 pass 的结果与在完整 pass 管道中运行相同的 pass 是相同的。这种方法对用户友好,因为 IR 形式可以手动创建,并且 IR 转换易于跟踪。

4.5 文档

        方言、Ops 和接口的文档是从其 ODS 描述生成的。除了单行摘要和更易读的描述外,生成的文档还包括参数和结果类型约束。由于相同的源代码用于验证代码和文档,因此文档更有可能与运行时行为保持同步。

4.6 验证器

        验证器用于强制执行 IR 的结构正确性和 Ops 的不变量,允许 passes 假设已验证的 IR 具有已检查的不变量,并且还用作调试工具。验证从检查 MLIR 整体的结构属性开始:类型必须完全匹配,值仅定义一次并遵循支配和可见性,符号名称在符号表中是唯一的,所有块以终止符 Ops 结束等。之后,应用单个 Op 和属性验证器。每个 Op 可以定义一组检查其有效性的结构和语义规则。例如,二元 Op 检查它有两个操作数,许多 Ops 只接受特定类型的值,许多需要附加特定的属性或区域。类似地,方言属性只能在特定 Ops 上允许或对其附加的 Ops 施加进一步的限制。例如,方言属性可以要求 Op 仅使用方言中定义的类型,即使 Op 本身更通用。验证失败被视为不变量违反并中止编译。

5 评估:MLIR 的应用

        MLIR 是一个旨在推广和推动广泛编译器项目的系统,因此我们的主要评估指标是展示它被采用并用于各种项目。我们提供了社区活动的总结,并详细描述了一些用例,以突出 MLIR 的通用性和可扩展性,并展示其在实现可定制性设计原则方面的表现。

        如今,MLIR 是一个不断发展的开源项目,社区涵盖了学术界和工业界。例如,关于在高性能计算(HPC)中使用 MLIR 的学术研讨会吸引了来自 16 所大学的个人,并涉及来自 4 个不同国家的 4 个国家实验室。MLIR 还得到了 14 家跨国公司的支持,在 LLVM 开发者会议上,超过 100 名行业开发者参加了关于 MLIR 的圆桌会议。社区的采用和参与是可用性和需求的代理指标。超过 26 种方言正在公共或私人开发中,跨不同公司的 7 个项目正在用 MLIR 替换自定义基础设施。我们认为这表明了对 MLIR 的真实需求,并认可了其可用性。

5.1 TensorFlow 图

        虽然其他讨论的表示对大多数编译器开发者来说很熟悉,但 MLIR 的一个关键用例是支持机器学习框架的开发。它们的内部表示通常基于具有动态执行语义的数据流图。TensorFlow 是此类框架的一个例子。其表示是一种高级数据流计算,其中节点是可以放置在各种设备上的计算,包括特定的硬件加速器。
        MLIR 用于 TensorFlow 中建模这种内部表示,并执行 图 1 中展示的用例的转换:从简单的代数优化到将图重新定位以在数据中心集群的硬件加速器上并行执行,从降低到适合移动部署的表示到使用 XLA 等工具生成高效的本机代码。图 7 展示了在 MLIR 中表示的 TensorFlow 图。

5.2 多面体代码生成

        MLIR 的最初动机之一是探索加速器的多面体代码生成。仿射方言是一种简化的多面体表示,旨在实现渐进式降低。虽然在此处全面探索设计点超出了本文的范围,但我们通过展示仿射方言的方面来展示 MLIR 的建模能力,并将仿射方言与过去的表示进行对比。

5.2.1 相似性

        MLIR 仿射方言在所有内存访问中操作于结构化多维类型。在默认情况下,这些结构化类型是单射的:通过构造保证不同的索引不会别名,这是多面体依赖分析的常见前提条件。
        仿射建模分为两部分。属性用于在编译时建模仿射映射和整数集,Ops 用于将仿射限制应用于代码。即,affine.for Op 是一个“for”循环,其界限表示为函数中需要不变的值的仿射映射。因此,循环具有静态控制流。类似地,affine.if 是受仿射整数集限制的条件。循环和条件的主体是使用 affine.load 和 affine.store 将索引限制为周围循环迭代器的仿射形式的区域。这使得精确的仿射依赖分析成为可能,同时避免了从有损的低级表示中推断仿射形式的需要。

5.2.2 差异

        与现有多面体框架的差异很多,我们可以将其分为四类:
(1) 丰富的类型:MLIR 结构化内存引用类型包含一个将缓冲区的索引空间连接到实际地址空间的布局映射。这种关注点的分离使得循环和数据转换更好地组合:数据布局的更改不会影响代码,也不会污染依赖分析。这样的转换混合以前已被探索,但并不常见。


(2) 抽象的混合:MLIR 中仿射循环的主体可以用类型化 SSA 值的操作来表示。因此,所有传统的编译器分析和转换仍然适用,并且可以与多面体转换交错。相反,多面体编译器通常完全抽象掉这些细节,使得多面体编译器难以操作,例如,向量类型。


(3) 更小的表示差距:多面体模型的一个关键特性是其能够在类型系统中表示循环迭代的顺序。在这个系统中,大量的循环转换直接组合,并且可以使用简单的数学抽象进行推理。然而,多面体转换需要提升到通常与原始表示截然不同的表示。此外,从转换后的多面体到循环的转换在计算上是困难的。基于 MLIR 的表示在低级表示周围保持高级循环结构,消除了提升的需要。


(4) 编译速度是 MLIR 的一个关键目标,但大多数现有多面体方法并未关注这一点。这些方法严重依赖于具有指数复杂度的算法:在整数线性规划上自动推导循环顺序,并在多面体扫描算法上将表示转换回循环。MLIR 采用的方法明确不依赖于多面体扫描,因为循环在 IR 中被保留。


        仿射方言的经验表明,它对广泛的代码生成项目是有用的,其开发是 MLIR 设计使之成为可能的重要探索。

5.3 Fortran IR (FIR)

        LLVM Fortran 前端“flang”目前正在由 NVIDIA/PGI 主导进行重大开发。与 Swift、Rust 和其他语言类似,flang 需要一个专门的 IR 以支持高性能 Fortran 代码库的高级转换,并使用 MLIR 来支持这些 Fortran 特定的优化。这些高级优化——高级循环优化、数组复制消除、调用专门化、去虚拟化——仅使用 LLVM 实现将是困难的。


        例如,FIR 能够将 Fortran 虚拟调度表建模为一级概念(见图9)。
        能够在结构化 IR 中建模编程语言的高级语义是非常强大的。例如,调度表的一级建模允许实现一个强大的去虚拟化 pass。虽然这可以用定制的编译器 IR 实现,但使用 MLIR 允许 flang 开发者将他们的工程资源集中在其领域的 IR 设计上,而不是重新实现基本基础设施。


        选择 MLIR 还解锁了其他非 Fortran 特定方言的可重用性:语言无关的 OpenMP 方言可以在 Fortran 和 C 语言前端之间共享。类似地,使用 OpenACC 目标异构平台在 MLIR 中变得可行,通过共享和重用面向 GPU 的方言和 passes。这很简单,因为 MLIR 专门设计用于支持可组合方言的混合。

5.4 领域特定编译器

        上述 MLIR 的应用在大型编译工作流中,但它也用于构建特定小型工作流的领域特定编译器。可重用和模块化的基础设施使这些专门路径成为可能,并且相对便宜。
        优化 MLIR 模式重写 MLIR 具有一个可扩展的模式重写系统,如第 4 节所述。除了静态声明的模式外,我们还有一些应用需要在运行时动态扩展重写模式,允许硬件供应商在驱动程序中添加新的降低。解决方案是将 MLIR 模式重写表达为 MLIR 方言本身,允许我们使用 MLIR 基础设施动态构建和优化高效的有限状态机(FSM)匹配器和重写器。这项工作包括在其他系统中看到的 FSM 优化,例如 LLVM SelectionDAG 和 GlobalISel 指令选择系统。
        晶格回归编译器 晶格回归是一种以快速评估时间和可解释性而闻名的机器学习技术。编译器的前身是使用 C++ 模板实现的。这允许使用元编程实现高性能代码,但在端到端模型上表达通用优化并不简单。这个特定的晶格回归系统用于拥有数百万用户的应用程序,因此性能改进至关重要。
        MLIR 被用作该专门领域的新编译器的基础,该编译器由专门的搜索方法驱动——实际上在编译期间解决了一个机器学习问题。最终的编译器通过投入 3 个人月的努力开发,结果在生产模型上实现了高达 8 倍的性能提升,同时在编译期间提高了透明度。

6 MLIR 设计的后果

        MLIR 的设计促进了新语言和编译抽象的建模,同时重用现有的通用抽象及其相关的编译方法。实际上,解决许多问题的方法是“添加新的 ops、新的类型”,可能收集到“一个新的方言”中。这对编译器工程来说是一个显著的设计转变。它带来了新的机会、挑战和见解。本节探讨其中的一些。

6.1 可重用的编译器 passes

        在一个 IR 中表示多个抽象层次的能力自然会产生编写跨多个抽象层次工作的 passes 的愿望。关于 MLIR 的一个常见问题是“当你有开放可扩展的操作和类型系统时,你如何编写编译器 pass?”虽然编译器 pass 总是可以以保守正确的方式处理未知构造,但我们的目标是生成高性能代码,因此我们需要在常见情况下做有用的事情。我们发现了四种

基本操作特性

        一些“基础”编译器 passes,如死代码消除和公共子表达式消除,依赖于非常简单的属性(如“无副作用”或“可交换”),我们将其定义为 Op 特性。ODS 中的操作定义允许操作的作者指定这些特性,passes 可以使用这些信息在许多不同的抽象域中保持适用性。

特权操作钩子

虽然某些特性可以用一个位来建模,但其他特性需要 C++ 代码来提供实现——例如常量折叠逻辑。MLIR 对适用于大量 passes 的某些钩子提供一级支持。这些钩子可以在每个操作的基础上实现,也可以在方言对象本身中实现。后一种方法对于 TensorFlow ops 的常量折叠等事情很方便,因为委托给现有逻辑很简单。

优化接口

MLIR 的一个主要目标是允许开放的可扩展性——不仅在操作和类型方面,还在转换方面。虽然规范化和常量折叠是关键操作,但有许多标准转换需要以某种方式进行参数化——例如,描述转换特定属性,实施成本模型等。

解决方案是一个称为“优化接口”的子系统。考虑 MLIR 的内联 pass:我们希望内联器能够在 TensorFlow 图、Flang 函数、函数式语言中的闭包等上工作——但内联器不知道调用点甚至被调用者是什么!内联器需要知道的核心特性是:

  • 是否可以将一个操作内联到给定区域;
  • 如何处理在内联后出现在块中间的终止操作。

为了了解这些属性,内联器 pass 在图 10 中定义了接口。各个操作和方言可以在 MLIR 中注册其对该接口的实现,以从通用内联器 pass 中受益。如果一个操作或方言未能提供接口,那么相应的优化 pass 将保守地处理该操作。此设计允许方言的实现者快速启动和运行,但随着时间的推移,通过在这些接口上投入更多的实现努力,从系统中获得更多的价值。

优化接口还为核心编译器提供了模块化的好处,因为方言特定的逻辑是在方言本身中实现的,而不是在核心转换中实现的。

方言特定 passes

最后,定义特定于特定方言的 passes 是有效且有用的,这些 passes 可以由其设计的方言中的操作的完整语义驱动。这些 passes 在 MLIR 系统中与在其他编译器系统中一样有用。例如,代码生成器希望根据特定机器约束或其他不适合更广泛框架的技巧进行自定义调度。这是新转换的简单而有用的起点,其中不需要泛化。

6.2 混合方言

       MLIR 的一个最深刻(但也最难理解)的方面是它允许并鼓励将来自不同方言的操作混合到一个程序中。虽然某些情况下这很容易理解(例如在同一模块中持有主机和加速器计算),但最有趣的情况发生在方言直接混合时——因为这启用了我们在其他系统中未见的整个类的重用。
       考虑第 5.2 节中描述的仿射方言。仿射控制流和仿射映射的定义独立于包含在仿射区域中的操作的语义。在我们的例子中,我们将仿射方言与表示简单算术的“标准”方言结合在一起,以目标无关的形式(如 LLVM IR),以及用于内部加速器的多个目标特定机器指令方言。其他人将其与其他问题域的抽象结合在一起。

6.3 互操作性

       我们的工作涉及与大量现有系统的互操作,例如,编码为协议缓冲区的机器学习图、包括 LLVM IR 在内的编译器 IR、专有指令集等。通常,表示有许多在现有系统背景下有意义的次优或不幸的决策,但 MLIR 的能力使得更具表现力的表示成为可能。
       因为导入器和导出器众所周知地难以测试(通常测试用例是二进制格式),我们希望确保它们的复杂性最小化。
       解决方案是定义一个尽可能直接对应于外部系统的方言——允许以简单和可预测的方式往返于该格式。一旦 IR 被导入到 MLIR 形式中,它可以使用所有 MLIR 基础设施进行这些转换的提升和降低,并允许这些转换像所有其他 MLIR passes 一样进行测试。

6.4 无偏见设计带来的新挑战

       虽然 MLIR 允许定义几乎任意的抽象,但它对应该做什么提供的指导很少:在实践中什么效果更好或更差?我们现在有一些工程师和研究人员将这些技术和技术应用于新问题领域的经验,并意识到编译器 IR 设计和抽象设计的“艺术”在编译器和语言领域并没有得到很好的理解——许多人在既定系统的约束内工作,但相对较少的人有机会自己定义抽象。

       这是一个挑战,但也是未来研究的另一组机会。更广泛的 MLIR 社区正在积累大量关于这些抽象设计权衡的专业知识,我们预计这将是一个随着时间推移的肥沃研究领域。

6.5 展望未来

       MLIR 的设计与其他编译器基础设施有很大不同,以至于我们仍在学习——即使在将其构建并应用于许多不同系统之后。我们相信仍有很多东西需要发现,并且需要几年的研究才能完全理解设计点并建立最佳实践。例如,树外方言的兴起,使用 MLIR 的源语言前端数量的增加,可能应用于抽象语法树,以及应用于结构化数据(如 JSON、协议缓冲区等),这些仍处于非常早期阶段,可能会揭示有趣的新挑战和机会。

7 相关工作

       MLIR 是一个与多个不同领域重叠的项目。虽然其组合基础设施提供了一个新颖的系统,但单个组件在文献中有类似物。有关直接与 IR 设计相关的参考和讨论,请参阅第 2 节。

       MLIR 是类似于 LLVM 的编译器基础设施,但 LLVM 在标量优化和同质编译方面取得了巨大成就,而 MLIR 旨在将丰富的数据结构和算法建模为一级值和操作,包括张量代数和算法、图表示以及异构编译。MLIR 允许混合和匹配优化,将编译 passes 分解为组件并重新定义降低、清理角色。这主要归功于模式重写基础设施,将完整的转换捕获为小的局部模式的组合,并控制在单个操作的粒度上应用哪些模式重写。自动扩展、形式化和验证重写逻辑将是重要的下一步。在后端方面,MLIR 的 DDR 类似于 LLVM 的指令选择基础设施,支持具有多结果模式和作为约束规范的可扩展操作。

       许多编程语言和模型解决硬件异构性。最初是同质编程模型,OpenMP 添加了对将任务和并行区域卸载到加速器的支持,基于早期的提案如 StarSs 和 OpenACC。C++ AMP、HCC 和 SyCL 利用传统的 Clang/LLVM 流程和现代 C++ 提供硬件加速的高级抽象。不幸的是,所有这些例子都很快将高级构造降低为对运行时执行环境的调用,依赖于主机语言(通常是 C++)中现有的优化来减轻抽象惩罚。很少有努力针对异构编译过程本身。扩展 LLVM IR 的并行中间表示解决了部分问题,但传统上专注于同质设置。迄今为止最雄心勃勃的努力可能是 Liquid Metal,具有共同设计的领域特定语言(DSL)和编译流程,将托管对象语义转换为静态、向量或可重构硬件;然而,其 Lime 编译器的大部分努力在于将圆形对象适配到方形硬件中。MLIR 提供了对高层语言的直接嵌入,通过可扩展的操作和类型集拥抱异构性,同时提供一个通用基础设施,以最大限度地重用不同目标之间的通用组件来逐步降低这些构造。

       解决语言异构性一直是元编程系统的长期承诺,特别是多阶段编程。轻量级模块化分期(LMS)是一个最先进的框架和运行时代码生成器,提供用于生成高效代码和在 Scala 中嵌入 DSL 的核心组件库。Delite 承诺为 DSL 开发者带来显著的生产力提升,同时支持并行和异构执行。我们认为这种方法是对 MLIR 的补充,提供了更高级别的抽象来嵌入 DSL 并通过通用元编程构造实现优化。

       更进一步进入语言语法,ANTLR 是一类解析器生成器,旨在使开发新编译器前端变得容易。MLIR 目前没有通用解析器生成,没有 AST 构建或建模功能。将 MLIR 与 ANTLR 这样的系统结合起来,可以从用户输入到代码生成创建可重用的编译器库。

       更狭义地由其在机器学习中的应用来定义,XLA、Glow 和 TVM 解决了类似的异构编译目标。然而,这些是从图形抽象开始并针对加速器的多维向量抽象的特定代码生成实例。所有这些都可以利用 MLIR 作为基础设施,利用通用功能,同时使用其当前的代码生成策略。类似地,Halide 和 TVM 的循环嵌套元编程技术,早期的循环嵌套元编程,以及完全自动化的流程如 PolyMage、Tensor Comprehensions、Stripe、Diesel、Tiramisu 及其底层多面体编译技术可以作为 MLIR 基于编译框架中的不同代码生成路径共存。序列化和互操作性格式,如 ONNX,通过提供一组不同框架可以映射到的通用 ops 来解决 ML 前端多样性的问题。ONNX 可以作为 MLIR 中的一个方言候选,其他 ops 可以降低到和从中降低。

8 结论与未来工作

       我们介绍了 MLIR,一个用于编译器构建的灵活且可扩展的基础设施。本文描述了 MLIR 的具体设计,展示了其在一系列重要领域的适用性,并描述了一些原创研究和工程影响。

       展望未来,我们渴望看到成熟的编译器社区(例如 Clang C 和 C++ 编译器)以及领域专家如何从引入更高级别、语言特定的 IR 中受益。我们也有兴趣看到 MLIR 是否能够启用教授编译器和 IR 设计艺术的新方法,并希望看到由此基础设施催化或加速的全新研究领域。

       未来方向 MLIR 正在追求多个未来方向。在 ML 和 HPC 领域,这些包括从具有符号形状的参考秩多态规范中推断高效的 Op 实现。它还涉及启用更广泛的数据结构(稀疏、图形)和程序转换,将符号推理(如自动微分和算法简化)与更传统的数据流和控制流优化结合起来。超越 ML 和 HPC,可以考虑 MLIR 在其他相关领域的适用性,如安全编译、安全关键系统、数据分析和图处理、关系查询优化等。

       回到通用语言的世界,一个明显缺失的前端是从 Clang 派生的 C++ 中级表示。比如,一个类似于 Swift 的 SIL 和 Rust 的 MIR 的“CIL”,将有助于优化当前需要从降低代码中重建的常见 C++ 习惯用法(例如,将 std::vector 视为数组而不是指针操作)。在 MLIR 中支持垃圾回收语言、更高阶和多态类型系统以及类型推断也是开放的挑战。

       在 LLVM 中探索并行性和并发构造一直很困难,主要是因为所需的更改是侵入性的且不易分层(例如,注入元数据并检查所有 passes 以确保元数据在抽象级别过低时传播,同时失去优化机会)。使用 MLIR,并行构造可以是一级操作,使用区域和并行习惯用法特定的验证。这将支持在降低到例如 LLVM 之前的更高级别的转换,在那里可以对已降低的代码执行常规转换。

       除了调试和测试,IR 的文本形式对教育也很有用。额外的工具可以展示高性能编译中的优化交互,可以为新学生揭开编译器的神秘面纱。IR 设计是开发新编译器或优化框架的一个组成部分,但许多本科编译器课程不涵盖 IR 设计。MLIR 为此类课程的新方法提供了探索的机会。

参考文献

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com

热搜词