本文来自知乎mackler,作者是清华计算机博士,其中是他对于DSA和AI软件栈问题的思考,文章很有洞察力,写得非常棒。

开个新系列,聊一聊专用架构和AI软件栈的问题。这个题目比较大,涉及面非常广,和之前的文章一样,我一般观点都比较激进,对于现有方法的缺点一般是毫不掩饰,各位看官请酌情食用,另外我也会抛出一些全新的解决思路。

这篇文章主要先聊硬件,AI软件栈的复杂性和根本性难题其实主要来源于硬件。如果硬件只是GPU,尤其是是Volta架构之前的GPU,其实AI软件栈根本不用这么复杂,所以我们先来讲讲硬件架构,也是我的老本行。

体系结构领域大的创新其实已经停滞了很多年了,虽然说这么多年硬件的提升一半是靠摩尔定律,一半是靠体系结构革新,作为一个做arch的,我原先也非常相信arch在其中的作用。但实际情况是,arch的“创新”其实很大程度依赖于工艺的演进。arch领域最近几年最常见的论调是摩尔定律要完蛋了,arch即将成为硬件性能的主要推动力,似乎未来架构要在工艺不变的情况下纯靠架构演进原地起飞。实际并不是,脱离了工艺的演进,架构带来收益非常困难!

架构的收益其实更多是工艺演进时,新约束下新tradeoff带来的超额收益。以前我也不明白这一点,因为理论上讲,架构只是电路逻辑,什么工艺下都可以做。但实际上工艺可以打开设计的空间,因为相同的面积功耗限制下,新工艺可以放更多的逻辑,我们就可以把原有架构的各个部分超级加倍,比如cache变大,发射数增大。但随之而来原先架构不突出的问题可能就变得明显,此时就需要调整架构,比如cache太大了就可以再加一级,利用率可以更高,从而获得超额的提升(相比简单根据工艺scale)。所以本质上讲,不同工艺带来的不同约束下,架构设计是有不同tradeoff选择的,所谓架构收益更多是这种新tradeoff带来的超额提升。但如果约束保持不变,其实最佳的tradeoff很快就收敛了,后面想继续靠arch压榨出更多性能就非常困难了。

实际上,靠持续天才般的创新来维持的商业模式根本没法持久,arch领域这种大的创新已经停滞很多年了,提出OoO机制的Sohi大佬之前就曾经表达过把2000年之后的ISCA论文全烧了对世界没啥影响,当然他是做CPU的乱序执行的,从单核CPU的视角确实如此,这种评价虽然非常激进,但也侧面说明了这方面的问题。

arch领域需要的不是天才式的创新,而是能持续数十年稳定提升算力的方法论。一次天才式的创新其实是很难持续的,硬件这么多年来指数级的持续性能提升靠的是类似于前面说的新工艺下重新tradeoff带来稳定超额收益的方法论,虽然说起来很low,但这才是硬件持续演进的动力,当然这种方法论也是有寿命的。以一个不太严谨的粗粒度划分,算力提升的方法论可以粗略分解成三个阶段:80年代以前是超标量、80年代到15年左右是并行,15年往后是专用架构。当然不是说所有产品都沿着这个年代划分,大家都在多条路径上挣扎,只是一个粗略的划分,像我前面举得例子其实就很贴近超标量的方法论,通过tradeoff来获得超额的收益,这种方法论在当今的CPU核的设计上同样在沿用,只是越来越吃力,像Intel之前著名的Tick-Tock商业模式也正是基于这样一套方法论。

在上个世纪80年代,单核提升确实遇到了不小的问题,激进的架构改动在当时的工艺下根本做不动,原有架构在当时的工艺下也基本发挥到极致了。所以就开始诉诸并行和向量化,当然这对软件的冲击是非常大的,软硬件接口其实是博弈了很多年,一方面要给硬件持续数十年“无脑”演进的空间,另一方面又不能给可编程性带来太大负担。并行做起来之后多核就成为了主流,比单核提升轻松,当然把代码难写(相比串行代码)的锅丢给了软件,慢慢也积累了大量并行代码的生态。此时多核的路就相对好走多了,也就没必要死扣单核性能了,毕竟要花更多的力气。而且在并行阶段还孕育出了SIMT和GPGPU这种为并行而生的编程模型与架构,软件也就朝着并行化的方向一路狂奔了,而CPU的单核架构已经很多年没有大的革新了。

其实从cuda和GPU的演进就可以看出,相比Intel的Tick-Tock挤牙膏,NVidia的性能持续演进之路要比Intel顺畅得多,因为SIMT编程模型白嫖了程序员提供的程序并行度最好的维度,GPU的演进只需要每一代产品用更多的晶体管不断提高流处理器的数量和并发度就可以让cuda生态的新老代码一起飞升。

