框架集成：MLIR工具链

可视化演示

展示MLIR工具链核心机制：多级中间表示逐层降级、优化/转换/折叠等工具链流水线、MLIR多级IR带来的性能提升，通过交互理解MLIR工具链原理。

工具链全景

傍晚六点，转运中心的传送带嗡嗡转动。收货口刚卸下的包裹贴着五花八门的面单，可几分钟后，它们都变成了同一套条形码，在传送带上匀速前进。你问调度员秘诀，他指了指头顶的大屏：“我们不是认识每一张面单，而是把'任何格式'先翻译成'中心格式'，再逐级翻译成'派送格式'。”MLIR 工具链，就是 AI 编译器世界里的这套传送系统。

“中心格式”就是 MLIR 的多级中间表示。之所以叫“多级”，是因为同一份程序可以同时以不同抽象等级存在：高层还能看出“这是一个卷积”，低层已经变成“读取这几块内存、做乘加、写回”。每一级被封装成一种 dialect（方言）。框架集成要做的第一件事，就是把 PyTorch、TensorFlow、ONNX 手里的计算图接进来，送上这条传送带。

整条工具链由三个环节组成：前端导入（frontend）把框架图转换为高层方言；中端优化由 mlir-opt 驱动一系列 pass，把高层方言逐级下降（lowering）到更低层；后端导出由 mlir-translate 对接 LLVM 等后端，最终生成机器码。这套设计理念可以概括为两点：可扩展——新硬件只需新增方言与转换规则；渐进式下降——每层只做自己擅长的事。这条“逐级下降”的链条可以用一句公式概括：

$$\displaystyle \text{graph} \;\rightarrow\; \text{tosa} \;\rightarrow\; \text{linalg} \;\rightarrow\; \text{affine} \;\rightarrow\; \text{llvm} \;\rightarrow\; \text{asm}$$

1. graph：框架侧的高层计算图，例如 PyTorch 导出的 torch 算子序列，此时还带着 Python 层语义。

$$\text{graph}$$

2. tosa：统一面单层，把不同框架的算子标准化成可移植的张量算子集合。

$$\text{tosa}$$

3. linalg：结构化线性代数层，算子变成带显式循环结构的张量运算，暴露出可优化的形状信息。

$$\text{linalg}$$

4. affine：循环结构显式化，借助仿射映射精确分析数组访问与依赖关系。

$$\text{affine}$$

5. llvm：下降到 LLVM 方言，操作与 LLVM IR 指令一一对应，是交给后端前的最后一站。

$$\text{llvm}$$

6. asm：最终产物——目标机器的汇编/机器码，由 LLVM 后端生成。

$$\text{asm}$$

从收货口到装车口，包裹每过一个环节就换一张更精确的标签。框架集成不是把某个框架“宠上天”，而是给所有框架一条同等的传送带——这正是 MLIR 工具链最迷人的地方。

Dialect方言体系

转运中心里并不只有一种标签。收货部用“大区编号”，仓储部用“库位层架”，派送部用“街道门牌”。同一个包裹，在不同部门眼里有不同的“方言”，但谁都不会搞混，因为每套标签都贴在同一个包裹上。

MLIR 的 dialect（方言）正是这样的标签体系。一个方言定义了一组合法的 Operation（操作）、Type（类型）和 Attribute（属性）。比如 func 方言描述函数骨架，arith 方言描述加减乘除，scf 方言描述结构化控制流，tosa 方言描述可移植张量算子，linalg 方言描述带有循环结构的线性代数运算，llvm 方言则直接对应 LLVM 指令。每一个方言，都是一套自洽的“词汇表”。

方言机制让 MLIR 变成乐高积木：任何团队可以自定义方言，再声明它与其他方言之间的转换规则，完全不必改动中心本身。以 affine 方言为例，它把数组访问限定为仿射映射——下标与地址之间必须是线性关系，这样编译器才能精确分析依赖、做循环变换。仿射映射的一般形式是：

$$\displaystyle f(d_0, d_1, \ldots, d_{k-1}) = \alpha_0 d_0 + \alpha_1 d_1 + \cdots + \alpha_{k-1} d_{k-1} + \beta$$

1. d0…dk-1：循环变量或数组下标，它们是仿射映射的输入参数。

$$d_0, d_1, \ldots, d_{k-1}$$

2. α0…αk-1：各下标的系数，对应内存布局中的 stride（步长），例如行优先时行下标系数等于列数。

