TVM User Tutorial -- Blitz Course to TensorIR

最新推荐文章于 2025-06-03 09:05:19 发布

姆克儿

最新推荐文章于 2025-06-03 09:05:19 发布

阅读量420

点赞数

CC 4.0 BY-SA版权

分类专栏： TVM文档翻译文章标签：人工智能深度学习 TVM

原文链接：https://tvm.apache.org/docs/tutorial/tensor_ir_blitz_course.html

TVM文档翻译专栏收录该内容

12 篇文章

订阅专栏

Author: Siyuan Feng

TensorIR 是一种特定领域语言，用于深度学习项目，有两个广泛的用途：

在各种硬件后端上实现转换和优化程序。
自动张力化程序优化的抽象。

import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np

IRModule

IRModule 是TVM的中心数据结构，它包含深度学习程序。它是研究IR变换和建模的基本对象。

这是IRModule的生命周期，可以从TVMScript创建。TensorIR调度原语和过程是转换IRModule的两种主要方式。此外，一个IRM模块上的一系列转换也是可以接受的。请注意，我们可以在任何阶段将IRModule打印到TVMScript。在所有转换和优化完成后，我们可以将IRModule构建为可运行的模块，以便部署到目标设备上。
基于TensorIR和IRModule的设计，我们能够创建一种新的编程方法：

用基于python AST的语法编写TVMScript程序。
使用python api转换和优化程序。
使用命令式转换API以交互方式检查和尝试性能。

构建IRModule

IRModule 可以通过编写 TVMScript 来创建，TVMScript 是TVM IR的一种圆形可折叠语法。
与通过Tensor表达式创建计算表达式（使用Tensor表达式与运算符一起工作）不同，Tensor允许用户通过TVMScript编程，TVMScript是一种嵌入python AST的语言。这种新方法可以编写复杂的程序，并对其进行进一步的调度和优化。
下面是向量加法的一个简单示例。

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        # We exchange data between function by handles, which are similar to pointer.
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # Create buffer from handles.
        A = T.match_buffer(a, (8,), dtype="float32")
        B = T.match_buffer(b, (8,), dtype="float32")
        for i in range(8):
            # A block is an abstraction for computation.
            with T.block("B"):
                # Define a spatial block iterator and bind it to value i.
                vi = T.axis.spatial(8, i)
                B[vi] = A[vi] + 1.0


ir_module = MyModule
print(type(ir_module))
print(ir_module.script())

Out:
<class 'tvm.ir.module.IRModule'>
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(8,), "float32"], B: T.Buffer[(8,), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i in T.serial(8):
            with T.block("B"):
                vi = T.axis.spatial(8, i)
                T.reads(A[vi])
                T.writes(B[vi])
                B[vi] = A[vi] + T.float32(1)

此外，我们还可以使用张量表达式DSL编写简单的运算符，并将它们转换为IRM模块。

from tvm import te

A = te.placeholder((8,), dtype="float32", name="A")
B = te.compute((8,), lambda *i: A(*i) + 1.0, name="B")
func = te.create_prim_func([A, B])
ir_module_from_te = IRModule({"main": func})
print(ir_module_from_te.script())

Out:
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(8,), "float32"], B: T.Buffer[(8,), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i0 in T.serial(8):
            with T.block("B"):
                i0_1 = T.axis.spatial(8, i0)
                T.reads(A[i0_1])
                T.writes(B[i0_1])
                B[i0_1] = A[i0_1] + T.float32(1)

构建并运行IRM模块

我们可以将IRModule构建成具有特定目标后端的可运行模块。

mod = tvm.build(ir_module, target="llvm")  # The module for CPU backends.
print(type(mod))

Out:
<class 'tvm.driver.build_module.OperatorModule'>

准备输入阵列和输出阵列，然后运行模块。

a = tvm.nd.array(np.arange(8).astype("float32"))
b = tvm.nd.array(np.zeros((8,)).astype("float32"))
mod(a, b)
print(a)
print(b)

Out:
[0. 1. 2. 3. 4. 5. 6. 7.]
[1. 2. 3. 4. 5. 6. 7. 8.]

转换IRModule

IRModule是程序优化的中心数据结构，可以通过调度进行转换。调度包含多个基本方法，用于以交互方式转换程序。每个原语都以特定的方式转换程序，以带来额外的性能优化。
在这里插入图片描述上图是优化张量程序的典型工作流程。首先，我们需要在由TVMScript或Tensor表达式创建的初始IRModule上创建一个调度。然后，一系列调度原语将有助于提高性能。最后，我们可以将其降低并构建成一个可运行的模块。
这里我们只是演示一个非常简单的转换。首先，我们在输入ir_模块上创建时间表。

sch = tvm.tir.Schedule(ir_module)
print(type(sch))

Out:
<class 'tvm.tir.schedule.schedule.Schedule'>

Tile the loop into 3 loops and print the result.

# Get block by its name
block_b = sch.get_block("B")
# Get loops surronding the block
(i,) = sch.get_loops(block_b)
# Tile the loop nesting.
i_0, i_1, i_2 = sch.split(i, factors=[2, 2, 2])
print(sch.mod.script())

Out:
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(8,), "float32"], B: T.Buffer[(8,), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0, i_1, i_2 in T.grid(2, 2, 2):
            with T.block("B"):
                vi = T.axis.spatial(8, i_0 * 4 + i_1 * 2 + i_2)
                T.reads(A[vi])
                T.writes(B[vi])
                B[vi] = A[vi] + T.float32(1)

我们还可以重新排列循环。现在我们将循环i_2移动到i_1的外部。

sch.reorder(i_0, i_2, i_1)
print(sch.mod.script())

Out:
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(8,), "float32"], B: T.Buffer[(8,), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0, i_2, i_1 in T.grid(2, 2, 2):
            with T.block("B"):
                vi = T.axis.spatial(8, i_0 * 4 + i_1 * 2 + i_2)
                T.reads(A[vi])
                T.writes(B[vi])
                B[vi] = A[vi] + T.float32(1)

转换为GPU程序

如果我们想在GPU上部署模型，线程绑定是必要的。幸运的是，我们还可以使用原语并进行增量转换。

sch.bind(i_0, "blockIdx.x")
sch.bind(i_2, "threadIdx.x")
print(sch.mod.script())

Out:
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(8,), "float32"], B: T.Buffer[(8,), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0 in T.thread_binding(2, thread="blockIdx.x"):
            for i_2 in T.thread_binding(2, thread="threadIdx.x"):
                for i_1 in T.serial(2):
                    with T.block("B"):
                        vi = T.axis.spatial(8, i_0 * 4 + i_1 * 2 + i_2)
                        T.reads(A[vi])
                        T.writes(B[vi])
                        B[vi] = A[vi] + T.float32(1)

绑定线程后，现在使用cuda后端构建IRModule。

ctx = tvm.cuda(0)
cuda_mod = tvm.build(sch.mod, target="cuda")
cuda_a = tvm.nd.array(np.arange(8).astype("float32"), ctx)
cuda_b = tvm.nd.array(np.zeros((8,)).astype("float32"), ctx)
cuda_mod(cuda_a, cuda_b)
print(cuda_a)
print(cuda_b)

Out:
[0. 1. 2. 3. 4. 5. 6. 7.]
[1. 2. 3. 4. 5. 6. 7. 8.]