Author: Tianqi Chen
https://docs.tvm.ai/tutorials/tensor_expr_get_started.html
Tensor Expression入门
这是TVM中Tensor表达语言的入门教程。TVM使用特定于域的张量表达式来进行有效的内核构造。
在本教程中，我们将演示使用张量表达式语言的基本工作流程。
from __future__ import absolute_import, print_function
import tvm
import numpy as np
# Global declarations of environment.
tgt_host="llvm"
# Change it to respective GPU if gpu is enabled Ex: cuda, opencl
tgt="cuda"
Vector Add Example
在本教程中，我们将使用向量添加示例来演示工作流程。
描述计算
作为第一步，我们需要描述我们的计算。TVM采用Tensor语义，每个中间结果表示为多维数组。用户需要描述生成Tensor的计算规则。
我们首先定义一个符号变量n来表示形状。然后我们定义两个占位符Tensor A和B，给定形状（n，）
然后我们用计算操作描述结果Tensor C. 计算函数采用张量的形状，以及描述张量的每个位置的计算规则的lambda函数。
在此阶段没有计算，因为我们只是声明应该如何进行计算。
n = tvm.var("n")
A = tvm.placeholder((n,), name='A')
B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C")
print(type(C))
输出：
<class 'tvm.tensor.Tensor'>
计算图 Schedule the Computation
虽然上面的行描述了计算规则，但我们可以以多种方式计算C，因为C轴可以以数据并行方式计算。TVM要求用户提供的计算描述，我们称之为一个schedule。
schedule是一组计算转换，它转换程序中的计算循环。
在我们构造schedule之后，默认情况下，schedule以行主要顺序以串行方式计算C.
for (int i = 0; i < n; ++i) {
C[i] = A[i] + B[i];
}
s = tvm.create_schedule(C.op)
我们使用split构造来分割C的第一个轴，这将原始迭代轴分成两次迭代的乘积。这相当于以下代码:
for (int bx = 0; bx < ceil(n / 64); ++bx) {
for (int tx = 0; tx < 64; ++tx) {
int i = bx * 64 + tx;
if (i < n) {
C[i] = A[i] + B[i];
}
}
}
bx, tx = s[C].split(C.op.axis[0], factor=64)
最后，我们将迭代轴bx和tx绑定到GPU计算网格中的线程。这些是GPU特定的构造，允许我们生成在GPU上运行的代码。
if tgt == "cuda" or tgt.startswith('opencl'):
s[C].bind(bx, tvm.thread_axis("blockIdx.x"))
s[C].bind(tx, tvm.thread_axis("threadIdx.x"))
Compilation
在完成指定Schedule后，我们可以将其编译为TVM函数。默认情况下，TVM编译成一个与类型无关的函数，可以从python端直接调用。
在下面的行中，我们使用tvm.build来创建一个函数。构建函数采用调度，函数的期望签名（包括输入和输出）以及我们要编译的目标语言。
编译fadd的结果是GPU设备功能（如果涉及GPU）以及调用GPU功能的主机包装器。fadd是生成的主机包装器函数，它包含对内部生成的设备函数的引用。
fadd = tvm.build(s, [A, B, C], tgt, target_host=tgt_host, name="myadd")
运行功能
编译的TVM函数公开了一个可以从任何语言调用的简洁C API。
我们在python中提供了一个最小的数组API，以帮助快速测试和原型设计。阵列API基于DLPack标准。
ctx = tvm.context(tgt, 0)
n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
fadd(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
检查生成的代码
您可以在TVM中检查生成的代码。tvm.build的结果是一个TVM模块。fadd是包含主机包装器的主机模块，它还包含用于CUDA（GPU）功能的设备模块。
以下代码获取设备模块并打印内容代码。
if tgt == "cuda" or tgt.startswith('opencl'):
dev_module = fadd.imported_modules[0]
print("-----GPU code-----")
print(dev_module.get_source())
else:
print(fadd.get_source())
日期：
-----GPU code-----
extern "C" __global__ void myadd_kernel0( float* __restrict__ C, float* __restrict__ A, float* __restrict__ B, int n) {
if ((((int)blockIdx.x) < (((n - 64) / 64) + 1)) && (((int)blockIdx.x) < (((n + 63) / 64) - 1))) {
C[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = (A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] + B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
C[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = (A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] + B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
}
}
}
}
注意
代码定制化
您可能已经注意到，A，B和C的声明都采用相同的形状参数，n。正如您将在打印的设备代码中找到的那样，TVM将利用此功能仅将单个形状参数传递给内核。这是一种专业化形式。
在主机端，TVM将自动生成检查参数中的约束的检查代码。因此，如果将具有不同形状的数组传递给fadd，则会引发错误。
我们可以做更多的专业化。例如，我们可以在计算声明中编写 n = tvm.convert(1024)而不是n = tvm.var(“n”)使生成的函数仅采用长度为1024的向量。
保存编译模块
除了运行时编译之外，我们还可以将已编译的模块保存到文件中，并在以后加载它们。我们称这个为提前编译。
以下代码首先执行以下步骤：
from tvm.contrib import cc
from tvm.contrib import util
temp = util.tempdir()
fadd.save(temp.relpath("myadd.o"))
if tgt == "cuda":
fadd.imported_modules[0].save(temp.relpath("myadd.ptx"))
if tgt.startswith('opencl'):
fadd.imported_modules[0].save(temp.relpath("myadd.cl"))
cc.create_shared(temp.relpath("myadd.so"), [temp.relpath("myadd.o")])
print(temp.listdir())
输出：
['myadd.so', 'myadd.ptx', 'myadd.o', 'myadd.tvm_meta.json']
注意
模块存储格式
CPU（主机）模块直接保存为共享库（.so）。可以有多种自定义格式的设备代码。在我们的示例中，设备代码存储在ptx中，以及元数据json文件中。它们可以通过导入实现单独加载和链接。
加载编译模块
我们可以从文件系统加载已编译的模块并运行代码。以下代码分别加载主机和设备模块并将它们重新链接在一起。我们可以验证新加载的函数是否有效。
fadd1 = tvm.module.load(temp.relpath("myadd.so"))
if tgt == "cuda":
fadd1_dev = tvm.module.load(temp.relpath("myadd.ptx"))
fadd1.import_module(fadd1_dev)
if tgt.startswith('opencl'):
fadd1_dev = tvm.module.load(temp.relpath("myadd.cl"))
fadd1.import_module(fadd1_dev)
fadd1(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
将所有内容打包到一个库中
在上面的示例中，我们分别存储设备和主机代码。TVM还支持将所有内容导出为一个共享库。我们将设备模块打包成二进制blob，并将它们与主机代码链接在一起。目前，我们支持Metal，OpenCL和CUDA模块的包装。
fadd.export_library(temp.relpath("myadd_pack.so"))
fadd2 = tvm.module.load(temp.relpath("myadd_pack.so"))
fadd2(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
注意，运行时API和线程安全
TVM的编译模块不依赖于TVM编译器。相反，它们仅依赖于最小运行时库。TVM运行时库包装设备驱动程序，并为编译的函数提供线程安全和设备无关的调用。
这意味着您可以从任何GPU上的任何线程调用已编译的TVM函数。
生成OpenCL代码
TVM为多个后端提供代码生成功能，我们还可以生成在CPU后端上运行的OpenCL代码或LLVM代码。
以下代码块生成OpenCL代码，在OpenCL设备上创建数组，并验证代码的正确性。
if tgt.startswith('opencl'):
fadd_cl = tvm.build(s, [A, B, C], tgt, name="myadd")
print("------opencl code------")
print(fadd_cl.imported_modules[0].get_source())
ctx = tvm.cl(0)
n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
fadd_cl(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy())
Summary
本教程使用向量添加示例介绍TVM工作流程。一般工作流程是
-
通过一系列操作描述您的计算。
-
描述我们如何计算使用schedule原函数。
-
编译到我们想要的目标函数。
-
（可选）保存稍后要加载的功能。
我们非常欢迎您查看其他示例和教程，以了解有关TVM中支持的操作，调度原语和其他功能的更多信息。
