TVM User Tutorial -- Optimizing Operators with Schedule Templates and AutoTVM

Authors: Lianmin Zheng, Chris Hoge

在本教程中，我们将展示如何使用 TVM 张量表达式 (TE) 语言编写调度模板，AutoTVM 可以搜索这些模板以找到最佳调度。 这个过程称为Auto-Tuning，它有助于自动化优化张量计算的过程。
本教程建立在前面关于如何使用 TE <tensor_expr_get_started> 编写矩阵乘法的教程的基础上。

自动调整有两个步骤。

• 第一步是定义搜索空间。
• 第二步是运行搜索算法来探索这个空间。
  在本教程中，您可以了解如何在 TVM 中执行这两个步骤。 整个工作流程由一个矩阵乘法示例说明。
  请注意，本教程不会在 Windows 或最新版本的 macOS 上运行。 要让它运行，您需要将本教程的主体包装在 if __name__ == "__main__": 块中。

安装依赖

要在 TVM 中使用 autotvm 包，我们需要安装一些额外的依赖项。

pip3 install --user psutil xgboost cloudpickle

为了让 TVM 在调优中跑得更快，建议使用 cython 作为 TVM 的 FFI。 在 TVM 的根目录下，执行：

pip3 install --user cython
sudo make cython3

现在回到python代码。 首先导入所需的包。

import logging
import sys

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

# the module is called `autotvm`
from tvm import autotvm

基本的TE矩阵乘法

回想一下使用 TE 进行矩阵乘法的基本实现。 我们把它写在这里，做一些改动。 我们将把乘法包装在 python 函数定义中。 为简单起见，我们将注意力集中在拆分优化上，使用定义重新排序的块大小的固定值。

def matmul_basic(N, L, M, dtype):

    A = te.placeholder((N, L), name="A", dtype=dtype)
    B = te.placeholder((L, M), name="B", dtype=dtype)

    k = te.reduce_axis((0, L), name="k")
    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")
    s = te.create_schedule(C.op)

    # schedule
    y, x = s[C].op.axis
    k = s[C].op.reduce_axis[0]

    yo, yi = s[C].split(y, 8)
    xo, xi = s[C].split(x, 8)

    s[C].reorder(yo, xo, k, yi, xi)

    return s, [A, B, C]

使用 AutoTVM 进行矩阵乘法

在前面的调度代码中，我们使用常量“8”作为tiling factor。 但是，它可能不是最好的，因为最好的平铺因子取决于真实的硬件环境和输入形状。
如果您希望调度代码在更广泛的输入形状和目标硬件上可移植，最好定义一组候选值并根据目标硬件上的测量结果选择最佳值。
在 autotvm 中，我们可以为这种值定义一个可调参数，或者一个“knob”。

一个基本的矩阵乘法模板

我们从一个示例开始，说明如何为拆分调度操作的块大小创建可调参数集。

# Matmul V1: List candidate values
@autotvm.template("tutorial/matmul_v1")  # 1. use a decorator
def matmul_v1(N, L, M, dtype):
    A = te.placeholder((N, L), name="A", dtype=dtype)
    B = te.placeholder((L, M), name="B", dtype=dtype)

    k = te.reduce_axis((0, L), name="k")
    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")
    s = te.create_schedule(C.op)

    # schedule
    y, x = s[C].op.axis
    k = s[C].op.reduce_axis[0]

    # 2. get the config object
    cfg = autotvm.get_config()

    # 3. define search space
    cfg.define_knob("tile_y", [1, 2, 4, 8, 16])
    cfg.define_knob("tile_x", [1, 2, 4, 8, 16])

    # 4. schedule according to config
    yo, yi = s[C].split(y, cfg["tile_y"].val)
    xo, xi = s[C].split(x, cfg["tile_x"].val)

    s[C].reorder(yo, xo, k, yi, xi)

    return s, [A, B, C]

在这里，我们对之前的调度代码进行了四次修改，并获得了一个可调的“模板”。 我们可以一一解释修改。

1. 使用装饰器将此函数标记为简单模板。

