本文翻译自Blitz Course to TensorIR — tvm 0.9.dev0 documentation
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调度原语和pass是转换IRModule的两种主要方式。同样,也可以在IRModule上做一系列转换。注意,我们可以在任意阶段打印一个IRModule到TVMScript。在完成所有的转换和优化之后,我们可以将IRModule构建为一个可运行模块,以便部署到目标设备上。
基于TensorIR和IRModule的设计,我们可以创建一种新的编程方法:
- 用TVMScript编写一个基于python-AST语法的程序。
- 使用python api转换和优化程序。
- 使用命令式转换API交互式地检查和测试性能。
创建IRModule
IRModule可以通过编写TVMScript来创建,这是TVM IR的一种双向转换的语法。
与通过张量表达式创建计算表达式不同,TensorIR允许用户通过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())
输出:
<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构建成具有特定目标后端的可运行模块。
mod = tvm.build(ir_module, target="llvm") # The module for CPU backends.
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进行转换。一个调度包含多个交互转换程序的原语方法。每个原语都以特定的方式转换程序,以带来额外的性能优化。
上图是优化张量程序的典型工作流。首先,我们需要在一个初始的IRModule上创建一个调度,这个IRModule可以使用TVMScript或张量表达式创建。然后,一系列的调度原语将有助于提高性能。最后,我们可以将其低级化并构建成一个可运行的模块。
这里我们只是演示一个非常简单的转换。首先,我们在输入ir_module上创建调度。
sch = tvm.tir.Schedule(ir_module)
print(type(sch))
输出:
<class 'tvm.tir.schedule.schedule.Schedule'>
将循环tile成3个循环并打印结果。
# 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())
输出:
# 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.]