谈谈对深度学习编译技术的一些思考

早上在朋友圈看到一个手工优化抠得很极致的朋友(哪位同学中枪了^-^)就最近蓝色同学发表的文章有一些comment,觉得有些意思,手痒痒打算写篇文章谈谈自己的一点理解。

我主要会从三个方面展开,1.对TVM技术栈的认识,2.TVM与XLA的对比,3.编译优化与手工优化的对比探讨。

  1. 对TVM技术栈的一些认识

我个人的认识,技术方面TVM整个系统设计里非常值得学习的是针对机器学习性能优化场景核心复杂性的拆解抽象以及整体的工程架构的组织,并不是在于编译技术深度(这里指的是经典编译技术)

具体来说,TVM的整个设计里,把优化问题定义为一个 design space exploration的问题,部分问题交给了系统框架设计(schedule/computation抽象的隔离以及auto-tuning的流程设计),部分问题交给了机器(auto-tuning过程),部分问题交给了用户(computation/schedule的撰写),部分问题交给了existing compiler toolchain(nvcc/LLVM),同时还提供了优化探索的扩展接口(tensorize/intrinsics),这里没有再讨论诸如TVM里关于runtime以及前端表述语言的内容,因为我个人认为这部分属于重要且消耗精力未必小,但并非整个系统设计里最core的问题,所以不再展开。如果我们看上面提到的每个子问题,会发现都有相关的工作,比如schedule/computation的抽象基于Halide IR(最近社区也已经在推进新的TVM IR的工作)进行扩展,auto-tuning的流程其实是一个经典的machine learning的问题(YY一下,包装一把,设计成一个Kaggle task完全是可能的),对existing compiler toolchain的leverage更不用说了。把这些子问题进行有机的组合,找到每个部分发挥的合适scope,从而对整个优化空间探索过程进行复杂性拆解,以一种现实的方式在优化任务里获得decent的performance是整个TVM系统设计里我最喜欢也最觉得应该学习的东西。一言以辟之,这是一个在system design方面我觉得很出色的工作,而不是单点的技术深度。而优秀的system design,再配合上整体不错的工程实现(比如python/C++的双向互调用和相对轻巧清晰的工程架构设计),以及对用户接口的考虑(topi的提供,用户tutorial的提供以及辅助工具的提供),使得其被社区接纳也方便得多。再加上创始团队非常活跃,对于接受外部贡献也很open。这些因素加总在一起,让TVM具备了现在的气候。

2. TVM与其他编译优化技术的对比

我们对XLA有过比较多的关注,所以我会把XLA作为深度学习编译优化技术的另一个体系来和TVM做一个对比审视。

XLA的整体设计理念,跟TVM存在一些比较明显的区别(以开源部分为例):

1). XLA目前打击的是非计算密集算子,TVM打击的是计算密集算子,优化目标不一样,导致其作法存在明显的差异;

2). XLA强调自动化,所以给定计算图,XLA会直接codegen出可执行码(XLA社区最近也有一些妥协,开始考虑让用户在计算图上标记XLA的compilation scope,从而应对缓解XLA可用性的问题,虽然我觉得这种作法有一些倒退,但如果配合上系统层面的设计,倒不失为一个务实的解法),而TVM则把一部分优化工作转移到用户身上(TVM开发者),这个设计理念的区别,我认为其影响还是比较深远的,因为涉及到了整个优化任务复杂性的拆解,XLA想在系统层面完成更多工作,TVM则认为可以把部分工作offload给用户,我认为这跟两个工具假定的目标用户存在差异有关。XLA假定的用户更倾向于是普通算法建模用户,而TVM目前的用户更多是DL引擎开发者(随着时间推移,我的观察和判断TVM会更多推广让普通算法建模用户的使用,但目前的用户还更多是具备引擎系统经验的开发者),这两类用户对于使用接口的容忍度存在比较明显的差异;

3). XLA整体的工程系统设计更为考究,也更为厚重,但不容易拆解出来以模块化的方式为外部使用,而TVM的设计相对更为轻巧,也比较容易以松耦合的方式被外部使用(比如TVM离线gen的kernel被集成到其他DL引擎框架里);

4). XLA在图优化方面,会有更为复杂专注的实现逻辑,而TVM在图优化方面的实现则相对简单得多;

5). 为了支持更为复杂的图优化,加上自动codegen可执行码的理念,XLA的codegen部分实现逻辑是比较复杂的,相比较而言TVM的codegen部分其实比较朴素直接,如果用技术语言来描述一下的话,TVM的codegen部分,更像是一个纯粹的1-1 mapping性质的visitor实现,而XLA的codegen则除了对IR DAG遍历以外,涉及到针对不同计算pattern的inter-op的codegen逻辑拼接,以及数据存取index的推导计算和复用优化等等。当然TVM的codegen也可能针对不同硬件,加入一些inter-op的graph pattern的处理逻辑,但并不影响主体的界定;

6).TVM是一个经典的machine learning-based system,在完成schedule/computation抽象以外,整个优化空间探索,转换成了一个data-driven的机器学习优化问题,这是一个轻巧,但也一力降十会的作法。XLA在这方面,因为是纯system guy的工作,所以比较实在,是以纯系统的方式来解决优化问题。但是除了机器学习的方式以外,改成heuristics的方式来进行优化空间探索是不是也可能获得相近的效果呢?我觉得这还是一个open的question。不过把历史数据使用起来,辅助指导优化过程的探索寻优,这个原则我是buy in的。

