MLIR编译框架：新一代AI编译器的基础设施

来源：互联网 2026-08-04 18:15:04

人工智能资讯合集

MLIR是一种高度模块化、开放可扩展的编译器基础设施，通过多级抽象表示和方言机制实现从高级框架到低级代码的逐步递降与优化，支持AI、高性能计算等多领域，并在操作、区域、基本块等结构基础上提供可复用的优化与转换能力。

编译器基础设施这个话题，近年来讨论热度一直不低。而MLIR的出现，可以说为整个编译器生态带来了一些不一样的思路。它是一个完全可扩展的编译器基础设施，最突出的标签就是高度模块化和开放的架构——操作集、属性集、类型系统全都是开放可扩展的，并且随着不同应用场景持续演进。这意味着，MLIR不仅能用来搭建AI编译器，在高性能计算、图形处理、量子计算这些领域，也同样适用。 MLIR的设计思想相当巧妙：它支持将高阶IR逐步递降到特定的目标形式，覆盖了从图形算法到低级代码生成的整个范围。在处理循环嵌套、数据布局这类典型的中级优化之外，它也能执行面向后端IR的低级调度和映射决策——比如把IR映射到专门的矢量指令、完成自动矢量化和软件流水线等。换句话讲，MLIR提供的多级抽象表示能力，不只是给出了一种新的表示思路，更是用统一的设计框架把多级抽象整合在一起，通过标准化的IR和一套内在一致的概念，去表示各种形式的计算和来自不同框架的数据流图。它为不同的问题域提供了可共享、可重用的基础设施，让开发者能在不同层次上插入优化、转换与分析，最终经由高级优化和并行化方法得到高性能目标代码。

MLIR中的操作与方言

程序的行为是由操作来描述的。MLIR里的操作集，覆盖范围从硬件指令到函数，再到AI模型的构建块，几乎涉及计算的方方面面。操作是MLIR的基本实体和最小代码单元，也是抽象和计算的核心单元。你可以把它想象成硬件指令或其它高度抽象操作的一种泛化表示。操作的运行方式是：它消费值，同时定义新值。每个操作会使用一组类型化的操作数，并产生一组类型化的结果。这些定值和使用的值，构成了SSA形式下不可变的数据单元。举个例子，在表达式“%3 = addi %0, %1”中，加法操作使用了%0和%1，定义了%3。当一个操作的输出结果连接到另一个操作的输入操作数时，就形成了SSA形式的数据流。若干个操作可以按顺序组成基本块，多个基本块又能组成区域——比如循环体或函数体。操作还带有属性，这些属性包含了编译时已知的信息。每个操作通过属性参数来描述自己的重要特征。和执行时才能获取的操作数值不同，属性的值在编译时就能确定下来并且固定不变。顺便提一句，MLIR中的每个值都有类型，类型里包含了该值在编译时已知的信息。 MLIR的操作是完全可扩展的，并且带有应用相关的语义，可以用来表示LLVM中所有核心IR结构——指令、函数、模块，全都不在话下。为了保证可扩展性，MLIR只把属性、类型、操作和区域这几种概念作为基本内建概念，其他概念都由它们衍生而来。比如，模块和函数实际上就是被定义为具有特定语义、符号名称属性，并且以区域为主体的操作。更进一步说，从for循环这样的高级结构到低级的硬件指令，都可以用操作来定义。模块化和重用性是MLIR的核心设计理念。它允许不同的组件以独立库的形式分别实现，开发者可以根据需要链接这些库——这和LLVM的设计理念高度一致，但又比LLVM更进一步地强调了灵活性与工程可扩展性。目前MLIR虽然在部分环节依赖于LLVM的支持库，但在架构上和LLVM IR是相互独立的。通过保持优化与转换逻辑与核心IR抽象的分离，MLIR实现了高度解耦的构建系统。这个构建系统可以独立选择编译链接特定的方言实现代码。 MLIR通过“方言”这个概念，支撑起了模块化、可扩展的设计以及中间表示的多样化。方言是IR对象的逻辑组合，在唯一的命名空间下，为不同应用领域中逻辑上相关的各种抽象（操作、属性和类型）提供了一种组合机制。你可以把MLIR方言笼统地理解为IR的逻辑层次。正是通过一系列可组合的方言来实现多阶段编译——这是MLIR区别于其他编译框架的重要特性。这个特性让开发者可以在面向不同异构硬件的编译栈上，复用公共的方法和组件。 MLIR包含一系列特定领域的方言，这些方言描述了给定应用场景下所支持的合法操作集，并根据各自的工作上下文与其他方言进行交互。不同的方言在代码生成的IR处理流程中扮演着不同角色。为了高效地与MLIR交互，开发者有时需要根据需求定义新的MLIR方言，为后续的分析和转换铺平道路。也就是说，MLIR除了提供若干内置方言来表达通用功能，还为开发者提供了开放基础设施，允许大家在不同粒度和抽象级别上定义新的方言、自定义类型、操作和属性。

MLIR中的常用方言简介

