用调度模板和 AutoTVM 优化算子#
在本教程中,展示了如何使用 TVM 张量表达式(TE)语言来编写调度模板,这些模板可以被 AutoTVM 搜索到,以找到最佳调度。这个过程被称为 自动调谐(Auto-Tuning),它有助于自动化优化张量的计算过程。
本教程建立在之前关于 如何使用 TE 编写矩阵乘法 的教程上。
自动调谐有两个步骤。
第一步是定义搜索空间。
第二步是运行搜索算法来探索这个空间。
在本教程中,学习如何在 TVM 中执行这两个步骤。整个工作流程通过矩阵乘法的例子来说明。
备注
本教程暂不能在 Windows 或最近版本的 MacOS 上运行。为了让它运行,你需要将本教程的主体包裹在一个 if __name__ == "__main__":
块中。
安装依赖项#
为了在 TVM 中使用 autotvm
包,需要安装一些额外的依赖项(也可安装 GPU 版本)。
pip3 install --user psutil xgboost cloudpickle
为了使 TVM 在 tuning 中运行得更快,建议使用 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
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)
# 调度
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” 作为平铺系数。然而,这可能不是最好的,因为最佳的平铺系数取决于实际的硬件环境和输入形状。
如果你想让调度代码在更大范围的输入形状和目标硬件上可移植,最好是定义一组候选值,并根据目标硬件上的测量结果挑选最佳值。
在 autotvm
中,可以定义可调整的参数,或者说是 “旋钮”,用于此类值。
基本的矩阵乘法模板#
以例子开始,说明如何为 split
调度操作的块大小创建可调度的参数集。
# 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. 依据配置进行调度
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 确定值来调度。
使用高级参数 API 的矩阵乘法模板#
在前面的模板中,手动列出了旋钮的所有可能值。这是定义空间的最底层的 API,它给出了要搜索的参数空间的明确列举。然而,TVM 还提供了另一组 API,可以使搜索空间的定义更容易、更智能。在可能的情况下,接受你使用这个更高级别的 API。
在下面的例子中,使用 ConfigSpace.define_split
来定义 split 旋钮。它将列举所有可能的方式来 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]
关于 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 中,我们编写了一个矩阵乘法模板,允许我们对分割调度中使用的块大小进行参数化。我们现在可以对这个参数空间进行搜索。下一步是选择一个调整器来指导对这个空间的探索。
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 提供的一些调谐器策略包括:
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, range_length=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")],
)
waiting for device...
device available
Get devices for measurement successfully!
No: 1 GFLOPS: 0.66/0.66 result: MeasureResult(costs=(0.40677344800000004,), error_no=MeasureErrorNo.NO_ERROR, all_cost=6.793164014816284, timestamp=1679468144.982454) [('tile_y', [-1, 512]), ('tile_x', [-1, 2])],None,19
No: 2 GFLOPS: 0.77/0.77 result: MeasureResult(costs=(0.3469925972,), error_no=MeasureErrorNo.NO_ERROR, all_cost=5.829528570175171, timestamp=1679468150.8401856) [('tile_y', [-1, 64]), ('tile_x', [-1, 2])],None,16
No: 3 GFLOPS: 7.75/7.75 result: MeasureResult(costs=(0.0346520524,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.8935184478759766, timestamp=1679468151.716168) [('tile_y', [-1, 512]), ('tile_x', [-1, 128])],None,79
No: 4 GFLOPS: 2.96/7.75 result: MeasureResult(costs=(0.0907275406,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.7036769390106201, timestamp=1679468153.448838) [('tile_y', [-1, 128]), ('tile_x', [-1, 8])],None,37
No: 5 GFLOPS: 10.52/10.52 result: MeasureResult(costs=(0.025506687,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.7988688945770264, timestamp=1679468154.142109) [('tile_y', [-1, 2]), ('tile_x', [-1, 128])],None,71
No: 6 GFLOPS: 13.43/13.43 result: MeasureResult(costs=(0.019993396200000003,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.5659496784210205, timestamp=1679468154.7467444) [('tile_y', [-1, 4]), ('tile_x', [-1, 512])],None,92
No: 7 GFLOPS: 4.19/13.43 result: MeasureResult(costs=(0.06404942579999999,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.2754650115966797, timestamp=1679468156.0508864) [('tile_y', [-1, 8]), ('tile_x', [-1, 2])],None,13
No: 8 GFLOPS: 4.52/13.43 result: MeasureResult(costs=(0.05936086380000001,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.1969263553619385, timestamp=1679468157.2878318) [('tile_y', [-1, 4]), ('tile_x', [-1, 2])],None,12
No: 9 GFLOPS: 1.68/13.43 result: MeasureResult(costs=(0.1596650536,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.795808792114258, timestamp=1679468160.1257796) [('tile_y', [-1, 8]), ('tile_x', [-1, 1])],None,3
No: 10 GFLOPS: 1.49/13.43 result: MeasureResult(costs=(0.1807046766,), error_no=MeasureErrorNo.NO_ERROR, all_cost=3.0970828533172607, timestamp=1679468163.2516549) [('tile_y', [-1, 128]), ('tile_x', [-1, 4])],None,27
调谐完成后,可以从日志文件中选择具有最佳测量性能的配置,并用相应的参数来编译时间表。还可以做快速验证,以确保时间表产生正确的答案。可以在 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
最后说明和总结#
在本教程中,展示了如何建立算子模板,让 TVM 搜索参数空间并选择优化的时间表配置。为了更深入地了解它的工作原理,建议在这个例子的基础上进行扩展,在 张量表达式入门 教程中演示的调度操作的基础上添加新的搜索参数。在接下来的章节中,将演示 AutoScheduler,这是 TVM 优化常见算子的方法,不需要用户提供自定义的模板。