极智AI | 谈谈 tvm ansor

  本文主要讨论一下 tvm ansor。

  参考论文:
极智AI | 谈谈 tvm ansor_第1张图片

  参考代码:

https://github.com/apache/tvm

文章目录

    • 1、深度学习系统堆栈
    • 2、相关的高性能张量生成工作
      • 2.1 auto tvm
      • 2.2 halide auto-scheduler
    • 3、tvm ansor 原理
      • 3.1 program sampling
        • 3.1.1 sketch generation
        • 3.1.2 random annotation
      • 3.2 performance fine-tuning
        • 3.2.1 evolutionary search
        • 3.2.2 learned cost model
      • 3.3 task scheduler
    • 4、性能数据

1、深度学习系统堆栈

  深度学习系统堆栈,前端是各种 AI 框架(pytorch、tf、mxnet、caffe … ),后端是各种硬件(NVGPU、ARMCPU、X86CPU、NPU … )
极智AI | 谈谈 tvm ansor_第2张图片

  对于深度学习编译器来说,就是把前端 AI 框架出来的网络 / 算子 code_gen / compiler 成对应 target 硬件后端高性能的机器码,而 tvm / halide 把这个compiler 的过程拆成计算(compute)和调度(optimized code)两个过程,理由很简单,这样解耦后只需要关注调度优化,因为调度是和硬件强相关的,而不用担心调度影响了精度。拿 dense_layer + relu 这个 subgraph 来说,halide 和 tvm 的计算表达可能是这样的:

# halide  
dense(o, b) += data(i, b) * weight(i, o) 
relu(o, b) = max(dense(o, b), 0.0) 

# tvm 
dense = compute(shape, lambda b, o: sum(data[b, i] * weight[o, i], i)) 
relu = compute(shape, lambda b, o: max(dense[b, o], 0.0))

极智AI | 谈谈 tvm ansor_第3张图片


2、相关的高性能张量生成工作

2.1 auto tvm

  做法是使用 template-guided 进行搜索,对于每个 op 使用 templates 去定义搜索空间。

  缺点是:

  (1)没有做到全自动:还需要大量的手动工作;

  (2)受限的搜索空间:不能达到更加好的优化性能;
极智AI | 谈谈 tvm ansor_第4张图片

2.2 halide auto-scheduler

  做法是基于 sequential construction 进行搜索,使用 beam search 算法搜索顺序生成。说一下 beam search 算法,beam search 是所谓的 束搜索,是对 greedy search(贪心搜索)的一个改进算,相对 greedy search 扩大了搜索空间,但远远不及穷举搜索指数级的搜索空间,是二者的一个折中方案。

  halide auto-scheduler 的缺点是:

  (1)因为它是顺序生成的,对于中间的算子优化来说是不连续不完整的,因此 cost model 没办法进行准确的预测,导致了错误的累积;

  (2)和 auto tvm 一样,搜索空间也是受限的;

  通过对以上两个调度搜索算法的讨论,得到三个痛点:(1)怎么样能够自动的获取一个很大的搜索空间;(2)如何让搜索更加的高效;(3)对于这么多的搜索任务我们怎么样去分配资源。对于这三个痛点,ansor 的解法是:对于(1)采用 hierarchical search space;对于(2)采用 sample complete programs and fine-tune them;对于(3)采用 task scheduler 程序对重要任务进行优先排序。


3、tvm ansor 原理

  先来看一下 tvm ansor 的逻辑栈,下面对主要的操作展开谈谈。
极智AI | 谈谈 tvm ansor_第5张图片

  前端接的还是 AI框架出来的 models,然后切分成很多 subgraphs 子图,对于每个 subgraph 都会进行如下的调度优化。

3.1 program sampling

  program sampling 的目标是自动地构建一个连续的大的搜索空间。这里要理解这个搜索空间到底是用于搜索个啥,我的理解是搜索计算的姿势或者说计算的调度方式,而不是计算本身。方法是使用 sketch generation 和 random annotation。其中 sketch 是 a few good high-level structures,annotaion 是 billions of low-level details,通俗一点理解 sketch 进行高层抽象,annotaion 进行底层映射,通过这两步操作就能构建了一个很大的搜索空间。

这个用于解决痛点(1)搜索空间大小问题。
在这里插入图片描述

