Author: Tianqi Chen
https://docs.tvm.ai/tutorials/tensor_expr_get_started.html

Tensor Expression入门

这是TVM中Tensor表达语言的入门教程。TVM使用特定于域的张量表达式来进行有效的内核构造。

在本教程中，我们将演示使用张量表达式语言的基本工作流程。

from __future__ import absolute_import, print_functionimport 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标准。

我们首先创建一个GPU上下文。
然后tvm.nd.array将数据复制到GPU。
fadd运行实际计算。
asnumpy（）将GPU阵列复制回CPU，我们可以使用它来验证正确性

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的向量。

保存编译模块

除了运行时编译之外，我们还可以将已编译的模块保存到文件中，并在以后加载它们。我们称这个为提前编译。

以下代码首先执行以下步骤：

它将已编译的主机模块保存到目标文件中。
然后它将设备模块保存到ptx文件中。
cc.create_shared调用编译器（gcc）来创建共享库

from tvm.contrib import cc
from tvm.contrib import utiltemp = 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 = 1024a = 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中支持的操作，调度原语和其他功能的更多信息。

TVM系列---1.开始使用Tensor Expression相关推荐

TVM：使用Tensor Expression (TE)来处理算子
TVM:使用Tensor Expression (TE)来处理算子在本教程中,我们将聚焦于在 TVM 中使用张量表达式(TE)来定义张量计算和实现循环优化.TE用纯函数语言描述张量计算(即每个表达式 ...
使用Tensor Expression张量表达式处理算子
使用Tensor Expression张量表达式处理算子这是TVM中Tensor表达语言的入门教程.TVM使用特定于域的张量表达式来进行有效的内核构造. 本文将演示使用张量表达式语言的基本工作流程. ...
【TVM系列二】TVM介绍
文章同步更新在公众号 AIPlayer,欢迎扫码关注,共同进步目录一.TVM的工作流程 1.整体流程 2.关键数据结构 3.Transformations 4.搜索空间和基于机器学习的转换 5.目 ...
TVM系列 #002_About VTA
About VTA [Text] The Versatile[通用的] Tensor Accelerator (VTA) is an extension of the TVM framework de ...
人脸表情系列：论文阅读——Facial Expression Recognition by De-expression Residue Learning
task是表情识别,将一张图像视为两部分组成:expressive component和neutral component.从一张图像中生成该个体的无表情图像称为de-expression learn ...
AI System 人工智能系统 TVM深度学习编译器 DSL IR优化计算图编译优化内存内核调度优化 DAG 图优化 DFS TaiChi 函数注册机 Registry
DSL 领域专用语言 TVM深度学习编译器 AI System 人工智能系统参考项目 TaiChi 三维动画渲染物理仿真引擎DSL TVM 深度学习DSL 密集计算DSL LLVM 模块化编译器编 ...
用动态实现扩展TVM
用动态实现扩展TVM Extending TVM with Dynamic Execution Outline ● Motivation for Dynamism ● Representing Dyn ...
TVM 各个模块总体架构
TVM 各个模块总体架构 Deploy Deep Learning Everywhere Existing Deep Learning Frameworks Limitations of Existi ...
TVM：使用 Schedule 模板和 AutoTVM 来优化算子
TVM:使用 Schedule 模板和 AutoTVM 来优化算子在本文中,我们将介绍如何使用 TVM 张量表达式(Tensor Expression,TE)语言编写 Schedule 模板,Aut ...

TVM系列---1.开始使用Tensor Expression