TVM：使用Tensor Expression (TE)来处理算子

在本教程中，我们将聚焦于在 TVM 中使用张量表达式（TE）来定义张量计算和实现循环优化。TE用纯函数语言描述张量计算（即每个表达式都没有副作用）。当在 TVM 的整体上下文中查看时，Relay 将计算描述为一组算子，并且其中每一个算子都可以表示为 TE 表达式，每个 TE 表达式获取输入张量并生成输出张量。
本文是TVM中 TE 语言的入门教程。TVM 使用领域专用（domain specific）的张量表达式来高效地构造内核。我们以两个使用 TE 语言的为例来演示基本工作流。第一个示例介绍了 TE 和带有向量加法的 schedule。第二个示例通过逐步优化矩阵与 TE 的乘法来扩展这些概念。这个矩阵乘法示例将作为未来涵盖更高级的 TVM 特性的教程的对比基础。

示例一：使用TE为CPU编写和调度向量加法

初始化 tvm环境

我们的第一个例子是使用 Python 来为向量加法实现一个 TE，然后是一个针对 CPU 的 schedule，我们从初始化 tvm 环境开始：

import tvm
import tvm.testing
from tvm import te
import numpy as np# 如果能够指定目标 CPU，那么将会得到更好地性能
# 如果用的是llvm，可以通过 `llc --version` 来查看 CPU 类型
# 可以通过查看 /proc/cpuinfo 来查看你的处理器可能支持的其他扩展，
# 比如，如果你的 CPU 有 AVX-512 指令集，那么你可以使用 `llvm -mcpu=skylake-avx512` 选项tgt = tvm.target.Target(target="llvm", host="llvm")

描述向量计算

我们首先描述向量加法计算。TVM 采用张量语义，每个中间结果表示为一个多维数组。我们需要描述规则来得到张量。我们首先定义一个符号变量 n 来表示形状。然后我们定义两个 placeholder 张量：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函数

te.compute方法的第二个参数是执行计算的函数。在本例中，我们使用一个匿名函数（也称为lambda函数）来定义计算，在本例中是对 a 和 B 的第 i 个元素的加法。

为计算创建一个默认的Schedule

虽然上面几行描述了计算规则，但我们可以用许多不同的方法计算 C 以适应不同的设备。对于具有多个 axis 的张量，您可以选择首先迭代哪个 axis ，另外计算可以跨不同的线程拆分。TVM要求用户提供一个 schedule，来描述应如何执行计算。TE 中的 schedule 操作可以更改循环顺序、跨不同线程拆分计算、将数据块分组在一起，以及其他操作。schedule 背后的一个重要概念是，它们只描述如何执行计算，因此相同 TE 的不同 schedule 一定会产生相同的结果。

在 TVM 中，我们可以创建一种朴素的 schedule ，按照行优先的顺序来计算 C。

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

s = te.create_schedule(C.op)

编译并验证默认的 schedule

通过 TE 表达式和 schedule，我们可以为目标语言和体系结构生成可运行的代码，在本例中是 LLVM 和 CPU 。我们向 TVM 提供 schedule、schedule 中的TE表达式列表、目标和主机，以及我们正在生成的函数的名称。输出的结果是可以直接从 Python 调用 type-erased 函数。

在下一行中，我们使用 tvm.build 创建一个函数。build 函数接受 schedule、函数所需的签名（包括输入和输出）以及我们要编译到的目标语言。

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

我们运行该函数，并将输出与 numpy 中的相同计算进行比较。编译后的 TVM 函数提供了一个简明的C API，可以被任何语言调用。我们首先创建一个设备（在本例中为CPU），这是一个 TVM 可以编译 schedule 的设备。在本例中，设备是LLVM CPU target。然后，我们可以在设备中初始化张量并执行自定义的加法操作。为了验证计算的正确性，我们可以将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 的速度差异，创建一个辅助函数来运行 TVM 生成代码的 profile。

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.000008
naive: 0.000006

使用并行性（paralleism）来优化 schedule

