本文翻译自Working with Operators Using Tensor Expression — tvm 0.9.dev0 documentation

在本教程中，我们将把注意力转向TVM如何使用张量表达式(TE)定义张量计算和实现循环优化。TE用纯函数语言描述张量计算(即每个表达式都没有副作用)。当把TVM看作一个整体时，Relay将一个计算描述为一组算子，每个算子可以表示为一个TE表达式，每个TE表达式接受输入张量并产生一个输出张量。

这是TVM中张量表达式语言的入门教程。TVM使用一个特定于域的张量表达式来有效地构造核心。我们将通过两个使用张量表达式语言的示例来演示基本工作流。第一个例子介绍了TE和使用向量加法的调度。第二部分对这些概念进行了扩展，逐步优化了使用TE的矩阵乘法。这个矩阵乘法示例将作为以后教程中介绍TVM的更高级特性的对比基础。

使用TE编写和调度CPU向量加法

让我们来看一个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 timeitnp_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 = 32768a = 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).meanprint("%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 = 4outer, 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             Performancenumpy    6.66109990561381e-06                     1.0naive              5.8282e-06      0.8749606044923809
parallel              6.0927e-06      0.9146687613655552vector    2.6245599999999998e-05    3.9401300643878434

代码说明：你可能已经注意到的，A, B和C的声明都有相同的形状参数n。TVM将利用这一点只传递一个形状参数给内核，正如你将在打印的设备代码中发现的那样。这是一种专门的形式。在主机端，TVM将自动生成检查代码，以参数的约束条件。因此，如果你将一个其他的数组传递给fadd，将会引发一个错误。我们可以做更多的专门化。例如，我们可以在计算声明中写n = tvm.runtime.convert(1024)，而不是n = te.var("n")。那么生成的函数将只接受长度为1024的向量。

我们已经定义、调度和编译了一个向量加法运算符，然后就可以在TVM运行时中执行它了。我们可以将操作符保存为一个库，然后使用TVM运行时加载它。

目标为GPU的向量加法

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 schedulen = 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 = 1024a = 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())

保存和加载编译后的模块

除了运行时编译之外，我们还可以将编译后的模块保存到一个文件中，稍后再将它们加载回来。

下面的代码首先执行以下步骤:

它将编译后的主机模块保存到一个目标文件中。
然后将设备模块保存到一个ptx文件中。
cc.create_shared调用编译器(gcc)来创建共享库

from tvm.contrib import cc
from tvm.contrib import utilstemp = 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编译了代码。

生成OpenCL代码

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 = 1024a = 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包含许多不同的调度原语:

split:根据定义的因子将指定的轴分割为两个轴。
tile:通过定义的因子在两个轴上分割计算。
fuse:融合一个计算的两个连续轴。
reorder:可以按指定顺序对一个计算的轴重新排序。
bind:可以将一个计算绑定到一个特定的线程，在GPU编程中很有用。
compute_at:默认情况下，TVM将计算函数最外层的张量，也就是根的张量。compute_at指定一个张量应该在另一个算子的计算的第一个轴上计算。
compute_inline:当标记为inline时，将展开计算，然后插入到需要张量的地址中。
compute_root:将计算移动到函数的最外层，也就是根。这意味着计算阶段将在进入下一阶段之前完全计算完成。

这些原语的完整描述可以在Schedule Primitives in TVM — tvm 0.9.dev0 documentation页面中找到。

Example2:使用TE手工优化矩阵乘法

现在，我们将考虑第二个更高级的示例，演示TVM如何只用18行python代码将一个常见的矩阵乘法运算加速18倍。

矩阵乘法是一种计算密集型运算。对于良好的CPU性能，有两个重要的优化:

提高内存访问的cache命中率。复杂的数值计算和热点内存访问都可以通过提高cache率来加速。这需要我们将原始内存访问模式转换为适合cache策略的模式。
SIMD(单指令多数据)，也称为矢量处理单元。在每个周期中，SIMD可以处理一小批数据，而不是只处理单个数。这就要求我们以统一的模式转换循环体中的数据访问模式，以便LLVM后端可以将其低级化为SIMD。

本教程中使用的技术是代码仓中提到的技巧的一部分。它们中的一些已经被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 funcc = 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).meanprint("%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)] = 0f32for (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)]))}}}
}

Optimization 1: 分块

提高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)]))}}}}}}
}

Optimization 2：向量化

