本文翻译自Working with Operators Using Tensor Expression — tvm 0.9.dev0 documentation
在本教程中，我们将把注意力转向TVM如何使用张量表达式(TE)定义张量计算和实现循环优化。TE用纯函数语言描述张量计算(即每个表达式都没有副作用)。当把TVM看作一个整体时，Relay将一个计算描述为一组算子，每个算子可以表示为一个TE表达式，每个TE表达式接受输入张量并产生一个输出张量。
这是TVM中张量表达式语言的入门教程。TVM使用一个特定于域的张量表达式来有效地构造核心。我们将通过两个使用张量表达式语言的示例来演示基本工作流。第一个例子介绍了TE和使用向量加法的调度。第二部分对这些概念进行了扩展，逐步优化了使用TE的矩阵乘法。这个矩阵乘法示例将作为以后教程中介绍TVM的更高级特性的对比基础。
让我们来看一个Python的例子，在这个例子中，我们将实现一个向量加法的TE，然后以CPU为target进行调度。首先初始化TVM环境。
import tvm
import tvm.testing
from tvm import te
import numpy as np如果能够识别并指定目标CPU，则可以获得更好的性能。如果你正在使用LLVM，可以用命令llc --version获取CPU类型的信息，并使用/proc/cpuinfo以获得你的处理器可能支持的其他扩展。例如，AVX-512指令的cpu，可以使用llvm -mcpu=skylake-avx512指定。
tgt = tvm.target.Target(target="llvm", host="llvm")我们来描述一个向量加法的计算。TVM采用张量语义，将每个中间结果表示为一个多维数组。用户需要描述生成张量的计算规则。我们首先定义一个符号变量n来表示形状。然后定义两个占位张量，A和B，具有给定的形状(n, )。然后我们用一个运算操作来描述结果张量C。在运算的定义中，输出符合指定的张量形状，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")lamada函数：te.compute方法的第二个参数是执行计算的函数。在这个例子中，我们使用了一个匿名函数，也称为lambda函数，来定义计算，在这个例子中是A和B的第i个元素的加法。
虽然上面几行描述了计算规则，但我们可以用许多不同的方式来计算C，以适应不同的设备。对于具有多个轴的张量，可以选择先迭代哪个轴，或者可以将计算拆分到不同的线程中。TVM要求用户提供一个调度，以描述如何执行计算。TE内的调度操作可以更改循环顺序，跨线程分割计算，以及在其他操作之间对数据块分组。调度背后的一个重要概念是，它们只描述了如何执行计算，因此对同一个TE，不同的调度将产生结果是相同的。
TVM允许您创建一个简单的调度，它将通过按行主顺序迭代来计算C。
for (int i = 0; i < n; ++i) {
C[i] = A[i] + B[i];
}s = te.create_schedule(C.op)通过TE表达式和调度策略，我们可以为我们的目标语言和体系结构生成可运行的代码，如当前示例中的LLVM和CPU。我们向TVM提供调度、调度的TE表达式列表、目标和主机，以及我们正在生成的函数名。输出的结果是一个可以在Python中直接调用的类型擦除函数。
接下来我们使用tvm.build创建函数。构建函数接受的参数包括调度策略、所需的函数签名(包括输入和输出)以及我们想要编译得到的目标语言。
fadd = tvm.build(s, [A, B, C], tgt, name="myadd")让我们运行这个函数，并将输出与使用numpy做相同计算得到的结果进行比较。编译后的TVM函数提供一个简明的C API，可以从任何语言调用它。我们首先创建一个设备，TVM可以将调度编译到这个设备(在本例中是CPU)上。在本例中，设备为LLVM CPU目标。然后，我们可以初始化设备中的张量，并执行自定义的加法运算。为了验证计算的正确性，我们可以将输出结果的c张量，与使用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())要比较这个版本与numpy的速度，可以创建一个helper函数来运行TVM生成的代码的概要文件。
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.000007
naive: 0.000006现在我们已经了解了TE的基本原理，接下来让我们更深入地了解调度的功能，以及如何使用它们来优化不同架构的张量表达式。调度是以一系列步骤，许多不同的方式对表达式做转换。当对TE中的表达式应用调度时，输入和输出保持不变，但在编译后，表达式的实现可能会改变。在示例中我们使用的默认调度，加法运算是串行的，但是很容易对它进行多线程并行化。我们可以对计算做并行调度操作。
s[C].parallel(C.op.axis[0])tvm.lower命令将根据对应的调度策略生成TE的中间表达（IR）。通过低级化不同调度策略的表达式，我们观察到调度对计算顺序的影响。我们使用标志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="myadd_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.000006现代cpu具备对浮点值执行SIMD操作的能力，我们可以利用这一点对计算表达式应用另一种调度。完成这一任务需要多个步骤:首先，我们必须使用分割调度原语将调度分割为内部和外部循环。内部循环通过向量化调度原语向量化，以利用SIMD指令，外部循环使用并行调度原语并行化。选择的分割因子为CPU上的线程数。
# 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))输出：
vector: 0.000026
@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.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) {
let cse_var_1: int32 = ((i.outer*4) + i.inner.s)
C[(cse_var_1*stride_2)] = (A[(cse_var_1*stride)] + B[(cse_var_1*stride_1)])
}
}
}
}现在我们可以比较不同的调度
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.66109990561381e-06 1.0
naive 5.8282e-06 0.8749606044923809
parallel 6.0927e-06 0.9146687613655552
vector 2.6245599999999998e-05 3.9401300643878434代码说明：你可能已经注意到的，A, B和C的声明都有相同的形状参数n。TVM将利用这一点只传递一个形状参数给内核，正如你将在打印的设备代码中发现的那样。这是一种专门的形式。在主机端，TVM将自动生成检查代码，以参数的约束条件。因此，如果你将一个其他的数组传递给fadd，将会引发一个错误。我们可以做更多的专门化。例如，我们可以在计算声明中写n = tvm.runtime.convert(1024)，而不是n = te.var("n")。那么生成的函数将只接受长度为1024的向量。
我们已经定义、调度和编译了一个向量加法运算符，然后就可以在TVM运行时中执行它了。我们可以将操作符保存为一个库，然后使用TVM运行时加载它。
TVM可以面向多种架构，下一个示例我们将编译一个GPU的向量加法
# If you want to run this code, change ``run_cuda = True``
# Note that by default this example is not run in the docs CI.
run_cuda = False
if run_cuda:
# Change this target to the correct backend for you gpu. For example: cuda (NVIDIA GPUs),
# rocm (Radeon GPUS), OpenCL (opencl).
tgt_gpu = tvm.target.Target(target="cuda", host="llvm")
# Recreate the schedule
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")
print(type(C))
s = te.create_schedule(C.op)
bx, tx = s[C].split(C.op.axis[0], factor=64)
################################################################################
# Finally we must bind the iteration axis bx and tx to threads in the GPU
# compute grid. The naive schedule is not valid for GPUs, and these are
# specific constructs that allow us to generate code that runs on a GPU.
s[C].bind(bx, te.thread_axis("blockIdx.x"))
s[C].bind(tx, te.thread_axis("threadIdx.x"))
######################################################################
# Compilation
# -----------
# After we have finished specifying the schedule, we can compile it
# into a TVM function. By default TVM compiles into a type-erased
# function that can be directly called from the python side.
#
# In the following line, we use tvm.build to create a function.
# The build function takes the schedule, the desired signature of the
# function (including the inputs and outputs) as well as target language
# we want to compile to.
#
# The result of compilation fadd is a GPU device function (if GPU is
# involved) as well as a host wrapper that calls into the GPU
# function. fadd is the generated host wrapper function, it contains
# a reference to the generated device function internally.
fadd = tvm.build(s, [A, B, C], target=tgt_gpu, name="myadd")
################################################################################
# The compiled TVM function exposes a concise C API that can be invoked from
# any language.
#
# We provide a minimal array API in python to aid quick testing and prototyping.
# The array API is based on the `DLPack <https://github.com/dmlc/dlpack>`_ standard.
#
# - We first create a GPU device.
# - Then tvm.nd.array copies the data to the GPU.
# - ``fadd`` runs the actual computation
# - ``numpy()`` copies the GPU array back to the CPU (so we can verify correctness).
#
# Note that copying the data to and from the memory on the GPU is a required step.
dev = tvm.device(tgt_gpu.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())
################################################################################
# Inspect the Generated GPU Code
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# You can inspect the generated code in TVM. The result of tvm.build is a TVM
# Module. fadd is the host module that contains the host wrapper, it also
# contains a device module for the CUDA (GPU) function.
#
# The following code fetches the device module and prints the content code.
if (
tgt_gpu.kind.name == "cuda"
or tgt_gpu.kind.name == "rocm"
or tgt_gpu.kind.name.startswith("opencl")
):
dev_module = fadd.imported_modules[0]
print("-----GPU code-----")
print(dev_module.get_source())
else:
print(fadd.get_source())除了运行时编译之外，我们还可以将编译后的模块保存到一个文件中，稍后再将它们加载回来。
下面的代码首先执行以下步骤:
from tvm.contrib import cc
from tvm.contrib import utils
temp = utils.tempdir()
fadd.save(temp.relpath("myadd.o"))
if tgt.kind.name == "cuda":
fadd.imported_modules[0].save(temp.relpath("myadd.ptx"))
if tgt.kind.name == "rocm":
fadd.imported_modules[0].save(temp.relpath("myadd.hsaco"))
if tgt.kind.name.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())模块存储格式：CPU(主机)模块直接保存为共享库(.so)。设备代码可以有多种自定义格式。在我们的例子中，设备代码存储在ptx中，以及一个元数据json文件中。它们可以各自独立的通过导入加载和链接。
我们可以从文件系统加载编译好的模块并运行代码。下面的代码分别加载主机和设备模块，并将它们链接在一起。我们可以验证新加载的函数是否正常工作。
fadd1 = tvm.runtime.load_module(temp.relpath("myadd.so"))
if tgt.kind.name == "cuda":
fadd1_dev = tvm.runtime.load_module(temp.relpath("myadd.ptx"))
fadd1.import_module(fadd1_dev)
if tgt.kind.name == "rocm":
fadd1_dev = tvm.runtime.load_module(temp.relpath("myadd.hsaco"))
fadd1.import_module(fadd1_dev)
if tgt.kind.name.startswith("opencl"):
fadd1_dev = tvm.runtime.load_module(temp.relpath("myadd.cl"))
fadd1.import_module(fadd1_dev)
fadd1(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())在上面的例子中，我们分别存储设备代码和主机代码。TVM还支持将所有内容导出为一个共享库。在内部，我们将设备模块打包成二进制块，并将它们与主机代码链接在一起。目前我们支持封装Metal, OpenCL和CUDA模块。
fadd.export_library(temp.relpath("myadd_pack.so"))
fadd2 = tvm.runtime.load_module(temp.relpath("myadd_pack.so"))
fadd2(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())运行时API和线程安全：
编译后的TVM模块不依赖于TVM编译器。相反，它们只依赖于一个最小的运行时库。TVM运行时库封装了设备驱动程序，并为编译后的函数提供了线程安全的和设备无关的调用。
这意味着你可以在任何GPU上的任何线程中调用编译过的TVM函数，前提是你已经为那个GPU编译了代码。
TVM为多种后端提供代码生成特性。我们也可以生成运行在CPU后端上的OpenCL代码或LLVM代码。
下面的代码块生成OpenCL代码，在OpenCL设备上创建数组，并验证代码的正确性。
if tgt.kind.name.startswith("opencl"):
fadd_cl = tvm.build(s, [A, B, C], tgt, name="myadd")
print("------opencl code------")
print(fadd_cl.imported_modules[0].get_source())
dev = tvm.cl(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_cl(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())TE调度原语：TVM包含许多不同的调度原语:
这些原语的完整描述可以在Schedule Primitives in TVM — tvm 0.9.dev0 documentation页面中找到。
现在，我们将考虑第二个更高级的示例，演示TVM如何只用18行python代码将一个常见的矩阵乘法运算加速18倍。
矩阵乘法是一种计算密集型运算。对于良好的CPU性能，有两个重要的优化:
本教程中使用的技术是代码仓中提到的技巧的一部分。它们中的一些已经被TVM抽象地自动应用，但是由于TVM的限制，它们中的一些不能自动应用。
我们首先收集numpy实现的矩阵乘法的性能数据。
import tvm
import tvm.testing
from tvm import te
import numpy
# The size of the matrix
# (M, K) x (K, N)
# You are free to try out different shapes, sometimes TVM optimization outperforms numpy with MKL.
M = 1024
K = 1024
N = 1024
# The default tensor data type in tvm
dtype = "float32"
# You will want to adjust the target to match any CPU vector extensions you
# might have. For example, if you're using using Intel AVX2 (Advanced Vector
# Extensions) ISA for SIMD, you can get the best performance by changing the
# following line to ``llvm -mcpu=core-avx2``, or specific type of CPU you use.
# Recall that you're using llvm, you can get this information from the command
# ``llc --version`` to get the CPU type, and you can check ``/proc/cpuinfo``
# for additional extensions that your processor might support.
target = tvm.target.Target(target="llvm", host="llvm")
dev = tvm.device(target.kind.name, 0)
# Random generated tensor for testing
a = tvm.nd.array(numpy.random.rand(M, K).astype(dtype), dev)
b = tvm.nd.array(numpy.random.rand(K, N).astype(dtype), dev)
# Repeatedly perform a matrix multiplication to get a performance baseline
# for the default numpy implementation
np_repeat = 100
np_running_time = timeit.timeit(
setup="import numpy\n"
"M = " + str(M) + "\n"
"K = " + str(K) + "\n"
"N = " + str(N) + "\n"
'dtype = "float32"\n'
"a = numpy.random.rand(M, K).astype(dtype)\n"
"b = numpy.random.rand(K, N).astype(dtype)\n",
stmt="answer = numpy.dot(a, b)",
number=np_repeat,
)
print("Numpy running time: %f" % (np_running_time / np_repeat))
answer = numpy.dot(a.numpy(), b.numpy())输出：
Numpy running time: 0.017980现在，我们使用TVM TE编写一个基本的矩阵乘法，并验证它产生的结果与numpy实现相同。我们还编写了一个函数来帮助我们测量调度优化的性能。
# TVM Matrix Multiplication using TE
k = te.reduce_axis((0, K), "k")
A = te.placeholder((M, K), name="A")
B = te.placeholder((K, N), name="B")
C = te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")
# Default schedule
s = te.create_schedule(C.op)
func = tvm.build(s, [A, B, C], target=target, name="mmult")
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)
def evaluate_operation(s, vars, target, name, optimization, log):
func = tvm.build(s, [A, B, C], target=target, name="mmult")
assert func
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)
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 = []
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="none", log=log)输出：
none: 3.230811让我们通过TVM低级函数来看看操作符和默认调度的中间表示。注意，这个实现本质上是一个简单的矩阵乘法实现，在A和B矩阵的索引上使用了三个嵌套循环。
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, [1048576], []),
B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}
buffer_map = {A_1: A, B_1: B, C_1: C}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {
for (x: int32, 0, 1024) {
for (y: int32, 0, 1024) {
C[((x*1024) + y)] = 0f32
for (k: int32, 0, 1024) {
let cse_var_2: int32 = (x*1024)
let cse_var_1: int32 = (cse_var_2 + y)
C[cse_var_1] = (C[cse_var_1] + (A[(cse_var_2 + k)]*B[((k*1024) + y)]))
}
}
}
}提高cache命中率的一个重要技巧是分块，精心安排内存访问在一个块内部，小范围内的相邻的内存访问将具有高度的内存局部性。在本教程中，我们选择块因子为32。这将导致一个block为一个32 * 32 *sizeof(float)的内存区域。这相当于4KB的cache大小，而L1缓存的参考缓存大小为32KB。
我们首先为C的运算创建一个默认调度，然后使用指定的块因子对其应用一个tile调度原语，调度原语以向量[x_outer, y_outer, x_inner, y_inner]的形式返回从最外层到最内层的循环顺序。然后我们得到操作输出的缩减轴，并使用因子4对其执行拆分操作。这个因素不会直接影响我们现在正在进行的分块优化，但在以后应用矢量化时会很有用。
现在计算已经被分块，我们可以重新排序计算，将缩轴运算放入计算的最外层循环中，帮助确保分块的数据仍然在缓存中。这样就完成了调度，我们可以构建并与原始调度性能做对比。
bn = 32
# Blocking by loop tiling
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
# Hoist reduction domain outside the blocking loop
s[C].reorder(xo, yo, ko, ki, xi, yi)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="blocking", log=log)输出：
blocking: 0.303535通过重新排序计算以利用cache，您应该会看到计算性能的显著改进。现在，打印内部表示，并将其与原始表示进行比较:
@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, [1048576], []),
B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}
buffer_map = {A_1: A, B_1: B, C_1: C}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {
for (x.outer: int32, 0, 32) {
for (y.outer: int32, 0, 32) {
for (x.inner.init: int32, 0, 32) {
for (y.inner.init: int32, 0, 32) {
C[((((x.outer*32768) + (x.inner.init*1024)) + (y.outer*32)) + y.inner.init)] = 0f32
}
}
for (k.outer: int32, 0, 256) {
for (k.inner: int32, 0, 4) {
for (x.inner: int32, 0, 32) {
for (y.inner: int32, 0, 32) {
let cse_var_3: int32 = (y.outer*32)
let cse_var_2: int32 = ((x.outer*32768) + (x.inner*1024))
let cse_var_1: int32 = ((cse_var_2 + cse_var_3) + y.inner)
C[cse_var_1] = (C[cse_var_1] + (A[((cse_var_2 + (k.outer*4)) + k.inner)]*B[((((k.outer*4096) + (k.inner*1024)) + cse_var_3) + y.inner)]))
}
}
}
}
}
}
}另一个重要的优化技巧是向量化。当内存访问模式是统一的时，编译器可以检测此模式并将连续内存传递给SIMD向量处理器。在TVM中，我们可以利用这个硬件特性，使用向量化接口来提示编译器这个模式。
在本教程中，我们选择对内层循环的行数据进行向量化，因为在之前的优化中，它已经对cache友好了。
# Apply the vectorization optimization
s[C].vectorize(yi)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="vectorization", log=log)
# The generalized IR after vectorization
print(tvm.lower(s, [A, B, C], simple_mode=True))输出：
vectorization: 0.336937
@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, [1048576], []),
B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}
buffer_map = {A_1: A, B_1: B, C_1: C}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {
for (x.outer: int32, 0, 32) {
for (y.outer: int32, 0, 32) {
for (x.inner.init: int32, 0, 32) {
C[ramp((((x.outer*32768) + (x.inner.init*1024)) + (y.outer*32)), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (k.inner: int32, 0, 4) {
for (x.inner: int32, 0, 32) {
let cse_var_3: int32 = (y.outer*32)
let cse_var_2: int32 = ((x.outer*32768) + (x.inner*1024))
let cse_var_1: int32 = (cse_var_2 + cse_var_3)
C[ramp(cse_var_1, 1, 32)] = (C[ramp(cse_var_1, 1, 32)] + (broadcast(A[((cse_var_2 + (k.outer*4)) + k.inner)], 32)*B[ramp((((k.outer*4096) + (k.inner*1024)) + cse_var_3), 1, 32)]))
}
}
}
}
}
}如果我们看一下上面的IR，我们可以看到内部循环行数据是向量化的，B被转换为PackedB(这可以通过内部循环的(float32x32*)B2部分来证明)。现在PackedB的遍历是顺序的。所以我们来看看A的访问模式。在当前中，A是逐列访问的，这不是缓存友好的类型。如果我们改变ki和内轴xi的嵌套循环顺序，A矩阵的访问模式将对缓存更加友好。
s = te.create_schedule(C.op)
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
# re-ordering
s[C].reorder(xo, yo, ko, xi, ki, yi)
s[C].vectorize(yi)
evaluate_operation(
s, [A, B, C], target=target, name="mmult", optimization="loop permutation", log=log
)
# Again, print the new generalized IR
print(tvm.lower(s, [A, B, C], simple_mode=True))输出：
loop permutation: 0.118288
@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, [1048576], []),
B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}
buffer_map = {A_1: A, B_1: B, C_1: C}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {
for (x.outer: int32, 0, 32) {
for (y.outer: int32, 0, 32) {
for (x.inner.init: int32, 0, 32) {
C[ramp((((x.outer*32768) + (x.inner.init*1024)) + (y.outer*32)), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (x.inner: int32, 0, 32) {
for (k.inner: int32, 0, 4) {
let cse_var_3: int32 = (y.outer*32)
let cse_var_2: int32 = ((x.outer*32768) + (x.inner*1024))
let cse_var_1: int32 = (cse_var_2 + cse_var_3)
C[ramp(cse_var_1, 1, 32)] = (C[ramp(cse_var_1, 1, 32)] + (broadcast(A[((cse_var_2 + (k.outer*4)) + k.inner)], 32)*B[ramp((((k.outer*4096) + (k.inner*1024)) + cse_var_3), 1, 32)]))
}
}
}
}
}
}另一个重要的技巧是数组包装。这个技巧是对数组的存储维度重新排序，以便在扁平化之后将某个维度上的连续访问模式转换为顺序模式。
如上图所示，在分块计算后，我们可以观察到B(扁平化后)的数组访问模式有规则但是不连续。我们期望经过一些转换后可以得到一个连续的访问模式。通过将数组a[16][16]重新排序为一个[16/4][16][4]数组，这样从打包的数组中获取相应的值时，B的访问模式将是顺序的。
为了做到这一点，结合B的新包装，我们必须启动一个新的默认调度。这里值得花点时间来评论一下:TE是一种用于编写优化算子的强大而富有表现力的语言，但它通常需要了解底层算法、数据结构和要编写面向的硬件目标。在本教程的后面部分，我们将讨论让TVM承担这个负担的一些选项。不管怎样，让我们继续执行新的优化计划。
# We have to re-write the algorithm slightly.
packedB = te.compute((N / bn, K, bn), lambda x, y, z: B[y, x * bn + z], name="packedB")
C = te.compute(
(M, N),
lambda x, y: te.sum(A[x, k] * packedB[y // bn, k, tvm.tir.indexmod(y, bn)], axis=k),
name="C",
)
s = te.create_schedule(C.op)
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
s[C].reorder(xo, yo, ko, xi, ki, yi)
s[C].vectorize(yi)
x, y, z = s[packedB].op.axis
s[packedB].vectorize(z)
s[packedB].parallel(x)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="array packing", log=log)
# Here is the generated IR after array packing.
print(tvm.lower(s, [A, B, C], simple_mode=True))输出：
array packing: 0.110371
@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, [1048576], []),
B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}
buffer_map = {A_1: A, B_1: B, C_1: C}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {
allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global {
for (x: int32, 0, 32) "parallel" {
for (y: int32, 0, 1024) {
packedB_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) + y)] = B[ramp(((y*1024) + (x*32)), 1, 32)]
}
}
for (x.outer: int32, 0, 32) {
for (y.outer: int32, 0, 32) {
for (x.inner.init: int32, 0, 32) {
C[ramp((((x.outer*32768) + (x.inner.init*1024)) + (y.outer*32)), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (x.inner: int32, 0, 32) {
for (k.inner: int32, 0, 4) {
let cse_var_3: int32 = ((x.outer*32768) + (x.inner*1024))
let cse_var_2: int32 = (k.outer*4)
let cse_var_1: int32 = (cse_var_3 + (y.outer*32))
C[ramp(cse_var_1, 1, 32)] = (C[ramp(cse_var_1, 1, 32)] + (broadcast(A[((cse_var_3 + cse_var_2) + k.inner)], 32)*packedB_1[(((y.outer*1024) + cse_var_2) + k.inner)]))
}
}
}
}
}
}
}到目前为止，我们所有的优化都集中在有效地访问和计算来自A和B矩阵的数据，以计算C矩阵。分块优化后，算子会逐块将结果写到C中，并且访问模式不是顺序的。我们可以通过使用顺序缓存数组来解决这个问题，使用cache_write、compute_at和unroll 组合来保存块结果，当所有块结果就绪时写入C 。
s = te.create_schedule(C.op)
# Allocate write cache
CC = s.cache_write(C, "global")
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
# Write cache is computed at yo
s[CC].compute_at(s[C], yo)
# New inner axes
xc, yc = s[CC].op.axis
(k,) = s[CC].op.reduce_axis
ko, ki = s[CC].split(k, factor=4)
s[CC].reorder(ko, xc, ki, yc)
s[CC].unroll(ki)
s[CC].vectorize(yc)
x, y, z = s[packedB].op.axis
s[packedB].vectorize(z)
s[packedB].parallel(x)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="block caching", log=log)
# Here is the generated IR after write cache blocking.
print(tvm.lower(s, [A, B, C], simple_mode=True))输出：
block caching: 0.110481
@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, [1048576], []),
B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}
buffer_map = {A_1: A, B_1: B, C_1: C}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {
allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global;
allocate(C.global: Pointer(global float32), float32, [1024]), storage_scope = global {
for (x: int32, 0, 32) "parallel" {
for (y: int32, 0, 1024) {
packedB_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) + y)] = B[ramp(((y*1024) + (x*32)), 1, 32)]
}
}
for (x.outer: int32, 0, 32) {
for (y.outer: int32, 0, 32) {
for (x.c.init: int32, 0, 32) {
C.global_1: Buffer(C.global, float32, [1024], [])[ramp((x.c.init*32), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (x.c: int32, 0, 32) {
let cse_var_4: int32 = (k.outer*4)
let cse_var_3: int32 = (x.c*32)
let cse_var_2: int32 = ((y.outer*1024) + cse_var_4)
let cse_var_1: int32 = (((x.outer*32768) + (x.c*1024)) + cse_var_4)
{
C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[cse_var_1], 32)*packedB_1[cse_var_2]))
C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 1)], 32)*packedB_1[(cse_var_2 + 1)]))
C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 2)], 32)*packedB_1[(cse_var_2 + 2)]))
C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 3)], 32)*packedB_1[(cse_var_2 + 3)]))
}
}
}
for (x.inner: int32, 0, 32) {
for (y.inner: int32, 0, 32) {
C[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] = C.global_1[((x.inner*32) + y.inner)]
}
}
}
}
}
}到目前为止，我们的计算只设计为使用单个核。几乎所有现代处理器都是多核的，并行计算可以使我们的计算受益。最后的优化是利用线程级并行。
# parallel
s[C].parallel(xo)
x, y, z = s[packedB].op.axis
s[packedB].vectorize(z)
s[packedB].parallel(x)
evaluate_operation(
s, [A, B, C], target=target, name="mmult", optimization="parallelization", log=log
)
# Here is the generated IR after parallelization.
print(tvm.lower(s, [A, B, C], simple_mode=True))输出：
parallelization: 0.143585
@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, [1048576], []),
B: Buffer(B_2: Pointer(float32), float32, [1048576], []),
C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}
buffer_map = {A_1: A, B_1: B, C_1: C}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {
allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global {
for (x: int32, 0, 32) "parallel" {
for (y: int32, 0, 1024) {
packedB_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) + y)] = B[ramp(((y*1024) + (x*32)), 1, 32)]
}
}
for (x.outer: int32, 0, 32) "parallel" {
allocate(C.global: Pointer(global float32), float32, [1024]), storage_scope = global;
for (y.outer: int32, 0, 32) {
for (x.c.init: int32, 0, 32) {
C.global_1: Buffer(C.global, float32, [1024], [])[ramp((x.c.init*32), 1, 32)] = broadcast(0f32, 32)
}
for (k.outer: int32, 0, 256) {
for (x.c: int32, 0, 32) {
let cse_var_4: int32 = (k.outer*4)
let cse_var_3: int32 = (x.c*32)
let cse_var_2: int32 = ((y.outer*1024) + cse_var_4)
let cse_var_1: int32 = (((x.outer*32768) + (x.c*1024)) + cse_var_4)
{
C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[cse_var_1], 32)*packedB_1[cse_var_2]))
C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 1)], 32)*packedB_1[(cse_var_2 + 1)]))
C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 2)], 32)*packedB_1[(cse_var_2 + 2)]))
C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 3)], 32)*packedB_1[(cse_var_2 + 3)]))
}
}
}
for (x.inner: int32, 0, 32) {
for (y.inner: int32, 0, 32) {
C[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] = C.global_1[((x.inner*32) + y.inner)]
}
}
}
}
}
}在仅用18行代码应用上述简单优化之后，我们生成的代码的性能就可以开始接近使用Math Kernel Library (MKL)的numpy了。因为我们一直在记录性能，所以我们可以比较结果。
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
none 3.2308110501000002 1.0
blocking 0.30353472839999995 0.09395000936083989
vectorization 0.3369373636 0.10428878643013527
loop permutation 0.1182884093 0.03661260515260981
array packing 0.11037113629999999 0.03416205237275754
block caching 0.11048060149999998 0.03419593401990512
parallelization 0.1435849393 0.044442382136694665请注意，web页面上的输出反映了非专有Docker容器上的运行时间，应该被认为是不可靠的。强烈建议您自己运行本教程，观察TVM所获得的性能增益，并仔细研究每个示例，以理解矩阵乘法运算的迭代改进。
如前所述，如何使用TE和调度原语进行优化可能需要一些底层架构和算法方面的知识。然而，TE被设计为可以探索潜在优化的更复杂算法的基础。通过本文对TE的介绍，我们现在可以开始探索TVM如何实现调度优化过程的自动化。
本教程提供了一个使用向量相加和矩阵乘法的TVM张量表达式(TE)工作流的演练。通常的工作流程是：
接下来的教程将扩展矩阵乘法示例，并展示如何构建矩阵乘法的通用模板和其他具有可调参数的运算，这些可调参数允许您针对特定平台自动优化计算。