而到了15年左右,并行的持续演进之路又遇到了一定的困难。AI的兴起对算力的需求已经不是翻倍的增长,而是数量级的增长。即使壕如NV激进地使用最新工艺,也很难靠并行满足算力地需求,各种新兴的AI芯片也开始用定制化电路来满足算力的需求。专用架构开始在各种芯片上出现苗头,AI芯片基本就是专用架构native的,NV的GPU也引入了TensorCore,Arm的SME、Intel的AMX都旨在用专用架构来获取性能提升。

arch领域新的方法论是DSA。图灵奖得主David Patterson大佬在17年的ISCA上提出体系结构的黄金时代,讲摩尔定律要完蛋了,以后算力进步得靠我们做arch的了,我们应该做DSA,后面又讲了很多RISC-V和芯片敏捷开发相关的内容。这些都是想构筑一套DSA的全新方法论,为未来数十年持续的算力提升铺路。很多人挺困惑DSA和RISC-V还有芯片敏捷开发是啥关系,其实如果是做高性能CPU,根本不需要敏捷开发,因为大的改动很少,基本架构非常稳定。敏捷开发其实是为了开发DSA指令。而RISC-V集齐开源实现则更多是为了减少重复造轮子,因为DSA要真正用起来,有很多脏活累活需要CPU协助,而这部分正是RISC-V希望补齐的。这也是David Patterson大佬为我们构筑的DSA演进路线的蓝图,他希望的是用RISC-V+DSA来持续提升算力,以后每年“无脑”算力提升的方法论就是去分析主流应用热点,然后设计DSA指令,嵌入到RISC-V生态,从而获得性能提升,确实不需要什么大的创新。不过这种方法论的整个故事构建是需要非常多大的创新叠加到一起才能拼成的,我这个系列其实要讨论的正是这部分。

DSA这件事其实也不是David Patternson大佬提出之后大家才开始投入的,早在之前就已经又很多人在实践这件事了,17年的时候早已经是AI加速器满天飞了。只不过大部分人是以一种从通用到专用一下子获得大额算力提升的角度来看待这件事的,而David Patterson考虑的则是专用到专用的迭代,如何创造一个持续提升算力的方法论。这个事情其实决定了AI加速器到底是历史上昙花一现的架构还是一个能长期留存下来的架构。因为大家可以想一下,我们开坑做第一款AI芯片的时候,我们的baseline是GPU,单纯从算力角度还是比较容易一下子超出同一代的GPU的。但一个成熟的商业芯片势必要按照一定的周期推陈出新,第二代、第三代芯片的baseline可就是上一代AI芯片了。第一代芯片你可以用一个systolic array专门处理矩阵乘,当然可以吊打通过SIMD+SMT实现SIMT的GPU芯片,第二代芯片你如何进一步用一种新的电路结构吊打上一代的专用电路呢?针对一个特定功能的专用电路,最佳数据通路其实一两代产品肯定收敛了,如果没收敛,这个芯片的架构师可以去面壁罚站了。

收敛了再往后怎么办?CPU/GPU这一类芯片可是持续演进了几十代的。从现有的实践来看,基本落入了两类:要么继续DSA走到黑,然后继续忍受我这个系列后面要讨论的DSA在软件上的各种缺点,直到能有方案解决这些问题,从此引领下一波arch演进路线;要么回头走老的路线,继续吃SIMT的收益,于是AI芯片成为昙花一现的产物,最终AI属于GPU。从目前的情况来看,大部分做AI加速器的团队都还没有能力去解决第一条路的困难,经过前几年的AI芯片投资热潮之后,现在投资也冷静下来了,改成投做GPGPU方案的AI芯片了,于是今年国内一下子冒出了一堆做GPGPU的创业公司。

不过走GPGPU路线,天花板也就在那儿,当然大家也不会傻傻的只吃SIMT收益,到头来肯定还是SIMT为主,但SIMT很难打过NV,所以DSA收益也有限的吃一点,DSA的吃法不同就可以在DSA上打性能,外宣上就可以吊打NV了。当然NV也不傻,当年一大波公司开始做AI芯片啃DSA的时候,NV适时加入TensorCore,也吞下一口DSA的收益,从此NV的GPU也变成SIMT+DSA混吃的模式,所以在这条路线上竞争某种意义上讲还是挺难的。

DSA路线最大的困难在软件。只要想吃DSA的收益,一定会带来软件上的巨大困难。吃的越多困难越大,吃得少一点也会引入很大的困难,除非只吃一次,当然只吃一次收益也就只有一次。即使像NV这样拥有强健cuda生态的公司,引入TensorCore一样需要承受这样的痛苦。cuda只能覆盖SIMT的软件问题,覆盖不了TensorCore,同样需要能解决DSA路线困境的软件方案。而以DSA为主的AI芯片,则更是会感受到软件上的巨大痛苦。

DSA指令的性能来源于固化了粗粒度的一整块计算,这些接口是需要通过ISA直接暴露给软件的,软件的主要痛苦则是来源于ISA的高度碎片化和不稳定性。不稳定性是软件最大的障碍,给一款芯片做适配,即使工程量再大,逼近芯片都固定了,总归是可以适配下来的。但不稳定性意味着适配的工程量无法被下一代芯片复用,那此时的适配工程量就是无穷无尽的了。应用是没法跟着芯片的迭代周期进行迭代的,生态就很难建立起来,于是很多人就诉诸编译器。

