TVM：使用 Schedule 模板和 AutoTVM 来优化算子

TVM：使用 Schedule 模板和 AutoTVM 来优化算子

在本文中，我们将介绍如何使用 TVM 张量表达式（Tensor Expression，TE）语言编写 Schedule 模板，AutoTVM 可以搜索通过这些模板找到最佳 Schedule。这个过程称为自动调整（Auto Tuning），它有助于自动优化张量计算的过程。

本教程需基于之前介绍的如何使用 TE 来写一个矩阵乘法的教程。

Auto Tuning 有两步：

  • 第一步是定义一个搜索空间
  • 第二步是运行相应的搜索算法来探索这个空间

本教程将展示如何在 TVM 中完成这两步，整个流程将以矩阵乘法为例。

安装依赖

要使用 TVM 中的 AutoTVM 包，需要安装这些额外的依赖：

pip install --user psutil xgboost cloudpickle

为了使 TVM 在 tuning 中运行得更快，建议使用 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.testingfrom tvm import autotvm

TE 实现基本的矩阵乘法

回想一下使用TE实现矩阵乘法的基本方法。我们把它写在这里，稍作修改。我们将把乘法封装在Python 函数定义中。简单起见，我们将把注意力集中在分割优化（split optimization）上，使用一个固定值来定义重排（reordering）的块的大小。

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)# scheduley, 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 作为 tiling 因子。但是，它可能不是最好的，因为最佳的tiling 因子取决于实际的硬件环境和输入的形状。

如果希望 Schedule 代码能够在更大范围的输入形状和目标硬件上移植，最好定义一组候选值，并根据目标硬件上的测量结果选择最佳值。

在 autotvm 中，我们可以定义一个可调参数，或为此类参数定义一个“knob”。

一个基本的矩阵乘法模板

我们的一个示例，介绍如何为 spliting schedule 操作的块大小创建一个可调参数集。

# 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)# scheduley, x = s[C].op.axisk = s[C].op.reduce_axis[0]# 2. get the config objectcfg = autotvm.get_config()# 3. define search spacecfg.define_knob("tile_y", [1, 2, 4, 8, 16])cfg.define_knob("tile_x", [1, 2, 4, 8, 16])# 4. schedule according to configyo, 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]

这里对于之前的 schedule 代码，我们有四处调整，从而可以得到一个可调节的 “模板”：

  1. 使用装饰器 @autotvm.template() 来将此函数标记为一个模板

  2. 获取配置项： cfg 可以视为本函数的一个参数，但是我们是通过一种不同的方式得到它的。有了这个参数，这个函数就不再是一个确定的 schedule。而是我们可以为这个函数传入不同的配置来得到不同的 schedules。一个像这样含有配置项的函数称为 “模板”。

    为了使得模板函数更加 compact，我们通过完成以下两件事情来在一个函数中定义参数搜索空间：

    • 通过一组值定义搜索空间。这是通过将 cfg 设置为一个 ConfigSpace 对象完成的。它会收集本函数中所有的可以调节的 knobs ，并从中建立一个搜索空间。
    • 根据此空间中的实体进行 schedule。这是通过将 cfg 设置为一个 ConfigEntity 对象完成的。当它是一个 ConfigEntity 时，它会忽略所有空间定义的API（即 cfg.define_XXX(...)）。相反，它会所有可调节的 knobs 保存确定的值，然后我们根据这些值进行 schedule。

    在 auto-tuning 时，我们将首先使用 ConfigSpace 对象调用此模板以构建搜索空间。然后，我们在构建空间中使用不同的 ConfigEntity 调用此模板，以获得不同的 schedule。最后，我们将度量由不同计划生成的代码，并选择最优的。

  3. 定义两个可调节的 knobs。第一个是具有 5 个可能值的 tile_y。第二个是 tile_x，具有相同的可能值列表。这两个 knob 是独立的，因此它们会有大小为 25=5x5 的搜索空间。

  4. 配置 knobs 被传递到 split schedule 操作，允许我们根据之前在 cfg 中定义的 5x5 确定值进行 schedule。

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

在前面的模板中，我们手动列出了 knobs 的所有可能值。这是定义空间的最低级别API，并提供要搜索的参数空间的显式枚举。但是，我们还提供了另一组API，可以使搜索空间的定义更简单、更智能。在可能的情况下，我们建议您使用此高级API。

