[MLC-05] 与机器学习框架的整合

第五讲讨论一个工程上绕不开的问题：用户已经在 PyTorch 等框架中写好模型，MLC 系统如何把它导入自己的 IR，而不是要求每个人重写模型。

1. 导入不是复制代码#

框架模型的价值不仅在 Python 函数，还包括模块层次、参数、算子语义和图拓扑。整合过程的目标是：

framework model
  -> capture graph and parameters
  -> map supported operators
  -> build Relax/TensorIR IRModule
  -> optimize and build for deployment

它需要明确哪些内容保留为高层算子、哪些内容降低为 TensorIR、哪些内容交给外部运行库。

2. Tensor Expression（TE）#

TE 是描述张量计算的 DSL。以矩阵乘法为例：

def te_matmul(A, B):
    k = te.reduce_axis((0, A.shape[1]), name="k")
    return te.compute(
        (A.shape[0], B.shape[1]),
        lambda i, j: te.sum(A[i, k] * B[k, j], axis=k),
        name="matmul",
    )

TE 关注“每个输出元素如何定义”，并可通过 te.create_prim_func 生成 TensorIR。它适合快速构造 shape 参数化的算子表达，但复杂控制流或端到端组合仍需要更高层 builder。

3. BlockBuilder#

直接手写完整 TVMScript 在大模型上很难维护。BlockBuilder 提供了编程式构造 IRModule 的接口：

emit_te      将 TE 计算引入 TensorIR
emit         在 Relax dataflow 中生成 binding
emit_func_output 结束函数并声明输出

它把“构造计算图”和“生成底层函数”放进同一套 API，适合 importer 根据外部图逐节点翻译。

4. PyTorch FX 导入#

PyTorch FX 能捕获一个 GraphModule，其中节点描述 placeholder、参数读取、函数调用、模块调用和输出。Importer 的基本结构通常是：

FX node
  -> resolve input variables and parameters
  -> dispatch a converter by op kind
  -> emit Relax expression
  -> store node-to-IR value mapping

例如 torch.nn.Linear 可以映射为高层 linear，也可以拆成 matmul + bias；选择取决于后端是否有更好的融合实现，以及后续图优化是否需要看见这些组成部分。

5. 导入边界#

一个健壮的 importer 不能假设所有 Python 都可翻译。常见边界包括：

• 动态控制流和数据依赖 shape；
• 未支持的自定义算子；
• Python side effect；
• 参数共享与权重布局；
• dtype、广播和语义差异；
• 高层算子应保留还是立即 lower。

因此导入器更像受控的语义翻译器，而不是通用 Python 编译器。支持范围应明确，fallback 路径也应清楚。

6. 学习笔记#

框架整合决定了编译器能否真正进入用户工作流。好的 IR 不是越低越好：导入过早地把一切降成循环，会失去图层的融合机会；保留过高又可能无法匹配目标硬件。TE、Relax 和 TensorIR 的分层正是在可移植性与可优化性之间提供不同落点。

参考#

• MLC 课程：Integration with Machine Learning Frameworks
• PyTorch FX