$$\alpha_i d_i$$

3. Σ 求和：把每一维的贡献累加，得到一维内存偏移的核心部分。

$$\sum_{i=0}^{k-1} \alpha_i d_i$$

4. β：常数偏移量，对应数组起始位置的修正（如偏置、padding）。

$$+ \beta$$

仓储部往货架上贴“库位层架”，派送部往包裹上贴“街道门牌”，标签虽多，但每套都精确服务于一个环节。MLIR 的高明之处，就是让这些标签在同一个包裹上共存，并且随时可以互相翻译。

核心工具命令

中控室里并排摆着三台机器。调度员每天的工作就围着它们转：第一台是“分拣程序控制器”，负责按顺序执行一道道分拣工序；第二台是“翻译机”，负责把外来单据翻译成内部格式；第三台是“模拟演练台”，在正式发车前先试跑一遍路线。

这三台机器对应 MLIR 工具链的三个核心命令：mlir-opt、mlir-translate、mlir-cpu-runner。mlir-opt 是主将，读取 .mlir 文件，按命令行参数指定的顺序执行 pass，输出优化后的 .mlir；mlir-translate 是翻译官，在 MLIR 内部表示与外部格式（LLVM IR、SPIR-V、Graphviz 文本等）之间互译；mlir-cpu-runner 是演练台，不经过完整编译链也能直接执行 MLIR 模块，适合快速验证。调试时还可以给 mlir-opt 加 --mlir-print-ir-after-all 或 --debug-only 这类调试选项，观察每一道工序之后的 IR 长相。

一条典型的命令形如 mlir-opt file.mlir --tosa-to-linalg --linalg-tile --convert-linalg-to-affine -o out.mlir，每个 flag 就是一个 pass。pass 的执行顺序决定了优化效果，这本质上是一个函数复合：

$$\displaystyle \mathrm{IR}_n = P_n \circ P_{n-1} \circ \cdots \circ P_1 (\mathrm{IR}_0)$$

1. IR0：mlir-opt 读入的初始中间表示，对应传送带起始处的原始包裹。

$$\mathrm{IR}_0$$

2. P1…Pn：按顺序注册的 pass，如 --tosa-to-linalg、--linalg-tile。

$$P_1, P_2, \ldots, P_n$$

3. ∘ 复合：表示“先执行 P1，再执行 P2”，相邻 pass 的输入输出首尾相连。

$$P_n \circ P_{n-1} \circ \cdots \circ P_1$$

4. IRn：最终产物，所有工序完成后的 IR，可直接交给翻译机或后端。

$$\mathrm{IR}_n$$

分拣工序必须“先分大区、再分街道、最后分楼栋”，顺序错了包裹就会迷路。mlir-opt 的存在，就是让你把这道顺序牢牢抓在自己手里。

框架前端集成

转运中心的收货口分成好几条车道：1 号口是顺丰（PyTorch），2 号口是京东（TensorFlow），3 号口是邮政小包（ONNX）。各家车道形状不同，但卸下来的包裹，最后都要先贴上一张行业通用的“统一面单”——TOSA。

这就是框架前端集成做的事。Torch-MLIR 项目负责把 PyTorch 模型经 torch 方言逐步合法化到 MLIR 高层方言；TensorFlow 官方通过 TOSA 或 MLIR-HLO 路径接入；ONNX 模型则借助 importer 导入。TOSA（Tensor Operator Set Architecture）是一套精心设计的“可移植张量算子”方言，像统一面单一样，让不同框架的算子先归一化，再进入 MLIR 的优化流水线。

从数据结构看，前端集成要完成一次“图到模块”的映射：计算图的每个节点（算子）变成一个 dialect 操作，每条边（张量）变成操作数或结果，整张图被包进 func 方言的函数体。可以记作：

$$\displaystyle \mathcal{G} = (\mathcal{V}, \mathcal{E}) \;\longmapsto\; \mathrm{module}\{\, \mathrm{func}\{\, \mathrm{op}_v : v \in \mathcal{V} \,\} \,\}$$

1. G=(V,E)：计算图，V 是算子节点，E 是张量数据边。

$$\mathcal{G} = (\mathcal{V}, \mathcal{E})$$

2. module：MLIR 的顶层容器，相当于一个编译单元，容纳多个函数。

$$\mathrm{module}$$

3. func：函数方言，把一张子图包装成可调用、可重用的函数。