在下面的示例中，我们使用 ConfigSpace.define_split 来定义拆分 knob。它将列举所有分割轴和构建空间的可能方法。

我们还有 ConfigSpace.define_reorder用于 reorder knob，ConfigSpace.define_annotation 用于展开、矢量化、线程绑定等注释。当高级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)# scheduley, x = s[C].op.axisk = 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 configyo, 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]

注意：More Explanation on cfg.define_split

更多关于 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个可能值。

在 schedule 期间，cfg[“tile_y”] 是一个 SplitEntity 对象。我们将外轴和内轴的长度存储在 cfg['tile_y'].size（包含两个元素的元组）中。在这个模板中，我们使用yo, yi=cfg['tile_y']] 来应用它。实际上，这相当于yo, yi=s[C].split（y，cfg[“tile\u y”].size[1]）或 yo, yi=s[C].split（y，npart=cfg[“tile\u y”].size[0]）

使用 cfg.apply API 的优点是，它使多级拆分（即当num_outputs>=3时）更加容易。

第二步：使用AutoTVM优化矩阵乘法

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

TVM 中的 Auto-tuners

调谐器的工作可以通过以下伪代码来描述：

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提供的一些调谐器策略包括：

  • tvm.autotvm.tuner.RandomTuner：按随机顺序枚举空间
  • tvm.autotvm.tuner.GridSearchTuner：按网格搜索顺序枚举空间
  • tvm.autotvm.tuner.GATuner：利用遗传算法进行空间搜索
  • tvm.autotvm.tuner.XGBTuner：使用基于模型的方法。训练 XGBoost 模型预测降低 IR 的速度，并根据预测选择下一批。

我们可以根据空间大小、时间预算和其他因素选择调谐器。例如，如果您的空间非常小（小于1000），gridsearch 调谐器或随机调谐器就足够了。如果您的空间级别为 10^9（这是CUDA GPU上conv2d 算子的空间大小），XGBoostTuner 可以更高效地探索并找到更好的配置。

开始 tuning

这里我们继续我们的矩阵乘法示例。首先，我们创建一个 tuning 任务。我们还可以检查初始化的搜索空间。在这种情况下，对于 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)

此处输出：

N, L, M = 512, 512, 512
task = autotvm.task.create("tutorial/matmul", args=(N, L, M, "float32"), target="llvm")
print(task.config_space)

然后我们需要定义如何测量生成的代码并选择调谐器。因为我们的空间很小，随机调谐器也可以。

在本教程中，我们只做了 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")],
)

调优 tuning 完成后，我们可以从日志文件中选择性能最好的配置，并使用相应的参数编译 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)

最后的注意事项与总结

在本教程中，我们展示了如何构建算子模板，以使得 TVM 能够对参数空间进行搜索并选择最优的的 schedule 配置。为了更深入地了解其工作原理，我们建议在此示例上进行扩展，可以根据使用张量表达式（TE）入门教程中演示的 schedule 操作向计划添加新的搜索参数。在接下来的部分中，我们将演示 AutoScheduler，这是一种TVM优化常用算子的方法，用户无需提供用户定义的模板。

Ref：
https://tvm.apache.org/docs/tutorial/autotvm_matmul_x86.html

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.mzph.cn/news/532664.shtml

如若内容造成侵权/违法违规/事实不符，请联系多彩编程网进行投诉反馈email:809451989@qq.com，一经查实，立即删除！

相关文章

TVM：使用 Auto-scheduling 来优化算子

TVM：使用 Auto-scheduling 来优化算子

TVM&#xff1a;使用 Auto-scheduling 来优化算子 在本教程中&#xff0c;我们将展示 TVM 的 Auto-scheduling 功能如何在无需编写自定义模板的情况下找到最佳 schedule。 与基于模板的 AutoTVM 依赖手动模板定义搜索空间不同&#xff0c;auto-scheduler 不需要任何模板。 用…
阅读更多...
C语言—sort函数比较大小的快捷使用--algorithm头文件下

C语言—sort函数比较大小的快捷使用--algorithm头文件下

