乐趣区

关于机器学习:TVM-学习资料用-Schedule-模板和-AutoTVM-优化算子

jiezi

残缺 TVM 中文文档，拜访→TVM 中文站
作者：Lianmin Zheng，Chris Hoge

本教程将展现如何用 TVM 张量表达式（TE）语言编写 schedule 模板，并通过 AutoTVM 对模板进行搜寻，从而找到最佳 schedule。这个主动优化张量计算的过程被称为 Auto-Tuning。

本教程基于后面的 TE 编写矩阵乘法教程 设立。

auto-tuning 包含两个步骤：

第一步：定义搜寻空间。
第二步：运行搜索算法来摸索这个空间。
通过本教程能够理解如何在 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

# 模块名叫 `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)

    # 调度
    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”作为循环切分因子，然而它可能不是最佳的。因为最佳的循环切分因子取决于实在的硬件环境和输出 shape。

如果心愿调度代码可能在更宽泛的输出 shape 和指标硬件上可移植，最好定义一组候选值，并依据指标硬件上的评估后果抉择最佳值。

autotvm 中能够为这种值定义一个可调参数，或者一个 “knob”。

根本矩阵乘法模板

以下示例将演示，如何为 split 调度操作的 block 大小创立一个可调的参数集。

# Matmul V1: 列出候选值
@autotvm.template("tutorial/matmul_v1")  # 1. 应用装璜器
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)

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

    # 2. 获取 config 对象
    cfg = autotvm.get_config()

    # 3. 定义搜寻空间
    cfg.define_knob("tile_y", [1, 2, 4, 8, 16])
    cfg.define_knob("tile_x", [1, 2, 4, 8, 16])

    # 4. 依据 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 参数使得这个函数不再是一个确定的 schedule。将不同的配置传递给这个函数，能够失去不同的 schedule。这种应用配置对象的函数称为“模板”。

为使模板函数更精炼，可在单个函数中定义参数搜寻空间：

  • 用一组值来定义搜寻空间。将 cfg 转为 ConfigSpace 对象，收集此函数中的所有可调 knob，而后从中构建一个搜寻空间。
  • 依据空间中的实体进行调度。将 cfg 转为 ConfigEntity 对象，当它被转为 ConfigEntity 后，会疏忽所有空间定义
    API（即 cfg.define_XXXXX(…)），但会存储所有可调 knob 的确定值，并依据这些值进行调度。

在 auto-tuning 的过程中，首先用 ConfigSpace 对象调用这个模板来构建搜寻空间，而后在构建的空间中用不同的 ConfigEntity 调用这个模板，来失去不同的 schedule。最初，咱们将评估由不同 schedule 生成的代码，而后抉择最佳的 schedule。

  1. 定义两个可调 knob。第一个是 tile_y，它有 5 个可能值。第二个是 tile_x，它和前者具备雷同的可能值。这两个 knob 是独立的，所以它们逾越大小为 25 = 5×5 的搜寻空间。
  2. 配置 knob 被传递给 split 调度操作，而后能够依据之前在 cfg 中定义的 5×5 确定值进行调度。

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

后面的模板手动列出了 konb 的所有可能值，它是用来定义空间的最底层 API，显示列出了要搜寻的参数空间。这里举荐应用另一组更高级的 API，它能够更简略、更智能地定义搜寻空间。

上面的示例用 ConfigSpace.define_split 来定义拆分 knob。它列举了所有可能的拆分 axis 和结构空间的办法。

同时，ConfigSpace.define_reorder 用于对 knob 从新排序，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)

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

    ##### 开始定义空间 #####
    cfg = autotvm.get_config()
    cfg.define_split("tile_y", y, num_outputs=2)
    cfg.define_split("tile_x", x, num_outputs=2)
    ##### 完结定义空间 #####

    # 依据 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 的长处是它使多级拆分（即当 num_outputs >= 3 时）变得更加简略。

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

第 1 步编写的矩阵乘法模板，可对拆分的 schedule 中的块大小进行参数化。通过第 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

调优器可采取不同的策略来打算下一批配置，包含：

  • 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 能够更无效地摸索并找到更好的配置。

开始调优

