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

本节学习如何使用TVM 张量表达式(TE)语言来编写调度模板，这些模板可以被autoTVM搜索到，以找到最佳调度。这个过程称为auto-Tuning,它有助于优化张量计算的自动化过程。

auto-tuning的步骤如下：

第一搜索空间
第二运行一个搜索算法来探索这个空间

Install dependencies

为在TVM使用autoTVM包，需要安装一些额外的依赖

pip3 install --user psutil xgboost cloudpickle

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

FFI:Foreign Function Interface,用一种编程语言写的程序能调用另一种编程语言写的函数

pip3 install --user cython

sudo make cython3

导入需要的库：

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

Basic Matrix Multiplication with TE

用一个python函数定义来包装TE矩阵乘法。为简单起见，将注意力放在分割优化上，使用一个固定值来定义重新排序的块大小

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]

Matrix Multiplication with AutoTVM

在上述就矩阵乘法代码中，使用常数“8”作为平铺系数(tiling factor)。这可能不是最好的，最佳的平铺系数取决于实际的硬件环境和输入形状。

如果想让调度代码在更大范围的输入形状和目标硬件上可移植，最好是定义一组候选值，并根据目标硬件上的测量结果挑选最佳值。

在 autotvm 中，我们可以定义一个可调整的参数，或者说是一个 “旋钮”，用于此类值。

# 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]

在这里，我们对之前的调度代码做了四项修改，得到了一个可调度的 “模板”。我们可以逐一解释这些修改：

使用装饰器将这个函数标记为一个简单的模板。
获取 config 对象。以把这个 cfg 看作是这个函数的一个参数,但可以以不同的方式获得它。有了这个参数，这个函数就不再是一个确定性的调度了。相反，我们可以向这个函数传递不同的配置，得到不同的调度。一个像这样使用配置对象的函数被称为 “模板”。

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

定义一个跨越一组数值的搜索空间。这是通过使 cfg 成为一个 ConfigSpace 对象来实现的。它将收集这个函数中的所有可调控旋钮，并从中建立一个搜索空间。

根据这个空间的一个实体来调度。这是通过使 cfg 成为一个 ConfigEntity 对象来实现的。当它是一个 ConfigEntity 时，它将忽略所有空间定义 API（即 cfg.define_XXXXX(...)）。相反，它将为所有可调度的旋钮存储确定的值，我们根据这些值来调度。

在自动调度过程中，我们将首先用 ConfigSpace 对象调用该模板来构建搜索空间。然后，我们在构建的空间中用不同的 ConfigEntity 调用该模板，以获得不同的调度。最后，我们将测量不同调度所产生的代码，并挑选出最好的一个。

定义两个可调度的旋钮。第一个是 tile_y，有 5 个可能的值。第二个是 tile_x，有相同的可能值列表。这两个旋钮是独立的，所以它们跨越了一个大小为 25=5x5 的搜索空间。
配置旋钮被传递给 split 调度操作，使我们能够根据我们先前在 cfg 中定义的 5x5 确定值来调度。

A Matrix Multiplication Template with the Advanced Parameter API

在前面的模板中，手动列出了一个旋钮的所有可能值。这是定义空间的最底层的 API，它给出了要搜索的参数空间的明确列举。然而，我们还提供了另一组 API，可以使搜索空间的定义更容易、更智能。在可能的情况下，我们接受你使用这个更高级别的 API。

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

还有ConfigSpace.define_reorder用于重新排序旋钮，以及 ConfigSpace.define_annotate 用于 unroll、矢量化、线程绑定等注释。当高级 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]

PS:关于 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 时）更容易。

Step 2: Use AutoTVM to Optimize the Matrix Multiplication

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

Auto-tuners in 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

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

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

可以根据你的空间大小、你的时间预算和其他因素来选择调谐器。例如，如果你的空间非常小（小于 1000），网格搜索调谐器或随机调谐器就足够好了。如果你的空间在的水平（这是 CUDA GPU 上 conv2d 运算器的空间大小），XGBoostTuner 可以更有效地探索并找到更好的配置。

Begin tuning

首先，创建一个tuning task，也可以检查初始化空间。在本例子中，对于一个512 * 512的矩阵乘法，空间大小玩10 * 10 = 100.注意这个task与搜索空间与所选的tuner无关。

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

)

然后，需要定义一个怎样衡量生产的代码和挑选的tuner。因为我们的空间很小，随机的tuner就可以了

在实践中，你可以根据你的时间预算做更多的试验。我们将把调谐结果记录到一个日志文件中。这个文件可以用来选择调谐器以后发现的最佳配置。

# 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")],

)

输出结果：

waiting for device...

waiting for device...

device available

device available

Get devices for measurement successfully!

Get devices for measurement successfully!

