TVM Compiler中文教程：使用TVM编写可调模板和使用自动调优器

使用TVM编写可调模板和使用自动调优器

这是TVM中auto-tuning模块的入门教程。

auto-tuning分两个步骤：第一步定义搜索空间；第二步是运行搜索算法来探索这个空间。在本教程中，你可以了解如何在TVM中执行这两个步骤。下面通过矩阵乘法展示auto-tuning的完整工作流程。

knob旋钮：可调值，搜索空间中的值

加载依赖库

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

pip3 install --user psutil xgboost

为了使TVM在调优中运行得更快，建议使用cython作为TVM的FFI。在TVM源码根目录运行如下命令：

pip3 install --user cython
sudo make cython3

现在开始python代码，首先导入包：

import logging
import sys

import numpy as np
import tvm

#导入autotvm模块
from tvm import autotvm

第一步：定义搜索空间

在本节中，我们将确定性TVM调度代码重写为可调调度模板。您可以将搜索空间定义的过程视为现有计划代码的参数化。

首先，我们展示如何在TVM中实现分块矩阵乘法。

#Matmul V0：恒定平铺tiling因子
def matmul_v0(N,L,M,dtype):
    A = tvm.placeholder((N, L), name='A', dtype=dtype)
    B = tvm.placeholder((L, M), name='B', dtype=dtype)
    
    #reduce轴
    k = tvm.reduce_axis((0,L), name='k')
    C = tvm.compute((N,M), lambda i,j: tvm.sum(A[i,k]*B[k,j],axis=k), name='C')
    s = tvm.create_schedule(C.op)
    
    #调度
    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]

参数化调度

在之前的调度代码中，我们使用常量“8”作为平铺因子。但是，它可能不是最好的，因为最佳的平铺因子取决于实际的硬件环境和输入形状。

如果你希望调度代码可以在更广泛的输入形状和目标硬件上移植，则最好定义一组候选值，并根据目标硬件上的测量结果选择最佳值。

在autotvm中，我们可以定义一个可调参数，或者为这种值定义“knob”。

#Matmul V1:列出候选值
@autotvm.template #1. 使用装饰器
def matmul_v1(N, L, M, dtype):
    A = tvm.placeholder((N, L), name='A', dtype=dtype)
    B = tvm.placeholder((L, M), name='B', dtype=dtype)

    k = tvm.reduce_axis((0, L), name='k')
    C = tvm.compute((N, M), lambda i, j: tvm.sum(A[i, k] * B[k, j], axis=k), name='C')
    s = tvm.create_schedule(C.op)
    
    #调度
    y, x = s[C].op.axis
    k = s[C].op.reduce_axis[0]
    
    #2. 获取autotvm配置对象
    cfg = autotvm.get_config()
    
    #3.定义搜索空间knob
    cfg.define_knob("tile_y", [1,2,4,8,16])
    cfg.define_knob("tile_y", [1,2,4,8,16])
    
    #4.根据配置来调度
    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.根据该空间中的实体定义搜索空间，2.根据该空间中的实体来调度。为此，我们将cfg设置为配置空间ConfigSpace或配置实体ConfigEntity对象。

   当它是ConfigSpace时，它将收集此函数中的所有可调旋钮knobs并构建搜索空间。当它是ConfigEntity时，它将忽略所有空间定义API（即cfg.define_XXXXX(...)）。相反，它存储所有可调旋钮的确定性值，并根据这些值进行调度。

   在auto-tuning期间，我们将首先使用ConfigSpace对象调用此模板函数来构建搜索空间。然后，我们使用不同的ConfigEntity对象调用此模板函数在构建的搜索空间中来获取不同的调度。最后，我们将测量不同调度生成的代码并选择最佳调度。

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

4. 根据cfg的确定值来调度。

使用更好的空间定义API

在上一个模板中，我们手动列出旋钮的所有可能值。这是定义空间的最低级API。但是，我们还提供了另一组API，以使空间定义更容易，更智能。建议使用这套高级API。

在下面示例中，我们使用ConfigSpace.define_split来定义分割旋钮。它将列举出分割轴和构造空间的所有可能方法。

我们还有用于重新排序的旋钮ConfigSpace.define_reorder和用于注释的ConfigSpace.define_annotate，如展开，矢量化，线程绑定。当高级API无法满足您的要求时，您始终可以回退使用低级API。

@autotvm.template
def matmul(N, L, M, dtype):
    A = tvm.placeholder((N, L), name='A', dtype=dtype)
    B = tvm.placeholder((L, M), name='B', dtype=dtype)

    k = tvm.reduce_axis((0, L), name='k')
    C = tvm.compute((N, M), lambda i, j: tvm.sum(A[i, k] * B[k, j], axis=k), name='C')
    s = tvm.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。在此模板中，将枚举所有可能的组合，这些组合可以将轴y拆分为两个轴，其长度为y。例如，如果y的长度是32并且我们想要使用32的因子将它分成两个轴，那么（外轴的长度，内轴的长度）对有6个可能的值，即（32,1） ，（16,2），（8,4），（4,8），（2,16）或（1,32）。它们只是tile_y的6个可能值。cfg.define_split("tile_y", y, num_outputs=2)