2. 获取配置对象：您可以将此 cfg 视为此函数的参数，但我们以不同的方式获取它。 有了这个参数，这个函数就不再是一个确定性的调度了。 相反，我们可以将不同的配置传递给这个函数并获得不同的调度。 像这样使用配置对象的函数称为“模板”。
   为了使模板函数更紧凑，我们可以做两件事来在单个函数中定义参数搜索空间。

   1. 跨一组值定义搜索空间。 这是通过使 cfg 成为 ConfigSpace 对象来完成的。 它将收集此函数中的所有可调旋钮并从中构建一个搜索空间。
   2. 根据该空间中的实体进行调度。 这是通过使 cfg 成为 ConfigEntity 对象来完成的。 当它是 ConfigEntity 时，它将忽略所有空间定义 API（即 cfg.define_XXXXX(...)）。 相反，它将存储所有可调旋钮的确定值，我们根据这些值进行调度。

   在自动调整期间，我们将首先使用 ConfigSpace 对象调用此模板来构建搜索空间。 然后我们在构建的空间中用不同的ConfigEntity 调用这个模板来得到不同的调度。 最后，我们将测量不同时间表生成的代码并选择最佳的。

3. 定义两个可调旋钮。 第一个是 tile_y，有 5 个可能的值。 第二个是 tile_x，具有相同的可能值列表。 这两个旋钮是独立的，因此它们跨越大小为 25 = 5x5 的搜索空间。

4. 配置旋钮被传递给拆分调度操作，允许我们根据之前在 cfg 中定义的 5x5 确定性值进行调度。

带有高级参数 API 的矩阵乘法模板

在前面的模板中，我们手动列出了旋钮的所有可能值。 这是定义空间的最低级别 API，并给出了要搜索的参数空间的显式枚举。 但是，我们还提供了另一组 API，可以使搜索空间的定义更简单、更智能。 在可能的情况下，我们建议您使用这个更高级别的 API
在以下示例中，我们使用 ConfigSpace.define_split 来定义拆分旋钮。 它将列举所有可能的分割轴和构造空间的方法。
我们还有用于重新排序旋钮的 ConfigSpace.define_reorder 和用于展开、矢量化、线程绑定等注释的 ConfigSpace.define_annotate。 当高级 API 无法满足您的要求时，您始终可以退回使用低级 API。

@autotvm.template("tutorial/matmul")
def matmul(N, L, M, dtype):
    A = te.placeholder((N, L), name="A", dtype=dtype)
    B = te.placeholder((L, M), name="B", dtype=dtype)

    k = te.reduce_axis((0, L), name="k")
    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")
    s = te.create_schedule(C.op)

    # schedule
    y, x = s[C].op.axis
    k = s[C].op.reduce_axis[0]

    ##### define space begin #####
    cfg = autotvm.get_config()
    cfg.define_split("tile_y", y, num_outputs=2)
    cfg.define_split("tile_x", x, num_outputs=2)
    ##### define space end #####

    # schedule according to config
    yo, yi = cfg["tile_y"].apply(s, C, y)
    xo, xi = cfg["tile_x"].apply(s, C, x)

    s[C].reorder(yo, xo, k, yi, xi)

    return s, [A, B, C]