编译器不是万能的。一定程度上很多架构师都喜欢画大饼用一个编译器衔接上下各种碎片化的前端和后端,似乎这样整个架构图就一切都舒爽了,啥都能自由支持,而且自动优化。但实际上编译器的能力是非常有限的,编译器的自动化能力很大程度依赖于IR的完备性,也就是图灵祖师爷给我们奠定的超强基础。而编译器的优化能力嘛,且不说AI编译器这种才几年探索的新兴事物,HPC领域卷了几十年了,要是编译器优化能力可以依赖那还有blas库啥事呢?而这两点其实在AI这一领域更是崩坏,所以这种架构图虽然吹牛一时爽,最后编译器其实是有多少人工就有多少自动化。

Patterson一定程度上尝试构建DSA的arch演进之路,但最大的障碍其实是软件。著名编译器llvm的创立者chris在21年ASPLOS上讨论了编译器的黄金时代,与patterson一唱一和。考虑到他本人也加入sifive投身到RISC-V事业中,他俩的黄金时代试图构建了一个关于DSA的完整故事:

  • 复用riscv做控制流处理大多数情况,用chisel快速搭建自己的dsa指令

  • 复用mlir的core dialect处理大多数情况,用mlir快速搭建编译器通路把新dsa指令用起来。

我对于整个故事在芯片这边的逻辑是非常支持的,但编译器这边的故事是深表怀疑的,而编译器这边的故事恰恰又是DSA路线能否走通的关键。对于从零搭建一个全新的dsa硬件而言,chris的这套是走的通的,而且会比现有的做法便捷一些,但这一套没法复制LLVM的成功,也无法支撑一个长期生态的发展,更没法撑起所谓的黄金时代。

其实这个道理和前面说的做AI芯片是一样的,从零做一款AI芯片,DSA的收益可以吃一大口,软件、编译器都可以完全针对这套新的ISA去定制。但要做一个持续多年的芯片,硬件上需要有持续可以啃的收益,软件上要做到老代码新硬件自动性能提升,而DSA的持续性能提升来自硬件接口的不稳定性,而这种不稳定性又会导致老代码适配新硬件非常困难,难以形成长期“无脑”的稳定路线,毕竟软件和硬件一样,不能指望每次新硬件兼容性出现灾难都指望一个“大创新”来解决,这在商业上是不可靠的。因此软硬件之间在DSA的路线上是存在一定结构性矛盾的。

不解决这个根本性的矛盾,AI芯片未来的发展只能一直主吃SIMT,但又不断尝试吃一些DSA的收益然后被软件扎得满嘴刺,反复横跳。

这一篇硬件方面聊的比较多,软件方面基本还是蜻蜓点水,下一篇可以从主要从AI软件栈角度聊一聊。

上一回聊了很多硬件方面的内容,总的来说希望给大家总结一下AI芯片目前面临的困境和选择,摆在面前的一条路是DSA,一条路是SIMT,这是一个非常困难的选择。

SIMT是NVidia已经趟过的老路,成熟可靠,但后劲不足,而且选择这条路意味着AI芯片最终收敛为GPU,GPU的所有天花板也能一眼望到头。在这条路上打败NV是非常困难的,在SIMT路线方面我算是个十足的N吹,NV在这方面的恐怖实力和把控力还是相当了得的。当然这篇不会太多聊SIMT,这方面我相信有很多非常详尽的分析。

DSA是目前看来arch未来的演进之路,但仍然处于探路的过程中,也依然有机会成为CPU和GPU之外的第三种独立的架构存在于历史长河中。即使是目前开始做GPGPU的创业公司,也必然需要通过DSA这一部分来增强产品竞争力,包括NV的TensorCore也算是走的DSA这条路,其实从这个角度来说NV在DSA方面也有一定布局,当然NV这部分布局还没像SIMT那样成熟,我们也还有很大的机会。

所以DSA or SIMT,我的选择当然是DSA!一定要把DSA啃下来才有希望打败NV。当然还是明确一下,吃DSA不代表不吃SIMT,更多是指通过DSA打开未来的演进空间,多打开一层空间才能真正超越。

而DSA能否成功,成败取决于软件!关键在于能不能设计出一个软硬件解耦合,硬件可以在一定范围内自由吃DSA收益的顶层设计。我一个原本做arch的人,用脚投票去做编译器也正是因为发现DSA能否走通其实不取决于硬件,而是软件。AI的软件栈也是这种顶层设计演化尝试和迭代的最前沿阵地,目前在软件上的尝试也有很多不同的路径,基本可以简单分类为两类。