另一个重要的优化技巧是向量化。当内存访问模式是统一的时，编译器可以检测此模式并将连续内存传递给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)]))}}}}}
}

Optimization 3:循环重排

如果我们看一下上面的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)]))}}}}}
}

Optimization 4:数组包装

另一个重要的技巧是数组包装。这个技巧是对数组的存储维度重新排序，以便在扁平化之后将某个维度上的连续访问模式转换为顺序模式。

如上图所示，在分块计算后，我们可以观察到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)]))}}}}}}
}

Optimization 5：通过cache优化块内存的写入

到目前为止，我们所有的优化都集中在有效地访问和计算来自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)]}}}}}
}

Optimization 6:并行化

到目前为止，我们的计算只设计为使用单个核。几乎所有现代处理器都是多核的，并行计算可以使我们的计算受益。最后的优化是利用线程级并行。

# 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             Performancenone      3.2308110501000002                     1.0blocking     0.30353472839999995     0.09395000936083989vectorization            0.3369373636     0.10428878643013527
loop permutation            0.1182884093     0.03661260515260981array packing     0.11037113629999999     0.03416205237275754block caching     0.11048060149999998     0.03419593401990512parallelization            0.1435849393    0.044442382136694665

请注意，web页面上的输出反映了非专有Docker容器上的运行时间，应该被认为是不可靠的。强烈建议您自己运行本教程，观察TVM所获得的性能增益，并仔细研究每个示例，以理解矩阵乘法运算的迭代改进。

最后注意事项及摘要

如前所述，如何使用TE和调度原语进行优化可能需要一些底层架构和算法方面的知识。然而，TE被设计为可以探索潜在优化的更复杂算法的基础。通过本文对TE的介绍，我们现在可以开始探索TVM如何实现调度优化过程的自动化。

本教程提供了一个使用向量相加和矩阵乘法的TVM张量表达式(TE)工作流的演练。通常的工作流程是：

通过一系列的运算来描述你的计算。
描述我们想要如何使用调度原语计算。
编译到我们想要的目标函数。
可选的，保存函数以便稍后加载。

接下来的教程将扩展矩阵乘法示例，并展示如何构建矩阵乘法的通用模板和其他具有可调参数的运算，这些可调参数允许您针对特定平台自动优化计算。

【TVM帮助文档学习】使用张量表达式处理算子相关推荐

【TVM帮助文档学习】使用TVMC编译和优化模型
本文翻译自Compiling and Optimizing a Model with TVMC - tvm 0.9.dev0 documentation 在本节中,我们将使用TVM命令行驱动程序TVM ...
python文档学习
文章目录 python文档学习 python解释器传入参数交互模式的提示符帮助信息基础知识 python保留字注释多行语句缩进 python数据类型与运算符变量与基本类型数字类型集 ...
Spring文档学习
Spring文档学习参考Spring Framework Documentation学习 1. IoC 容器 1.1 容器实例化 <beans><import resource=& ...
EasyUI文档学习心得
概述 jQuery EasyUI 是一组基于jQuery 的UI 插件集合,它可以让开发者在几乎完全不需要CSS以及复杂的JS代码情况下完成美观且功能强大的Web界面. 本文主要说明一些如何利用Eas ...
UHS-II文档学习
#UHS-II文档学习 UHS-II 引脚分配 SD4.0与SD3.0相比多了底下一排引脚[10-17],UHS-II Mode使用的是三组差分信号分别为:D0.D1和RCLK差分信号.RCLK差分信 ...
Django 4.0文档学习（一）
本系列文章基于Django4.0版本官方网站文档学习使用开发工具为pycharm > python -m django --version 4.0 文章目录编写你的第一个 Django 应用 ...
Spring Framework 5.3文档学习（一）
Spring Framework 5.3文档学习(一) Overview 1.What We Mean by "Spring" 2. History of Spring and t ...
使用Tensor Expression张量表达式处理算子
使用Tensor Expression张量表达式处理算子这是TVM中Tensor表达语言的入门教程.TVM使用特定于域的张量表达式来进行有效的内核构造. 本文将演示使用张量表达式语言的基本工作流程. ...
python定义函数的组成部分有_Python文档学习笔记（4）--定义函数
定义函数关键字 def 引入函数的定义.其后必须跟有函数名和以括号标明的形式参数列表.组成函数体的语句从下一行开始,且必须缩进. 执行一个函数会引入一个用于函数的局部变量的新符号表. 因此,在函数 ...

【TVM帮助文档学习】使用张量表达式处理算子