我们已经说明了 TE 的基本原理，现在让我们更深入地了解 schedule 的作用，以及它们如何用于优化不同体系结构的张量表达式。schedule 是应用于表达式的一系列步骤，用于以多种不同方式对其进行转换。当一个 schedule 应用于TE中的一个表达式时，输入和输出保持不变，但在编译时，表达式的实现可能会改变。在默认 schedule 中，这个张量加法是串行运行的，但该操作其实是很容易在所有处理器线程之间并行。我们可以将我们的操作并行调度到计算中：

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

tvm.lower 命令将生成 TE 的中间表示（IR）以及相应的 schedule 。通过在执行不同的 schedule 操作时 lowing 表达式，我们可以看到 schedule 对计算顺序的影响。我们使用标志 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)])}
}

TVM现在可以在独立的线程上运行这些块。我们在执行并行操作的情况下编译并运行这个新的 schedule：

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.000005

使用矢量化（vectorization）来优化 schedule

现代 CPU 能够对浮点数进行 SIMD 操作，我们可以对计算表达式使用另一个 schedule 来利用这一点。实现这一点需要多个步骤：首先，我们必须使用 split scheduling 原语将 schedule 拆分为内部循环和外部循环。内部循环可以使用向量化来使用使用向量化调度原语的 SIMD 指令，然后外部循环可以使用并行调度原语进行并行化。选择分割因子作为CPU上的线程数。

注：SIMD，全称 Single Instruction Multiple Data，单指令多数据流，能够复制多个操作数，并把它们打包在大型寄存器的一组指令集。

# 由于我们需要修改之前例子中的并行操作，因此这里要重建 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")s = te.create_schedule(C.op)# factor 的选择需要适合你的线程数，这取决于架构，
# 建议将此系数设置为等于可用CPU核心数。
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.000016
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, [n: int32], [stride: int32], type="auto"),C: Buffer(C_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_1)] = ((float32*)A_2[(((i.outer*4) + i.inner.s)*stride)] + (float32*)B_2[(((i.outer*4) + i.inner.s)*stride_2)])}}}
}

对比不同的 schedule

下面我们来对比以下之前提到的不同 schedule：

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    7.98278022557497e-06                     1.0naive              5.9189e-06      0.7414584684465222
parallel    4.9771999999999995e-06    0.6234920490550659vector    1.6127399999999997e-05    2.0202735819196875

注意：Code Specialization

代码专门化

正如我们所看到的，A、B 和 C 的声明都采用相同的形状参数 n。TVM将利用这一点，只向 kernel 传递一个 shape 参数，我们在打印的设备代码中找到它。这是专门化化的一种形式。

在 host 端，TVM 将自动生成检查代码，以检查参数中的约束。因此，如果将具有不同形状的数组传递到 fadd 中，将引发错误。

我们可以做更多的专门化。例如，我们可以在计算声明中写入n=tvm.runtime.convert（1024）而不是 n=te.var（“n”）。生成的函数将只获取长度为1024的向量。

我们已经定义、调度并编译了一个向量加法运算符，然后可以在 TVM Runtime 执行它。我们可以将算子保存为库，稍后可以使用 TVM Runtime 加载该库。

针对GPU的矩阵加法（可选）

在介绍保存与加载自定义算子库的方法之前，我们先来看一下如何针对 GPU 做矩阵加法。

TVM能够针对多种体系结构。在本例，我们将针对GPU中矢量加法的编译。

