菜鸟AI
菜鸟AI AI提示词 · 教程 · 资讯
首页>其他资讯

TVM TensorIR创建完全教程：从零到一掌握深度学习编译器开发

2026-06-11阅读 0热度 0
Tensor

Apache TVM 0.21.0 正式发布，中文文档已同步更新，内容覆盖 TensorIR 等新特性。

【TVM教程】TensorIR 创建

Apache TVM 作为面向 CPU、GPU 及专用加速器的深度学习编译栈，此次 0.21.0 版本在 TensorIR 创建路径上做了关键升级——无论是通过 TVMScript 还是 Tensor Expression，都更加直接且灵活。以下逐一展开。

利用 TVMScript 构建 TensorIR 函数

若要快速创建 TensorIR 函数，TVMScript 是最直接的方式。这种 Python 方言能让你用接近数学表达的形式，直接描述底层的计算与调度逻辑，省去额外转换步骤。

基础写法

以《理解 TensorIR 抽象》中的 mm_relu 为例，下面展示标准的 ir_module 与 TVMScript 写法：

import numpy as np
import tvm
from tvm.script import ir as I
from tvm.script import tir as T


@I.ir_module
class MyModule:
    @T.prim_func
    def mm_relu(
        A: T.Buffer((128, 128), "float32"),
        B: T.Buffer((128, 128), "float32"),
        C: T.Buffer((128, 128), "float32"),
    ):
        Y = T.alloc_buffer((128, 128), dtype="float32")
        for i in range(128):
            for j in range(128):
                for k in range(128):
                    with T.block("Y"):
                        vi = T.axis.spatial(128, i)
                        vj = T.axis.spatial(128, j)
                        vk = T.axis.reduce(128, k)
                        T.reads(A[vi, vk], B[vk, vj])
                        T.writes(Y[vi, vj])
                        with T.init():
                            Y[vi, vj] = T.float32(0)
                        Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i in range(128):
            for j in range(128):
                with T.block("C"):
                    vi = T.axis.spatial(128, i)
                    vj = T.axis.spatial(128, j)
                    T.reads(Y[vi, vj])
                    T.writes(C[vi, vj])
                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

语法糖让代码更紧凑

当编写复杂计算时，TVM 提供了几个实用语法糖来简化代码：

  • 用 T.grid 替代多层 for 循环
  • 借助 T.axis.remap 精简 block 迭代器注解
  • 当读写信息可从主体推导时，可省略 T.reads 和 T.writes 声明
@I.ir_module
class ConciseModule:
    @T.prim_func
    def mm_relu(
        A: T.Buffer((128, 128), "float32"),
        B: T.Buffer((128, 128), "float32"),
        C: T.Buffer((128, 128), "float32"),
    ):
        Y = T.alloc_buffer((128, 128), dtype="float32")
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

验证两个模块是否结构等价：

print(tvm.ir.structural_equal(MyModule, ConciseModule))

输出：

True

利用 Python 变量动态化 TVMScript

TVMScript 无法由 Python 直接运行，但可与 Python 变量联动。例如，用 Python 变量指定 TensorIR 的形状与数据类型，使脚本更灵活。

# Python 变量
M = N = K = 128
dtype = "float32"


# 使用 TVMScript 定义的 IRModule
@I.ir_module
class ConciseModuleFromPython:
    @T.prim_func
    def mm_relu(
        A: T.Buffer((M, K), dtype),
        B: T.Buffer((K, N), dtype),
        C: T.Buffer((M, N), dtype),
    ):
        Y = T.alloc_buffer((M, N), dtype)
        for i, j, k in T.grid(M, N, K):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    Y[vi, vj] = T.cast(T.float32(0), dtype)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(M, N):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype))

验证结构等价性：

print(tvm.ir.structural_equal(ConciseModule, ConciseModuleFromPython))

输出：

True

支持动态形状的 TensorIR 函数

静态形状适用于固定场景，但实际模型输入尺寸常变化。TVMScript 同样支持动态形状：只需将参数声明为 T.handle，再通过 T.match_buffer 绑定动态尺寸即可。