在调度期间，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时）更容易。

第二步：搜索空间

在第1步中，我们通过将旧的调度代码扩展到模板来构建搜索空间。下一步是选择一个调优器来探索这个空间。

在TVM中使用自动调优器

可以通过伪代码来描述调优器的工作。

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

当取出下一批配置时，调优器可以采取不同的策略。我们在autotvm中提供四种不同策略的调优器。

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

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

开始调优

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

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

输出：

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

然后我们需要定义如何测量生成的代码并选择一个调优器。由于我们的空间很小，RandomTuner就可以了。

我们在本教程中只进行了10次试验来进行演示。在实践中，您可以根据您的时间预算进行更多试验。我们将调优结果记录到日志文件中。此文件可用于以后获得最佳配置。

#日志配置（用于将调优日志打印到屏幕）
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))

#开始调优，将日志记录到文件`matmul.log`
tuner = autotvm.tuner.RandomTuner(task)
tuner.tune(n_trial=10,
           measure_option=measure_option,
           callbacks=[autotvm.callback.log_to_file('matmul.log')])

输出：

Get devices for measurement successfully!
No: 1   GFLOPS: 9.16/9.16       result: MeasureResult(costs=(0.0293049456,), error_no=0, all_cost=1.421377420425415, timestamp=1560502458.516948)       [('tile_y', [64, 8]), ('tile_x', [128, 4])],,None,23
No: 2   GFLOPS: 1.63/9.16       result: MeasureResult(costs=(0.1650307972,), error_no=0, all_cost=3.7540791034698486, timestamp=1560502461.7132463)     [('tile_y', [32, 16]), ('tile_x', [256, 2])],,None,14
No: 3   GFLOPS: 18.28/18.28     result: MeasureResult(costs=(0.0146827124,), error_no=0, all_cost=1.6729228496551514, timestamp=1560502462.497204)      [('tile_y', [32, 16]), ('tile_x', [1, 512])],,None,94
No: 4   GFLOPS: 8.93/18.28      result: MeasureResult(costs=(0.030057214999999998,), error_no=0, all_cost=2.0133767127990723, timestamp=1560502463.5245106)     [('tile_y', [4, 128]), ('tile_x', [32, 16])],,None,47
No: 5   GFLOPS: 19.67/19.67     result: MeasureResult(costs=(0.0136481436,), error_no=0, all_cost=1.7570922374725342, timestamp=1560502464.2963889)     [('tile_y', [64, 8]), ('tile_x', [1, 512])],,None,93
No: 6   GFLOPS: 18.90/19.67     result: MeasureResult(costs=(0.014203376200000001,), error_no=0, all_cost=1.6420013904571533, timestamp=1560502465.0690885)     [('tile_y', [16, 32]), ('tile_x', [4, 128])],,None,75
No: 7   GFLOPS: 3.65/19.67      result: MeasureResult(costs=(0.0735679294,), error_no=0, all_cost=2.345797538757324, timestamp=1560502466.8066292)      [('tile_y', [1, 512]), ('tile_x', [64, 8])],,None,39
No: 8   GFLOPS: 8.85/19.67      result: MeasureResult(costs=(0.030322758999999998,), error_no=0, all_cost=1.4428813457489014, timestamp=1560502467.836355)      [('tile_y', [128, 4]), ('tile_x', [128, 4])],,None,22
No: 9   GFLOPS: 11.23/19.67     result: MeasureResult(costs=(0.0239063918,), error_no=0, all_cost=0.9399979114532471, timestamp=1560502469.8295026)     [('tile_y', [16, 32]), ('tile_x', [32, 16])],,None,45
No: 10  GFLOPS: 17.75/19.67     result: MeasureResult(costs=(0.015127030999999999,), error_no=0, all_cost=0.8141782283782959, timestamp=1560502470.6222835)     [('tile_y', [256, 2]), ('tile_x', [2, 256])],,None,81

最后，我们从缓存文件应用最佳性能结果并检查其正确性。我们可以直接在autotvm.apply_history_best上下文中调用matmul函数。当我们调用此函数时，它将使用其参数查询调度上下文，并使用相同的参数获取最佳配置。

#从log文件应用最佳性能
with autotvm.apply_history_best('matmul.log'):
    with tvm.target.create('llvm'):
        s, arg_bufs = matmul(N,L,M,'float32')
        func = tvm.build(s, arg_bufs)
#检查正确性
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)

#调用调优的matmul函数
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.asnumpy(), rtol=1e-2)