初探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）

我们随后开始看一看TE的基础内容，我们深入了解下调度是做什么的，以及怎么在目标平台上优化一个张量表达式。调度是一系列对于张量表达式的变形，当我们应用一个调度时，表达式的输入输出不变，仅仅在编译过程中的实现方法会改变。在张量加法算法中，加法是顺序被执行，因此比较容易使用处理器的多线程做并行化优化。代码这样子的：

s[C].parallel(C.op.axis[0])

tvm.lower()函数可以生成TE在对应优化调度下的IR（intermediate representation）。通过lower表达式的方式，我们可以查看不同调度对计算顺序的影响。我们可以用一个标志位simple_mode=true来输出一个可读的类C的伪代码：

print(tvm.lower(s, [A, B, C], simple_mode=true))

会看到下面的输出：

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),
             B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, n) "parallel" {
    C_2[(i*stride)] = ((float32*)A_2[(i*stride_1)] + (float32*)B_2[(i*stride_2)])
  }
}

现在就可以在不同的线程上并行运行这段表达式了，重新编译后可以再运行下这个并行化的表达式，看看效果：

fadd_parallel = tvm.build(s, [A, B, C], tgt, name="myadd_parallel")
fadd_parallel(a, b, c)

tvm.testing.assert_allclose(c.numpy(), a.numpy(), b.numpy())

evaluate_additiong(fadd_parallel, tgt, "parallel", log=log)

看看效果：

parallel: 0.000032

好像是比以前快了，但是没有到预期，因为128线程，只加速了两倍。猜测是因为n太小， 主要的耗时在访存上了，因此多线程加速起不到应有的效果。

使用向量化的优化调度

现在的各种CPU DSP都支持SIMD的向量化操作，我们可以用一个调度，把SIMD用上的。在TVM中，我们需要两个步骤来完成simd的调度：1.把循环拆分成内层和外层，内层是可以用单一simd指令完成的，外层是用于多线程完成的。由于要拆分循环，所以需要设定一个拆分的单位，以适配CPU的线程数和simd的宽度。

# Recreate the schedule, since we modified it with the parallel operation in
# the previous example
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")

s = te.create_schedule(C.op)

# This factor should be chosen to match the number of threads appropriate for
# your CPU. This will vary depending on architecture, but a good rule is
# setting this factor to equal the number of available CPU cores.
factor = 4

outer, inner = s[C].split(C.op.axis[0], factor=factor)
s[C].parallel(outer)
s[C].vectorize(inner)

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

evaluate_addition(fadd_vector, tgt, "vector", log=log)

print(tvm.lower(s, [A, B, C], simple_mode=True))

看一下vectorize后的伪代码吧：

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),
             B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i.outer: int32, 0, floordiv((n + 3), 4)) "parallel" {
    for (i.inner.s: int32, 0, 4) {
      if @tir.likely((((i.outer*4) + i.inner.s) < n), dtype=bool) {
        C_2[(((i.outer*4) + i.inner.s)*stride)] = ((float32*)A_2[(((i.outer*4) + i.inner.s)*stride_1)] + (float32*)B_2[(((i.outer*4) + i.inner.s)*stride_2)])
      }
    }
  }
}

这时的耗时为：

vector: 0.000062

比不用simd还慢了。。。。

对比几种优化调度

上面尝试了很多调度后，我们可以对比一下：

baseline = log[0][1]
print("%s\t%s\t%s" % ("Operator".rjust(20), "Timing".rjust(20), "Performance".rjust(20)))
for result in log:
    print(
        "%s\t%s\t%s"
        % (result[0].rjust(20), str(result[1]).rjust(20), str(result[1] / baseline).rjust(20))
    )

Operator                  Timing             Performance
               numpy    6.184579997352557e-06                    1.0
               naive    4.635100000000001e-06     0.7494607559420615
            parallel             3.20865e-05       5.188145357281386
              vector    6.325680000000001e-05      10.22814807587232

我们在定义A/B/C的时候，给了他们一个相同的shape，就是n。tvm会利用这个相同的shape的特性做出些有效的优化。但是在调用优化库的时候，tvm会自动生成shape检测代码，如果给传入的向量不是这个shape，就会报错。也有些其他的方法来避免这些错误的出现，比如在定义运算的阶段，我们使用n=tvm.runtime.convert(1024)来代替n=te.var("n")，这样编译出的函数就会仅仅对1024内的长度做向量化。

使用tvm，我们能够定义、优化、编译一个向量加法的算子，并且在tvm的运行时组件上可以运行。这个算子可以被保存为一个库，在tvm的运行时组件中加载他。