使用Tensor Expression张量表达式处理算子

使用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中支持的操作，调度原语和其它功能的更多信息。

 

参考文献：

https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html

https://blog.csdn.net/xxradon/article/details/98213214