AI Compiler是什么?
Homura 我摸到了!

为了让更多人对AI compiler有个了解,在此对这两者的区别和联系做一个科普,也因此本文以科普区别为主,不会深入。这篇文章一直想写,也算是对我去年工作中所学到的一部分东西的总结,但是硬是咕咕咕到了现在,最后选择了假期结束前把这一篇赶出来以提前适应上班状态,避免假期太强的假期综合症。个人水平有限,如有偏颇之处欢迎联系我指正

本文将从两方面讲述内容,首先是AI compiler是什么,都在做什么,其次是和传统compiler的异同。为了让读者能更好的理解内容,所需的背景知识我会尽可能的在文中做注解

AI compiler是做什么的

将各种框架训练产生的模型文件进行编译,生成目标平台的代码。从这个角度来看是和传统compiler是非常类似的,但是模型文件更像一个DSL(ldomain-specific language)

编译流程

先从解析输入开始,按照编译的流程来讲解各个阶段的异同。而这些异同大都是由于ai compiler输入的特殊性质导致的

编译对象

首先要编译的对象就有很大的不同。传统compiler则是编译的语言源代码,而AI compiler编译的是各种各样的模型,编译对象的不同导致了后面的各种处理大相径庭。

先来科普一下模型的组成:模型中包含了一个计算图以及各种数据,而计算图又是由许多算子构成的。每一个算子代表了一种计算

我个人觉得模型也可以算是一种DSL,从模型的输出向上看相当于一个expr,而每个op结点相当于一些特定函数。

解析输入方式不同

而传统语言源代码需要经过各种的parse,手写parser更是非常费力不讨好的一件事情,尽管现在的parser generator技术比较成熟,parser写起来依然是非常麻烦的。

对于AI compiler来说需要支持各种各样的模型的解析。这里不需要写什么复杂的parser了,像onnx会提供一个文件,可以通过protobuf解析这个文件生成对应的解析模型的源代码,直接调用生成的解析模型的源码中的函数即可。由于这个原因,你也不用担心不使用特定的语言进行编写还要自己做解析的工作,极端的讲,哪怕有一天要换语言做你也不需要担心解析的过程。在这个过程中更多的是将模型的数据取出来,放入设计好的IR中。

但是对于ai compiler来说你需要支持各种模型,如果只是支持某一种格式的模型是远远不会有用户的,这里不像传统compiler只需要支持自己语言的parser就可以了。目前主流的框架大致有三类,pytorch、onnx以及tensorflow,这三类有着各自的模型格式,而三者都有一定的用户群体,框架的支持程度对于用户来讲是一个非常关键的点。

高层IR设计

上面提到了ai compiler需要支持多种格式的模型,而不同格式模型的算子定义又是有许多差异。想要做到兼容各种格式的模型又是一个非常麻烦的问题,假设你设计了一套对应了算子的高层IR,可能需要对输入的模型中的算子前后添加一些操作,使其达到等效于你所选择的这个算子的实现。
用常规编程语言的例子就是C语言中要做到类似于成员函数调用通常会在这个函数的第一个成员传入结构体的指针。但是实际上麻烦的事情更多,很多算子甚至不能在不同框架转换,有的能转换也非常复杂,而关于这个问题本文就不深入探索了。

而传统编程语言的高层IR(通常为ast)相对简化很多,高层不需要考虑兼容与转换的问题。

优化

将输入读取进来后要做的事情当然是优化,编译器不仅要能够正确的生成对应平台的可执行代码,还要尽可能保证性能。

这个方面可讲的实在太多,不同层面的IR的优化方式又是各不相,而我所知的也比较有限(很多地方没有参与,但是有一些了解,我想做一个简单的科普还是不成问题),就挑ai编译器讲一下通常都有哪些方面的优化(简单提及概念),都是做什么用的,为什么要有这样的优化。

图优化

  1. 图节点合并:这个想法非常自然,只要减少了节点数量那么计算所需要的时间自然也会减少许多
  2. 更换顺序:有的时候更换节点顺序后一些节点就可以自然的合并
  3. 还有很多优化是基于算子自身定义的,这些在此就不提及了,本质上都是为了减少计算

fuse

将多个节点融合到一个子图中,直接影响到后面的tiling、调度、buffer分配,这是比较常见的一个步骤,因为本质上是为优化服务因此放到了这里。

tiling :数据切分与重排

tiling这边我没有实际参与过,所以我只能大概讲一下在编译到ai加速器上的情况下我的理解。

对于ai加速器来说,通常只会适应某一些满足条件的数据大小的计算,而实际给加速器的数据大小则是各种各样的,因此需要将数据切分到适应加速器的结构。而大部分情况图上的每个节点所需要数据的大小则是已知的,因此可以提前切分好数据进行计算。(也有节点的数据大小不固定的情况,这里暂且不谈)