而除了上述的区别以外,XLA和TVM我认为也存在一些可以相互借鉴的地方:

1). schedule/computation分离的作法,同样可以参考,实际上我们目前在XLA里已经加入了一些理念相似的尝试,让XLA在纯粹的codegen之外,多了一重抽象tuning的层次,也确实在优化效果上取得了一些收益;

2). auto-tvm对优化空间探索的作法,在XLA里可以参考,这也是machine learning for system思想的应用了;

3). XLA里关于复杂计算图优化的处理,在TVM里是可以借鉴的,特别是XLA里为了支持复杂图优化在codegen方面的设计,同样可以为TVM所借鉴,这里面有着一些蛮精巧的考虑,也是比较hard core的内容;

4). TVM在设计上的可扩展性的考虑,我觉得是XLA可以参考的,这样在系统本身不具备某些优化能力的情况下,可以通过用户的外部扩展来进行支持,比如tensorize和intrinsics的思想,但我们确实得承认,tensorize和intrinsics的思想是那种一但想到,相对并不难实现的作法(因为跟context的依赖较小),而XLA在打击非计算密集算子时偏重较复杂计算子图的特质导致加入外部扩展机制的支持很可能并不像加入tensorize/intrinsics支持那样straightforward;

5).TVM在用户使用便利性以及模块化方面的考虑,是XLA可以借鉴的,怎样在整个系统接入使用存在复杂性的同时,能够先把其中一部分功能以便于外界使用的方式拆解出来?比如在XLA里,针对特定的fusion pattern,直接将fuse后形成的结果gen出一个离线plan文件,甚至直接生成一个custom op以及对应的fusion pattern描述,便于外部集成调用。

3. 深度学习编译优化技术与手工优化的对比

在我看来,深度学习编译优化与手工优化并不是互斥关系,而是互补关系。在一个完整的系统里,应该既有深度学习编译优化,也有手工优化,让各自解决其适合解的问题。

1). 编译优化适合解决给定策略,涉及较多routine性质tedious work的问题。比如我们知道loop unrolling会可能带来性能收益是一个策略问题,但是按什么样的strides来unroll,是一个trial-and error的问题。以及我们都知道对于GEMM进行分tile计算可以提升计算访存比,这是一个策略问题,但是给定一款硬件,按什么尺寸,在什么维度上分tile则是一个trial-and-error的问题。这类问题,适合采取编译优化的手段来解,也往往是编译优化能够在生产效率上显著优于手工优化的地方;

2). 手工优化适合那种不容易精确形式化描述成一个清晰策略,带有一定非逻辑思维的直觉性质(由我们认识规律的能力水平决定)的问题,往往涉及到全局性质优化的问题具备这种性质。比如最近我们针对TensorCore在进行手工优化,会在访存pattern上进行精细的优化,以期最大可能将计算与访存overlap,就会发现涉及到精细的访存排布,至少基于目前我们对TVM schedule描述的理解,如果只是基于TVM显式提供的schedule,不去手写TVM IR,是不容易表达出来的。实际上TVM在设计上提供了Tensorize/intrinsics的接口,也是在一定程度上需要将手工优化的经验嵌入到优化流程里,但也并不是所有的手工优化都可以基于目前的tensorize/intrinsics机制来完成扩充的。这也算是计算机系统里经典的leaky abstraction的一个例子。

3).手工优化是可能向编译优化迁移的。比如通过扩展编译引擎的内核,来加入对应的支持。比如我们最近在TVM社区里针对NV GPU TensorCore提供了基于graph/IR pass/codegen模块改造的作法,能够做到用户完全无感,自动完成TensorCore kernel优化生成的效果,而社区的另一个相关工作作法则是需要显式提供TensorCore相关intrinsics的描述,将一部分工作offload到用户层。这算是一个手工优化,向编译优化层次迁移的示例。

4).总会有些优化在编译优化层面完成会存在事倍功半的效果,这类优化我们就应该考虑either是通过手工优化扩充,或是通过提供pre-defined library,甚至runtime强化的方式来进行协同优化,而不是什么优化都往编译层面压。反过来也一样。手工优化可以在极限场景下找到非常精细的性能优化空间,但是并不是所有的手工优化所探索的性能空间都复杂精细到编译优化不能支持的程度。找到不同技术适合的土壤就好。之前跟NV的同学沟通,他们针对TensorCore kernel的支持,考虑采取设计若干个小的recipe,recipe提供可定制的可能,再进行拼装组合来实现不同尺寸的GEMM kernel,这种作法,包括CUTLASS的设计思想,在我看来,都具备了一定的将手工优化的经验向编译优化层次转移的味道,只是程度不同而已。

最后,还是做个小广告,这里所讨论的很多内容都跟我们平时工作所关注讨论的内容有关,欢迎感兴趣的同学联系,这是一个非常有趣,也充满想像空间的领域,邮件地址muzhuo.yj@alibaba-inc.com。

发布于 2019-10-19

文章被以下专栏收录