使用 TE 来处理运算符

张量表达式（TE）指的是一种纯函数语言来描述张量计算（这意味着每一个表达式都没有副作用）。当在TVM整体环境下来看，Relay 用一系列操作符来描述一个计算过程。这一系列操作符中的每一个操作符可以都被表示为一个TE表达式，并且这每一个TE表达式都采用张量输入，然后生成一个张量输出。
调度是应用于将一个表达式通过多种不同的方式进行变换的一系列步骤。当调度应用于TE中的表达式时，输入和输出保持不变，但当编译时表达式的实现可能会发生变化。

为 CPU 用 TE 的形式写一个向量加法和调度

默认调度

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

# 指定 target
tgt = tvm.target.Target(target="llvm -mcpu=skylake-avx512",host="llvm -mcpu=skylake-avx512")

# 描述计算
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 表达式来表示计算方法

# 为计算创建默认的调度,不同的调度将会生成不同的代码,此处，默认的调度为：
"""
for (int i = 0; i < n; ++i) {
  C[i] = A[i] + B[i];
}
"""
s = te.create_schedule(C.op)

# 编译并且评估默认调度
# tvm.build 来创建一个函数。build函数接收调度，所需的函数签名（包括输入和输出）还有我们想要编译成的目标语言。
fadd = tvm.build(s,[A,B,C],tgt,name="vec_add")

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())
print("Done.")

性能分析

import timeit

np_repeat = 100
np_running_time = timeit.timeit(
    setup="import numpy\n"
    "n = 100000\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 = 100000
    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.000035
naive: 0.000026

并行调度

在默认的调度中，这种张量加运算是串行的，但很容易跨所有处理器线程并行话。我们可以将并行调度操作应用于我们的计算。

# 更新调度
s[C].parallel(C.op.axis[0])

# tvm.lower 生成 TE 的中间表示，以及对应的调度
# simple_mode=True 标志返回可读的 C 风格语句
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出:

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

TVM现在可以在独立的线程中运行这些块。让我们编译和运行这个新的应用了并行操作的调度：

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

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

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

输出：

parallel: 0.000050

向量化调度

现代CPU还能够对浮点值执行SIMD操作，我们可以通过将另外一个调度应用到我们的表达式计算中以利用这一点。实现这个需要多个步骤：首先，我们必须使用split调度原语将调度拆分为内循环和外循环。内循环可以使用矢量化调度源语来使用SIMD指令进行矢量化，外循环可以使用并行调度源语并行化。

# 由于之前更新过调度，所以这里重新创建调度
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)

# 这个参数根据 CPU 核的具体核数量来指定
factor = 8

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))

输出：

vector: 0.000658

综合分析不同的调度性能结果：

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))
    )

</