@I.ir_module
class DynamicShapeModule:
    @T.prim_func
    def mm_relu(a: T.handle, b: T.handle, c: T.handle):
        # 动态形状定义
        M, N, K = T.int32(), T.int32(), T.int32()

        # 使用动态形状绑定输入缓冲区
        A = T.match_buffer(a, [M, K], dtype)
        B = T.match_buffer(b, [K, N], dtype)
        C = T.match_buffer(c, [M, N], dtype)
        Y = T.alloc_buffer((M, N), dtype)
        for i, j, k in T.grid(M, N, K):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    Y[vi, vj] = T.cast(T.float32(0), dtype)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(M, N):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype))

测试动态形状运行时推理：

def evaluate_dynamic_shape(lib: tvm.runtime.Module, m: int, n: int, k: int):
    A = tvm.runtime.tensor(np.random.uniform(size=(m, k)).astype("float32"))
    B = tvm.runtime.tensor(np.random.uniform(size=(k, n)).astype("float32"))
    C = tvm.runtime.tensor(np.zeros((m, n), dtype="float32"))
    lib(A, B, C)
    return C.numpy()


# 只需编译一次
dyn_shape_lib = tvm.compile(DynamicShapeModule, target="llvm")
# 支持不同的输入形状
print(evaluate_dynamic_shape(dyn_shape_lib, m=4, n=4, k=4))
print(evaluate_dynamic_shape(dyn_shape_lib, m=64, n=64, k=128))

输出：

[[1.6744074  1.8393843  0.9076001  0.32640088]
 [1.3455076  1.5298209  0.75502187 0.32371795]
 [1.9979694  2.221868   1.0828729  0.43582058]
 [1.7054784  1.8512932  0.89285195 0.34154552]]
[[30.544813 29.938599 33.654526 ... 29.934391 30.73088  25.106636]
 [30.644558 31.062693 32.34803  ... 29.584583 32.756992 25.280499]
 [33.73643  33.23441  34.2736   ... 34.284283 35.100815 27.748833]
 ...
 [31.313179 30.462463 30.996958 ... 28.831778 32.279408 25.663143]
 [33.129818 31.630735 33.334507 ... 29.682335 32.925854 26.043703]
 [32.44726  30.645096 33.926357 ... 29.750242 32.810432 25.420698]]

通过 Tensor Expression 生成 TensorIR

多数开发者无需直接处理 TensorIR 底层细节，更倾向用简洁 API 描述计算。此时 Tensor Expression（TE）是最佳选择。

TE 是面向计算的领域特定语言，提供类似表达式的 API。需要留意，TE 仅负责计算描述；其传统调度能力已被 TVM Unity 栈中的 TensorIR 调度替代，因此本节聚焦表达式构建。

构建静态形状函数

仍以 mm_relu 为例，展示 TE 的创建方式：

from tvm import te

A = te.placeholder((128, 128), "float32", name="A")
B = te.placeholder((128, 128), "float32", name="B")
k = te.reduce_axis((0, 128), "k")
Y = te.compute((128, 128), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name="C")

te.compute 签名为 te.compute(output_shape, fcompute)，fcompute 定义每个位置元素 Y[i, j] 的计算逻辑：

lambda i, j: te.sum(A[i, k] * B[k, j], axis=k)

该 lambda 等价于矩阵乘法累加：Yi,j = Σk Ai,k × Bk,j。计算定义完成后，结合输入输出参数即可生成 TensorIR 函数。本例期望的函数签名包含两个输入 (A, B) 和一个输出 (C)：

te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
TEModule = tvm.IRModule({"mm_relu": te_func})
TEModule.show()

输出：

# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
        T.func_attr({"tir.noalias": True})
        # with T.block("root"):
        Y = T.alloc_buffer((128, 128))
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
                T.reads(A[v_i, v_k], B[v_k, v_j])
                T.writes(Y[v_i, v_j])
                with T.init():
                    Y[v_i, v_j] = T.float32(0.0)
                Y[v_i, v_j] = Y[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                v_i, v_j = T.axis.remap("SS", [i, j])
                T.reads(Y[v_i, v_j])
                T.writes(C[v_i, v_j])
                C[v_i, v_j] = T.max(Y[v_i, v_j], T.float32(0.0))

构建动态形状函数

TE 动态形状创建与静态版本基本一致，只需将张量形状设为符号变量：

# 定义符号变量
M, N, K = te.var("m"), te.var("n"), te.var("k")
A = te.placeholder((M, N), "float32", name="A")
B = te.placeholder((K, N), "float32", name="B")
k = te.reduce_axis((0, K), "k")
Y = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
C = te.compute((M, N), lambda i, j: te.max(Y[i, j], 0), name="C")

dyn_te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
DynamicTEModule = tvm.IRModule({"mm_relu": dyn_te_func})
DynamicTEModule.show()

输出：

# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def mm_relu(var_A: T.handle, var_B: T.handle, var_C: T.handle):
        T.func_attr({"tir.noalias": True})
        m, n = T.int32(), T.int32()
        A = T.match_buffer(var_A, (m, n))
        k = T.int32()
        B = T.match_buffer(var_B, (k, n))
        C = T.match_buffer(var_C, (m, n))
        # with T.block("root"):
        Y = T.alloc_buffer((m, n))
        for i, j, k_1 in T.grid(m, n, k):
            with T.block("Y"):
                v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k_1])
                T.reads(A[v_i, v_k], B[v_k, v_j])
                T.writes(Y[v_i, v_j])
                with T.init():
                    Y[v_i, v_j] = T.float32(0.0)
                Y[v_i, v_j] = Y[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
        for i, j in T.grid(m, n):
            with T.block("C"):
                v_i, v_j = T.axis.remap("SS", [i, j])
                T.reads(Y[v_i, v_j])
                T.writes(C[v_i, v_j])
                C[v_i, v_j] = T.max(Y[v_i, v_j], T.float32(0.0))

以上分别通过 TVMScript 与 Tensor Expression 展示了 TensorIR 的完整创建过程。静态与动态、简洁与显式，开发者可根据实际场景灵活选择——这正是 TVM 架构设计的精妙之处。

上一篇亿风投教父终极警告：软件不再卖给人，270亿Cursor陷困境 下一篇Triton教程：triton_language.cos函数使用与实战指南
免责声明

本网站新闻资讯均来自公开渠道，力求准确但不保证绝对无误，内容观点仅代表作者本人，与本站无关。若涉及侵权，请联系我们处理。本站保留对声明的修改权，最终解释权归本站所有。

相关阅读

更多
其他资讯06-11
TVM TensorIR创建完全教程：从零到一掌握深度学习编译器开发

TVM0 21 0版本更新了TensorIR创建方式：TVMScript支...
其他资讯06-11
TVM教程：深入理解TensorIR抽象核心原理

TensorIR是TVM中的张量程序抽象，通过显式定义块轴属...
其他资讯05-28
TensorMesh获Nvidia等投资，AI内存优化方案深度评测

大语言模型推理的算力成本，尤其是GPU资源的消耗，始...
其他资讯05-09
2025年Tensor G7芯片深度测评：Pixel 12系列2nm工艺性能前瞻

谷歌下一代TensorG7芯片已在开发中，代号“Lajolla”...
其他资讯05-05
谷歌Pixel 11系列手机配置曝光：Tensor G6芯片、放弃三星基带

谷歌Pixel 11系列核心信息泄露：Tensor G6、新相机...
其他资讯04-29
Google Pixel 11 系列曝光 Tensor G6 芯片规格首次泄露

Google Pixel 11系列前瞻：Tensor G6芯片关键信息...

最新教程

Stable Diffusion WebUI整合包下载与模型放置全指南HunyuanVideo安装失败排查指南：依赖、显存与工作流问题解决Runway官网入口与使用指南：下载注册及常见问题全解析Notion AI新手入门指南：从下载到模板设置的完整教程GitHub Copilot安装指南：JetBrains插件市场一键配置与激活全流程2026年ComfyUI安装与配置终极指南：从零部署到高效出图全流程解析CogVideoX安装包获取与部署指南：从下载到剪辑机配置的完整教程2024图像识别实战精选：基于EasyDL的完整案例解析与测评

最新资讯

Harness vs Loop Engineering：2025 DevOps工具深度对比EMR Spark Relational Cache跨集群同步精选指南Hermes技能测评：一条命令省10分钟，效率提升300%瑞幸Skill实测：点一杯咖啡需要消耗多少Token云HIS前后端分离源码实践：SpringBoot+Bootstrap+jQueryAI智能体开发技术方案深度对比与推荐阿里云ESSD-AutoPL云盘性能测评与计费对比AI开发微信小程序教程：5分钟快速入门指南
欢迎回来 登录或注册后，可保存提示词和历史记录
登录后可同步收藏、历史记录和常用模板
注册即表示同意服务条款与隐私政策