一类是希望找一个稳定的中间层IR,这个IR可以是硬件层面的virtual ISA,可以是library层面的一层稳定的API,也可以是一个算子库,甚至是深度学习框架。希望在剧烈变化的前端和后端之间找到一个稳定不变的层来解耦合,在变化中求稳的办法只有依靠完备性,像CPU基础指令集一样稳定,但CPU基础指令集的完备性是来源于图灵祖师爷带给我们的,但这完备性仅限于细粒度IR。但细粒度IR是没法向下对接DSA硬件的,因为DSA硬件的指令是粗粒度的,否则拿不到DSA的收益。

这方面的尝试一直都有,深度学习框架早期从caffe那种粗粒度的layer模式转向相对细粒度的operator转变就是一种尝试,不过现在主流的框架都有上千个算子了,还在不断演进,我们也永远可以设计出框架写不出来的情况。后面出现的ONNX和NNVM算是构建一套标准化算子的尝试,NNVM算是浅尝辄止,可能发现了这方面的问题,于是转向另一条路,开坑TVM搞编译器去了,而ONNX坚持到今天,算子集合短短几年已经更新了十几代了,而且至今没能做到两个主流框架pytorch和tensorflow的全覆盖,可以说是一个标准的“试图用一个大一统的框架统一现在的n个框架,最后导致出现了n+1个框架”的典范了。MLIR一定程度上放弃了这种完备性的追求,转而躺平让大家随意在不同层之间打洞写conversion,既然躺平了其实也就放弃了去解决这方面的问题,硬件更新之后,相应的编译器通路也需要跟着一起更新,这茫茫多的conversion其实和适配DSA硬件的工作量是一个规模的。当然MLIR也引入了core dialect来尝试增加复用,不过core dialect当中相对high-level的也一样不具备完备性,low-level一些的完备性好很多,但又是大家为DSA硬件打洞需要跳过的层次,所以虽然chris在编译器黄金时代的演讲中讲了很多解决碎片化的问题,尤其提到了要从O(前端x后端)变成O(前端+后端)的工程量,但MLIR实际解决的是编译器基础设施的模块化和复用,而不是IR的复用,没有改变O(前端x后端),只是让这个大O的常数变小了一些。碎片化还是存在,只是换了个形式,变成MLIR dialect的碎片化。

既然说到MILR,就顺便进一步展开来讲讲MLIR。

MLIR本身更接近一个库的定位而非生态!其实很多关于MLIR的争论都来源于定位的差异,MLIR是一个还不错的编译器库,具备丰富的编译器基础组件,可以减少重复造轮子。很多支持MLIR的人主要的论点也是在实践上使用MLIR构建自己的编译通路多么多么方便。那么库和生态主要区别是什么呢?简单概括就是使用者的收益主要来源于其他使用者贡献的代码还是这个项目本身的代码。像LLVM就是一个标准的生态,新的编程语言可以复用其他后端用户提供的后端适配代码,新的硬件可以复用所有前端使用者贡献的不同编程语言代码。对于生态而言,用户越多每个用户都能越爽,即使LLVM项目本身不再更新了,新的用户加入仍然可以使所有参与者受益,因此生态也需要大家共同来建设,这种建设不是给LLVM本身提交pr,而是围绕LLVM建自己的编程语言或者适配后端硬件。而库则不同,库的收益主要来源于库这个项目本身的代码提供的功能,而MLIR套用生态的判断标准显然不是一个合格的生态设计,从目前公开的一些实践来看,大多数参与者贡献的dialect和conversion都是在自己的闭环内,甚至tf团队自己的iree也是如此,大多算子基本bypass了所有core dialect。那么此时新的使用者其实很难复用其他人的闭环dialect,大多还是在前端框架的dialect、后端的dialect之间搭建自己的闭环,失去了这种不同参与者的互动,也就失去了大家一起来contribute建设这个生态的意义,毕竟你的大多数dialect往往也只有你自己用,无法带来用的人越多大家都爽的优点。大家建立自己闭环不是不配合建生态,毕竟high-level的dialect没有完备性,大家的需求也不一样,导致了大家都有对自己更有利的dialect设计和取舍。

那么再回到DSA路线的问题,一个新的前端/后端想要接入这个生态,需要做的是给前端可能出现的各种类型的算子打造通路到达后端硬件,MLIR灵活的设计允许不同dialect混用,所以可以不用考虑某一个特定的dialect能不能吃下所有的前端并传递给后端(当然可能也不存在这样的dialect)但如果前端有很多呢?作为一个后端厂商,我肯定希望所有前端用户我都能支持,毕竟NV的cuda没有说我们只支持tensorflow不支持pytorch或者说只支持主流的框架,小众的框架不支持对吧。

对于各种新技术,个人认为还是要透过现象看本质,AI编译器和芯片的发展虽然日新月异,但更多的还是在一些已有的路径上反复横跳,或者把问题换了个形式呈现出来,实际进展还是异常缓慢的。

