AI 编译器入门：TVM 与 MLIR 的原理、差异与实战- 惊觉

1. 为什么需要 AI 编译器

深度学习模型从论文走向生产，中间隔着一道巨大的鸿沟：同一份模型，要跑在 NVIDIA GPU、AMD GPU、ARM CPU、手机 NPU、甚至各种自研加速器上。如果每个后端都手写一套算子实现，工程量会呈指数级爆炸。

传统做法是「框架 + 算子库」：PyTorch 定义计算图，后端调用 cuDNN、MKL、NNAPI 等库。这套方案在通用硬件上很成熟，但遇到新硬件、新算子、新量化格式时，往往要等厂商补齐库支持，迭代周期长。

AI 编译器要解决的核心问题就是：把高层计算图自动、可优化地映射到多样化的硬件上。它把「写算子」变成「写编译规则」，让一次模型定义可以生成多后端的机器码。

一个典型的 AI 编译流程可以抽象为：

在这条链路上，TVM 和 MLIR 是当前最具代表性的两套基础设施，但它们的定位并不相同。

2. TVM：端到端的深度学习编译栈

2.1 TVM 的设计哲学

TVM 由陈天奇等人提出，最初的目标是「让深度学习框架的模型能高效跑在各种硬件上」。它的核心创新是计算与调度分离（compute / schedule separation）。

在 TVM 中，你只描述「算什么」，不描述「怎么算」：

import tvm
from tvm import te

# 计算：C[i, j] = sum_k A[i, k] * B[k, j]
n = 1024
A = te.placeholder((n, n), name="A")
B = te.placeholder((n, n), name="B")
k = te.reduce_axis((0, n), name="k")
C = te.compute(
    (n, n),
    lambda i, j: te.sum(A[i, k] * B[k, j], axis=k),
    name="C",
)

这段代码只声明了矩阵乘的数学语义，没有任何循环结构。真正的性能来自调度：

s = te.create_schedule(C.op)

# 分块，提升缓存命中率
bn = 32
i_outer, i_inner = s[C].split(C.op.axis[0], factor=bn)
j_outer, j_inner = s[C].split(C.op.axis[1], factor=bn)
k_outer, k_inner = s[C].split(k, factor=bn)

# 调整循环顺序，让内层访问连续
s[C].reorder(i_outer, j_outer, k_outer, i_inner, j_inner, k_inner)

# 向量化最内层
s[C].vectorize(j_inner)

# 多线程并行
s[C].parallel(i_outer)

同一份计算定义，换一套调度就能适配 CPU、GPU 或加速器。这种「一份计算、多份调度」的模式，是 TVM 能覆盖大量后端的关键。

2.2 TVM 的编译流水线

TVM 的完整链路大致分为四层 IR：

层级	名称	作用
高层	Relay / Relax	计算图表示，承载图级优化
张量级	TIR（Tensor IR）	显式循环与调度，承载算子级优化
底层	TIR 降级后的低级 IR	接近机器语义
目标	LLVM IR / CUDA C / 汇编	交给后端编译器

图级优化包括算子融合、常量折叠、死代码消除、布局变换等。算子级优化则依赖调度原语（split、reorder、tile、vectorize、parallel、bind）以及自动调优。

2.3 AutoTVM 与 Ansor：自动搜索调度

手写调度门槛高、调参耗时。TVM 提供了自动调优：

AutoTVM：基于模板的参数搜索，人工给出调度骨架，工具搜索分块大小、展开因子等参数。
Ansor（Auto-scheduler）：无需模板，自动生成调度空间并搜索，适合复杂算子。

from tvm import autotvm

@autotvm.template("tutorial.matmul")
def matmul_template(n, bn):
    A = te.placeholder((n, n), name="A")
    B = te.placeholder((n, n), name="B")
    k = te.reduce_axis((0, n), name="k")
    C = te.compute((n, n), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))

    s = te.create_schedule(C.op)
    cfg = autotvm.get_config()
    cfg.define_split("tile_i", C.op.axis[0], num_outputs=2)
    cfg.define_split("tile_j", C.op.axis[1], num_outputs=2)
    # ... 应用配置并返回
    return s, [A, B, C]

调优结果会存成日志，编译时直接查表，避免线上重复搜索。

2.4 TVM 的适用场景

TVM 的优势在于端到端：从模型导入（ONNX、PyTorch）到多后端代码生成，一条链路打通。它特别适合：

需要把模型部署到非主流硬件（自研 NPU、嵌入式芯片）；
需要针对特定 shape 做极致性能调优；
需要统一管理多后端的推理服务。