sort函数 一般情况下要将一组数从的大到小排序或从小到大排序&#xff0c;要定义一个新的函数排序。 而我们也可以直接使用在函数下的sort函数&#xff0c;只需加上头文件&#xff1a; #include<algorithm> using namespace std;sort格式&#xff1a;sort(首元素地址&…
阅读更多...
散列的使用

散列的使用

散列 散列简单来说&#xff1a;给N个正整数和M个负整数&#xff0c;问这M个数中的每个数是否在N中出现过。 比如&#xff1a;N&#xff1a;{1,2,3,4}&#xff0c;M{2,5,7}&#xff0c;其中M的2在N中出现过 对这个问题最直观的思路是&#xff1a;对M中每个欲查的值x&#xff0…
阅读更多...
关于C++中的unordered_map和unordered_set不能直接以pair作为键名的问题

关于C++中的unordered_map和unordered_set不能直接以pair作为键名的问题

关于C中的unordered_map和unordered_set不能直接以pair作为键名的问题 在 C STL 中&#xff0c;不同于有序的 std::map 和 std::set 是基于红黑树实现的&#xff0c;std::unordered_map 和 std::unordered_set 是基于哈希实现的&#xff0c;在不要求容器内的键有序&#xff0c…
阅读更多...
AI编译器与传统编译器的联系与区别

AI编译器与传统编译器的联系与区别

AI编译器与传统编译器的区别与联系 总结整理自知乎问题 针对神经网络的编译器和传统编译器的区别和联系是什么&#xff1f;。 文中提到的答主的知乎主页&#xff1a;金雪锋、杨军、蓝色、SunnyCase、贝壳与知了、工藤福尔摩 笔者本人理解 为了不用直接手写机器码&#xff0…
阅读更多...
python学习1：注释\变量类型\转换函数\转义字符\运算符

python学习1：注释\变量类型\转换函数\转义字符\运算符

python基础学习 与大多数语言不同&#xff0c;python最具特色的就是使用缩进来表示代码块&#xff0c;不需要使用大括号 {} 。缩进的空格数是可变的&#xff0c;但是同一个代码块的语句必须包含相同的缩进空格数。 &#xff08;一个tab4个空格&#xff09; Python语言中常见的…
阅读更多...
Python、C++ lambda 表达式

Python、C++ lambda 表达式

Python、C lambda 表达式 lambda函数简介 匿名函数lambda&#xff1a;是指一类无需定义标识符&#xff08;函数名&#xff09;的函数或子程序。所谓匿名函数&#xff0c;通俗地说就是没有名字的函数&#xff0c;lambda函数没有名字&#xff0c;是一种简单的、在同一行中定义函…
阅读更多...
python 学习2 /输入/ 输出 /列表 /字典

python 学习2 /输入/ 输出 /列表 /字典

python基础学习第二天 输入输出 xinput("输入内容") print(x)input输出&#xff1a; eval :去掉字符串外围的引号&#xff0c;按照python的语法执行内容 aeval(12) print(a)eval输出样式&#xff1a; 列表 建立&#xff0c;添加&#xff0c;插入&#xff0c;删去…
阅读更多...
Linux、Mac 命令行快捷键

Linux、Mac 命令行快捷键

Linux、Mac 命令行快捷键 Linux 命令行编辑快捷键&#xff0c;参考了好多个&#xff0c;应该算是比较全的了&#xff0c;Linux 和 Mac 的都有&#xff0c;笔者本人比较常用的也已经红色标出来了&#xff0c;如有错误或遗漏&#xff0c;欢迎留言指出。 光标移动及编辑&#xff…
阅读更多...
Python 命令行传参

Python 命令行传参

Python 命令行传参 说到 python 命令行传参&#xff0c;可能大部分人的第一反应就是用 argparse。的确&#xff0c;argparse 在我们需要指定多个预设的参数&#xff08;如深度学习中指定模型的超参数等&#xff09;时&#xff0c;是非常有用的。但是如果有时我们只需要一个参数…
阅读更多...
快速排序 C++

快速排序 C++

快速排序 C 本文图示借鉴自清华大学邓俊辉老师数据结构课程。 快速排序的思想 快速排序是分治思想的典型应用。该排序算法可以原地实现&#xff0c;即空间复杂度为 O(1)O(1)O(1)&#xff0c;而时间复杂度为 O(nlogn)O(nlogn)O(nlogn) 。 算法将待排序的序列 SSS 分为两个子…
阅读更多...
Linux命令行下感叹号的几个用法