MLIR方言提供了多个抽象级别，这让我们能够完成那些在单一抽象级别下很难执行的转换和优化。举个例子，Torch、TOSA、MHLO等上层方言大多和AI领域相关，它们通常负责导入高层AI框架的图并处理算子集合，捕获算子语义后再将其转换为统一的表示形式，消除不同框架之间的不一致性，从而简化AI模型导入MLIR环境的过程和代码生成任务。在统一表示的基础上，MLIR可以执行各种转换和优化，比如算子融合、内存布局优化、并行化等等。这些优化有助于生成高效的低层代码，并在不同硬件平台上跑得更快。再往中间层看，Affine、SCF（Structured Control Flow）、Linalg、Tensor、Vector等方言主要负责张量和缓冲区的操作及其他转换流程。中间层方言通常支持多样化的上层输入方言——比如TOSA、MHLO都可以作为Linalg方言的上层输入。在这些中间层方言里，Affine方言使用了多面体编译技术，主要用来表达和优化带约束的循环和内存访问模式。它提供了强大的循环变换能力，比如循环拆分、循环融合、循环交换等。 SCF方言则用于表示IR中不受仿射分类规则约束的各种控制流结构，包括循环和条件。例如，scf.for循环操作接受整数值作为循环的上下限或步长，但不支持仿射映射。这些结构化的控制流有助于在高层次上进行控制流分析和优化，让代码生成更简单。 Linalg方言负责表达和优化与线性代数相关的张量操作。比如矩阵乘法、卷积等操作可以通过Linalg方言描述，并利用分块操作把大规模计算问题分解成小规模的子问题——一方面提升数据局部性和缓存命中率，另一方面更好地将计算映射到硬件资源，达到在目标硬件上高效执行的优化目标。张量表示抽象的值类型数据序列。Tensor方言专注于张量操作的抽象，提供了处理多维数组的操作，可以在高层次上对张量进行操作，为后续的优化和转换打下基础。使用张量和Tensor方言操作，能提高算法开发者的效率。而MemRef方言则代表较低级别的缓冲区访问，用于构建与物理内存之间的桥梁。Vector方言提供了向量化的操作抽象，可以高效地进行数据并行优化，比如向量化和SIMD化。总结一下：Affine、SCF、Math、Linalg等方言主要用于表示计算结构和控制流；Tensor、MemRef、Vector等方言则主要负责表示数据负载。往底层看，NVVM、ROCDL等平台方言和LLVM方言，通过与其他编译器（尤其是传统编译器和领域专用编译器）后端兼容的中间表示，实现了MLIR与不同外部编译器后端之间的对接。针对特定硬件的底层方言，比如NVGPU、AMDGPU，则实现了对特定硬件特性的抽象。这些方言可以看作是MLIR中间表示在导出到外部编译器后端之前的过渡阶段。 MLIR完成优化和转换后，通常会输出LLVM IR，然后借助LLVM框架调用不同的处理器后端完成代码生成。LLVM方言提供了LLVM IR指令和类型到MLIR的直接映射。考虑到编译过程后续需要LLVM进一步处理，作为叶方言的LLVM方言，能够简化MLIR和LLVM IR两种表示之间的翻译过程。一个值得注意的优点是：MLIR把域内的映射问题和其他问题限制在各自领域的方言内处理。这意味着数据结构之间的转换可以做到尽可能简单，得到正确转换结果的概率也更高。

MLIR的高层结构

