实在之路
黑格尔💾Ansor论文笔记
type
status
date
slug
summary
tags
category
icon
password
AI summary
Background
Template guided
比如TVM, FlexTensor
编写模板，定义一些可调参数.
TVM存在的问题:
- 需要人工参与，存在工作量
- 编写模板需要专家级别的知识
- 手工编写的模板通常描述整个优化搜索空间 FlexTensor的问题：
- 模板只描述单个op，多个op fuse时，能力就弱了。
Sequential construction based
- 预定义的一系列独立的固定顺序的 搜索空间/步骤
- 每一步关注一个 主题/子搜索空间(tiling，...)， 每一步只保留 topK个选项
- 用learned cost model 局限:
- 中间的步骤的结果是incomplete program， cost model很难评估起性能，可信度不高
- 固定顺序的 子搜索空间/步骤 限制了整体的搜索空间的规模， 因为本质上子搜索空间之间很难做到正交，相互之间会影响。
- 有限的 子搜索空间 很难说考虑全了，也限制了整体搜索空间的大小.
Design
- Programmer Sampler
- Perf tuner
- task scheduler
Programmer Sampling
维特根斯坦: 我的语言的界限意味着我的世界的界限
The search space an algorithm explores determines the best programs it can find.
现有的构造搜索空间方法的局限：
- 手工写模板，不能覆盖整个搜索空间，因为人通常考虑不全
- 早期基于incomplete program的激进剪枝策略
Sketch Generation
定义状态 , S代表当前已经生成的sketch， i代表当前处理的node的index。其中node的index按照拓朴序逆序编排.
初始状态为：
递归应用规则:
最终状态: 当 i = 0时，表示状态为最终状态。
Predicates
应用规则需要满足条件，下面是一些谓词：
IsStrictInliable: 是否elem-wise
HasDataReuse: 是否为计算密集，计算中是否有大量数据复用
HashFusibileConsumer: node i 只有一个consumer，并且可以fuse
HasMoreReductionParallel: 空间维度没有什么可并行性， reduce维度有更多可并行性。比如:
空间维度指的是决定Ouput shape的那些维度
值得注意的是，这些谓词的值是通过自动分析得到的，通过分析 计算定义-数学表达式 的读写模式。
怎么自动分析得到的？
Derivation rules
rule 1: Skip: 跳过当前node
rule 2: Always Inline: 融合
rule 3: Multi-level Tiling: 进行多级tiling，cpu上采用"SSRSR"的方式，S代表空间维度，R代表reduce维度。比如: , SSRSR 把 扩展到:
rule4: Multi-level Tiling with Fusion: 把consumer融合到tiling后的计算中
rule 5: Add Cache Stage: 对于可融合的计算密集node， 添加cache node, 先将多个output cache到cache block中，cache block满了后，一次写出。 添加后可应用 "Multi-level Tiling with Fusion"。
rule 6: Reduction Factorization: 切分reduce轴，提供更多的并行度
Ansor也支持用户自定义规则.
Example
GPU support
multi-level tiling 从 SSRSRS 切换到 SSSRRSRS
两个新rule:
- 针对shared memory的cache规则
- 跨线程的reduction
Performace Tunning
使用进化算法
1. 初始轮次通过随机采样从sketch得到大量sample, 以下叫作 后代。
2. 对于每一轮的后代，依据cost model，选出一批后代，然后实际用硬件去跑，同时将得到的结果反过来训练cost model.
- 这里并不是出依据cost model的结果选出top k， 而是一种概率选择，每个采样被选中的概率和其cost-model的结果成正相关。
3. 对于每一轮，进化算法使用上一轮选出的结果以及随机采样作为这一轮的初始人口，然后进行变异和交叉，产生下一代。
4. 总共的轮次是一个固定的数，全部轮次结束后，选最好的。
变异指对sample进行一些预定义的修改，具体有：
- Tile size mutation: 随机选择一个tile loop，将其tile size除以一个随机因子，拆分成两个loop。
- Parallel mutation: 随机选择一个被标记了parallel的loop，将其与相邻loop合并或者除以一个随机因子拆分成两个loop。
- Program mutation: 随机选择一个program(比如
auto_unroll_max_step=N)，然后将其随机修改为另一个合法的值.
- Computation location mutation: 随时选择一个flexible node(没有被multi-level tiled的节点，比如conv layer的padding), 随机将其的attach point改为另一个合法的node
老实说，没理解这一点
- Node-based crossover: 杂交。Ansor中，program的基因被定义为：rewriting steps, 也就是从naive program开始的一系列规则变换以及随机标注。因此可以对两个program的基因进行杂交。为了避免随机杂交导致的非法program，Ansor的做法是将杂交的粒度限定为DAG中的node这个级别，因为不同node的rewrite step会有更少的依赖。然后Anosr也会对杂交的程序进行依赖验证，保证正确性。
TVM是固定的参数空间，而Ansor除此之外，还会有结构性的变换。
Learned Cost Model
巨大的搜索空间，意味着需要一个准备且高效的cost model。一个模型可以在多个不同硬件后端上复用，只需要用对应的数据训练就行了。
tensor program基本都是多个嵌套的循环，可分为循环语句，以及循环内部的计算赋值语句。
cost model直接预测的是循环内部的计算赋值语句的cost，一个program的cost可以约等于所有循环内部的计算赋值语句的cost的和。
使用加权平方误差作为损失函数。
优化DNN时，一般需要计算cost model的program不会超过30000，数据量很小，训练起来很开，所以选择每次都重新训练一个，而不是增量更新。
增量不是更优？
Task Scheduler
一个DNN可以被划分成多个子图(据说采用TVM的子图划分规则),其中一些子图，即使花时间去优化，对整体的性能提升也没撒提升，可能的原因：
- 本身不是瓶颈，比如只在整网中出现了一次，且数据量不大
- 这个子图本身已经很难优化 Task Scheduler就是要去找出当前的热点子图，然后为其分配一个时间单位，进行优化。 目标函数: ,其中 为子图i在DNN整网中出现的次数，t为了每个子图花的时间组成的向量，即 , 表示为子图i分配的时间。 对于每一轮，我们有一个当前的状态t向量，表示当前各个子图已经消耗了时间资源，我们的要选出一个子图x,为其分配一个时间单元，并且使得目标函数的收益最高。 具体而言，通过计算f对于所有 的偏导，然后选择值最小的那个子图。
Prev
torch代码笔记
Next
TVM论文笔记
Loading...