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

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标准。

• 我们首先创建一个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 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中支持的操作，调度原语和其他功能的更多信息。