$$\mathrm{func}$$

4. op_v：把图节点逐个映射为方言操作，边则映射为操作数或结果。

$$\mathrm{op}_v : v \in \mathcal{V}$$

面单格式再乱，只要进了统一面单通道，后面的传送带就对它一视同仁。前端集成把“各家快递”变成“统一包裹”，MLIR 工具链的价值，正是从这一刻开始积累。

转换下降机制

面单统一之后，真正的重头戏才开始：把“这是一箱易碎玻璃器皿，送往朝阳区”翻译成“送往 XX 街道 YY 号，轻拿轻放，置 3 号货架”。越接近派送，指令就越具体。方言下降（dialect lowering）做的就是这件事——从“描述做什么”逐级降到“描述怎么做”。

MLIR 的下降由两个核心机制驱动：RewritePattern（重写模式）和 ConversionDriver（转换驱动）。RewritePattern 声明“看到什么样的操作，就替换成什么样的操作序列”；ConversionDriver 则按合法化规则（legality）反复应用模式，直到所有操作都落在目标方言内。合法化规则告诉驱动：转换目标（conversion target）里哪些操作是“天生的好公民”（already legal），哪些必须被转换。

一个高层操作往往展开成多个低层操作，这种一对多映射是下降的基本形态：

$$\displaystyle \text{Op}(x_1, x_2, \ldots, x_n) \;\Longrightarrow\; \text{op}_1(\cdots),\; \text{op}_2(\cdots),\; \ldots,\; \text{op}_m(\cdots)$$

1. Op(x1,…,xn)：源方言中的一个高层操作，带若干操作数。

$$\text{Op}(x_1, x_2, \ldots, x_n)$$

2. ⟹：模式重写符号，表示匹配成功后触发整体替换。

$$\Longrightarrow$$

3. op1…opm：替换产物，一组合法的低层操作，共享或重排原操作数。

$$\text{op}_1(\cdots),\; \text{op}_2(\cdots),\; \ldots,\; \text{op}_m(\cdots)$$

4. m：展开系数，由模式设计者决定，体现“一词多义”的拆分逻辑。

$$m$$

从“送往朝阳区”到“送到 3 号货架”，每一步翻译都在增加可执行性、减少歧义。转换下降机制就是这条翻译流水线的“语法规则”，保证任何一句话都不会翻出歧义。

优化策略

面单落地后，分拣员开始琢磨怎么装车更省。大件拆成小件，方便塞进车厢缝隙——这是 tiling（循环分块）；几个小包捆成一捆，少贴几张签——这是 fusion（算子融合）；一辆大车一次拉走一条街的货——这是 vectorization（向量化）；按派送路线重新码放货架——这是 memory layout（内存布局）。

这些优化大多发生在 linalg/affine 等结构化方言层，因为此时循环和访存还“看得见”。循环分块把一个大循环切成小块，让数据块能留在 CPU 缓存里，避免反复访问内存；算子融合把 element-wise 链合并成一个算子，减少中间张量的写出和读入；向量化让一条 SIMD 指令同时处理多个数；内存布局则选择行优先、列优先或分块布局，匹配硬件的访存习惯。

分块的本质是对循环空间的重参数化。一个双重循环 (i, j) 被改写为四重循环：

$$\displaystyle i = i_0 \times T_i + i_1 \quad (0 \le i_1 < T_i), \qquad j = j_0 \times T_j + j_1 \quad (0 \le j_1 < T_j)$$

1. i, j：原始大循环的循环变量，遍历整个数据空间。

$$i,\; j$$

2. i0, j0：外层循环变量，表示当前处理第几个数据块。

$$i_0,\; j_0$$

3. Ti, Tj：分块尺寸，决定每块处理多少数据，是调优的关键旋钮。

$$T_i,\; T_j$$

4. i1, j1：内层循环变量，表示块内的局部偏移。

$$i_1,\; j_1$$

5. 0 ≤ i1 < Ti：约束条件，保证块内偏移不越界，块与块不重叠。

$$0 \le i_1 < T_i$$

装车方案千变万化，但原则只有一条：让搬运工少跑路、让车厢不空载。MLIR 把优化放在结构化方言层做，正因为此时“车厢和路线”还清晰可见，编译器才有机会施展这些装车智慧。

代码生成链路