No: 1   GFLOPS: 8.17/8.17       result: MeasureResult(costs=(0.0328701162,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.6089231967926025, timestamp=1658363130.4992497)  [('tile_y', [-1, 16]), ('tile_x', [-1, 16])],None,44

No: 1   GFLOPS: 8.17/8.17       result: MeasureResult(costs=(0.0328701162,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.6089231967926025, timestamp=1658363130.4992497)  [('tile_y', [-1, 16]), ('tile_x', [-1, 16])],None,44

No: 2   GFLOPS: 7.45/8.17       result: MeasureResult(costs=(0.0360509632,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.695908784866333, timestamp=1658363131.2425098)   [('tile_y', [-1, 8]), ('tile_x', [-1, 16])],None,43

No: 2   GFLOPS: 7.45/8.17       result: MeasureResult(costs=(0.0360509632,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.695908784866333, timestamp=1658363131.2425098)   [('tile_y', [-1, 8]), ('tile_x', [-1, 16])],None,43

No: 3   GFLOPS: 25.86/25.86     result: MeasureResult(costs=(0.010381651799999999,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.3583974838256836, timestamp=1658363131.6205783)  [('tile_y', [-1, 32]), ('tile_x', [-1, 128])],None,75

No: 3   GFLOPS: 25.86/25.86     result: MeasureResult(costs=(0.010381651799999999,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.3583974838256836, timestamp=1658363131.6205783)  [('tile_y', [-1, 32]), ('tile_x', [-1, 128])],None,75

No: 4   GFLOPS: 30.35/30.35     result: MeasureResult(costs=(0.0088443436,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.31720852851867676, timestamp=1658363131.9254267) [('tile_y', [-1, 1]), ('tile_x', [-1, 64])],None,60

No: 4   GFLOPS: 30.35/30.35     result: MeasureResult(costs=(0.0088443436,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.31720852851867676, timestamp=1658363131.9254267) [('tile_y', [-1, 1]), ('tile_x', [-1, 64])],None,60

No: 5   GFLOPS: 9.35/30.35      result: MeasureResult(costs=(0.028700182600000002,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.7562305927276611, timestamp=1658363132.5526962)  [('tile_y', [-1, 16]), ('tile_x', [-1, 128])],None,74

No: 5   GFLOPS: 9.35/30.35      result: MeasureResult(costs=(0.028700182600000002,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.7562305927276611, timestamp=1658363132.5526962)  [('tile_y', [-1, 16]), ('tile_x', [-1, 128])],None,74

No: 6   GFLOPS: 28.15/30.35     result: MeasureResult(costs=(0.0095366542,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.29946398735046387, timestamp=1658363132.9198403) [('tile_y', [-1, 8]), ('tile_x', [-1, 512])],None,93

No: 6   GFLOPS: 28.15/30.35     result: MeasureResult(costs=(0.0095366542,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.29946398735046387, timestamp=1658363132.9198403) [('tile_y', [-1, 8]), ('tile_x', [-1, 512])],None,93

No: 7   GFLOPS: 19.80/30.35     result: MeasureResult(costs=(0.0135571626,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.30517101287841797, timestamp=1658363133.3043413) [('tile_y', [-1, 512]), ('tile_x', [-1, 512])],None,99

No: 7   GFLOPS: 19.80/30.35     result: MeasureResult(costs=(0.0135571626,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.30517101287841797, timestamp=1658363133.3043413) [('tile_y', [-1, 512]), ('tile_x', [-1, 512])],None,99

No: 8   GFLOPS: 5.75/30.35      result: MeasureResult(costs=(0.0467114476,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.8827676773071289, timestamp=1658363134.2409544)  [('tile_y', [-1, 8]), ('tile_x', [-1, 32])],None,53

No: 8   GFLOPS: 5.75/30.35      result: MeasureResult(costs=(0.0467114476,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.8827676773071289, timestamp=1658363134.2409544)  [('tile_y', [-1, 8]), ('tile_x', [-1, 32])],None,53

No: 9   GFLOPS: 27.67/30.35     result: MeasureResult(costs=(0.0097023816,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.38503575325012207, timestamp=1658363134.582773)  [('tile_y', [-1, 2]), ('tile_x', [-1, 128])],None,71

No: 9   GFLOPS: 27.67/30.35     result: MeasureResult(costs=(0.0097023816,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.38503575325012207, timestamp=1658363134.582773)  [('tile_y', [-1, 2]), ('tile_x', [-1, 128])],None,71

No: 10  GFLOPS: 2.62/30.35      result: MeasureResult(costs=(0.1026008602,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.708867073059082, timestamp=1658363136.3652241)   [('tile_y', [-1, 8]), ('tile_x', [-1, 1])],None,3

No: 10  GFLOPS: 2.62/30.35      result: MeasureResult(costs=(0.1026008602,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.708867073059082, timestamp=1658363136.3652241)   [('tile_y', [-1, 8]), ('tile_x', [-1, 1])],None,3

tuning完成后，从日志文件中选出就v有最佳测量性能的配置，并用相应的参数来编译schedule.也可以做一个快速验证，以确保schedule产生正确的答案。可以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)

输出结果：

Finish loading 10 records

Finish loading 10 records

matmul.log内容如下：