# 本段代码默认不运行，如果想要运行的话，请将 ``run_cuda = True``run_cuda = False
if run_cuda:# 这里的 target 需要根据自己的 GPU 类型修改：# NVIDIA：cuda# Radeon：rocm# OpenCL：opencltgt_gpu = tvm.target.Target(target="cuda", host="llvm")# 重建 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)################################################################################# 最后，我们必须将迭代轴bx和tx绑定到GPU计算网格中的线程。# 朴素的 schedule 对GPU无效，这些是允许我们生成在GPU上运行的代码的特定构造。s[C].bind(bx, te.thread_axis("blockIdx.x"))s[C].bind(tx, te.thread_axis("threadIdx.x"))####################################################################### 编译# -----------# 在指定完 schdule 之后，我们可以将其编译成一个 TVM 函数。默认情况下，TVM编译成一个 type-erased 函       # 数，可以从python端直接调用该函数。# 在下一行中，我们使用 tvm.build 来创建一个函数。build 函数采用 schedule、函数所需的签名（包括输如和输出）以及我们要编译到的目标语言。# 编译 fadd 的结果是一个GPU设备函数（如果涉及GPU）以及一个调用 GPU 函数的 host wrapper。fadd是生成的主机包装函数，它在内部包含对生成的设备函数的引用。fadd = tvm.build(s, [A, B, C], target=tgt_gpu, name="myadd")################################################################################# 编译过的 TVM 函数会有一个简洁的 C API，它可以被任意的语言调用## 我们提供一个 Python 的最小的数组 API 来帮助快速的测试和原型化# 该数组 API 是基于 `DLPack <https://github.com/dmlc/dlpack>`_ 标准.## - 我们首先创建一个 GPU 设备# - 然后 tvm.nd.array 从 GPU 拷贝数据# - ``fadd`` 运行真正的计算# - ``numpy()`` 从 GPU 数组拷贝回 CPU (这样我们就能验证正确性了).## 请注意，将数据复制到 GPU 上的内存和从中复制数据是必需的步骤。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())################################################################################# 检查生成的 GPU 代码# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~# 我们可以检查在 TVM 中生成的代码，tvm.build 的结果是一个 TVM 模块。fadd 是一个 host 模块其中包含         # host wrapper 的 host module，它同样包含一个CUDA（GPU）设备模块## 下面的代码取得设备模块并打印内容代码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())

此处输出：

['myadd.o', 'myadd.so']

注意：Module Storage Format

模块存储格式

CPU（Host）模块直接保存为共享库（.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 还支持将所有内容导出为一个共享库。在 hood 下，我们将设备模块打包成二进制blob，并将它们与主机代码链接在一起。目前我们支持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())

注意：Runtime API and Thread Safety

运行时API与线程安全

TVM 的编译模块并不依赖于 TVM 编译器。它们只依赖于最小 Runtime Library。TVM Runtime Library 包装设备驱动程序，并向编译函数提供线程安全和设备无关调用。

这意味着我们可以从任何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 Scheduling Primitives

TE 调度原语

TVM 包括许多不同的调度原语：
split：按定义的因子将指定轴拆分为两个轴。
tile：平铺将按定义的因子沿两个轴分割计算。
fuse：融合一次计算的两个连续轴。
reorder：可以将计算轴重新排序为定义的顺序。
bind：可以将计算绑定到特定线程，在GPU编程中很有用。
compute_at：默认情况下，TVM将在函数的最外层或根计算张量。compute_at指定应在另一个运算符的第一个计算轴上计算一个张量。
compute_inline：当标记为inline时，计算将展开，然后插入到需要张量的地址中。
compute_root：将计算移动到函数的最外层或根。这意味着，在进入下一个阶段之前，将对计算阶段进行完全计算。

可以在Schedule primitives 文档页面中找到这些原语的完整描述。

示例二：用TE手动优化矩阵乘

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

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

提高内存访问的缓存命中率。高缓存命中率可以加速复杂的数值计算和热点内存访问。这要求我们将源内存访问模式转换为适合缓存策略的模式。
SIMD（单指令多数据），也称为矢量处理单元。在每个循环中，SIMD 都可以处理一小批数据，而不是处理单个值。这要求我们以统一模式转换循环体中的数据访问模式，以便LLVM 后端可以将其 lower 到 SIMD。

本教程中使用的技术是这个仓库中提到的技巧的一部分。其中一些已被 TVM 抽象自动使用，但由于 TVM 的一些约束，有一些无法自动使用。

准备工作和性能baseline

我们首先采集 numpy 实现的矩阵乘的数据：