另一类是尝试做编译器,这里说的编译器特指算子编译器,负责自动算子生成,关于图编译器以及AI软件栈的分层暂且不讨论,因为我们现在关注的还是DSA的可编程性问题,至关重要的点就是如何用好DSA指令而不影响可编程性。算子编译器方面目前主要的探索无非两类,一类是Halide/TVM这种尝试解耦计算和调度的,编译算法层面没有太多新的东西;一类是基于多面体编译技术的进行调度优化的。当然还有大量两者都涉及的众多编译器方案,在这里我不会展开具体的编译器进行分析,还是透过这些编译器来看后面这两类具体的技术。

Halide这一类核心在于解耦计算和调度,然后定义了一系列调度原语把调度优化的空间给建模出来,然后计算写一遍,调度可以手工tune或者通过一些算法来搜索,这样似乎可以一个算子部署到不同硬件上,包括dsa指令也是建模到调度空间里了。但是我们回过头来想一想计算和调度解耦之后,可以复用的其实还是计算,但计算本身的工程量很小,聊胜于无,真正的大头还是调度,调度还是极其碎片化的。调度不仅和硬件强相关,也和计算的形式强相关,因此写调度的时候还是要针对特定计算和特定后端硬件来写调度,仍然是一个O(前端*后端)的碎片化状况。

TVM在实际实践中主要的作用还是进一步降低了算子优化的门槛,某种程度上也算是推动了整个行业的发展。现在大大小小的AI编译器团队非常多,其实真正搞编译器算法的人非常少也非常难招,而写算子的门槛则相对低很多,招人相对容易得多,很多团队围绕TVM进行工具链开发建设,可以招大量写算子的人来覆盖碎片化的工程量,从产业实践上来讲也算是一条不错的路径。MLIR其实也是类似的路径,不过是从写调度换成写lower通路的各种pass,都比较容易堆人力,本质上还是尝试用人力战胜复杂性。

但降低门槛只能降低起步阶段的工程量!复杂性的可怕之处随着这些路径演进的深入会逐渐展现出来,用TVM和MLIR更多的是降低了起步阶段的门槛和工程量,让大家能快速写一写性能还可以的算子,在对接的业务越来越多,业务灵活性需求越来越高,追求算子覆盖率和性能调优的阶段,这个工程量和直接用native工具写kernel已经没有什么本质区别了。

当然TVM这一类工作是给了空间做autoTVM的,但是调度空间其实非常碎片化,一般也只是一些启发式的通用搜索算法,算法可以设计得相对具有统一性和简洁性,但简洁的后果必然就是效率不行,很难搜到高效的解。对调度空间做比较细致的建模则需要直面搜索空间的碎片化,工作量一样惊人,比如我们可以设计一个专门针对tensor core在某一类pattern下的自动调优算法,但是我们需要多少这种算法才能覆盖主要场景呢?而且这一类工作门槛又比较高,很难招到足够的人去面对这种碎片化。

既然说到这种自动化的代码生成,就和多面体编译一起说了,多面体编译算是整个编译领域形式化的一个集大成者,90年代有过一段高歌猛进的阶段,但是局限性依然很多,比如tiling的自动求解其实就做不到,而tiling影响的则是体系结构最核心的瓶颈,也就是缓存的利用,像Pluto算法其实也只是是辅助tiling,所以多面体编译后面慢慢就没落了。现在虽然又活跃起来了,其实依然是新瓶装旧酒。

多面体编译是一种建模方法,不是一个具体的算法,针对具体问题需要设计具体的算法去搜索和求解,当然过去几十年的研究已经积累了一些算法,都已经封装到一些成熟的多面体编译库里,比如ISL。很多AI编译器其实也只是在调这些库来完成一部分搜索调优工作。

不过,多面体编译现有的大多数算法还是针对通用硬件的,比如我们在CPU和GPU这一类硬件上生成更高效的调度。当硬件架构转向DSA时,多面体编译的适用性也会大幅下滑。当然这些倒是可以通过开发新算法来解决的,比如上面说的,我们可以设计一个专门针对tensor core在某一类pattern下的自动调优算法,设计新的多面体编译算法其实更是难上加难,但方法上是可行的。

不过回想一下我们最开始讨论的DSA核心困难是什么,是DSA本身的不稳定性和碎片化,这种编译算法一样会面临碎片化的问题,当硬件开发新的DSA指令时,我们需要精通多面体编译建模的专家设计出相应的编译算法,才能充分把新的后端指令用好,这种做法显然比TVM那一类堆人写算子的做法更加难scale,编译器专家堆都堆不起来,人太少了,这是一条需要编译算法持续创新的商业路径,是非常不稳定也不可靠的。

要想基于编译器设计一个针对DSA的可持续商业化路径,同时减少工程复杂性,我们还是要透过这些纷繁的具体技术选型去看整个栈从上到下有哪些问题,以及整个编译器领域总体的方法论以及能力边界。这篇在细节的技术上讲的有点多了,也主要给大家讨论一下编译器在解决DSA持续演进之路上面临的困境,篇幅原因,下一篇我们再透过这些细节去讨论编译器本身的方法论和这些困境的本质以及DSA时代编译器可以做些什么。