3.1.1 sketch generation

  先贴一些论文里的描述:
极智AI | 谈谈 tvm ansor_第6张图片

  做法大致流程是这样:先出数学表达 -> 再出程序表达 -> 根据程序表达和数学表达出有向无环图描述计算节点和计算逻辑 -> 根据设计好的 sketch gen rules 来生成 sketch 表达(比如对于CPU,有 “SSRSRS”的循环展开结构)。所以 sketch 的整个流程可以用如下这个 dense_layer + relu 的例子展示,可以看到最后生成的 sketch 草图里有很多类似 TILE_I0 的参数,这些是超参,是需要调优搜索的值,也是 random annotation 的工作之一,所以最后生成的 sketch 草图不止一个,往往是一个草图列表。
极智AI | 谈谈 tvm ansor_第7张图片

  还有这个例子也是一样:
极智AI | 谈谈 tvm ansor_第8张图片

3.1.2 random annotation

  annotaion 是获得 sketch 草图的后续动作,annotaion 操作是 randomly 选择一个sketch,randomly 填充 tile sizes,randomly parallelize 外部循环,randomly vectorize 内部循环,randomly change 计算节点位置以对 tile stucture 做微调等。还是接着上面的 dense_layer + relu 的例子展示:
极智AI | 谈谈 tvm ansor_第9张图片

3.2 performance fine-tuning

  接着 program sampling 后的操作就到了 performance tuner 了。这里这要包含搜索算法设计和代价评估模型的设计,使 ansor 能够更加高效的进行搜索和收敛。

  这个用于解决痛点(2)搜索效率问题。

3.2.1 evolutionary search

  random annotation 的随机生成/抽样不能保证性能,ansor 采用类似遗传算法的进化搜索算法来进一步对 annotation 抽样程序进行进化搜索。主要涉及 基因突变和基因交叉。

  (1)突变:随机改变 tile size;随机改变 parallel/unroll/vectorize 因子和粒度;随机改变 computation location;

  (2)交叉:从算子计算有向无环图的角度解释,如下:
在这里插入图片描述

3.2.2 learned cost model

  预测每个 non-loop innermost statement 的得分,可能比较抽象,上一些论文里的解释
极智AI | 谈谈 tvm ansor_第10张图片

  再上个例子进行说明,如下,statement B 和 statement C 就是所谓的 non-loop innermost statement 了,这个比较直观一点。
极智AI | 谈谈 tvm ansor_第11张图片

  每个 non-loop innermost statement 提取的特征包括:使用的 cache lines、使用的 memory、reuse distance、arithmetic intensity 等

3.3 task scheduler

  ansor 中的 task scheduler 会对重要任务进行优先排序。

  这个是用于解决痛点(3)多任务调度资源分配问题。

  我们拿 resnet50 来说,在 ansor 中会把它切成 29 个独立的 subgraphs。如果按照传统的任务调度方式,一般会顺序地对这些 subgraphs 进行调度,就像这样。
在这里插入图片描述

  ansor 中的 task scheduler 会将时间进行切片并对重要的子图优先排序,就像这样,这种优化很像推理里的 pipeline 优化。
在这里插入图片描述


4、性能数据

  下面给出一些测试的性能数据,论文里也有:

  (1) 单算子

  测试平台是 Intel-Platinum 8124M (18 cores):
极智AI | 谈谈 tvm ansor_第12张图片
极智AI | 谈谈 tvm ansor_第13张图片

  (2) subgraph

  测试平台是 ‘’@C’’ for Inter CPU (8124M)、‘’@G’’ for NVIDIA (V100)。
在这里插入图片描述
极智AI | 谈谈 tvm ansor_第14张图片

  (3) 整网

  测试平台是 Intel CPU (8124M)、NVIDIA GPU (V100)、ARM CPU (A53)。
在这里插入图片描述
极智AI | 谈谈 tvm ansor_第15张图片

  从上面给出的数据可以看出,不管是单算子、subgraph、整网来说,ansor 的表现都很棒。


  以上讨论了我目前对 tvm ansor 的理解。


  收工~



扫描下方二维码即可关注我的微信公众号【极智视界】,获取更多AI经验分享,让我们用极致+极客的心态来迎接AI !

你可能感兴趣的:(极智AI,编译器,人工智能,深度学习,搜索引擎优化)