对于ai来说计算很多数据都是多维数组,而实际计算很多又是多层循环,常规的数据计算方式相对低效,而许多数据又是编译期间固定的,所以需要重新以一个高效的方式重排数据。

tiling这个过程可以说是对性能影响最大的部分之一,相信大部分人都看过那个经典的按行遍历与按列遍历二维数组的例子,不仅如此,还会牵扯到计算单元的利用率以及数据传输的带宽利用。

调度

将计算图序列化,计算出一个合适的算子执行顺序

如果要涉及多设备还要尽可能做到多设备之间减少依赖,同时要考虑到数据在不同设备之间传输的带宽

buffer分配

数据是通常以一个tensor为单位(最简单的说法tensor就是一个多维数组,但是这样并不确切,但是理解这里足以),而一个tensor通常存在一个buffer之中,运算的时候从buffer中取数据

通过合理的分配方式减少运算中内存的使用,其中牵扯到计算buffer的生命周期,什么时候可以及时释放掉这块buffer,什么时候可以重用已有的buffer等等

生成代码与运行

在这方面其实都是差不多的,ai compiler还经常会生成一些ai加速器用的代码。对于ai加速器来说更多是几条配置指令加一条计算指令来执行特定的算子(上面提到的数据切分重排也和这个问题有非常大的联系)

关于生成产物的运行,和传统compiler相同也是有两类选择

  1. 生成一个直接可以执行的程序
  2. 生成类似于字节码的东西供另一个运行时的程序读取并且执行

编译时间敏感度

传统compiler还是对编译时间比较敏感的,因此导致了一些算法必须选择一个较好解,而最优解是需要很长时间的。

对于ai compiler编译时间的敏感程度相对较小,而且如果开启量化需要跑量化矫正集那根本无法控制时间(也因此需要高性能的evaluator)。而且对目标执行速度要求高,因此有更多的时间去搜索更好的解,相对于炼丹来说这点时间洒洒水啦

AI compiler特有的部分

量化

这里的量化并不是指量化交易,而是指一种将浮点数转换为定点数的计算。在ai中通常使用float进行计算,为了缩小数据大小通常会将float量化到int8,而最后还会转换回float输出(这个则是反量化过程)。如果只是做常规的数据类型转换那一定会有很大的精度损失,因此需要各种量化的算法来尽可能减少这一影响。

而量化通常需要统计数值范围,并且使用这个范围来算出一个适当的量化参数,而这个范围我们需要通过在编译期间实际执行整个模型来得到,这个时候我们就需要一个evaluator来执行。

evaluator与kernel

我觉得这里可以视为以模型和参数作为输入的解释器,对于模型来说最小单位是一个算子,那么我们就需要添加每一个算子所对应的实现,又称为kernel。

而kernels的实现不仅要正确,还要尽可能的高效。原因有如下两条

  1. 这直接影响到开启量化后的编译时间。
  2. 对于ai加速器来说只能够加速特定的算子,而其他算子依然会使用cpu来执行。cpu上的算子执行如果要利用这些kernels的话那它们的性能也是非常重要。

关于第二条,这只是做法之一,实际上加速器加速不到的算子也有很多的实现方式,在这里只是提及有这种实现方式不进行评价好与坏,所以仅供参考。

利用传统编译器的技术来做ai compiler

常量折叠、寄存器分配等技术都是可以从传统编译器来借鉴的。

一个非常典型的利用传统编译器技术的莫过于TVM(最有名的开源ai compiler)。其中的高层IR(Relay)直接利用了lambda calculus

既然我们知道如何做control flow(lambda calculus),为啥不直接用lambda calculus当IR呢?这就是relay了。(当然,传统DL compiler能做的还是一样,但是没啥好讲的(maybe sized tensor?but sized tensor is boring))。
选取了lambda calculus为ir以后,由于这上面的研究很多,我们实现需求复杂的任务比其他框架简单得多 - 因为我们只需要照抄经典compiler算法。

原文链接:https://www.zhihu.com/question/331611341/answer/875630325

不过只是这些当然还不够,需要探索更多专为ai相关的技术才能做好ai compiler。

最后

当了解了一些知识以后,就很难从一个完全不知道的视角去讲述了,所以写完本文我也难以把握哪里是相关知识较少的人看不明白的地方。读者如果能够通过本文了解ai compiler大致是什么样子的话那是再好不过了,如果读完本文对ai compiler产生了兴趣也欢迎进入这个行业和我一起摸爬滚打

  • 本文标题:AI Compiler是什么?
  • 本文作者:Homura
  • 创建时间:2022-02-05 10:38:18
  • 本文链接:https://homura.live/2022/02/05/Other/what-is-ai-compiler/
  • 版权声明:本博客所有文章除特别声明外,均采用 BY-NC-SA 许可协议。转载请注明出处!