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

爱被打了一巴掌 2022-01-16 02:51 393阅读 0赞

### 文章目录 ###

*  使用TVM编写可调模板和使用自动调优器
 *   *  加载依赖库
     *  第一步：定义搜索空间
     *   *  参数化调度
         *  使用更好的空间定义API
     *  第二步：搜索空间
     *   *  在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)