但 TVM 也有短板：自定义算子接入需要理解 TIR 与调度体系，学习曲线陡峭；图级 IR 与算子级 IR 之间的边界相对固定，跨层优化不够灵活。

3. MLIR：可复用的编译器基础设施

3.1 MLIR 要解决什么问题

MLIR（Multi-Level IR）由 LLVM 社区发起，出发点不是「做一个 AI 编译器」，而是「做一套可扩展的编译器基础设施」。

它观察到：不同领域（机器学习、量子计算、硬件描述、数据库查询）都需要 IR，但各自需要的抽象层级不同。如果每个领域都从零造 IR，会重复实现大量通用能力（类型系统、Pass 管理、模式匹配、诊断）。

MLIR 的答案是：提供一套 IR 元框架，让各领域定义自己的方言（Dialect）。

3.2 核心概念：Dialect、Operation、Pass

MLIR 中一切皆 Operation。一个 Operation 由名字、操作数、结果、属性、区域（Region）组成。方言是一组相关 Operation 的集合。

// 一个简单的 MLIR 片段
func.func @matmul(%A: tensor<128x256xf32>, %B: tensor<256x64xf32>) -> tensor<128x64xf32> {
  %C = linalg.matmul ins(%A, %B : tensor<128x256xf32>, tensor<256x64xf32>)
                     outs(%init : tensor<128x64xf32>) -> tensor<128x64xf32>
  return %C : tensor<128x64xf32>
}

这里出现了 func、linalg 等方言。MLIR 生态中常见的方言包括：

方言	抽象层级	典型用途
`func`	函数	函数定义与调用
`arith`	标量运算	加减乘除、比较
`tensor` / `memref`	张量 / 内存	值语义与缓冲语义
`linalg`	结构化线性代数	矩阵乘、卷积等
`scf`	结构化控制流	for / while / if
`affine`	仿射循环	多面体优化
`gpu`	GPU 抽象	线程块、共享内存
`llvm`	LLVM IR	对接 LLVM 后端

Pass 是 MLIR 的优化单元，可以按需组合：

void buildPipeline(mlir::PassManager &pm) {
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(mlir::createCSEPass());
  pm.addNestedPass<mlir::func::FuncOp>(mlir::createLinalgElementwiseOpFusionPass());
  pm.addPass(mlir::createConvertLinalgToLoopsPass());
  pm.addPass(mlir::createConvertSCFToCFPass());
}

3.3 渐进式降级

MLIR 最核心的思想是渐进式降级（progressive lowering）：从高层方言逐步转换到低层方言，每一层只做该层该做的事。

这种分层让优化可以「在合适的抽象层级做合适的事」：算子融合在 linalg 层做，循环变换在 affine 层做，寄存器分配交给 LLVM。相比一次性从高层 IR 直接降到机器码，可维护性和可扩展性都更好。

3.4 MLIR 的适用场景

MLIR 适合作为编译器框架使用：

构建领域专用编译器（AI、DSL、硬件描述）；
在已有编译器中引入多层 IR，做跨层优化；
复用社区成熟的方言与 Pass，减少重复造轮子。

代价是：MLIR 本身不提供端到端的模型导入与部署能力，你需要自己搭建从框架到 MLIR、再从 MLIR 到目标硬件的完整链路。

4. TVM 与 MLIR 的对比

两者经常被拿来比较，但严格说它们不在同一层面：TVM 是「AI 编译器产品」，MLIR 是「编译器基础设施」。

维度	TVM	MLIR
定位	端到端深度学习编译栈	可扩展编译器基础设施
核心抽象	Relay/Relax + TIR	Dialect + Operation
计算与调度	显式分离，调度原语丰富	通过方言与 Pass 表达
自动调优	AutoTVM / Ansor 成熟	需自行集成或借助外部工具
模型导入	内置 ONNX / PyTorch 前端	需自行实现或借助 torch-mlir
多后端	内置 CPU/GPU/加速器后端	依赖方言与下游后端
学习曲线	调度体系较陡	概念多但分层清晰
扩展方式	新增调度与后端	新增方言与 Pass

一个常见的误解是「MLIR 会取代 TVM」。实际上两者可以互补：TVM 的 Relax 已经引入 MLIR 风格的架构，部分后端也借助 MLIR 生态；而基于 MLIR 的 torch-mlir、IREE 等项目，则在 MLIR 之上重建了端到端能力。

选择建议：

想快速把模型部署到多种硬件，优先考虑 TVM 或 IREE；
想自研编译器、需要深度定制 IR 与优化流程，优先考虑 MLIR；
团队已有 LLVM 积累，MLIR 的迁移成本更低。

5. 实战：一个最小可运行的 TVM 示例

