成都网站建设哪家公司好,网站快速收录提交,兰州网站建设小程序,在线室内设计网站完整 TVM 中文文档#xff0c;访问→TVM 中文站 作者#xff1a;Lianmin Zheng#xff0c;Chris Hoge
本教程将展示如何用 TVM 张量表达式#xff08;TE#xff09;语言编写 schedule 模板#xff0c;并通过 AutoTVM 对模板进行搜索#xff0c;从而找到最佳 schedule。…完整 TVM 中文文档访问→TVM 中文站 作者Lianmin ZhengChris 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 sysimport numpy as np
import tvm
from tvm import te
import tvm.testing# 模块名叫 autotvm
from tvm import autotvmTE 的基本矩阵乘法
回想一下用 TE 进行矩阵乘法的基本实现下面做一些改变。将矩阵乘法放在 Python 函数定义中。简单起见重点关注拆分的优化将重新排序的块大小设为固定值。
def matmul_basic(N, L, M, dtype):A te.placeholder((N, L), nameA, dtypedtype)B te.placeholder((L, M), nameB, dtypedtype)k te.reduce_axis((0, L), namek)C te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axisk), nameC)s te.create_schedule(C.op)# 调度y, x s[C].op.axisk 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), nameA, dtypedtype)B te.placeholder((L, M), nameB, dtypedtype)k te.reduce_axis((0, L), namek)C te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axisk), nameC)s te.create_schedule(C.op)# 调度y, x s[C].op.axisk 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]下面将对前面的调度代码作出四个修改然后得到一个可调的“模板”。一一解释这些修改 使用装饰器将此函数标记为简单模板。 获取 config 对象将 cfg 视为此函数的参数但我们以另外的方式获取它。cfg 参数使得这个函数不再是一个确定的 schedule。将不同的配置传递给这个函数可以得到不同的 schedule。这种使用配置对象的函数称为“模板”。
为使模板函数更精炼可在单个函数中定义参数搜索空间
用一组值来定义搜索空间。将 cfg 转为 ConfigSpace 对象收集此函数中的所有可调 knob然后从中构建一个搜索空间。根据空间中的实体进行调度。将 cfg 转为 ConfigEntity 对象当它被转为 ConfigEntity 后会忽略所有空间定义 API即 cfg.define_XXXXX(…)但会存储所有可调 knob 的确定值并根据这些值进行调度。
在 auto-tuning 的过程中首先用 ConfigSpace 对象调用这个模板来构建搜索空间然后在构建的空间中用不同的 ConfigEntity 调用这个模板来得到不同的 schedule。最后我们将评估由不同 schedule 生成的代码然后选择最佳的 schedule。 定义两个可调 knob。第一个是 tile_y它有 5 个可能值。第二个是 tile_x它和前者具有相同的可能值。这两个 knob 是独立的所以它们跨越大小为 25 5x5 的搜索空间。 配置 knob 被传递给 split 调度操作然后可以根据之前在 cfg 中定义的 5x5 确定值进行调度。
带有高级参数 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), nameA, dtypedtype)B te.placeholder((L, M), nameB, dtypedtype)k te.reduce_axis((0, L), namek)C te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axisk), nameC)s te.create_schedule(C.op)# 调度y, x s[C].op.axisk s[C].op.reduce_axis[0]##### 开始定义空间 #####cfg autotvm.get_config()cfg.define_split(tile_y, y, num_outputs2)cfg.define_split(tile_x, x, num_outputs2)##### 结束定义空间 ###### 根据 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_outputs2) 枚举了所有可能的组合以 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, npartscfg[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 configsmeasure this batch of configs on real hardware and get resultsct 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 可以更有效地探索并找到更好的配置。
开始调优
下面继续矩阵乘法的示例。首先创建一个调优任务然后检查初始的搜索空间。下面示例中是 512x512 的矩阵乘法空间大小为 10x10100。注意任务和搜索空间与选择的调优器无关。
N, L, M 512, 512, 512
task autotvm.task.create(tutorial/matmul, args(N, L, M, float32), targetllvm)
print(task.config_space)输出结果
ConfigSpace (len100, space_map0 tile_y: Split(policyfactors, product512, num_outputs2) len101 tile_x: Split(policyfactors, product512, num_outputs2) len10
)然后定义如何评估生成的代码并且选择一个调优器。由于我们的空间很小所以随机调优器就可以。
本教程只做 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(builderlocal, runnerautotvm.LocalRunner(number5))# 用 RandomTuner 开始调优, 日志记录到 matmul.log 文件中
# 可用 XGBTuner 来替代.
tuner autotvm.tuner.RandomTuner(task)
tuner.tune(n_trial10,measure_optionmeasure_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_noMeasureErrorNo.NO_ERROR, all_cost0.638512134552002, timestamp1657225928.6342561) [(tile_y, [-1, 1]), (tile_x, [-1, 256])],None,80
No: 2 GFLOPS: 2.30/8.48 result: MeasureResult(costs(0.1165478966,), error_noMeasureErrorNo.NO_ERROR, all_cost2.0105199813842773, timestamp1657225930.6636436) [(tile_y, [-1, 4]), (tile_x, [-1, 8])],None,32
No: 3 GFLOPS: 11.82/11.82 result: MeasureResult(costs(0.0227097348,), error_noMeasureErrorNo.NO_ERROR, all_cost0.5589795112609863, timestamp1657225931.7059512) [(tile_y, [-1, 64]), (tile_x, [-1, 32])],None,56
No: 4 GFLOPS: 1.66/11.82 result: MeasureResult(costs(0.1616202114,), error_noMeasureErrorNo.NO_ERROR, all_cost2.6911513805389404, timestamp1657225934.9635096) [(tile_y, [-1, 1]), (tile_x, [-1, 4])],None,20
No: 5 GFLOPS: 3.65/11.82 result: MeasureResult(costs(0.073561817,), error_noMeasureErrorNo.NO_ERROR, all_cost1.3051848411560059, timestamp1657225936.3988533) [(tile_y, [-1, 256]), (tile_x, [-1, 16])],None,48
No: 6 GFLOPS: 1.85/11.82 result: MeasureResult(costs(0.1452834464,), error_noMeasureErrorNo.NO_ERROR, all_cost2.5179028511047363, timestamp1657225938.961955) [(tile_y, [-1, 512]), (tile_x, [-1, 4])],None,29
No: 7 GFLOPS: 0.87/11.82 result: MeasureResult(costs(0.30933780240000003,), error_noMeasureErrorNo.NO_ERROR, all_cost5.067087888717651, timestamp1657225944.589149) [(tile_y, [-1, 512]), (tile_x, [-1, 2])],None,19
No: 8 GFLOPS: 10.53/11.82 result: MeasureResult(costs(0.025489421,), error_noMeasureErrorNo.NO_ERROR, all_cost0.5452830791473389, timestamp1657225945.1592515) [(tile_y, [-1, 4]), (tile_x, [-1, 64])],None,62
No: 9 GFLOPS: 1.58/11.82 result: MeasureResult(costs(0.16960762680000002,), error_noMeasureErrorNo.NO_ERROR, all_cost2.8109781742095947, timestamp1657225948.0900776) [(tile_y, [-1, 2]), (tile_x, [-1, 2])],None,11
No: 10 GFLOPS: 2.42/11.82 result: MeasureResult(costs(0.11083148779999999,), error_noMeasureErrorNo.NO_ERROR, all_cost1.8757600784301758, timestamp1657225950.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(), rtol1e-4)输出结果
Finish loading 10 records总结
本教程展示了如何构建算子模板使得 TVM 能够搜索参数空间并选择优化的调度配置。为了更深入地了解其工作原理推荐基于 :ref: 张量表达式入门 tensorexpr_get_started 教程中演示的调度操作向调度添加新的搜索参数。接下来的章节将演示 AutoScheduler它是TVM 中一种优化常用算子的方法同时无需用户提供自定义的模板。
