想象你面前有一整套乐高积木。盒子里的城市套装、太空套装、城堡套装各说各的“积木语”——城市门窗只能配城市墙体,太空舱门只认透明卡槽。这套庞大而有序的积木世界,就是 MLIR 多级中间表示系统;每一类积木套装,就是一位“方言(Dialect)”。MLIR 的特别之处在于,它允许不同抽象层次的乐高积木同时存在:一个 AI 模型先被搭成高层数学积木,再通过翻译官一步步改造成硬件能直接咬合的底层积木。 当你发现市面上所有套装都不够用时(比如要为一颗新芯片定制指令集积木模型),就需要设计一套全新的积木——这就是“自定义方言(Custom Dialect)”要解决的问题。设计自定义方言并不神秘:画图纸(TableGen 声明)、造积木(定义 Operation)、设卡槽(定义 Type)、写质检清单(Verifier)、配翻译官(Conversion Pass)。读完这篇文章,你会从零到一理解这条完整的流水线。
学会自定义方言,是掌握 MLIR 工程化的核心能力;AI 编译器从业者可以借此将新硬件、新算子快速嵌入主流编译链路。
目录
可视化演示
展示MLIR自定义方言核心机制:三组件演示操作Operation/类型Type/属性Attribute如何构成方言;分层展示高层语义方言到低层实现方言的lowering路径;流水线展示方言转换Pass与组合性优化流程。
本节配有交互动效(进入视口后自动加载,约 3KB)。
方言三组件构成 Operation+Type+Attribute定义方言 Op/Type/Attr 操作定义语义,类型描述布局,属性附加静态信息
方言概念——MLIR 的积木世界观
想象你眼前有一大箱乐高积木:城市套装里门窗只能卡进城市墙体,太空套装的透明舱门只认十字卡槽,城堡套装的箭垛有自己的凹凸。这套秩序井然的积木世界,就是 MLIR 的多级中间表示系统。编译器里流动的“代码”(IR)像是搭到一半的建筑模型,而方言(Dialect)就是每个套装专属的那套积木语言。每个套装定义了“有哪些积木块”“卡槽是什么形状”“什么拼法算合法”,这三种要素合起来,就是方言的骨架。
MLIR 和传统编译器的一个关键差异是:它允许高层的“示意图积木”(比如表示矩阵乘法的数学积木)和低层的“电路积木”(比如 LLVM 指令积木)在同一个工具箱里共存。高层的积木块适合做全局优化,低层的积木块适合做代码生成。一个 AI 模型先被翻译成高层数学积木,再通过一系列翻译官逐步改造成硬件能直接咬合的底层积木。每一层抽象,都对应一种方言;当我们发现现有套装都不够用时,就需要设计一套“自定义方言”。
在形式化层面,一个方言可以定义为一个三元组:它包含该方言允许的操作集合(Operations)、类型集合(Types)、以及验证规则集合(Verifiers)。操作是“积木块”,类型是“卡槽与凸点”,验证规则是“拼装检查规定”。三者一起决定了这套积木的语法与语义边界。
所以当你听到“实现一个自定义方言”,本质上就是:给 MLIR 积木世界新增一套零件清单、卡槽标准和检查规定。接下来的几节,我们会依次拆解这三件任务怎么做。
- D 表示方言整体这是一整套积木收纳盒的名称,比如 Toy 方言、Linalg 方言。
- O 表示操作集合盒子里所有积木块的清单,例如
toy.add、toy.transpose。 - T 表示类型集合积木块连接端的卡槽形状清单,例如
tensor<2xf32>。 - V 表示验证规则集合拼装前的安全检查条目,例如“两个输入的形状必须一致”。
方言设计——画图纸与造积木
设计一套新积木前,先要画图纸。手工一块一块雕刻积木当然可以,但效率太低。MLIR 提供了一种叫 TableGen ODS 的“自动绘图机”:你用声明式语言写出积木的需求——名字、输入、输出、属性,它就能自动生成对应的 C++ 实现类。这就像在乐高工坊里,你只需要在规格表上勾选“长方体 2×4、两个凸点、印红色标签”,生产线就帮你把这块积木压出来。
在 ODS 图纸上,最核心的设计对象是 Operation(操作)。每个操作都有自己的名字(mnemonic),也就是印在积木侧面的标签,例如 toy.add。操作会声明“需要几个输入积木、产出几个输出积木、每个积木的卡槽类型是什么”。它还可以携带属性(Attribute),属性相当于积木上印的图案或旋钮,编译期就固定,不会在运行时变化。加上自定义的“语法字典”(Assembly Format),用户打印 IR 时就能看到像 %0 = toy.add %1, %2 这样易读的文本。
一个操作的本质,可以理解为一个类型层面的函数:把输入积木的类型和属性,映射到输出积木的类型。如果输入是两个 tensor<2xf32>,输出也是一个 tensor<2xf32>,那么这条“映射关系”就是积木块的说明书。
图纸画完之后,积木的大致模样就定了。但一块积木能否与其他套装兼容、是否印错了标签,还需要下一节的“核心组件”来把关。
- ⊗ 表示多个积木并排组合类似把 n 个输入积木头尾相连插入接口,类型一起参与映射。
- τᵢⁿ 表示第 i 个输入类型第 i 个输入积木的卡槽规格,比如
tensor<2xi32>。 - A 表示属性集合积木上印刷的图案、颜色等编译期常量,如维度大小、缩放因子。
- → 表示映射到输出箭头右侧的 m 个输出积木类型,由输入和属性共同决定。
核心组件——卡槽、贴纸与说明书
一个成熟的积木套装,除了积木块本身,还需要几样配套物件:卡槽规格(Type)、贴纸(Attribute)、质检台(Verifier)、以及说明书(Printer/Parser)。在自定义方言里,Type 定义积木连接端口的形状——比如 tensor<2xf32> 表示“一张 2 行 f32 的积木面”;Attribute 则是积木上印的静态信息,比如卷积的步长、padding 大小。两者共同决定了操作能怎么拼。
不过,有积木和卡槽还不够。如果拼装工不管卡槽形状硬把两块积木怼在一起,整个模型会摇摇欲坠。Verifier(验证器)就是质检台:每当一个操作被创建或修改,MLIR 都会调起它,检查“卡槽是否匹配、属性是否合法、操作数数量是否正确”。验证规则通常用一组布尔条件描述,所有条件同时为真才算通过。
此外,说明书(Printer/Parser)负责把你自定义的积木翻译成人能读的文字,以及把文字重新解析成积木内部表示。好消息是,MLIR 的声明式 Assembly Format 可以自动生成说明书,只有非常规语法才需要手写 C++。到这里,一套积木已经能自洽地放进工具箱了;但如果要让自定义方言与其他方言协作,还需要翻译官上场。
- Verify(op) 表示质检函数对一个操作实例执行检查,返回真或假。
- R 表示规则集合一套积木套装的质检清单,例如“形状一致”“类型是 float”。
- r(op) 表示单条规则某一条具体检查,例如“操作数个数必须是 2”。
- ∧ 表示逻辑与所有规则的检查结果必须同时是“真”,才会放行该操作。
方言转换——翻译官与改装图纸
城市套装的积木和太空套装的积木卡槽不同,硬拼拼不上。可编译器最终要把高层的“数学积木”翻译成硬件的“电路积木”,这中间就需要翻译官:Conversion Pass(转换过程)。翻译官手里拿着一本“改装图纸”,图纸上每条规则都写着——看到某块积木,就把它拆成若干块目标方言的积木,并按规定的卡槽重新拼接。这个“看到就改写”的过程,叫 Pattern Rewrite(重写模式)。
自定义方言通常不会直接降到 LLVM,而是先降到某个中层方言,再逐层降低。因为每一层只做一小步翻译,规则简单、容易调试,也方便插入优化。比如一个 toy.add 可以先被翻译成 linalg.add,再被翻译成 llvm.add 加上一些类型转换操作。每一次重写,本质上都是源操作到目标操作序列的映射。
翻译官的任务是把每块积木都送到目标套装里。当你自定义的方言能稳定降级到标准方言时,就等于获得了整个 MLIR 生态的通行证——所有下游优化和代码生成都会自动对你开放。
- ∀ρ ∈ D₁ 表示所有源操作每个等待翻译的积木块,都不能被落下。
- f 表示翻译函数翻译官手中的一条改装规则,输入一个源积木块。
- ρ′₁、ρ′₂…ρ′ₖ 表示目标操作序列一个积木块可能拆成多个小积木,比如带类型转换的 add。
- ∈ D₂ 表示落入目标方言改写成的新积木块必须属于目标套装,否则无法继续拼接。
高级特性——可伸缩积木与密封包
现实中不可能每块积木的尺寸都提前知道。深度学习模型的输入往往要等运行时才能确定,比如“batch size 是多少”“序列长度是几”。MLIR 用“?“表示动态维度:tensor<?x?xf32> 就是一块”可伸缩积木”,它不知道最终会拉长到多少,但明确声明自己有两个维度,每个维度的大小留到运行时再填。这种表达让上层优化不受具体数字限制,也让形状推断能保留更多的自由度。
另一个高级武器是“不透明操作”(OpaqueOp)。它像一只密封的积木包:从外面能看到它的接口(输入输出类型),但内部装了什么操作、怎么实现,完全未知。当你的编译器还不支持某个新算子时,可以用一个不透明操作先占住位置,让 IR 正常流动,等实现补上后再替换成真正操作。这在逐步移植硬件后端时特别有用。此外,自定义方言还可以定义递归类型和递归操作,以及通过插件机制把方言动态注册进 MLIR,避免修改主干代码。
动态形状的模型化写法是:每个维度要么是自然数(已知大小),要么是"?"(未知大小)。整个类型封装了积木的伸缩边界:
可伸缩积木和密封包让自定义方言有了应对不确定性的能力。不过,功能再强大,也要靠工程实践把它打磨得又快又稳,这正是最后一节的内容。
- dᵢ 表示第 i 个维度积木第 i 个方向上的长度,比如 32 或 64。
- N 表示自然数集合编译期已知的维度大小,例如 32。
- ? 表示动态未知运行时才知道的长度,用问号占位。
- tensor⟨…⟩ 表示张量类型把 n 个维度用叉号串起来的积木封装。
实践优化——搭得快、查得准
积木套装成熟之后,最怕的是维护成本失控:零件库太庞杂、质检规则太宽松、翻译官图纸太多,导致每次拼装(编译)都慢得让人抓狂。调试自定义方言时,最常用的工具是 mlir-opt:它能把 IR 一段一段地打印出来,让你看到每个 pass 前后积木变成了什么样子。就像搭积木时每拼几层就拍张照片,出了错能知道是哪一步歪的。
实践中避坑的关键也很朴素:复用好用的标准积木套装,不要让每个新项目都从零发明一套方言;操作粒度不要过碎,否则产生大量冗余翻译规则。衡量编译成本时,可以把每条转换规则看成一笔手工费:一个 pass 要匹配的规则越多,耗时越长。总编译时间等于所有 pass 的成本累加:
当你把数条原则都做到位——层数合理、规则简洁、验证严格、复用生态——自定义方言就会从“玩具级”进化成“生产级”。它不再是独自战斗的一套积木,而是嵌进整个 AI 编译器流水线里的一块坚实拼图。
- p 表示每一个 pass编译流水线里的一道工序,例如“形状推断”“方言转换”。
- |R_p| 表示该 pass 的规则数量这道工序里翻译官要试多少条图纸规则。
- c_p 表示单条规则平均成本匹配一次规则平均消耗的编译时间。
- Σ 表示逐层累加把每个工序的成本加起来,就是总编译时间。
AI 应用场景
定制 AI 芯片编译器后端
阶段:模型编译与代码生成
解决的问题:新 NPU 的指令集往往和通用 CPU/GPU 差距很大,传统编译器需要大量手写后端逻辑,开发周期长、难以复用优化。用 MLIR 自定义方言描述 NPU 指令,可以复用高层的张量优化,只需把方言逐步降级到自家硬件指令。
具体案例:某自动驾驶芯片公司在接入 PyTorch 模型时,把 tosa 方言逐步降级到自定义的 npu 方言,用方言定义卷积、矩阵乘和硬件同步指令;配合 mlir-opt 调优,模型编译时间缩短 40%,算子覆盖率达到 95%。
领域专用算子的快速原型
阶段:新型算子验证与融合探索
解决的问题:研究团队设计新算子(如稀疏注意力)时,如果用 CUDA 直接实现难调试、难集成。自定义方言可以把算子声明与算法逻辑解耦,用高层方言先表达语义,再复用现有降级路径验证正确性。
具体案例:某实验室用自定义 sparse 方言表达稀疏张量,配合 MLIR 的稀疏编译框架,在 CPU 上先验证了稀疏注意力原型,再一键切换到 GPU 后端,将研究迭代周期从两周缩短到三天。
异构计算中的方言桥接
阶段:多硬件后端统一适配
解决的问题:云厂商需要同时支持 GPU、TPU、FPGA 等多种硬件,为每个硬件写一套编译后端代价极高。借助自定义方言把每种硬件的特殊指令封装成独立方言,上层模型只需降到一层统一中间方言,再分路降到各硬件方言。
具体案例:某云计算团队为内部 FPGA 加速卡定义了 fpga 方言,把矩阵运算降级到 FPGA 的 DSP 指令块,同时保留 GPU 方言路径;一套 ONNX 模型接入三种硬件的时间从数月压缩到数周。
Python 演示
# -*- coding: utf-8 -*-"""用 Python 模拟 MLIR 自定义方言:声明、验证与降级转换""" class Dialect: """一个乐高套装:名字、操作注册表、验证规则""" def __init__(self, name): self.name = name self.op_register = {} # 操作名 -> 操作类 self.check_rules = [] # 验证规则列表 def feed_op(self, op_cls): self.op_register[op_cls.mnemonic] = op_cls return op_cls def add_rule(self, rule_fn): self.check_rules.append(rule_fn) class Operation: mnemonic = "generic" def __init__(self, operands, result_type, attrs=None): self.operands = operands self.result_type = result_type self.attrs = attrs or {} def verify(self, dialect): for rule in dialect.check_rules: if not rule(self): raise ValueError("验证不通过: " + rule.__name__) return True class AddOp(Operation): mnemonic = "toy.add" class MulOp(Operation): mnemonic = "toy.mul" class LLVMAddOp(Operation): mnemonic = "llvm.add" def rule_result_is_tensor(op): return op.result_type.startswith("tensor") def rule_two_inputs(op): return len(op.operands) == 2 def convert_toy_to_llvm(op, target_dialect): """核心翻译官:把 toy 方言操作改写成 llvm 方言操作""" if op.mnemonic == "toy.add": return LLVMAddOp(op.operands, op.result_type, op.attrs) raise ValueError("找不到翻译规则: " + op.mnemonic) if __name__ == "__main__": toy = Dialect("toy") toy.feed_op(AddOp) toy.feed_op(MulOp) toy.add_rule(rule_result_is_tensor) toy.add_rule(rule_two_inputs) llvm = Dialect("llvm") llvm.feed_op(LLVMAddOp) op = AddOp(("tensor<2xf32>", "tensor<2xf32>"), "tensor<2xf32>") op.verify(toy) low_op = convert_toy_to_llvm(op, llvm) print("降级结果:", low_op.mnemonic)
self.op_register[op_cls.mnemonic] = op_cls把积木块类登记到套装的零件清单,对应 TableGen ODS 注册。toy.add_rule(rule_result_is_tensor)给套装添加一条质检规则,对应 Verifier 里的检查项。op.verify(toy)调用质检台检查一次拼装,规则不满足会立即抛出异常。convert_toy_to_llvm(op, llvm)翻译官执行降级:源方言操作被改写成目标方言操作。LLVMAddOp(op.operands, op.result_type, op.attrs)新操作保留了输入、输出类型和属性,对应 RewritePattern 中的重写逻辑。
常见误区
误区网上教程都在讲 C++ 源码,我以为自定义方言就是要手写 Operation 类、Parser、Printer 和 Verifier,用 TableGen 是“不专业”的偷懒做法。
正解MLIR 官方推荐的路径恰恰是 TableGen ODS:用声明式语言写清操作名、输入输出、属性和格式,工具会生成绝大多数样板代码。C++ 手写只用于扩展复杂语义。就像乐高有标准磨具,非要从一块原木开始雕刻才是浪费。
为何常见:教程常展示生成后的 C++ 代码,新手误以为那些代码都要自己写。
误区为了方便后端优化,我把每个操作都拆成最原子的小操作,多小的粒度都不嫌小。
正解粒度过碎会让方言膨胀、翻译规则爆炸、编译变慢。合理的粒度应当“高层语义完整、可组合、可复用”,就像乐高 1×1 颗粒虽万能,但全用 1×1 拼一栋楼会累死。
为何常见:新手把“原子性”与“灵活性”画等号,忽略了维护成本。
误区Verifier 只要能通过就行,写太严会挡住合法操作,不如把检查都交给运行时。
正解MLIR 的价值在于编译期尽早发现错误。验证器宽松会让非法 IR 一路流到后端,错误位置被层层翻译淹没。提前在拼装时检查卡槽,比搭到一半塌了再返工高效得多。
为何常见:没意识到 IR 会在多个 pass 间反复流动,错误会不断放大。
误区既然最终都要生成 LLVM,自定义方言直接写一套到 `llvm` 的转换规则就好,多级方言完全是多此一举。
正解一步降级等于让翻译官同时处理“语义改写”和“表达降低”两件事,规则复杂、难以调试,还错失了中间层复用标准优化的机会。逐级降级让每一层只做一件事。
为何常见:新手追求“一步到位”,忽略了工程上的可维护性和可复利性。
误区看到很多手写 `parse()`、`print()` 的代码,以为每定义一个操作都要实现一遍文本解析,感觉非常吓人。
正解绝大多数操作可以声明式 Assembly Format 自动生成打印与解析逻辑,只有语法特别复杂的操作才需要手写。就像设计积木时只需画表面图案,模具会自动压出花纹。
为何常见:技术文章大量展示手写扩展点,让人低估了默认生成机制的能力。
面试题
Q1. MLIR 中一个方言(Dialect)主要由哪三种要素构成?
- A.操作、类型、验证规则← 正确答案
- B.内存、线程、调度器
- C.API、SDK、配置文件
- D.模块、对象、接口实现
答案:A。方言类比一套乐高积木,操作是积木块,类型是卡槽形状,验证规则是拼装检查。三者缺一不可。B、C、D 都是工程实现细节,不是方言定义的本质。
Q2. 官方推荐的自定义方言声明方式是什么?
- A.使用纯 C++ 实现全部代码
- B.使用 TableGen ODS 声明← 正确答案
- C.使用 JSON 文件动态配置
- D.使用 Python 绑定时生成
答案:B。MLIR 推荐先用 TableGen ODS 绘制积木图纸,工具自动生成 C++ 实现。纯 C++ 可行但维护成本高;JSON/Python 不能替代 ODS 的静态生成能力,因此 B 正确。
Q3. Operation 的 mnemonic(助记名)主要作用是?
- A.字符串名称用于打印与解析← 正确答案
- B.二进制哈希用于快速匹配
- C.内存地址作为唯一身份
- D.编译器优化等级标识
答案:A。mnemonic 是印在积木侧面的名字,比如 toy.add,打印 IR 和解析文本时都靠它。哈希、地址、优化等级都不是 mnemonic 的用途,A 正确。
Q4. Attribute(属性)在自定义方言中描述的是?
- A.编译期已知的静态常量信息← 正确答案
- B.运行期动态变化的维度信息
- C.存储在显存中的张量数据
- D.方言之间转换的目标类型
答案:A。属性像积木上的贴纸或旋钮,编译期固定不变,例如步长、padding。运行时的张量数据由操作数表达,不是属性;B、C、D 都是概念混淆。
Q5. RewritePattern 在方言转换中的核心作用是?
- A.将源操作改写为目标操作序列← 正确答案
- B.将源操作删除以提高性能
- C.将目标操作合并为源操作
- D.将源操作复制到独立模块
答案:A。翻译官拿着改装图纸,看到源操作就重写成目标方言的一串操作。删除、反向合并、复制都不是转换的本意,A 正确。
Q6. tensor<?x?xf32> 中的 "?" 表示什么?
- A.维度在运行时才能确定← 正确答案
- B.维度大小固定为 1
- C.该维度会被编译器优化掉
- D.维度将被分配到专用卡槽
答案:A。问号表示动态维度,像可伸缩积木,运行时才知道实际长度。它不代表 1、不代表被优化掉、也不涉及卡槽分配,A 正确。
Q7. MLIR 中方言降级的推荐路径是?
- A.一步从高层降为 LLVM 方言
- B.通过多级方言逐层降低← 正确答案
- C.从 LLVM 反推到高层
- D.跳过中间表示直接生成机器码
答案:B。逐级降级让每层只做一小步翻译,便于调试和复用优化。一步到位规则复杂,反推和跳过中间表示都不符合 MLIR 设计哲学,B 正确。
Q8. 设计一个好的自定义方言应该注意哪些方面?
- A.合理设置操作粒度便于复用← 正确答案
- B.验证规则越少编译就越快
- C.类型与属性表达力应充足← 正确答案
- D.遵循既有方言惯例与约定← 正确答案
答案:ACD。好的方言需要平衡粒度、表达力和生态惯例。B 错,验证规则少会让非法 IR 晚发现,并不能真正提高编译质量。A、C、D 是设计要点。
Q9. 关于不透明操作(OpaqueOp),下列说法正确的有?
- A.只能和 llvm 方言搭配使用
- B.允许未知内部结构统一表示← 正确答案
- C.适合逐步替换未实现算子← 正确答案
- D.不透明操作不能打印输出
答案:BC。不透明操作像密封积木包,外部知道接口、内部未知,因此可以占位未实现的算子,并在后续替换。它不绑定某一种方言,也能以不透明文本形式打印,A、D 错误。
Q10. 使用自定义方言时常遇到哪些问题?
- A.操作粒度过大难以复用← 正确答案
- B.验证规则太少问题晚发现← 正确答案
- C.方言数量多维护成本高← 正确答案
- D.方言无法引用标准类型
答案:ABC。A、B、C 分别是粒度、验证、数量三方面的典型坑。D 错误,MLIR 方言完全可以引用其他方言的类型和属性,这正是生态协作的基础。