上面持续矩阵乘法的示例。首先创立一个调优工作，而后查看初始的搜寻空间。上面示例中是 512×512 的矩阵乘法，空间大小为 10×10=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 次试验进行演示。实际上能够依据本人的工夫估算进行更多试验。调优后果会记录到日志文件中。这个文件可用于抉择之后发现的调优器的最佳配置。

# 记录 config（为了将 tuning 日志打印到屏幕）logging.getLogger("autotvm").setLevel(logging.DEBUG)
logging.getLogger("autotvm").addHandler(logging.StreamHandler(sys.stdout))

评估配置有两个步骤：构建和运行。默认用所有 CPU core 来编译程序。而后顺次进行评估。为了缩小方差，对 5 次评估后果取平均值。

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

# 用 RandomTuner 开始调优, 日志记录到 `matmul.log` 文件中
# 可用 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...
device available
Get devices for measurement successfully!
No: 1   GFLOPS: 8.48/8.48       result: MeasureResult(costs=(0.0316434228,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.638512134552002, timestamp=1657225928.6342561)        [('tile_y', [-1, 1]), ('tile_x', [-1, 256])],None,80
No: 2   GFLOPS: 2.30/8.48       result: MeasureResult(costs=(0.1165478966,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.0105199813842773, timestamp=1657225930.6636436)       [('tile_y', [-1, 4]), ('tile_x', [-1, 8])],None,32
No: 3   GFLOPS: 11.82/11.82     result: MeasureResult(costs=(0.0227097348,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.5589795112609863, timestamp=1657225931.7059512)       [('tile_y', [-1, 64]), ('tile_x', [-1, 32])],None,56
No: 4   GFLOPS: 1.66/11.82      result: MeasureResult(costs=(0.1616202114,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.6911513805389404, timestamp=1657225934.9635096)       [('tile_y', [-1, 1]), ('tile_x', [-1, 4])],None,20
No: 5   GFLOPS: 3.65/11.82      result: MeasureResult(costs=(0.073561817,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.3051848411560059, timestamp=1657225936.3988533)        [('tile_y', [-1, 256]), ('tile_x', [-1, 16])],None,48
No: 6   GFLOPS: 1.85/11.82      result: MeasureResult(costs=(0.1452834464,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.5179028511047363, timestamp=1657225938.961955)        [('tile_y', [-1, 512]), ('tile_x', [-1, 4])],None,29
No: 7   GFLOPS: 0.87/11.82      result: MeasureResult(costs=(0.30933780240000003,), error_no=MeasureErrorNo.NO_ERROR, all_cost=5.067087888717651, timestamp=1657225944.589149)  [('tile_y', [-1, 512]), ('tile_x', [-1, 2])],None,19
No: 8   GFLOPS: 10.53/11.82     result: MeasureResult(costs=(0.025489421,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.5452830791473389, timestamp=1657225945.1592515)        [('tile_y', [-1, 4]), ('tile_x', [-1, 64])],None,62
No: 9   GFLOPS: 1.58/11.82      result: MeasureResult(costs=(0.16960762680000002,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.8109781742095947, timestamp=1657225948.0900776)        [('tile_y', [-1, 2]), ('tile_x', [-1, 2])],None,11
No: 10  GFLOPS: 2.42/11.82      result: MeasureResult(costs=(0.11083148779999999,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.8757600784301758, timestamp=1657225950.0266354)        [('tile_y', [-1, 4]), ('tile_x', [-1, 4])],None,22

调优实现后，可从日志文件中抉择具备最佳评估性能的配置，并用相应参数来编译 schedule。疾速验证 schedule 是否产生了正确的后果，可间接在 autotvm.apply_history_best 上下文中调用 matmul 函数，它会用参数查问调度上下文，而后可用雷同的参数获取最优配置。

# 从日志文件中利用历史最佳
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)

# 验证正确性
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

总结

本教程展现了如何构建算子模板，使得 TVM 可能搜寻参数空间，并抉择优化的调度配置。为了更深刻地理解其工作原理，举荐基于 :ref: 张量表达式入门 <tensorexpr_get_started> 教程中演示的调度操作，向调度增加新的搜寻参数。接下来的章节将演示 AutoScheduler，它是 TVM 中一种优化罕用算子的办法，同时无需用户提供自定义的模板。

退出移动版