最后一公里，包裹被交到末端快递员手里。快递员不关心“易碎”“保鲜”这些业务词，他只认两件事：送到哪个小区，几号楼几单元。LLVM 就是那位快递员——它只认最底层的指令，但胜在路熟、覆盖广。

代码生成链路把 llvm 方言翻译成 LLVM IR，再由 LLVM 后端生成目标机器码，最后链接运行时完成交付。这条链路有两种打开方式：JIT（Just-In-Time，运行时编译再执行）适合快速验证和动态形状场景；AOT（Ahead-Of-Time，预先编译成可执行文件或 .so）适合追求稳定部署的生产环境。mlir-cpu-runner 就是典型的 JIT 路径。交付后还可以借助 perf 或 MLIR 的 timing pass 做性能分析，量化每一步优化到底值多少。

到达 LLVM 层后，所有看似“高级”的数组访问都必须退化为赤裸裸的地址计算。行优先布局下，二维下标 (d0, d1) 的内存地址是：

$$\displaystyle \mathrm{addr}(d_0, d_1) = \mathrm{base} + (d_0 \times S_0 + d_1 \times S_1) \times \mathrm{size}$$

1. d0, d1：数组的逻辑下标，例如卷积输出特征图的行和列。

$$d_0,\; d_1$$

2. S0, S1：各维度的 stride，行优先时 S0 = 每行元素数，S1 = 1。

$$S_0,\; S_1$$

3. size：单个元素的字节数，float32 为 4，int8 为 1。

$$\mathrm{size}$$

4. base：数组在内存中的起始地址，由运行时分配或静态指定。

$$\mathrm{base}$$

从框架到机器码，包裹终于走完了整条传送带：统一面单、逐级细化、装车优化、最后交给只认门牌的快递员。MLIR 工具链的价值，就是让这一切“对任何快递公司、任何小区”都成立。

AI 应用场景

Python 演示

# -*- coding: utf-8 -*-# 一个极简的 MLIR 工具链模拟器:# 用字符串形式的 IR 演示 方言下降 与 算子融合 两个 pass class Op:    def __init__(self, name, args, result):        self.name = name        self.args = list(args)        self.result = result     def __repr__(self):        return "%" + self.result + " = " + self.name + "(" + ", ".join(self.args) + ")"  LOWER_MAP = {    "tosa.add": "linalg.add",    "tosa.mul": "linalg.mul",    "tosa.relu": "linalg.relu",}  def lower_tosa(module):    lowered = []    for op in module:        if op.name in LOWER_MAP:            lowered.append(Op(LOWER_MAP[op.name], op.args, op.result))        else:            lowered.append(op)    return lowered  def fuse_elemwise(module):    fused = []    i = 0    while i < len(module):        op = module[i]        if op.name == "linalg.add" and i + 1 < len(module):            nxt = module[i + 1]            if nxt.name == "linalg.relu" and nxt.args[0] == op.result:                fused.append(Op("linalg.add_relu", op.args, nxt.result))                i += 2                continue        fused.append(op)        i += 1    return fused  def run_pipeline(module, passes):    current = module    for pass_name, pass_fn in passes:        current = pass_fn(current)        print("---- pass:", pass_name)        for op in current:            print(op)    return current  if __name__ == "__main__":    src = [        Op("tosa.add", ["a", "b"], "c"),        Op("tosa.relu", ["c"], "d"),        Op("tosa.mul", ["d", "e"], "f"),    ]    print("== 原始 TOSA 模块 ==")    for op in src:        print(op)    print()    run_pipeline(src, [("lower_tosa", lower_tosa), ("fuse_elemwise", fuse_elemwise)])

1. class Op:定义 IR 的最小单元：操作名、操作数、结果。
2. LOWER_MAP = {一张“方言映射表”，模拟从 tosa 到 linalg 的下降规则。
3. def lower_tosa(module):第一个 pass：把高层 tosa 算子逐一下降为 linalg 算子。
4. if op.name in LOWER_MAP:模式匹配：命中下降规则的操作被替换，未命中的原样保留。
5. def fuse_elemwise(module):第二个 pass：把 add 与紧随的 relu 合并成 add_relu。
6. if nxt.name == "linalg.relu" and nxt.args[0] == op.result:检查数据依赖，确保后面的 relu 确实消费 add 的结果。
7. def run_pipeline(module, passes):模拟 mlir-opt：按注册顺序逐个执行 pass。
8. if __name__ == "__main__":入口，构建一个小模块并用两张 pass 串联。

常见误区

面试题