Linux命令行下感叹号的几个用法

Linux命令行下 " ! " 的几个用法 ! 在大多数编程语言中表示取反的意思&#xff0c;但是在命令行中&#xff0c;他还有一些其他的神奇用法。熟练掌握这些用法&#xff0c;可以大大提高我们日常命令行操作的效率。 1 执行历史命令 !! ! 在命令行中可以用来执行历史…
阅读更多...
三地址码简介

三地址码简介

三地址码简介 三地址码&#xff08;Three Address Code&#xff09;是一种最常用的中间语言&#xff0c;编译器可以通过它来改进代码转换效率。每个三地址码指令&#xff0c;都可以被分解为一个四元组&#xff08;4-tuple&#xff09;的形式&#xff1a;&#xff08;运算符&am…
阅读更多...
llvm与gcc

llvm与gcc

llvm与gcc llvm 是一个编译器&#xff0c;也是一个编译器架构&#xff0c;是一系列编译工具&#xff0c;也是一个编译器工具链&#xff0c;开源 C11 实现。 gcc 相对于 clang 的优势&#xff1a; gcc 支持更过语言前端&#xff0c;如 Java, Ada, FORTRAN, Go等gcc 支持更多地 …
阅读更多...
攻防世界web新手区解题 view_source / robots / backup

攻防世界web新手区解题 view_source / robots / backup

1**. view_source** 题目描述&#xff1a;X老师让小宁同学查看一个网页的源代码&#xff0c;但小宁同学发现鼠标右键好像不管用了。 f12查看源码即可发现flag 2. robots 题目描述&#xff1a;X老师上课讲了Robots协议&#xff0c;小宁同学却上课打了瞌睡&#xff0c;赶紧来教教…
阅读更多...
python参数传递*args和**kwargs

python参数传递*args和**kwargs

python参数传递*args和**kwargs 和* 实际上真正的Python参数传递语法是 * 和 ** 。*args 和 **kwargs 只是一种约定俗成的编程实践。我们也可以写成 *vars 和 **kvars 。就如同其他常规变量的命名一样&#xff0c; args 和 kwargs 只是一种习惯的名称。 *args 和 **kwargs 一…
阅读更多...
听GPT 讲Rust源代码--src/tools(25)

听GPT 讲Rust源代码--src/tools(25)

File: rust/src/tools/clippy/clippy_lints/src/methods/suspicious_command_arg_space.rs 在Rust源代码中&#xff0c;suspicious_command_arg_space.rs文件位于clippy_lints工具包的methods目录下&#xff0c;用于实现Clippy lint SUSPICIOUS_COMMAND_ARG_SPACE。 Clippy是Ru…
阅读更多...
Java一次编译，到处运行是如何实现的

Java一次编译，到处运行是如何实现的

Java一次编译&#xff0c;到处运行是如何实现的 转自&#xff1a;https://cloud.tencent.com/developer/article/1415194 &#xff08;排版微调&#xff09; JAVA编译运行总览 Java是一种高级语言&#xff0c;要让计算机执行你撰写的Java程序&#xff0c;也得通过编译程序的…
阅读更多...
JIT（动态编译）和AOT（静态编译）编译技术比较

JIT（动态编译）和AOT（静态编译）编译技术比较

JIT&#xff08;动态编译&#xff09;和AOT&#xff08;静态编译&#xff09;编译技术比较 转自&#xff1a;https://www.cnblogs.com/tinytiny/p/3200448.html Java 应用程序的性能经常成为开发社区中的讨论热点。因为该语言的设计初衷是使用解释的方式支持应用程序的可移植…
阅读更多...
python解释器

python解释器

python解释器 计算机编程语言 本部分参考自&#xff1a;https://zhuanlan.zhihu.com/p/141212114 从计算机编程语言说起&#xff0c;它主要分为三类&#xff1a;机器语言、汇编语言、高级语言。 机器语言是一种计算机可以直接识别并执行的二进制指令集。由于其可以直接交给…
阅读更多...
最新文章

Copyright @ 2022~2023