import tvm
import tvm.testing
from tvm import te
import numpy# 矩阵的尺寸：
# (M, K) x (K, N)
# 你可以自己试一些不同的尺寸，有时候 TVM 的优化结果会好于含 MKL 的numpy
M = 1024
K = 1024
N = 1024# tvm 中默认的数据类型
dtype = "float32"# 与之前一样，这里可以根据自己的处理器及其是否支持某些指令集来改变 targettarget = tvm.target.Target(target="llvm", host="llvm")
dev = tvm.device(target.kind.name, 0)# 随机生成一些 tensor 用于测试
a = tvm.nd.array(numpy.random.rand(M, K).astype(dtype), dev)
b = tvm.nd.array(numpy.random.rand(K, N).astype(dtype), dev)# 重复实验，得到 numpy 的矩阵乘实现的 baseline
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.009308

现在，我们用 TVM TE 编写一个基本矩阵乘法，并验证它产生的结果与numpy实现相同。我们还编写了一个函数，它将帮助我们度量进度优化的性能。

# 使用 TE 实现的 TVM 的矩阵乘
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")# 默认 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.109406

让我们看一下使用 TVM lower 函数的算子和默认 schedule 的中间表示 IR。请注意，该实现本质上是矩阵乘法的简单实现，在 A 和 B 矩阵的索引上使用三个嵌套循环。

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, [1024, 1024], []),A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}buffer_map = {A_1: A, B_1: B, C_1: C} {for (x: int32, 0, 1024) {for (y: int32, 0, 1024) {C_2[((x*1024) + y)] = 0f32for (k: int32, 0, 1024) {C_2[((x*1024) + y)] = ((float32*)C_2[((x*1024) + y)] + ((float32*)A_2[((x*1024) + k)]*(float32*)B_2[((k*1024) + y)]))}}}
}

优化1：blocking阻塞

提高缓存命中率的一个重要技巧是阻塞，在这种情况下，我们可以构造内存访问，使块内部是具有高内存局部性的小邻域。在本教程中，我们选择块因子 32。这会使得一个块填充内存的 32*32*sizeof（float）区域。这对应于 4KB 的缓存大小，和一级缓存 32KB 的参考缓存大小。

我们首先为 C 操作创建一个默认的调度，然后使用指定的块因子对其应用一个 tile 调度原语，调度原语以向量 [x_-outer，y_-outer，x_-inner，y_-inner] 的形式返回从最外层到最内层的循环顺序。然后，我们得到操作输出的缩减轴，并使用因子4对其执行拆分操作。这个因素不会直接影响我们现在正在进行的阻塞优化，但在以后应用矢量化时会很有用。
现在操作已被阻塞，我们可以对计算进行重新排序，将简化操作放入计算的最外层循环中，从而帮助确保被阻塞的数据仍保留在缓存中。这就完成了 schedule，我们可以构建和测试与原始 schedule 相比的性能。

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.291928

通过重新排序计算以利用缓存，我们可以看到计算性能的显著提高。现在，打印内部表示并将其与原始表示进行比较：

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 = {A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}buffer_map = {A_1: A, B_1: B, C_1: C} {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_2[((((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) {C_2[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] = ((float32*)C_2[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] + ((float32*)A_2[((((x.outer*32768) + (x.inner*1024)) + (k.outer*4)) + k.inner)]*(float32*)B_2[((((k.outer*4096) + (k.inner*1024)) + (y.outer*32)) + y.inner)]))}}}}}}
}

优化2： vectorization矢量化

另一个重要的优化技巧是矢量化。当内存访问模式一致时，编译器可以检测到该模式并将连续内存传递给 SIMD 向量处理器。在TVM中，我们可以利用这个硬件特性，使用矢量化接口来提示编译器这个模式。

在本教程中，我们选择对内部循环行数据进行矢量化，因为它已经是我们之前优化中的缓存友好型数据。