上篇讲了很多DSA软件栈所面临各种现实的问题,这篇主要想从一个更宏观的角度讨论一下编译器的本质方法论和摆在我们面前的选择。

不过在展开本篇内容前,还是要先明确一下我这个系列里讲的DSA到底指什么,因为专用架构、包括异构这些词本身确实太宽泛了,不加以限定容易造成各种误解,如果放到几年前说异构编程很多人第一反应还是GPU编程吧。这个系列讲的DSA都是指狭义的DSA指令,也就是用硬件电路直接实现一些粗粒度计算。从软件视角看,就是提供了一系列高性能的粗粒度原子函数。这种方式背后对应的是一条潜在的硬件持续演进路线:通过不断将当前主流应用的热点计算特征提取出来,硬化成硬件指令来持续提升性能,具体可以看这个系列的第一篇。而我们这里要讨论的编译器也是如何克服这条路径的困难,让这条潜在的路线成为现实。

实际上具体的芯片往往不会只遵循一条演进路径,所以我们不能说某一款芯片是DSA,另一款芯片是SIMT,无论是CPU也好,GPU也好,还是眼花缭乱的AI芯片,基本都是不同路径的大杂烩,差别只是不同路径的比例而已。实际上目前硬件的演进路径从大面上分也就三条:超标量(乱序、预取、分支预测、…)、并行(向量、SIMT、多核、…)、专用(也就是DSA)。像前一阵Hotchips很热闹,看起来玩出了花,实际上更多是16年左右那波ai芯片在专用这条路上因为软件没太走通,硬件妥协了一波,回到并行的路上借着wafer-scale和chiplet打开的空间上扑腾,这种软硬件的妥协未来注定还会持续好几轮,关于近期这波硬件抛出来的方案,我们以后再聊……(这个系列果然好大的坑)我们这个系列关注的更多是专用指令这条路径如何走通。

回到DSA编译器的问题上,我想先讨论一个非常灵魂的问题,编译器的本职工作是什么?说到做编译器,很多人第一反应都是性能优化。性能优化确实很重要,但并不是编译器的本职工作,更多是锦上添花的作用。之所以大家第一反应很多都是性能优化,是因为传统编译器在解决本职工作上已经做得非常棒,没有什么搞头了。但这件事在AI编译器上还不太成立。

编译器的本职工作是编译!听起来很废话,但这个确实是DSA编译器目前一个巨大的误区。编译器首先要解决的问题是如何将硬件无关的用户代码自动转换成硬件指令。大家想想编译器最早出现的时候是为了解决什么问题,是解决大家不想写汇编指令的问题,是为了把与具体硬件无关的高级编程语言转换成具体硬件代码,而优化则是在编译的问题解决得足够好的情况下锦上添花的功能。

实际上,即使编译器优化做的不够好(实际上也是),只要能用高级编程语言,大多数程序员也可以去做一定的代码优化工作,但如果让广大程序员抛弃编译器开始写汇编,哪怕有一个特别牛逼的汇编优化器能把程序员写的汇编进行调优,但不支持高级语言,软件行业也直接炸了。同样,软件行业在编程语言上早已玩出花了,但是如果让硬件行业直接在芯片上支持各种高级语言的特性,硬件也一样炸锅。所以,编译器在整个软件栈中占据一席的主要原因必然是因为其本职工作——解耦合!

解耦合其实很多人都会讲,一方面是让程序员不用直面硬件的细节,同样让硬件不用直面软件的复杂性。

最最容易被忽视的是让编译器不用直面软硬件组合的复杂性!实际上如果只是考虑不让程序员直面硬件细节,目前的AI软件栈并没有放弃这一点,至少目前大多数深度学习开发者面对的都是深度学习平台提供的计算图抽象,基本不用考虑后端硬件的细节。硬件这边也相对比较难受的,需要做很多妥协来提高可编程性,一定程度上束缚了沿着DSA这条路演进的空间。但最难受的还是DSA编译器自己,目前编译器需要直面软硬件组合的复杂性。

传统编译器其实没有这样的问题,软件可以基于语言的机制玩得非常花,比如C++的模板元编程,可以玩出乱七八糟的黑魔法,而编译器并不需要直面乱七八糟的黑魔法,只需要将模板的机制做好即可,硬件这边同样如此,编译器将代码转换到类汇编的表示之后,剩下的工作量已经非常小了,也给硬件自己争取竞争力保留了一定的自由度和空间。

现在的DSA编译器主流的策略其实都是做非常厚重的编译器软件栈,内部包含了很多层级的转换和优化,工作量极大。其实这样做,一定程度上费力不讨好,因为一方面优化的压力全部压在了编译器上,另一方面,过于厚重的软件栈也屏蔽了其他层在性能优化出力的空间,其实这不是一个好的协作模式。

软件的分层核心不光要解耦,更需要合理的分赃!刚刚说的还是解耦的问题,其实一个合理的生态一方面要让每个层面的人只需要面对自己的任务,还有一件重要的事情是让每个层面的人都能在自己这一层出力来改善整个系统。

