[Python知识库]初探TVM--使用Tensor Engine来编写算子

什么是TE

就像在题目中写的那样，TE就是Tensor Engine的简称，其实就是用这些接口来定义一个计算算子是干什么的。可以转到tvm的一个大体介绍里面再看看。

通过这个教程，我们会学习TVM中，怎样通过TE定义一个张量计算算子。TE使用纯粹的功能性的语言描述一个张量计算，在这里是TE是没有边际效应（其实这个side effect一直不好理解，我觉得就是说换了硬件平台的话，也不应该对它产生影响）的。从tvm整体上看，relay是用来描述一组算子组成的计算图，TE是用来描述计算图中的每个算子节点，也就是说TE中可以认为是接受一个或多个输入张量，然后输出一个或多个输出张量。

在这个TE的入门的学习文章中，TVM用几个领域专用的张量计算来完成高效的算子构建。这里会通过两个算子来展示使用TE的技巧。第一个例子是个向量加法的例子，通过它来理解使用TE和TE的调度。然后我们会把这些调度概念一步一步拓展到一个compute bound的矩阵乘算子的例子上面去。同时，矩阵乘会被用到其他的TVM的操作当中，这个例子是一个非常好的入门的例子。

第一个例子：用TE写一个CPU的向量加法

我们先在Python的TE里面实现一个简单的向量加法吧，在中间加入一些基于CPU优化的优化调度试试。写Python的时候，还是老样子，首先import需要的依赖库。

import tvm 
import tvm.testing
import numpy as np
from tvm import te

还要再提一个说了三四次的东西了，就是如果把CPU的型号信息告知llvm，它会做出更多基于CPU特性的优化，比如被Linus疯狂吐槽的avx512。可以通过llc --version来查看CPU的版本信息，也可以通过cat /proc/cpuinfo来查看本台CPU的拓展新能包。例如，再查看过后可以使用llvm -mcpu=skylake-avx512来使能你的向量指令集。很遗憾，我用的AMD CPU，没有avx512. 不过CPU的型号是znver2，好像更高级一些啊。

tgt = tvm.target.Target(target="llvm", host="llvm")

描述一个向量计算

描述一个向量加法的计算。tvm支持向量语义，可以把每一层级的中间结果表达成为多维数组形式。用户需要描述生成向量的计算规则。我们可以先定义一个变量n来表示向量的形状，然后定义两个张量（tensor）占位符，A和B，他们都有相同的形状(n,)。然后我们可以用一个compute操作来描述结果张量c。compute操作定义了一种计算，使得输出遵循一个特定的张量形状，并且在张量的每个元素上都做一个使用lambda表达式定义的计算。因为n此时是一个变量，所以A B和C被定义了一个一致的形状。此时还没有实际的计算发生，我们仅仅声明了计算是什么。

n = te.var('n')
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i]+B[i], name="C")

lambda表达式（lambda functions）
te.compute方法的第二个参数是一个执行计算的函数，在这个例子中间，我们使用了一个匿名函数的方式，也就是一个lambda表达式，本例中是期望对于A和B的第i个元素做加法。

给计算创造一个默认的优化调度

虽然上面的几行代码定义了计算规则，我们仍然可以用不同的方式计算C以适应不同的设备。对于一个有多个维度的张量，你可以选择比如：1.首先计算的维度；2.哪些计算可以被拆分到不同的线程中。TVM要求使用者提供一套调度来描述计算是怎样被执行的。TE中的调度操作可以改变循环顺序，拆分计算到不同线程，以及对数据分块。一个重要的概念是，调度仅仅描述计算怎样被执行，也就是说，一个tensor expression上面应用不同的调度，应该产生同样的计算结果。

TVM允许创造一个基本的调度以row major的方式去计算C。

for (int i = 0; i < n; i++)
{
    C[i] = A[i] + B[i];
}

s = te.create_schedule(C.op)

编译并且评估默认的调度

使用TE的表达式和一个优化调度，我们可以生成一个在目标架构和语言的可执行代码，现在就是LLVM和X86 CPU。我们给TVM提供：调度，TE表达式的列表，目标机器和主机，函数的名称。这个输出结果是一个无类型的函数，并且可以在python端直接被调用。

用下面一行代码，我们可以通过tvm.build生成一个函数。build函数拿到这个调度，函数签名（包含输入输出张量），以及目标语言。

fadd = tvm.build(s, [A, B, C], tgt, name="myadd")

我们可以运行这个函数，并且把输出结果和用numpy写的代码的结果作比较。编译后的TVM函数暴露出一个简明的C接口，这个结果可以通过任何语言调用。首先我们需要定义一个device，TVM可以在这个device上编译出相应的优化调度。当前我们的device是LLVM CPU。然后我们可以在device上初始化张量并且执行这个加法操作。我们通过对比TVM函数的输出结果和numpy的输出结果来验证正确性。

dev = tvm.device(tgt.kind.name, 0)

n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)
fadd(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

可以通过一个helper函数来profile这个TVM生成的代码，以此来和numpy函数比较一下运行速度。

import timeit

np_repeat = 100
np_running_time = timeit.timeit(
    setup="import numpy\n"
    "n = 32768\n"
    'dtype = "float32"\n'
    "a = numpy.random.rand(n, 1).astype(dtype)\n"
    "b = numpy.random.rand(n, 1).astype(dtype)\n",
    stmt="answer = a + b",
    number=np_repeat,
)
print("Numpy running time: %f" % (np_running_time / np_repeat))


def evaluate_addition(func, target, optimization, log):
    dev = tvm.device(target.kind.name, 0)
    n = 32768
    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
    b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)
    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)

    evaluator = func.time_evaluator(func.entry_name, dev, number=10)
    mean_time = evaluator(a, b, c).mean
    print("%s: %f" % (optimization, mean_time))

    log.append((optimization, mean_time))


log = [("numpy", np_running_time / np_repeat)]
evaluate_addition(fadd, tgt, "naive", log=log)

输出：

Numpy running time: 0.000010
naive: 0.000010

不得不说，还挺慢。。。再把架构指定一下吧：

tgt = tvm.target.Target(target="llvm -mcpu=znver2", host="llvm")

这次快了点儿：

Numpy running time: 0.000007
naive: 0.000005

使用并行化优化调度（paralleism）