# 应用矢量化的优化方式
s[C].vectorize(yi)evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="vectorization", log=log)# 矢量化之后生成的 IR
print(tvm.lower(s, [A, B, C], simple_mode=True))

此处输出：

vectorization: 0.331263
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, [1024, 1024], []),A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}buffer_map = {A_1: A, B_1: B, C_1: C} {for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) {C_2[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) {C_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)] = ((float32x32*)C_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.inner*1024)) + (k.outer*4)) + k.inner)], 32)*(float32x32*)B_2[ramp((((k.outer*4096) + (k.inner*1024)) + (y.outer*32)), 1, 32)]))}}}}}
}

优化3：Loop Permutation循环置换

如果我们看一下上面的 IR，我们可以看到内环行数据被矢量化，B 被转换成 PackedB（这在内环的（float32x32）B2部分中很明显）。PackedB 的遍历现在是顺序的。因此，我们将研究 A 的访问模式。在当前 schdule中，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
)# 再一次打印新生成的 IR
print(tvm.lower(s, [A, B, C], simple_mode=True))

此处输出：

loop permutation: 0.113750
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, [1024, 1024], []),A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}buffer_map = {A_1: A, B_1: B, C_1: C} {for (x.outer: int32, 0, 32) {for (y.outer: int32, 0, 32) {for (x.inner.init: int32, 0, 32) {C_2[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) {C_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)] = ((float32x32*)C_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.inner*1024)) + (k.outer*4)) + k.inner)], 32)*(float32x32*)B_2[ramp((((k.outer*4096) + (k.inner*1024)) + (y.outer*32)), 1, 32)]))}}}}}
}

优化4：Array Packing数组打包

另一个重要技巧是数组打包。此技巧是对阵列的存储维度重新排序，以便在展平后将特定维度上的连续访问模式转换为序列模式。

如上图所示，在阻塞计算后，我们可以观察到 B 的阵列访问模式（平坦后），它是规则的但不连续的。我们希望经过一些转换后，我们可以得到一个连续的访问模式。通过将[16][16]数组重新排序为[16/4][16][4]数组，在从压缩数组中获取相应值时，B 的访问模式将是顺序的。

为了实现这一点，我们必须从一个新的默认 schedule 开始，考虑到 B 的新 wrapper。花点时间对此进行讨论是值得的：TE 是一种用于编写优化算子的功能强大的表达性语言，但它通常需要一些底层算法、数据结构，以及您正在编写的硬件 target。在本教程的后面，我们将讨论让 TVM 承担这一负担的一些选择。不管怎样，让我们继续新的优化 schedule。

# 我们要轻微地重写算法
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)# 这里是数组打包之后生成的 IR
print(tvm.lower(s, [A, B, C], simple_mode=True))

此处输出：

array packing: 0.224114
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, [1024, 1024], []),A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}buffer_map = {A_1: A, B_1: B, C_1: C} {allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global {for (x: int32, 0, 32) "parallel" {for (y: int32, 0, 1024) {packedB[ramp(((x*32768) + (y*32)), 1, 32)] = (float32x32*)B_2[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_2[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) {C_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)] = ((float32x32*)C_2[ramp((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.inner*1024)) + (k.outer*4)) + k.inner)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) + (k.outer*128)) + (k.inner*32)), 1, 32)]))}}}}}}
}

优化5：Optimizing Block Writing Through Caching通过缓存优化块写入

到目前为止，我们所有的优化都集中在高效地访问和计算来自 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.224213
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, [1024, 1024], []),A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}buffer_map = {A_1: A, B_1: B, C_1: C} {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[ramp(((x*32768) + (y*32)), 1, 32)] = (float32x32*)B_2[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[ramp((x.c.init*32), 1, 32)] = broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (x.c: int32, 0, 32) {C.global[ramp((x.c*32), 1, 32)] = ((float32x32*)C.global[ramp((x.c*32), 1, 32)] + (broadcast((float32*)A_2[(((x.outer*32768) + (x.c*1024)) + (k.outer*4))], 32)*(float32x32*)packedB[ramp(((y.outer*32768) + (k.outer*128)), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] = ((float32x32*)C.global[ramp((x.c*32), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.c*1024)) + (k.outer*4)) + 1)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) + (k.outer*128)) + 32), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] = ((float32x32*)C.global[ramp((x.c*32), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.c*1024)) + (k.outer*4)) + 2)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) + (k.outer*128)) + 64), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] = ((float32x32*)C.global[ramp((x.c*32), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.c*1024)) + (k.outer*4)) + 3)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) + (k.outer*128)) + 96), 1, 32)]))}}for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C_2[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] = (float32*)C.global[((x.inner*32) + y.inner)]}}}}}
}

