<jupyter_start><jupyter_text>(tir_blitz)= TensorIR 的突击课程**作者**: [<NAME>](https://github.com/Hzfengsy)TensorIR 是用于深度学习程序的特定域语言，有两个广泛的目的：- 在各种硬件后端进行程序变换和优化的实现。- 用于自动张量化程序优化的抽象。<jupyter_code>import tvm
from tvm.script.parser import ir_module
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np<jupyter_output><empty_output><jupyter_text>IRModuleIRModule 是 TVM 的中心数据结构，它包含深度学习程序。它是 IR 变换和模型构建的基本关注对象。```{image} https://daobook.github.io/tvm-web-data/images/design/tvm_life_of_irmodule.png:width: 85%```这是 IRModule 的生命周期（life cycle），它可以从 TVMScript 创建。TensorIR 调度原语（primitive）和传递（pass）是变换 IRModule 的两种主要方式。另外，对 IRModule 进行一系列的变换也是可以接受的。请注意，可以在 **任何** 阶段向 TVMScript 打印 IRModule。在所有变换和优化完成后，可以将 IRModule 构建为可运行的模块，以部署在目标设备上。基于 TensorIR 和 IRModule 的设计，能够创建新的编程方式：1. 用 TVMScript 写基于 Python-AST 语法的程序。2. 用 python api 变换和优化程序。3. 通过命令式的变换 API，交互式地检查和尝试性能。 创建 IRModuleIRModule 可以通过编写 TVMScript 来创建，TVMScript 是 TVM IR 的可圆润化（round-trippable）的语法。与通过 [张量表达式](tutorial-tensor-expr-get-started) 创建计算表达式不同，TensorIR 允许用户通过 TVMScript（嵌入式 python AST 的语言）来编程。这种新方法使得编写复杂的程序并进一步调度和优化它成为可能。<jupyter_code>@ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        # 通过句柄（handle）在函数之间交换数据，此操作类似于指针（pointer）。
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # 从句柄创建 buffer。
        A = T.match_buffer(a, (8,), dtype="float32")
        B = T.match_buffer(b, (8,), dtype="float32")
        for i in range(8):
            # block 是计算的抽象。
            with T.block("B"):
                # 定义 spatial block 迭代器，并将其绑定到值 i。
                vi = T.axis.spatial(8, i)
                B[vi] = A[vi] + 1.0


ir_module = MyModule
print(type(ir_module))
print(ir_module.script())<jupyter_output><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)<jupyter_text>此外，还可以使用张量表达式 DSL 来编写简单的算子，并将其转换为 IRModule。<jupyter_code>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())<jupyter_output># 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)<jupyter_text>构建和运行 IRModule我们可以将 IRModule 构建为具有特定目标后端的可运行模块。<jupyter_code>mod = tvm.build(ir_module, target="llvm")  # 用于 CPU 后端的模块
print(type(mod))<jupyter_output><class 'tvm.driver.build_module.OperatorModule'><jupyter_text>准备好输入 array 和输出 array，然后运行该模块。<jupyter_code>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)<jupyter_output>[0. 1. 2. 3. 4. 5. 6. 7.]
[1. 2. 3. 4. 5. 6. 7. 8.]<jupyter_text>转换 IRModuleIRModule 是程序优化的中心数据结构，它可以通过 `Schedule` 进行转换。调度包含多个原语方法，以交互式地转换程序。每个原语都以某些方式改造程序，以带来额外的性能优化。上面的图片是优化张量程序的典型工作流程。首先，需要在由 TVMScript 或 Tensor Expression 创建的初始 IRModule 上创建调度。然后，一连串的调度原语将有助于提高性能。最后，我们可以将其降低并构建为可运行的模块。这里只演示了非常简单的变换。首先，在输入的 `ir_module` 上创建调度。<jupyter_code>sch = tvm.tir.Schedule(ir_module)
print(type(sch))<jupyter_output><class 'tvm.tir.schedule.schedule.Schedule'><jupyter_text>将该循环分为 3 个循环，并打印结果。<jupyter_code># 按 name 获取 block
block_b = sch.get_block("B")
# 获取 block 周围的 loops
(i,) = sch.get_loops(block_b)
# 平铺（tile）循环嵌套。
i_0, i_1, i_2 = sch.split(i, factors=[2, 2, 2])
print(sch.mod.script())<jupyter_output># 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)<jupyter_text>也可以重新调度循环的顺序。现在将循环 `i_2` 移到 `i_1` 的外面。<jupyter_code>sch.reorder(i_0, i_2, i_1)
print(sch.mod.script())<jupyter_output># 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)<jupyter_text>转化为 GPU 程序如果想在 GPU 上部署模型，线程绑定是必要的。幸运的是，也可以使用原语并做增量变换。<jupyter_code>sch.bind(i_0, "blockIdx.x")
sch.bind(i_2, "threadIdx.x")
print(sch.mod.script())<jupyter_output># 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)<jupyter_text>绑定线程后，现在用 `cuda` 后端构建 IRModule。<jupyter_code>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)<jupyter_output>[0. 1. 2. 3. 4. 5. 6. 7.]
[1. 2. 3. 4. 5. 6. 7. 8.]
{'fraction_non_alphanumeric': 0.15808170515097691, 'fraction_numerical': 0.026198934280639432, 'mean_word_length': 3.158153846153846, 'pattern_counts': {'":': 13, '<': 36, '<?xml version=': 0, '>': 41, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '155348', 'license': 'Apache-2.0', 'n_tokens_mistral': 3281, 'n_tokens_neox': 3357, 'n_words': 513, 'path': 'xinetzone/docs/tutorial/tensor_ir_blitz_course.ipynb', 'repo_name': 'xinetzone/tvm', 'size': '6761'}