More Explanation on cfg.define_split
在此模板中，cfg.define_split("tile_y", y, num_outputs=2) 将枚举所有可能的组合，这些组合可以将轴 y 分成两个轴，其中两个轴的长度为 y 的因子。 例如，如果 y 的长度是 32，我们想用 32 的因子将它分成两个轴，那么 (外轴长度，内轴长度) 对有 6 个可能的值，即 (32, 1) , (16, 2), (8, 4), (4, 8), (2, 16) 或 (1, 32)。 这些都是 tile_y 的 6 个可能值。
在调度期间，cfg["tile_y"] 是一个 SplitEntity 对象。 我们将外轴和内轴的长度存储在 cfg['tile_y'].size （具有两个元素的元组）中。 在这个模板中，我们使用 yo, yi = cfg['tile_y'].apply(s, C, y) 来应用它。 实际上，这相当于 yo, yi = s[C].split(y, cfg["tile_y"].size[1]) 或 yo, yi = s[C].split(y, nparts=cfg[' tile_y"].size[0]) 使用 cfg.apply API 的优点是它使多级拆分（即当 num_outputs >= 3 时）更容易。

第 2 步：使用 AutoTVM 优化矩阵乘法

在第 1 步中，我们编写了一个矩阵乘法模板，允许我们参数化拆分计划中使用的块大小。 我们现在可以对这个参数空间进行搜索。 下一步是选择一个调谐器来指导这个空间的探索。

TVM 中的 Auto-tuners

调谐器的工作可以通过以下伪代码来描述

ct = 0
while ct < max_number_of_trials:
    propose a batch of configs
    measure this batch of configs on real hardware and get results
    ct += batch_size

在提出下一批配置时，调谐器可以采取不同的策略。 TVM 提供的一些调谐器策略包括：

• tvm.autotvm.tuner.RandomTuner：以随机顺序枚举空间
• tvm.autotvm.tuner.GridSearchTuner：以网格搜索顺序枚举空间
• tvm.autotvm.tuner.GATuner：使用遗传算法搜索空间
• tvm.autotvm .tuner.XGBTuner：使用基于模型的方法。 训练一个 XGBoost 模型来预测降低 IR 的速度，并根据预测选择下一批。

您可以根据您的空间大小、您的时间预算和其他因素来选择调谐器。 例如，如果您的空间非常小（小于 1000），则 gridsearch 调谐器或随机调谐器就足够了。 如果您的空间在 10^9 级别（这是 CUDA GPU 上的 conv2d 算子的空间大小），XGBoostTuner 可以更有效地探索并找到更好的配置。

开始调优

在这里，我们继续我们的矩阵乘法示例。 首先我们创建一个调优任务。 我们还可以检查初始化的搜索空间。 在这种情况下，对于 512x512 方阵乘法，空间大小为 10x10=100 请注意，任务和搜索空间与选择的调谐器无关。

N, L, M = 512, 512, 512
task = autotvm.task.create("tutorial/matmul", args=(N, L, M, "float32"), target="llvm")
print(task.config_space)

Out:
ConfigSpace (len=100, space_map=
   0 tile_y: Split(policy=factors, product=512, num_outputs=2) len=10
   1 tile_x: Split(policy=factors, product=512, num_outputs=2) len=10
)

然后我们需要定义如何测量生成的代码并选择一个调谐器。 由于我们的空间很小，随机调谐器就可以了。
我们在本教程中仅进行 10 次试验以进行演示。 在实践中，您可以根据自己的时间预算进行更多试验。 我们会将调整结果记录到日志文件中。 该文件可用于选择调谐器稍后发现的最佳配置。

# logging config (for printing tuning log to the screen)
logging.getLogger("autotvm").setLevel(logging.DEBUG)
logging.getLogger("autotvm").addHandler(logging.StreamHandler(sys.stdout))

测量配置有两个步骤：构建和运行。 默认情况下，我们使用所有 CPU 内核来编译程序。 然后我们依次测量它们。 为了帮助减少方差，我们进行了 5 次测量并对它们进行平均。

measure_option = autotvm.measure_option(builder="local", runner=autotvm.LocalRunner(number=5))

# Begin tuning with RandomTuner, log records to file `matmul.log`
# You can use alternatives like XGBTuner.
tuner = autotvm.tuner.RandomTuner(task)
tuner.tune(
    n_trial=10,
    measure_option=measure_option,
    callbacks=[autotvm.callback.log_to_file("matmul.log")],
)

调优完成后，我们可以从日志文件中选择具有最佳测量性能的配置，并使用相应的参数编译调度。 我们还快速验证时间表是否产生了正确的答案。 我们可以直接在 autotvm.apply_history_best 上下文下调用函数 matmul。 当我们调用这个函数时，它将使用它的参数查询调度上下文，并使用相同的参数获取最佳配置。

# apply history best from log file
with autotvm.apply_history_best("matmul.log"):
    with tvm.target.Target("llvm"):
        s, arg_bufs = matmul(N, L, M, "float32")
        func = tvm.build(s, arg_bufs)

# check correctness
a_np = np.random.uniform(size=(N, L)).astype(np.float32)
b_np = np.random.uniform(size=(L, M)).astype(np.float32)
c_np = a_np.dot(b_np)

c_tvm = tvm.nd.empty(c_np.shape)
func(tvm.nd.array(a_np), tvm.nd.array(b_np), c_tvm)

tvm.testing.assert_allclose(c_np, c_tvm.numpy(), rtol=1e-4)

最后的说明和总结

在本教程中，我们展示了如何构建允许 TVM 搜索参数空间并选择优化的调度配置的算子模板。 为了更深入地了解其工作原理，我们建议扩展此示例，根据张量表达式入门 <tensor_expr_get_started>_ 教程中演示的调度操作向调度添加新的搜索参数。 在接下来的部分中，我们将演示 AutoScheduler，这是一种 TVM 优化常用算子的方法，无需用户提供用户定义的模板。