想象你是高校里负责“社团登记”的教务老师。新学期,几十个新社团涌来申请注册:围棋社、机器人兴趣小组、区块链研讨社……每个社团都有自己的行话、自己的成员分类、自己的活动章程。如果有人问“机器人社到底有哪些活动”,你会去翻那本总登记册。编译器世界也是这样——MLIR(多级中间表示)就是一座容纳各种“语言社团”的大学,方言(Dialect)就是这些社团。传统编译器只有一套固定 IR,想加一个新算子往往要改框架核心;MLIR 用一套方言注册机制,让任何团队都能写一张“新社团登记表”,把新的操作、类型、属性和转换规则登记到统一基础设施里。这个机制看似只是几个 API 调用,却决定了 MLIR 为什么能成为当代 AI 编译器的地基。本文就用“社团登记处”的故事,带你从头到尾捋一遍 MLIR 方言注册机制。
对 AI 编译工程师来说,掌握方言注册机制等于拿到了“编译器扩展的钥匙”:你能把自己的 AI 算子、自定义 IR 和优化 Pass 像插件一样无缝装进 LLVM/MLIR 生态,这是 AI 芯片编译器岗位的核心技能。
目录
可视化演示
展示MLIR方言注册流程、延迟加载策略和方言间降级转换Pass,理解多级中间表示的模块化扩展机制。
本节配有交互动效(进入视口后自动加载,约 3KB)。
方言注册流程 DialectRegistry中心节点管理所有方言注册 每个方言通过initialize()注册操作/类型/属性 注册覆盖率 = 已注册方言 / 总方言数 方言在注册后即可参与编译流水线
方言基础概念:多级IR的“方言家族”
在一所大学里,要成立一个围棋社,就得有自己的活动定义:何时摆棋盘、怎样算打吃、社员的“段位”怎么表示。这些规则聚在一起,就成了围棋社的“内部语言”。MLIR 把这种“内部语言”抽象成方言:一个方言不是零散的几个类,而是一整套自包含的操作、类型、属性和接口,像社团章程一样能独立被人读懂。
传统编译器通常只有一套固定 IR,想表达神经网络算子、GPU 指令或自定义硬件行为,往往要往核心数据结构里“塞”新字段,牵一发动全身。MLIR 则允许同一份 IR 里同时存在 tensor、arith、linalg、affine 等方言,各自负责不同抽象层级的表达,就像校园里同时存在围棋社、机器人社和话剧社,各有各的语言。
方言的自包含性很关键:每种方言都清楚地规定“我的操作叫什么、我的类型是什么、我的属性有哪些”,不需要依赖外部的隐藏约定。命名空间(如 arith.*、tensor.*)相当于社团的正式 ID,避免同名操作互相冲突。没有注册机制,这些方言只是头文件里的类;有了注册,编译器才真正“认识”它们,并让它们进入后续的调度和优化。
你看,社团登记册最前面的几页,写的就是“这个社团包含哪些东西”。也正是这套四元组,让 MLIR 敢于说自己的方言体系是模块化、可插拔的。
- 方言符号
\mathcal{D}代表一个完整的方言,比如arith或你的自定义方言。 - 四元组括号
\langle \cdots \rangle表示把四类元素打包成一个整体,像一个社团档案袋。 - 操作集合
\mathcal{O}是所有 Operation 的集合,对应社团可发起的“活动”。 - 类型集合
\mathcal{T}是数据类型的集合,对应会员的“身份标签”,如整数、浮点、张量。 - 属性集合
\mathcal{A}是 Attribute 的集合,对应社团档案上的“备注信息”,如对齐方式、步长。 - 接口集合
\mathcal{I}是方言接口的集合,对应社团对外承诺的“合作窗口”。
注册中心机制:DialectRegistry 总名册
学校社团登记处有一本总名册:哪个社团叫什么、负责人是谁、何时注册的,全在册子上。MLIR 里的 DialectRegistry 就是这本总名册。它不是某个“方言类”,而是一个专门保存“方言登记信息”的注册表,负责在编译器启动阶段收集所有可能用到的方言。
实际开发中,我们会在初始化阶段创建 mlir::DialectRegistry,然后通过 registry.insert<MyDialect>() 把自定义方言类放进去。接着把注册表交给 MLIRContext,MLIR 才知道“世界上存在哪些方言”。这个过程很像开学前把所有社团登记表装订成册,再交给学生事务处。
DialectRegistry 的设计也考虑了性能:它支持按需加载。一个方言只有在被真正引用时才会被实例化,而不是启动时就全部创建。这就像社团登记处只保存名册,真正要办活动时才会联系社团负责人。MLIRContext 会维护“已经加载”的方言集合,保证同一个方言在上下文中只有一个实例,避免重复和冲突。
简单说,一个方言要参与编译,第一步必须“在册”。很多初学者直接在 MLIRContext 里 new 一个方言对象,却忘了先把方言放进注册表,编译时就会得到“方言未注册”的错误。记住:总名册是先决条件。
- 判定函数
\rho(\cdot)是一个布尔函数,回答“方言 D 是否可加载”。 - 目标方言
D是待查询的方言,比如mydialect::MyDialect。 - 注册表集合
\mathcal{R}是DialectRegistry中已注册方言的集合。 - 中括号判定
[D \in \mathcal{R}]表示“如果 D 在登记册里,结果为真,否则为假”。
方言类开发:继承 Dialect 基类
成立一个正式社团,要填“社团章程”:名称、宗旨、活动范围、成员类型。MLIR 中定义自定义方言,对应要做的事是写一个继承自 mlir::Dialect 的 C++ 类。这个基类提供了方言的通用骨架,你要在构造函数里告诉框架“我叫什么名字”,并把该方言的操作、类型、属性一一添加进去。
在构造函数里,我们通常会调用 addOperations<...>()、addTypes<...>()、addAttributes<...>()。如果方言还有自定义接口,比如可以用来解析或优化的 DialectInterface,也会在初始化时通过 addInterfaces<...>() 注册。可以把这个过程理解成:交上去一张登记表,表上不仅写了社团名,还贴上了“可举办活动清单”“会员身份样本”和“对外接口说明”。
传统写法里,遍历操作列表需要手写大量样板代码。MLIR 用宏生成机制减轻负担:比如定义 #define GET_OP_LIST 后包含生成的 MyDialectOps.cpp.inc,就能自动把所有操作类汇总成一个列表。这个列表再交给 addOperations 注册,代码量大大下降,也让方言类开发从“体力活”变成“填空”。
本质上,方言类的实现就是把“社团宣言”翻译成 C++ 代码。写好 MyDialect : public Dialect 并不难,难的是理解这个类不是独立存在的,它最终要被放回 MLIRContext 里才有意义。
- 注册函数
\mathrm{registerDialect}是方言注册过程的抽象,对应insert与load的完整动作。 - 四个参数
N, O, T, A分别是方言名、操作、类型和属性列表。 - 映射符号
\mapsto表示“经过注册后变成”,强调这是一个放入环境的动作。 - 上下文目标
\mathrm{Context}是MLIRContext,注册后的方言会被它持有和调度。
操作注册技术:给社团添加活动
社团不能只挂着名号不办活动。每个操作(Operation)就像一条活动规则:活动有哪些输入、产生哪些输出、有什么前置条件。MLIR 中,操作是 IR 的基本组成单元,比如 tensor.extract、arith.addi、linalg.matmul 都是不同的 Operation。
自定义操作时,最推荐的做法是用 ODS(Operation Definition Specification)描述操作。ODS 是一种表格驱动描述语言,你可以声明操作名、参数、结果、验证逻辑,甚至自定义汇编语法。写完 ODS 后,MLIR 的 mlir-tblgen 工具会生成对应的 C++ 类和解析代码,再通过宏把操作列表注入方言构造函数。
注册操作时,操作类必须具备语义验证能力。verify() 函数会在 IR 被创建或读取时检查操作是否合法,比如张量形状是否匹配、属性是否必需。这是社团活动规则中的“安全员”:防止不合规的 IR 流入后续优化阶段。一个只注册但没有验证逻辑的 Operation,就像一场没有规则的乱跑活动,在大型编译流水线里非常危险。
如果把方言看成社团,那么 Operation 就是一场场活动。注册活动的时候,不仅要写活动名,还要规定怎么检查参与者、怎么产出成果——这就是 ODS 加 verify 的职责。
- 操作符号
\mathrm{op}代表一个具体操作,如mydialect.add。 - 操作数集合
\mathit{operands}是输入操作数,相当于活动的“原材料”。 - 输入下括号
\underbrace{\mathit{operands}}_{\text{inputs}}强调操作数作为输入的语义。 - 映射箭头
\to表示操作执行时发生的转换。 - 结果集合
\underbrace{\mathit{results}}_{\text{outputs}}是操作完成后产出的新值或新副作用。
类型与属性注册:身份证与档案卡
社团里的会员有不同身份:普通社员、干事、副社长。MLIR 方言也要定义自己的类型(Type),比如一个自定义硬件方言里可能有 npu.tensor 这种带特定布局的张量类型。属性(Attribute)则更像档案卡上的备注:在 IR 中存放编译期常量或附加信息,例如张量的步长、量化参数。
自定义 Type 和 Attribute 都需要注册进方言,否则 IR 里一旦出现它们,编译器会因为“不认识”而拒绝处理。注册类型用 addTypes<MyType>(),注册属性用 addAttributes<MyAttr>()。除此之外,类型和属性通常还需要实现 parse 和 print,用于文本形式的 IR 读写,保证 .mlir 文件能正常输出和反序列化。
MLIR 内部对类型和属性采用“唯一化”存储(uniquing):相同结构的类型或属性在上下文中只有一个共享实例。这就像办社团身份证:相同段位相同姓名的成员只需登记一次,后面所有人都引用同一个副本。这既节省内存,也让比较操作变得非常快。
成员没有身份证,社团就无法确认身份;IR 里的类型和属性没有注册,编译器也无法确认它们是否合法。所以,类型与属性注册不是“可选项”,而是方言落地的硬性要求。
- 类型实例
T表示方言中出现的某个具体类型,如npu.tensor。 - 类型集合
\mathcal{T}_D表示方言 D 已注册的类型集合。 - 属性实例
A表示某个具体的属性,如量化步长属性。 - 属性集合
\mathcal{A}_D表示方言 D 已注册的属性集合。
转换Pass与加载:改造活动并接入流水线
社团登记完,不能只躺在名册里,还得能和其他社团协同。编译流水线里的 Pass(遍)就是对 IR 做扫描和改写的“改造工程”。一个自定义方言如果只能创建,却不能降级到标准方言,就无法真正完成编译。因此,注册方言的同时,往往要注册配套的转换 Pass。
MLIR 中,Pass 通过 PassManager 按顺序执行。自定义方言通常会注册 RewritePattern(重写模式),描述“看到什么 IR 就替换成什么 IR”。方言降级(Dialect Conversion)则通过 ConversionTarget 声明哪些方言是合法的目标,再调用 applyPartialConversion 或 applyFullConversion,把高层方言逐步翻译到低层方言,最终变成 LLVM IR 或硬件指令。
动态加载是这套机制的放大器:因为方言注册表可以在 MLIRContext 构造前注入,第三方库就能以插件形式注册新方言。AI 芯片公司在主分支中集成 MLIR,再通过插件加载自家 npu 方言,既不改动上游框架,又能复用标准优化。这就像学校允许外部专业人士凭“客座社团”身份入驻,快速带来新活动方案。
登记在册只是拿到“社团资格”,真正让方言为 AI 编译器所用的是 Pass 和降级机制。没有 Pass,方言就像一座孤岛;有了注册和降级,孤岛才通过桥梁接入编译大陆。
---
- 流水线整体
\mathrm{Pipeline}代表从输入 IR 到输出 IR 的完整变换链。 - 单个 Pass
P_i表示流水线中的第 i 个 Pass,比如降级 Pass、优化 Pass。 - 复合运算符
\circ表示“先执行后一个,再执行前一个”的复合关系,在编译器中表现为按顺序执行。 - 省略号
\cdots表示流水线可以继续延伸,Pass 数量由编译选项决定。
AI 应用场景
AI芯片自定义算子方言集成
阶段:编译器后端算子实现与代码生成
解决的问题:AI 芯片厂商的 NPU 有特殊指令集,比如二维卷积、量化计算。传统做法是在 LLVM 后端逐一加 intrinsic,代码侵入严重。利用 MLIR 方言注册机制,可创建 npu 方言,注册 npu.conv、npu.quant 等操作和属性,再通过注册降级 Pass 把 linalg 或 tosa 方言翻译到 npu 方言,所有扩展都像插件一样挂在主框架外。
具体案例:某 AI 芯片工具链把 ONNX 模型导入 MLIR,经过 tosa → linalg → npu 的方言降级链,最终生成 NPU 驱动指令。添加新算力扩展时,只需在新版本插件中注册更多操作和模式,主分支代码几乎不需要改动。
学术框架中的硬件加速方言接入
阶段:前端算子表达与中间优化
解决的问题:研究团队要快速验证一个新的稀疏矩阵加速算法,希望把稀疏表达直接融合进 MLIR,又不破坏现有社区方言。通过注册自包含方言和对应转换模式,团队能在隔离环境中实验,评估不同抽象层级的性能收益。
具体案例:学院实验室为 GPU Tensor Core 训练场景注册 sparse_gpu 方言,定义稀疏张量类型和归约操作,并注册一个转换 Pass 把 linalg.generic 变成专用内核调用,最终在 A100 上获得 20% 推理加速。
多产品线共享编译框架的插件化扩展
阶段:编译流水线构建与产品定制
解决的问题:大型公司内部有多个 AI 加速器产品线,每个产品都有自己的方言。方言注册机制允许各团队向同一个 MLIR 框架注入自己的方言和 Pass,同时避免命名冲突和重复加载,使公共优化库得以复用。
具体案例:云厂商的编译服务在启动时读取产品配置文件,通过 DialectRegistry 加载不同硬件方言插件。用户提交一个模型,服务端自动按目标设备选择注册表,然后运行统一流水线。
Python 演示
"""迷你方言注册表:模拟MLIR DialectRegistry的核心行为.""" class Dialect: def __init__(self, name): self.name = name self.operations = {} self.types = set() self.attributes = set() self.transforms = [] def register_operation(self, op_name, validator=None): self.operations[op_name] = validator or (lambda *args, **kwargs: True) return self def register_type(self, type_name): self.types.add(type_name) return self def register_attribute(self, attr_name): self.attributes.add(attr_name) return self def add_transform(self, func): self.transforms.append(func) return self def info(self): return { "name": self.name, "ops": sorted(self.operations), "types": sorted(self.types), "attrs": sorted(self.attributes), "transforms": [f.__name__ for f in self.transforms], } class DialectRegistry: def __init__(self): self._registry = {} def insert(self, dialect): if dialect.name in self._registry: raise ValueError("dialect already registered: " + dialect.name) self._registry[dialect.name] = dialect def load(self, name): if name not in self._registry: raise KeyError("dialect not registered: " + name) return self._registry[name] def names(self): return sorted(self._registry.keys()) def npu_conv_validator(shape_pair): return isinstance(shape_pair, tuple) and len(shape_pair) == 2 if __name__ == "__main__": npu = Dialect("npu") npu.register_operation("npu.conv", npu_conv_validator) npu.register_operation("npu.quant") npu.register_type("npu.tensor") npu.register_attribute("npu.strides") npu.add_transform(lambda mod: mod + " -> lowered") registry = DialectRegistry() registry.insert(npu) loaded = registry.load("npu") print(loaded.info())
DialectRegistry 和 Dialect 类,模拟 MLIR 方言注册表“先注册、后加载、可查重”的核心行为,帮助理解 C++ 版 registry.insert<MyDialect>() 和 context.loadDialect<...>() 的关系。class Dialect:定义方言类,对应 MLIR 中继承mlir::Dialect的自定义方言。self.operations = {}用字典保存操作名和验证函数,对应addOperations<...>()。def register_operation(self, op_name, validator=None):注册一个 Operation,并允许附加验证规则。def add_transform(self, func):给方言添加一个转换函数,模拟后续的 Pass。class DialectRegistry:构建总登记册,对应 MLIR 的DialectRegistry。def insert(self, dialect):把方言放入登记册,并检查重名,对应registry.insert<MyDialect>()。def load(self, name):从登记册取出方言,对应context.loadDialect<MyDialect>(),未注册会报错。registry.insert(npu)演示自定义 AI 芯片方言npu完成登记。loaded.info()打印方言包含的操作、类型、属性和转换信息,验证注册是否成功。
常见误区
误区我觉得把同一个方言在 `MLIRContext` 里反复注册能提高稳定性,反正多注册几次编译器也能忍。
正解一个方言在同一个 MLIRContext 中只会有一个实例。重复注册可能导致断言失败或未定义行为,因为上下文会维护“已加载方言”的唯一集合。正确做法是通过 DialectRegistry 登记一次,后续按需加载;即使多个 Pass 使用同一方言,也不需要重复注册。
为何常见:很多初学者把“注册”理解成“安装”,觉得像装软件一样多装几次无害,忽略了 MLIR 中方言实例与上下文是一一绑定的关系。
误区Pass 运行时发现缺一个方言,直接在 Pass 里创建一个并塞给 `MLIRContext`,以为这样能动态补救。
正解方言注册应该在 MLIRContext 构造之前通过注册表完成,或在上下文构造之后显式 loadDialect。在 Pass 执行阶段动态注册新方言会破坏编译流水线的静态结构,导致无法规划转换目标和验证 IR 合法性。
为何常见:计算机里“运行时加载库”的观念太强了,但 MLIR 的方言注册阶段有明确的时机窗口,不是任何时刻都能插入新语言。
误区手写 C++ 的 Operation 子类更直观,ODS 那些自动生成代码反而不透明,所以应该完全绕开 ODS。
正解MLIR 官方推荐使用 ODS 声明操作,因为它自动生成解析、打印、验证、文档和 Python 绑定等大量代码。手写 Operation 并非不可能,但很容易遗漏标准化接口,导致后续无法参与方言转换和模式匹配。ODS 不是可选的装饰,而是保证方言可维护性的关键工具。
为何常见:新手看到 ODS 文件里全是 TableGen 语法,觉得陌生难懂,宁可回到熟悉的 C++ 手工实现。
误区反正注册完编译器就能识别,`parse`/`print` 先空着,以后用到文本输出再补。
正解MLIR 的 IR 文本格式是生态基石,类型和属性没有实现打印解析,就无法在 .mlir 文件中表达,测试用例和调试器也会失效。注册类型/属性时,至少要实现 parse、print 和验证逻辑,才能保证 round-trip 完整。
为何常见:很多教学例子只展示 addTypes 一行代码,导致学习者忽略了解析打印这些“幕后工作”。
误区只要我把方言 insert 进 Registry,编译器就会自动帮我优化并生成代码。
正解注册只解决“编译器认识你”的问题,不解决“编译器会优化你”的问题。要真正接入流水线,还必须编写 Pass、RewritePattern 和 ConversionTarget,把高层方言逐步降级到可执行的低层方言。没有这些转换,方言就像只会登记、不会办活动的社团,最终会在编译后端卡住。
为何常见:看到“注册机制”四个字,容易误以为注册是全部工作,而忽略了 MLIR 庞大的 Pass 基础设施才是真正搬砖的地方。
面试题
Q1. 在 MLIR 中,哪个组件扮演“方言总登记册”的角色?
- A.PassManager:管理优化Pass运行
- B.DialectRegistry:维护方言登记表← 正确答案
- C.OpBuilder:帮助构建操作节点
- D.ModuleOp:表示顶层IR模块
答案:B。DialectRegistry 专门保存方言注册信息,相当于社团登记总名册。PassManager 负责执行 Pass,OpBuilder 负责构建操作,ModuleOp 代表 IR 模块,它们都不负责方言登记。
Q2. 自定义方言类通常需要继承哪个基类?
- A.mlir::Op操作基类
- B.mlir::Type类型基类
- C.mlir::Dialect方言基类← 正确答案
- D.mlir::Attr属性基类
答案:C。方言的“社团章程”必须继承 mlir::Dialect,才会获得添加 Operations、Types、Attributes 和 Interfaces 的能力。Op、Type、Attr 都是方言内部成员,不是方言本身的基类。
Q3. 把自定义方言放进注册表的典型调用是?
- A.context.loadAllAvailable()
- B.registry.insert<MyDialect>()← 正确答案
- C.op.setDialect<MyDialect>()
- D.pass.addRequired<MyDialect>()
答案:B。registry.insert<MyDialect>() 把方言类插入 DialectRegistry,相当于在总名册上登记。context.loadAllAvailable() 是加载所有已注册方言,不是插入操作;C 和 D 不是标准注册方式。
Q4. 描述自定义操作的语法和验证规则,推荐使用哪种表格驱动方式?
- A.ODS操作描述规范← 正确答案
- B.JSON配置规范文件
- C.YAML配置规范文件
- D.XML配置规范文件
答案:A。MLIR 使用 ODS(Operation Definition Specification)声明操作,通过 mlir-tblgen 自动生成 C++ 代码。JSON、YAML、XML 都不是 MLIR 官方的操作描述方案。
Q5. 方言降级时,用来声明哪些方言在目标中是合法的对象是?
- A.ConversionTarget转换目标← 正确答案
- B.RewriterBase重写器基类
- C.DialectRegistry方言注册表
- D.MLIRContext全局上下文
答案:A。ConversionTarget 告诉转换框架哪些操作与方言可以保留,哪些需要被改写,就像活动方案里标出哪些社团可以直接参加联合演出。
Q6. 让一个已注册方言被 MLIRContext 实际加载的调用是?
- A.pass.addRequiredDialect
- B.context.loadDialect<MyDialect>← 正确答案
- C.module.appendDialect
- D.builder.insertDialect
答案:B。context.loadDialect<MyDialect>() 显式加载方言到上下文中,使它可被后续 IR 和 Pass 使用。其他选项不是 MLIR 的标准方言加载 API。
Q7. 自定义 Attribute 注册时最关键的要求是?
- A.必须包含唯一id且自动注册
- B.必须继承Type且实现克隆
- C.必须实现解析打印并注册库← 正确答案
- D.必须与每个Operation配对使用
答案:C。属性要能参与 IR 文本的读写,就必须实现解析与打印,并通过 addAttributes 注册进方言。Attribute 不继承 Type,也不需要与 Operation 一一配对。
Q8. 一个 MLIR 自定义方言通常登记哪些要素?
- A.自定义Operation操作← 正确答案
- B.自定义Type类型体系← 正确答案
- C.自定义Attribute属性← 正确答案
- D.自定义Module模块名
答案:ABC。方言四元组包含操作、类型、属性和接口。模块名是 IR 顶层单元的概念,不属于方言注册内容,因此 D 不正确。
Q9. 方言转换 Pass 通常涉及以下哪些组件?
- A.ConversionTarget目标设置← 正确答案
- B.RewritePattern重写规则← 正确答案
- C.applyPartialConversion执行← 正确答案
- D.OpVerifier操作验证器
答案:ABC。转换 Pass 需要用 ConversionTarget 声明合法目标,用 RewritePattern 描述改写规则,最后用 applyPartialConversion 或 applyFullConversion 执行转换。操作验证器是常规 IR 验证组件,不是方言转换专用组件。
Q10. 关于方言动态加载,下列说法正确的是?
- A.未注册的方言不会参与编译← 正确答案
- B.相同方言可以重复注册多次
- C.注册表可在上下文构造前注入← 正确答案
- D.插件机制支持运行时扩展方言← 正确答案
答案:ACD。方言必须先注册再加载,否则上下文不认识它;DialectRegistry 可在 MLIRContext 构造前注入,第三方插件也利用这一机制做运行时扩展。同一个方言在上下文中只会有一个实例,重复注册会引发冲突,因此 B 错误。