前面提到过，在MLIR中，基本块被定义为一串不含控制流的顺序操作。由基本块构成的控制流图在MLIR中被组织成区域，控制可以从区域中的一个块流向后续块。MLIR支持递归结构，区域可以附加到一个操作上，该操作定义控制如何流入和流出这些区域——也就是说，操作可能包含块的关联区域，这些区域里又包含一系列块，每个块可能又包含一系列操作，而这些操作也可能包含其他区域，从而让IR在任意多个级别上实现嵌套。每个块以终止操作结束，控制流会转移到该终止操作的后继块。图1总结了MLIR嵌套结构的组成。图1. MLIR的高层嵌套结构以下示例代码说明了MLIR的关键概念，包括操作、操作数和结果、属性、区域、基本块等的基本使用方法。 ``` %value_definition = "dialect.operation"(%value_use) {attribute_name = #attr_kind<"value">} ({^block(%block_argument: !argument_type):"dialect.other_operation"() : () -> ()}) : (!operand_type) -> !dialect.result_type ``` 在MLIR中，操作是基本单位，由唯一的字符串（如tf.Conv2d、x86.repmovsb等）标识，可以接收多个操作数和返回多个结果，支持属性字典、后继（successor），以及零个或多个嵌套的区域。上面这段代码中的dialect.operation，就属于dialect方言，它接受一个操作数%value_use，返回一个类型为!dialect.result_type的结果。这个操作还包含一个命名属性attribute_name，使用了#attr_kind<"value">的语法，这是MLIR中通用的属性机制。操作的表达范围非常广——既可以表示函数定义（比如func.func操作）、函数调用（func.call操作）、缓冲区分配（memref.alloc操作）这样的高层抽象，也可以表示目标无关算术操作（如arith.addi操作）、特定硬件平台的底层指令（如nvvm.mma.sync操作）等底层逻辑。MLIR中的pass是实现操作转换和优化的核心机制。借助pass及其管理框架，MLIR能够模块化、渐进式地对多层次的操作进行变换和优化，并且支持开发者灵活地插入和定制新的转换逻辑。上面代码中的dialect.operation，通过大括号{}引入了一个区域。区域是一个有序的基本块列表，其语义由包含该区域的操作（也就是dialect.operation操作）来定义。区域必须包含在操作中，没有名称或地址，也没有类型或属性。函数体可以看作是区域的一个例子。在函数体中，位于基本块末尾的终止操作必须跳转到不同的块，或者从函数返回，返回参数的类型要和函数签名的结果类型匹配。同样，函数参数必须与区域参数的类型和数量匹配。区域中不同基本块之间的引用或跳转，必须局限在所属区域之内——也就是说，某个基本块的终止操作不能跳转到不在同一区域内的其他基本块。而某个值的可见性取决于它是否能够被后续操作引用。因此，区域的跳转限制意味着：一个值在某区域中定义，就只能被该区域内部的操作访问。这样一来，值的作用域自然被限制在区域内。另外，MLIR默认允许区域内的操作引用外层区域中已存在的值，只要这些值本来就是封闭操作的合法操作数。如果想彻底隔离内层区域、阻止它引用外部值，可以为封闭操作加上OpTrait::IsolatedFromAbove特征（trait）。上面代码中，dialect.operation操作的区域里包含一个基本块^block。基本块是MLIR中控制流的基本单元，包含一组顺序执行的操作序列，只有一个入口和一个出口，没有显式的分支结构，并且必须以终止操作结束。终止操作负责实现基本块之间的控制流分支。区域中的第一个基本块被称为入口块。入口块的块参数同时也是该区域的参数，这些参数绑定的值由包含该区域的操作语义决定。对于区域中其他基本块的块参数，其绑定值由控制流终止操作决定——比如分支跳转操作，这些操作会把该基本块作为后继块之一。在下面的函数@simple中，入口块^bb0的参数%a和%cond的值，由func.func操作定义和绑定。而块^bb0作为入口块支配区域中的其他基本块，这意味着在其他基本块内可以引用%a和%cond，它们的作用域是整个区域，这也体现了入口块参数兼具区域参数的特点。对于区域中其他基本块的块参数，比如块^bb3的参数%c，它的绑定值来自前驱块^bb1和^bb2的cf.br跳转操作（也是终止操作）。块^bb3将参数%c和%a一起传递给块^bb4，但%a不是^bb3的参数，而是直接引用定义在^bb0的%a。 MLIR利用这种基本块结构和参数传递方法，替代了传统SSA通过复杂的显式PHI节点来合并可能来自多个前驱基本块的变量值的机制，使得IR表达更简洁。 ``` func.func @simple(i64, i1) -> i64 { ^bb0(%a: i64, %cond: i1): cf.cond_br %cond, ^bb1, ^bb2 ^bb1: cf.br ^bb3(%a: i64) ^bb2: %b = arith.addi %a, %a : i64 cf.br ^bb3(%b: i64) ^bb3(%c: i64): cf.br ^bb4(%c, %a : i64, i64) ^bb4(%d : i64, %e : i64): %0 = arith.addi %d, %e : i64 return %0 : i64 } ``` 回到前面的代码块，基本块^block中嵌套了另一个操作dialect.other_operation。这个操作也可以拥有一个或多个区域，从而实现MLIR中间表示的递归层次化结构。总结一下：操作是语义与结构的基础单位，也是IR层级结构中的根节点——它定义自身的语义，同时承载一个或多个区域；区域扮演结构性作用域的角色，本质是一个有序的基本块列表，用于封装控制流结构并限定值的可见性；基本块则是基本的控制流单元，包含按顺序排列的一组操作，这些操作又可以进一步嵌套区域。在MLIR的中间表示中，另一种重要的结构关系是值在操作之间传递与引用。MLIR中的值，要么是从控制流前驱传入的块参数，要么是某个操作的输出结果。值的使用者是其他操作，通过操作的操作数来引用。图2展示了MLIR通过值的使用-定值链，在操作、基本块和操作数之间建立的数据流连接结构。图中直角灰色框表示值的定义，可以来自块参数或某个操作的结果；直角白色框表示操作的操作数，是对已定义值的引用；虚线箭头表示操作数使用值的方向。图2. MLIR中值的使用与定值基本块的每个块参数就是一个值。基本块中的每个操作，要么引用块参数，要么引用前驱操作的结果作为操作数。每个操作可以返回多个值，并被其他操作使用，从而形成清晰的使用-定值链，便于分析、优化pass追踪值的来源和使用位置。每个值记录了它的第一个使用（OpOperand），并且可以从这个使用出发，通过OpOperand的nextUse指针遍历后续所有OpOperand，或者通过back指针向前回溯。显然，MLIR中值的使用是通过这种典型的双向链表结构来组织的。

侠游戏发布此文仅为了传递信息，不代表侠游戏网站认同其观点或证实其描述