怎么理解这件事呢?还是用传统软件栈来理解这件事,当一个程序性能不好的时候,程序员可以在算法和代码实现层面做优化,从而改善性能;编译器同样可以改善编译优化算法来改善性能;芯片也可以针对应用在指令上的主要瓶颈进行相应的数据通路的优化。

进一步讲,编译器虽然控制了代码生成,也做了很多优化,但程序员不需要知道编译器内部代码生成细节,但仍然可以控制程序的很多优化机会。所以编译器做太厚,一方面承担了很多痛苦,但另一方面其实也抢夺了其他层的优化机会。这样的缺点在于,当前端用户发现某些程序性能很差的时候,除了骂编译器做得太烂、催促编译器解决之外,有没有什么办法能在当下通过改代码来控制程序的执行调度流程,从而改善性能,而不必等编译器更新。同样,当硬件觉得编译器/上层框架在特定情况下优化的很烂的时候,可以做一些改变来把硬件性能发挥出来。

同样,上面的这些优化机会一样要在解耦的基础上,程序员能做的优化大多是算法层面或者一些通用的编程模型层面(例如并行模型等),具体硬件相关的优化可能就知识包袱太重了,硬件能做的优化也是如此,要避免直接面对软件复杂性。

此外,目前编译器、框架、算子库等软件栈概念在深度学习领域都融合到一起了,做深度学习编译器的人,可能是主要写算子库的,可能是主要做循环优化算法的,可能是主要做计算图优化的,也可能是主要做分布式调度的。大家聚在一起,关注的点各不相同,虽然说的都是编译优化,想的则是完全不同的东西。而一个好的软件栈是要让各个领域的人划分好边界,在自己关注的优化空间内自由发挥,各司其职去促进整个系统进步,而不是盲目扩张和内卷。算子编译器和算子库一定程度上就是一种竞争性的关系,算子编译器做得越强,覆盖的长尾算子就越多,算子库做得越全面,未覆盖的算子就越少,因为两者竞争的是同一块优化空间,结果就是算子编译器团队也会写很多手动优化塞进编译器里,算子库团队也会尝试自动化很多优化策略来扩充算子库。实际上这一块优化空间,编译器抢起来很难,算子库覆盖起来也很累。本该解耦合共同协作的各方很容易变成相互抢活的内卷。

铺垫了这么多,接下来进入本篇的重点,编译器到底擅长什么?我也想借此从编译器的方法论角度来讨论一下dsa相关的软件栈目前为啥这么痛苦。

编译器的手段其实也就几种

最常见的是term rewrite,匹配程序IR中的某种特征,并按照预设的模板进行替换,从而达到对程序进行变形的目的。有了这种机制我们就可以通过增加term rewrite rule来不断覆盖各种手工变换代码的经验。这种方法扩展性很好,遇到各种corner case很容易添加rule来workaround。

当然这种方式在变换空间较大的时候会变得很低效,可能需要无穷无尽的rule才能覆盖。这时候我们就需要一些形式化的方法来建模和变换程序,其中polyhedral model算是其中的集大成者,定义了一个很大的affine loop调度变换空间。

之前曾经有人把编译的方法论类比成滑雪,从高层次抽象不断lower到低层次抽象,然后可以在每一个层次内部做一些等价变换进行优化,在每个的层级进行合适的优化,最后到底层变成硬件指令。从完备性的角度来看,越底层的表示越完备,越上层的完备性越差,因为上层表示的基本要素一定是底层表示的某种粗粒度的结构,比如for循环和if-else分支,到底层都是jump指令,jump按照某种结构去组织就形成了for循环和if-else分支,但在底层要重建上层这种结构的语义就会比较困难,因为lower下来再进行一定的等价变换进行优化,上层的结构就很难重新发掘出来了。所以大家一直在讲lowering很容易,rising或者叫lifting很难。

山顶半山腰这些高层级IR的完备性比较差,程序的变换形式也相对离散,一般会用term rewrite的方式覆盖,比如各种计算图编译器,会包含各种图融合的算法。而山底细粒度IR完备性好,而且变换空间大,尤其是loop变换,比较容易设计各种形式化的方法。

dsa的问题在于,指令的粒度较大,并不是山脚那种细粒度图灵完备的标量指令。那有人可能就想,那是不是在半山腰,我们从山顶lower到半山腰就好了。但问题这种粗粒度表达通常是缺乏完备性的,两组不同的粗粒度算子集可能不一定能轻松实现相互转换。所以这个问题并不是软件在山顶,硬件指令在半山腰,而是硬件指令在另一座山的半山腰。要从软件打通到硬件指令,往往需要在两座山之间架桥,这种架桥往往只能基于term rewrite的方式来覆盖,需要很大的工程量来实现这一点,而且换一下硬件指令,这座桥还需要重新调整和翻修。

