关于人工智能:TVM-学习资料TensorIR-快速入门

41次阅读

共计 4758 个字符,预计需要花费 12 分钟才能阅读完成。

本篇文章译自英文文档 Blitz Course to TensorIR
作者是 Siyuan Feng。更多 TVM 中文文档可拜访→TVM 中文站

TensorIR 是深度学习畛域的特定语言,次要有两个作用:

  • 在各种硬件后端转换和优化程序。
  • 主动 tensorized 程序优化的形象。
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 创立。转换 IRModule 的两种次要办法是 TensorIR 的 schedule 原语转换和 pass 转换。此外,也可间接对 IRModule 进行一系列转换。留神,能够在任何阶段将 IRModule 打印到 TVMScript。实现所有转换和优化后,可将 IRModule 构建为可运行模块,从而部署在指标设施上。

基于 TensorIR 和 IRModule 的设计,可创立一种新的编程办法:

  1. 基于 Python-AST 语法,用 TVMScript 编写程序。
  2. 应用 Python API 转换和优化程序。
  3. 应用命令式转换 API 交互检查和进步性能。

创立 IRModule

IRModule 是 TVM IR 的一种可往返语法,可通过编写 TVMScript 来创立。

与通过张量表达式创立计算表达式(应用张量表达式操作算子)不同,TensorIR 容许用户通过 TVMScript(一种嵌在 Python AST 中的语言)进行编程。新办法能够编写简单的程序并进一步调度和优化。

上面是向量加法的示例:

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        # 咱们通过 T.handle 进行数据交换,相似于内存指针
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # 通过 handle 创立 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"):
                # 定义一个空间(可并行)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())

输入后果:

<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 编写简略的运算符,并将它们转换为 IRModule。

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())

输入后果:

# 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)

构建并运行 IRModule

可将 IRModule 构建为特定 target 后端的可运行模块。

mod = tvm.build(ir_module, target="llvm")  # CPU 后端的模块
print(type(mod))

输入后果:

<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)

输入后果:

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

转换 IRModule

IRModule 是程序优化的外围数据结构,可通过 Schedule 进行转换。schedule 蕴含多个 primitive 办法来交互地转换程序。每个 primitive 都以特定形式对程序进行转换,从而优化性能。

上图是优化张量程序的典型工作流程。首先,用 TVMScript 或张量表达式创立一个初始 IRModule,而后在这个初始 IRModule 上创立 schedule。接下来,应用一系列调度原语来进步性能。最初,咱们能够将其升高并构建成一个可运行模块。

下面只演示了一个简略的转换。首先,在输出 ir_module 上创立 schedule:

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

输入后果:

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

将嵌套循环开展成 3 个循环,并打印后果:

# 通过名字获取 block
block_b = sch.get_block("B")
# 获取突围 block 的循环
(i,) = sch.get_loops(block_b)
# 开展嵌套循环
i_0, i_1, i_2 = sch.split(i, factors=[2, 2, 2])
print(sch.mod.script())

输入后果:

# 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())

输入后果

# 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())

输入后果:

# 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)

输入后果:

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

下载 Python 源代码
下载 Jupyter Notebook

以上就是该文档的全部内容,查看更多 TVM 中文文档,请拜访→TVM 中文站

正文完
 0