使用调度模板和AutoTVM优化算子

本文翻译自Optimizing Operators with Schedule Templates and AutoTVM — tvm 0.9.dev0 documentation

在本教程中，我们将展示如何使用TVM张量表达式(TE)语言来编写调度模板，AutoTVM可以搜索这些模板来找到最佳的调度。这个过程被称为自动调优，它有助于优化张量计算的过程的自动化。

本教程建立在上一篇关于如何使用TE编写矩阵乘法的教程的基础上。

自动调优有两个步骤。

• 第一步是定义搜索空间。
• 第二步是运行搜索算法来探索这个空间。

在本教程中，您可以学习如何在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函数。为了简单起见，我们将把注意力集中在split优化上，使用一个固定的值来定义重新排序的块的大小。 

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因子。然而，它可能不是最优的，因为最优的tiling银子取决于真实的硬件环境和输入形状。

如果您希望调度代码能够跨更广泛的输入形状和目标硬件进行移植，那么最好定义一组候选值，并根据目标硬件上的测量结果选择最佳值。

在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. 获取一个config对象:你可以把这个cfg作为这个函数的参数，但是我们可以通过不同的方式获取它。有了这个参数，这个函数就不再是确定性的调度了。相反，我们可以将不同的配置传递给这个函数，从而得到不同的调度。使用这种配置对象的函数称为“模板”。

为了使模板函数更紧凑，我们可以做两件事来在单个函数中定义参数搜索空间。

• 在一组值中定义搜索空间。这可以通过将cfg设置为一个ConfigSpace对象来实现。它将收集此函数中的所有可调旋钮，并从中构建一个搜索空间。
• 根据此空间中的某一实体进行调度。这是通过将cfg设置为一个ConfigEntity对象来实现的。当它是ConfigEntity时，它将忽略所有空间定义API(即cfg.define_XXXXX(…))。相反，它将存储所有可调旋钮的确定性值，我们将根据这些值进行调度。

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

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

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

使用高级参数API的矩阵乘法模板

在前面的模板中，我们手工列出了一个旋钮的所有可能值。这是定义空间的最低级别API，并给出要搜索的参数空间的显式枚举。但是，我们还提供了另一组API，可以使搜索空间的定义更简单、更灵活。在可能的情况下，我们建议您使用这个高级API

在下面的例子中，我们使用ConfigSpace.define_split来定义一个split旋钮。它将枚举所有可能的方式来切分一个轴，并构建空间。

我们也有ConfigSpace.define_reorder用于重新排序旋钮，configspace . define_annotation用于unroll、vectorization、线程绑定等注释。当高级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]

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的优点是它使得多级split(即当num_outputs >= 3时)更容易。

使用AutoTVM优化矩阵乘法 

前面我们写了一个矩阵乘法模板，它允许我们参数化split调度中使用的块大小。我们现在可以对这个参数空间进行搜索。下一步是选择一个调优器来指导这个空间的探索。 

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)，则网格搜索调优器或随机调优器就足够了。如果你的空间是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)

输出：

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搜索参数空间并选择优化调度配置的算子模板。为了更深入地理解它是如何工作的，我们建议对这个例子进行扩展，根据张量表达式入门教程中演示的调度操作，向调度中添加新的搜索参数。在接下来的小节中，我们将演示AutoScheduler, TVM使用它来优化通用算子，而不需要用户提供用户定义的模板。