这种工程量使得也有一批人热衷于建schedule空间然后在空间里搞最优化求解或启发式搜索的方法,这个空间可以是polyhedral model的调度变换空间,也可以是halide这种手工定义的一套变换集合。这一套方法对于传统的细粒度编译还是行之有效的,但在dsa这里会面临一定问题。

如何把dsa指令纳入到调度搜索空间里?

一种办法对具体的dsa指令进行建模,并建模到搜索空间里,这种方式本质上还是和dsa指令强耦合,工作量可能未必比用term rewrite覆盖来得小。而且建模这种门槛高的方式可能都很难找到足够多的人来完成这部分工作量。

另一种办法是做一种机制把dsa指令以一直通用的方式引入到搜索空间中,例如tvm的tensorize原语。但采用这种方式你会发现,要找到合法的tensorize本身已经非常困难了。现在大多数编译寻优搜索算法还是在一个相对稠密的合法解空间中,尝试找到一直性能更好的最优解。而dsa指令的引入会把调度空间中的合法解变得非常稀疏,找到一个合法解本身就已经非常困难了,现有的编译寻优算法其实是不太好使的。

这种状况其实也是路径依赖导致的,过去编译器关注的都是通用处理器上的代码优化问题,合法调度在解空间里非常稠密,只要解决依赖性分析的问题,搜索空间里基本都是合法调度,关键在于调度后程序的性能。

而体系结构核心瓶颈是访存,因此调度的核心在于优化访存。访存一方面取决于代码执行顺序,一方面取决于片上缓存如何分配,这两个问题也是过去编译优化两个最关键的问题,即auto-schedule和auto-tiling。可以说,在过去这么多年HPC领域的编译技术发展下,auto-schedule的问题被polyhedral model部分解决,而auto-tiling的问题形式化后往往会变得极其非线性,所以一直以来都是老大难问题。

而DSA指令的引入,其实将搜索空间中的合法解变得稀疏化了!此时找到合法解本身就是一件很难的事情,合法解集合随着DSA指令的粒度变大而变小。算是在新的历史阶段引入的新问题,即auto-tensorize的问题。新问题其实也带来了前面auto-schedule和auto-tiling问题的解决新思路。毕竟如果DSA指令是一个非常大颗粒的函数,用户代码要变成能用这样一个指令执行的可行变换是非常少的,甚至可能是唯一的。此时优化问题可能已经完全退化了,毕竟只有唯一的可行解,找到即是胜利,而且是专用指令,性能优化实际上是由DSA指令的物理实现完成的。

换个角度想,原来的搜索空间中,其实大多数都是硬件上不怎么高效的解。现在DSA指令相当于直接把这个问题换了个形式,DSA指令实际上是标记了那些形式的变换是可以被硬件高效执行的,这部分其实一直都是很稀疏的,而编译器现在只要找到这部分解,就可以一定程度上拿到DSA指令提供的性能,当然编译器在可行解的集合里进一步求最优化。优化问题实际上是变成了软硬件协同完成了,DSA指令优化一部分(manual-tiling和manual-schedule),编译器优化一部分(auto-tiling和auto-schedule)。而且这种协同是以编译器和后端解耦的方式完成的,后端粒度越大,DSA指令优化的部分占比越高,编译器优化的压力就越小,当然覆盖的应用程序也就越少,而后端粒度越小,DSA指令优化的部分占比就越少,编译器优化的压力就越大,当然覆盖的应用程序就越多。所以实际上一个具体的后端可以构建层次化的DSA指令,包含不同粒度的DSA指令,从而在通用性和高性能方面达到兼顾。此外,广义上来讲,DSA指令不仅可以包括芯片上实际的物理指令,也可以是高性能库提供的算子,硬件指令和手写库可以共同构建一套层次化的后端模板集合。

这种软硬件的解耦与协同的前提,是编译器可以解决auto-tensorize的问题,在调度空间种自动寻找极其稀疏的可行解。

当然,寻找可行解的开放问题是非常困难的,这个问题朝着最开放的方式去形式化,其实就是SAT的问题:即软件写一段代码,硬件指令的功能也用一段代码描述,要判断这个硬件指令和软件的这段代码功能是否等价,基本是标准的SAT问题。当然更一般的情况是要判断软件代码中是否包含这条指令的等价功能片段,更加变态。反过来讲,这个问题也可以很简单,如果我们要求软件代码和硬件功能描述保持字面上完全一致,这个问题做一个简单的IR匹配就可以实现,但软件编写的灵活性就没了,这个灵活性其实就在于各种代码的等价变换。

所以,从另一个角度来看,DSA编译器的关键是要提供这种代码等价变换的灵活性,通过一系列通用算法逐渐覆盖各种等价变换的域。将手工/硬件优化的宝贵算子应用到更加广泛的应用中去,而后端指令/算子则通过覆盖更多优化场景来分担编译器性能优化的压力。这篇抛出了比较多的观点,也是希望抛砖引玉,引起大家更多的思考和讨论。篇幅原因,下一篇我们再具体展开这种新的生态可能性探讨以及DSA这条硬件持续演进之路的走法。