/ MLC chentianqi
[MLC-03] End-to-End Model Execution and Relax
Notes for lesson three: using Relax graphs, call_tir, dataflow blocks, and parameter binding to assemble TensorIR functions into an executable model.
上一讲处理的是单个元张量函数;第三讲回答的是另一个问题:这些函数怎样按模型拓扑连接,并与权重、运行库和 build 流程组成一次端到端执行。
1. 两层 IR 的分工#
课程使用的一个重要分层是:
Relax
负责:计算图、函数调用、数据流、参数与端到端组合
|
| call_tir
v
TensorIR
负责:单个张量函数的循环、缓冲区与底层实现Relax 并不试图替代 TensorIR。它要表达的是“哪个算子接哪个算子、哪个值能被优化器看见、哪个函数应该调用哪个底层实现”。
2. call_tir 是图与程序之间的边#
一个 Relax 函数可将高层值传入已有的 TensorIR primitive function:
@R.function
def main(x: R.Tensor((1, 784), "float32"), w: R.Tensor((784, 10), "float32")):
with R.dataflow():
lv0 = R.call_tir("linear", (x, w), out_sinfo=R.Tensor((1, 10), "float32"))
gv0 = R.nn.relu(lv0)
R.output(gv0)
return gv0这里的 call_tir 很像 ABI 边界:图层知道输入输出的 tensor 信息和被调用函数名,但不需要了解 linear 内部是怎样 tiled、vectorized 或映射线程的。
3. Dataflow block#
dataflow block 用来区分函数内部的临时值和需要输出到外部作用域的值:
temporary value: 只在 block 内使用,可被局部优化
output value: 离开 block 后仍被引用,构成图的可见边界这个边界利于编译器做 DCE、融合、重写和内存规划。不是所有中间张量都应被当作长期可见对象;明确的数据流语义让优化器有更多安全空间。
4. 参数绑定#
一个模型函数通常接收输入和权重:
部署期常常希望把固定权重 绑定到 IRModule,让入口变成只接受运行时输入 的函数:
before binding: main(x, w1, b1, w2, b2)
after binding: main(x)这样做的好处包括减小调用接口、便于常量折叠、让运行时加载器统一管理权重,并为后续内存规划和序列化提供明确对象。
5. 与运行库混合#
不是每个函数都必须由 TensorIR 重新生成。端到端模型可以混合:
Relax graph
-> TensorIR matmul kernel
-> runtime library function
-> TensorIR elementwise kernel
-> external packed function这很符合真实部署:高性能库、手工 kernel 和编译生成代码会长期共存。MLC 的任务是让它们在一个统一模块中有明确调用边界和可组合的构建流程。
6. 对 serving 的启发#
LLM serving 的一次请求也可看作端到端执行图:embedding、attention、MLP、采样、缓存操作与外部通信不必都落在同一层 IR。高层图负责组织依赖,底层程序负责执行算子,运行时负责资源与生命周期。
本讲的重点不是某个 API,而是两层 IR 的职责边界:图层表达组合,程序层表达实现。