下面用 TVM 编译一个向量加法，并验证结果。示例基于 TVM 0.15 及以上版本。

import numpy as np
import tvm
from tvm import te

n = 1024
A = te.placeholder((n,), name="A", dtype="float32")
B = te.placeholder((n,), name="B", dtype="float32")
C = te.compute((n,), lambda i: A[i] + B[i], name="C")

s = te.create_schedule(C.op)
# 按 64 分块并向量化
i_outer, i_inner = s[C].split(C.op.axis[0], factor=64)
s[C].vectorize(i_inner)

# 编译为 LLVM 后端
target = tvm.target.Target("llvm")
fadd = tvm.build(s, [A, B, C], target=target, name="vec_add")

# 在 CPU 上执行
ctx = tvm.cpu(0)
a_np = np.random.rand(n).astype("float32")
b_np = np.random.rand(n).astype("float32")
a_tvm = tvm.nd.array(a_np, ctx)
b_tvm = tvm.nd.array(b_np, ctx)
c_tvm = tvm.nd.array(np.zeros(n, dtype="float32"), ctx)

fadd(a_tvm, b_tvm, c_tvm)

np.testing.assert_allclose(c_tvm.numpy(), a_np + b_np, rtol=1e-5)
print("TVM 向量加法结果正确")

这段代码展示了 TVM 的完整闭环：声明计算 → 编写调度 → 指定目标 → 编译执行 → 数值验证。把 target 换成 cuda，同一份计算就能生成 GPU 代码（需相应调整调度）。

6. 实战：一个最小可运行的 MLIR 示例

MLIR 的入门通常从 mlir-opt 开始。安装 LLVM/MLIR 后，可以写一个 .mlir 文件并用 Pass 管道处理。

// add.mlir
func.func @add(%a: tensor<4xf32>, %b: tensor<4xf32>) -> tensor<4xf32> {
  %c = arith.addf %a, %b : tensor<4xf32>
  return %c : tensor<4xf32>
}

用 mlir-opt 做规范化与 CSE：

mlir-opt add.mlir \
  --canonicalize \
  --cse \
  -o add_opt.mlir

如果要观察降级过程，可以把 linalg 逐步降到循环：

mlir-opt matmul.mlir \
  --convert-linalg-to-loops \
  --convert-scf-to-cf \
  --convert-func-to-llvm \
  --reconcile-unrealized-casts \
  -o matmul_lowered.mlir

每一步都会输出不同抽象层级的 IR，这正是 MLIR 渐进式降级思想的可视化体现。

如果需要自己写 Pass，可以用 C++ 或 Python 绑定：

#include "mlir/Pass/Pass.h"
#include "mlir/IR/PatternMatch.h"

struct MyPass : public mlir::PassWrapper<MyPass, mlir::OperationPass<mlir::func::FuncOp>> {
  void runOnOperation() override {
    auto func = getOperation();
    // 遍历 Operation，做你的变换
    func.walk([](mlir::Operation *op) {
      // ...
    });
  }
};

7. 选型与学习路径

给不同背景的读者一些具体建议：

如果你是算法工程师，想把模型部署到新硬件：从 TVM 的 Relax 或 IREE 入手，先跑通官方教程里的端到端示例，再逐步深入调度与后端。

如果你是编译器工程师，想自研 AI 编译器：从 MLIR 的 Toy Tutorial 开始，理解 Dialect 定义、Pass 管理与降级流程，再研究 torch-mlir、IREE 等真实项目。

如果你在做推理框架：两者都值得了解。TVM 的图优化与自动调优经验可以直接借鉴；MLIR 的分层 IR 设计则适合作为长期架构参考。

学习资源方面，TVM 官方文档与「TVM 深度学习编译器」系列文章质量较高；MLIR 官方文档的 Toy Tutorial 是最友好的入门材料，配合 LLVM 社区的演讲视频效果更好。

8. 总结

AI 编译器的本质，是在「模型语义」与「硬件执行」之间建立可自动化的映射关系。TVM 用计算与调度分离、自动调优和端到端链路，解决了「一份模型跑多硬件」的工程问题；MLIR 用方言体系和渐进式降级，解决了「编译器如何可扩展、可复用」的架构问题。

两者不是替代关系，而是不同层次的答案。理解 TVM，你会知道一个可用的 AI 编译器长什么样；理解 MLIR，你会知道如何把它拆解、重组、扩展到新的领域。把这两条线都走一遍，再回头看任何一个推理框架或芯片工具链，都会清晰很多。

转载自 CSDN-专业IT技术社区

原文链接：https://blog.csdn.net/Maxwell_Newton/article/details/169466172

AI 编译器入门：TVM 与 MLIR 的原理、差异与实战