优化6：Parallelization并行化

# 并行
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
)# 这里是并行化之后的 IR
print(tvm.lower(s, [A, B, C], simple_mode=True))

此处输出：

parallelization: 0.067949
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, [1024, 1024], []),A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], []),B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], [])}buffer_map = {A_1: A, B_1: B, C_1: C} {allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global {for (x: int32, 0, 32) "parallel" {for (y: int32, 0, 1024) {packedB[ramp(((x*32768) + (y*32)), 1, 32)] = (float32x32*)B_2[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[ramp((x.c.init*32), 1, 32)] = broadcast(0f32, 32)}for (k.outer: int32, 0, 256) {for (x.c: int32, 0, 32) {C.global[ramp((x.c*32), 1, 32)] = ((float32x32*)C.global[ramp((x.c*32), 1, 32)] + (broadcast((float32*)A_2[(((x.outer*32768) + (x.c*1024)) + (k.outer*4))], 32)*(float32x32*)packedB[ramp(((y.outer*32768) + (k.outer*128)), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] = ((float32x32*)C.global[ramp((x.c*32), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.c*1024)) + (k.outer*4)) + 1)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) + (k.outer*128)) + 32), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] = ((float32x32*)C.global[ramp((x.c*32), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.c*1024)) + (k.outer*4)) + 2)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) + (k.outer*128)) + 64), 1, 32)]))C.global[ramp((x.c*32), 1, 32)] = ((float32x32*)C.global[ramp((x.c*32), 1, 32)] + (broadcast((float32*)A_2[((((x.outer*32768) + (x.c*1024)) + (k.outer*4)) + 3)], 32)*(float32x32*)packedB[ramp((((y.outer*32768) + (k.outer*128)) + 96), 1, 32)]))}}for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C_2[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] = (float32*)C.global[((x.inner*32) + y.inner)]}}}}}
}

矩阵乘例子的总结

在仅用 18 行代码应用上述简单优化之后，我们生成的代码就可以得到与使用数学内核库（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.1094061458                     1.0blocking     0.29192816779999997     0.09388550549895809vectorization            0.3312631714     0.10653583220302389
loop permutation            0.1137497149    0.036582456445468314array packing            0.2241142794     0.07207623221003798block caching     0.22421289339999997     0.07210794694763607parallelization            0.0679485881    0.021852593361526892

请注意，以上的输出反映的是非独占 Docker 容器上的运行时间，因此并不可靠。强烈建议您自己运行本教程，观察 TVM 实现的性能增益，并仔细阅读每个示例，以了解矩阵乘法运算的迭代改进。

总结

如前所述，如何使用 TE 和调度原语应用优化可能需要一些底层架构和算法的知识。然而，TE 设计为更复杂的算法是为了可以搜索潜在的优化。有了本 TE 简介中的知识，我们现在可以开始探索 TVM 如何自动化进度优化过程。

本教程提供了使用向量加法和矩阵乘法示例的TVM张量表达式（TE）工作流演练。一般的工作流程是：

通过一系列操作描述您的计算。
描述我们希望如何计算和使用调度原语。
编译到我们想要的目标函数。
保存要稍后加载的函数（可选）。

接下来的教程将扩展矩阵乘法示例，并展示如何使用可调参数构建矩阵乘法和其他操作的通用模板，这些参数使得我们能够自动优化特定平台的计算。

Ref：

https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html