破解国产生态困境,为AI芯片搭建编译之桥。
作者 |  GACS
9月14日~15日,2023全球AI芯片峰会(GACS 2023)在深圳南山圆满举行。在9月15日AI大算力芯片专场上,中科加禾创始人崔慧敏分享了主题为《编译技术从学术研究到产业落地》的演讲。
崔慧敏认为,大模型让国产AI芯片面临的生态困境更加凸显,编译器成为一大破解之道。简单来说,编译器是用于将机器语言翻译为人类语言的工具,在AI芯片产业链中,相当于AI芯片与AI用户之间的“桥梁”。
在GPU时代,编译器的功能比较简单,且很长时间里都是英伟达一家独大。随着AI时代到来,编译器的内涵和外延进一步扩大,涉及了图、张量、指令、图算融合等不同层面的编译。
大模型时代,编译器能够促进芯片性能提升、跨平台优化及安全性增强。
在性能方面,编译技术可以发掘更激进的融合策略,进一步全局对计算、同步进行重排优化,为模型推理侧带来2~3倍的性能提升
在跨平台方面,基于机器模型概念,崔慧敏团队研发的PPOpenCL编译器可支持一份代码跨平台优化,达到接近平台调优版本最优效果。
在安全方面,编译技术能以很低的开销来实现软件TEE(可信执行环境),增强安全性。
目前,中科加禾正聚焦搭建对不同芯片都适用的工具集,通过组件化服务助力AI芯片完成生态迁移等目标,并计划在2023年12月推出跨平台大模型推理引擎
以下为崔慧敏的演讲实录:
大家好!我是中科加禾的创始人崔慧敏,我们是一家年轻的公司,8月刚刚成立,今天我们会分享我们为什么从学术走向产业来。我本人在学术圈做了20多年,一直在做编译的事情,现在确实发现在生态里有非常多编译可以做的事情和机会,希望大家可以开放地跟我们讨论,我们的产品定义和产品发布都在很开放的讨论中。
今天我的报告主要分为三个部分:一是传统时代和现代时代的编译有什么不一样的地方;其次我会分享我们这么多年,特别是近几年,在AI编译方面的探索,我们来讨论AI编译在生态的潜力非常大,跟大家分享我们的一些工作;最后稍微跟大家介绍一下我们在产业界落地的想法。
01.
AI时代,编译的内涵和外延被不断扩大
对传统的编译器,大家都比较清楚了,它是作为从高级的编程语言到底层芯片的桥梁,把上层人类理解的语言翻译成机器理解的语言。在这个过程中,编译器的角色就是能翻译对、翻译好、翻译得特别高效,在传统的CPU上基本上是这样的期望。
在CPU时代,编译器相对比较稳定,基础架构以GCC、LLVM为代表,有两个基础设施,指令集(ISA)比较收敛,现在RISC-V的出现又对编译器提出了新的机会。对上层应用来说,语言不停地演进,但底层的编程语言相对比较稳定。
所以在CPU时代,编译器的存在感比较低,除了芯片厂商和两大社区以外,更多的人在日常工作和科研中不太会碰到编译器。对编译的科研人员来说,中间这个时间相对来说比较平淡。
到了GPU的时代,我从读博士的时候CUDA刚刚出现,我们去接收CUDA,再在上面做编译优化,也过了非常长的时间。它刚出来的时候不像现在这么火,当时推广都面临非常多的问题,做了很大的努力。到现在,CUDA已经形成了事实上的标准,这个标准包括了它的API、编程语言,底下有丰富的编译器。围绕它,我们整个AI的生态和市场主导都是这样被定义出来的。
从CPU到GPU,编译器干的活儿已经比原来多很多了,大家意识到编译这个事情很重要。所以在CUDA上层,大家做了很多编译的工作,包括怎么从C自动生成CUDA、怎么能自动地把编译成CUDA的代码编译得更好等工作。这个时候,编译的存在感就已经比CPU时代更高一些了。
现在到AI的时代,整个编译的内涵和外延都被不停地扩大。这张图是大家经常讲到的AI框架和底层编译器以及中间算子的关系。
用户用上层像PyTorch这样的一些接口来写网络的流图,PyTorch负责把数据流图用框架的引擎执行起来,变成一个个算子,跑在底下各个不同的芯片上。在这个过程中,上层的框架怎么更好地跟底下的算子适配、怎么能生成更快的底层算子、整个一套性能怎么发挥得很好,从产业界和学术界,大家提出了非常多的新的编译概念,比如图层面编译、张量层面编译、底层指令层面编译。尤其是芯片厂商,一直在做指令层面的编译,学术界和工业界也都在做。
近些年大家又在提图和算子的抽象,其实一定程度上影响了芯片算力的发挥,所以现在大家又在提图算融合编译、多层次编译等新的思想。
另一个方面,由不同的芯片所带来的生态问题,也是大家一直在强调的。我这里列了几个代表性的厂商,也是我们合作比较紧密的厂商的软件栈架构。我们在跟他们合作过程中发现整个生态的问题里,芯片厂商付出了非常大的人力,做了非常好的规划,但是从芯片厂商到用户中间还有比较大的Gap,我们公司就想来做中间的桥梁。
02.
发掘更激进的编译融合策略,
可将推理侧性能提升2~3倍
下面来分享我们之前在编译的科研方面所做的一些探索。(围绕)编译的潜力,我们从三个方面都进行了很多探索。
首先是它的老本行性能,它能帮助芯片把更好的算力发挥出来,可以让一个应用在上面跑得更快,即使在英伟达的平台上,大家的优化也一直往前提,在英伟达的基础软件上可以做得更好。
另外一方面,现在很多芯片厂商确实这么做的,软件和硬件之间的CO-DESIGN(协同设计),可以辅助结构设计人员来寻找更好的架构探索,这是我们之前做的一个学术研究的探索。
编译在安全性方面也有很大的潜力,我们也稍微分享一下我们之前的工作。
今天的重头还是性能提升问题,最早的时候大家在图方面、算子方面做了很多优化。因为我们是做编译的,我们有一个最朴素的出发点,就是当你看到的信息越全面,所能做的优化就越彻底,效果就越好
我们当时出发点就是能不能把整个流图信息和指令信息拿到一起来看,这样可以拿到更多优化的信息,看到更广的范围,做更多全局最优的决策。所以我们做了这样一个框架,能把图和张量信息,以及张量和指令的信息结合在一起,做图算融合的优化。这样就把原来图和张量层面,以及张量和指令层面的结合更紧密,我们真的做了一个图、张量、指令全层次的编译优化器
这样的好处是什么呢?之前大家在AI编译里面做了非常多算子融合的技术,这些融合会从基于规则、模式来做融合的决策,最终大家看到的范围就是融合做得好,性能会得到非常显著的提升
我们的思路是,融合这件事情在编译里面是非常经典的问题,只不过在编译里不是叫算子融合,它叫循环融合或者代码融合。而且在编译里面,我们可以通过编译分析去寻找最优的融合策略,这个最优,我们可能考虑了通信量、访存量、计算量,编译能全局地来做融合。
我们基于这样一个非常朴素的出发点,做出来了后来看起来是非常复杂的系统,做出来效果非常好。当时我们用的是Bert里的一个Transformer模块,经过优化以后,可以把整个下半部分框起来的灰色部分融成一个kernel,这里面包含了十几个算子,可以像(上图)右边这样全局排布它的计算,包括我们在重新调整它的同步模式,这样就可以实现非常好的效果。
这是我们给出来的例子,我们把九个算子融合到一个,可以取得右边这样的时间加速,不管是在执行时间上还是执行指令上,我们都得到了非常好的优化效果。
所以我们的一个结论就是:利用深度的编译技术,可以发掘更多更激进的融合策略,而且这个融合是跟模型没有关系的,不是说这个模型过来能做,换一个模型不能做了,因为我是全局分析它的最优,拿过来这个模型的流图以后去看到底怎么融合才能做到最好。当然进一步融合以后,我们就能做很多大家会做的事情,全局对计算、同步进行重排优化,所以我们跟TritonRT相比可以带来2~3倍的性能提升
我们两篇的论文,一个是今年的MLSys’ 23,已经发布了,一个是ASPLOS’ 24,是明年年初会发,我们正在准备最终的版本
效果还是非常明显的。从另一方面讲,编译最擅长的是能做一套优化,在不同平台都适用,这也是做编译的人一直都在做的。
就像大家熟悉的,在做CPU编译的时候,里面都有一个机器模型,把这个东西定义好以后,就不需要在每一个平台上重新做一遍指令调度,重新做一遍分配,一切都可以自动适配这个模型。
我们开始这个工作非常早,大概2010年初的时候就开始,那会还是OpenCL,大家抱有很大的希望,大家原来期望同一份代码在不同的平台上都能运行,这是OpenCL最初的初衷,但我们运行的时候发现,同一份代码在不同平台跑的时候,性能表现差距非常大。真的拿一份程序换到新的平台去跑,其实还需要重新来写一份,这就导致平台的可迁移性没有原来预想的那么好。
所以我们做了一个叫做PPOpenCL的编译器,最核心的是把机器模型抽象得更好,当时还是针对不同的CPU和GPU,还是英伟达和AMD的GPU以及国内一些超算平台,我们做了很好的抽象,把机器模型的概念加入到OpenCL的编译器以后,就能实现一份代码在不同平台上都能达到这个平台上手工调得最优的效果
这是我们机器模型的概念,所以我们在整个产品的规划里也把机器模型作为核心理念。机器模型的好处是可以做一个优化,在不同的平台上都可以适配得很好。
下面是我们的另外一个例子,这是针对大家都非常在乎的一些动态形状的矩阵层,我们知道不同形状的矩阵层,即使在英伟达的平台上,它的性能差异也会非常大。针对这个,我们用编译的思路把机器模型构建好,这里我会提到机器模型是编译优化的灵魂,把这个灵魂构建好以后,依据机器模型所做出的优化具备良好的跨平台性质
这张图展示的是我们在英伟达和华为昇腾上的效果。黑色的是我们的性能,越高越好,我们的性能可以比平台自己的库或者用户自己写的实现,在不同形状上都取得比较好的效果。这是编译在跨平台优化方面的特殊优势。
前面讲的是大家最关心的编译对性能的贡献都来自哪些,一个来自对全局的分析,能够分析做更深的优化,另外一个优化是能够跨平台的,这样适配不同机器的时候可以做得更好
我们也做了很多其他方面的探索,这个纯粹跟大家分享一下,我们在公司的运营可能也不太会做这方面。
一方面是软硬件协同设计,我们当时在CPU上能加一个比较宽的SIMD单元,到底加到多宽、加了以后到底是什么样的,在这中间做一个探索。
传统来说,每个SIMD单元是每个核私有的,这会出现众所周知的利用率低,没有那么高的并行度,或者计算很快结束了,另外一个核也用不起来。所以我们做了一套能弹性共享,不同的核能有软件及时感知CUDA行为的变化,到底能用多少向量部件,我就给你多少向量部件,你不用的,别人就可以来征用。
这个过程中,大家就想到可能会错,所以我们用编译器和硬件来协同,让它做到可以及时感知软件行为的变化,按需调整资源的使用量,这样就可以用得更好。编译和硬件的设置里可以有非常良好的交互,共同来定义指令级,定义一些外围编程的语言、跟操作系统怎么衔接等技术。
我们还做了另外一个非常有意思的工作,就是安全性增强工作。比较典型的是我们在CPU上有TEE的环境去保证数据安全,但是在GPU上没有这个机制,现在的GPU也在加这个机制了,但更早的GPU上没有TEE的环境,这样对于AI里面一个非常经典的场景,就是模型放上去,怎么保证模型不会被其他人拿走。
所以,我们通过软件来做了一个TEE的环境,当然涉及到对操作系统的改动,对底下Hypervisor的改动,对中间不同驱动的改动。
对编译来说,对核心的就是编译约定了一个程序的行为,我们定义了什么叫做程序的行为是良好的,定义好以后用编译写了一个验证器,证明当你给我的输入满足什么条件的时候,我的程序行为就能够保证它是良好的。
这样我们通过driver、VM,通过编译的验证器,一起来保证当我要启动一个kernel的时候,我会插入一个检查,去检查这个输入的条件跟原来预设的条件是不是匹配,如果匹配,就可以直接放心地执行,这个代码已经经过我的编译器和验证器一起验证过,就可以保证它的行为是良好的,这个良好就是它不会越界,不会有一些飞掉的控制流等特征。
所以编译器的能力,比原来经典的提升程序性能或者只是个“翻译”,还是要高很多。
03.
为芯片企业搭建桥梁,
降低英伟达生态迁移门槛
接下来分享一下我们在产业界想做的事情,我们的落地,包括我们为什么来做这件事情。
我们觉得中国的市场会是一个非常多的不同芯片共存的市场,从用户到这些芯片的厂商之间,大家都在一步步地构建自己的桥梁或者梯子,大家现在可能都构建得非常辛苦。
我们希望在大家构建梯子的过程中,能够有自己的位置,我们可以来搭建一些对不同芯片都适用的工具集。这个工具集现阶段更多地是帮助大家把英伟达的一些生态迁移到我们自己芯片的生态上
在这个过程中,我们做编译的思路是非常好地构建一个模块化的基础平台,这个基础平台可以理解成是梯子的木板,大家可以用我们提供的木板,来更方便地构建不管是用户侧还是芯片侧的工具。当然我们也希望后面的大桥能变得更像现代化的大桥,这样就接近英伟达的生态环境,这是我们的一个理想。
我们要做的就是多层次化、模块化优化的转换工具,我们的目标是构建中间一个大的平台,在这个平台上,向上会接入更多的应用,包括一些CUDA的程序。这就涉及到模块里要有非常多的前端。
我们这个模块中间会是一个多层次的表示,基于MLIR来做的,包括图层面的中间表示、张量层面的中间表示、指令层面的中间表示,我们把这些中间表示放在一起,不同的模块会使用到不同的中间表示,中间表示可以实现不同层次之间的转换,所以这个中间表示上我们做的这些分析和优化的模块,就是中间画出来黄色的模块。
另外,我们也会跟不同的芯片厂家做对接,把中间表示高效地映射到底下的芯片,这个映射基本是跟芯片方提供的底层编程语言或者一些编程接口做映射。我们希望我们这个平台上有非常多模块化的组件,通过这些组件可以根据用户需求,帮用户定制出他需要的工具。这是我们的目标。
我们现在规划了需要做的100多个模块,我举几个比较典型的来说。这里会提到我们为什么需要使用不同的多层次的中间表示,因为多层次的中间表示每一层上会携带不同的信息,不同的分析会用到不同的信息,我们把这些信息都放在里面,这是当时MLIR设计的初衷。
这样从左到右我们编译的粒度越来越细,中间表示的层次也越来越高,我们上层的应用就可以翻译到框架上面来,所以大家可以看到,我们这个还是面向推理侧做更多优化为主的平台。
给大家举几个比较典型的、我们已经做的模块,也欢迎大家跟我们联系,我们可以帮大家部署。
一个是细粒度依赖分析,我们拿CUDA的程序过来,可以对它里面的计算进行非常细粒度的分析,除了CUDA还包括TVM的输入,我们比TVM自己的依赖分析做出来的粒度更精细,这里面包括从上一个数据到下一个数据的使用和依赖,它是线程内的,还是一个block内的,还是全局的,这些信息过来后就有助于我们后面到底怎么做融合,怎么做切分,这是一个比较基础的依赖分析的模块,如果大家有兴趣可以直接拿来用。
中间是刚才讲Bert优化时提到的一点,把它融合到一起后,可以做一个代码的优化,指令层面的优化,这个我们现在主要在英伟达上做的是指令级同步优化,通过插入一定量的同步指令,可以让它的计算顺序进行重构。
我们也有基于多面体模型的访存行为分析与提取的模块,在深度学习里大家也会用多面体来做优化,它有它自己的优势,比如它可以找最优。我们这个就是帮大家把现在LVM上的IR,把每一个线程、每一个block、每一个kernel的内存访问信息总结出来,提供给第三方,然后你可以在这上面去接基于多面体的优化,或者我们自己也有基于多面体的优化部分,这样的话我们可以找一个全局最优的融合策略。
另外在层次化机器模型构建方面,左边是GPU机器模型的构建,右边是NPU机器模型的构建,可以看出是非常粗粒度的,再往下会对每一个芯片构建自己独立的机器模型,这样机器模型就可以用在上面的优化和转换工具里面去。
这是现在有一些模块可以拿出来给大家分享的。
04.
结语:首款大模型推理引擎今年见,
全局统一优化显著提升推理性能
我们近期的规划产品,第一个是跨平台大模型推理引擎
我们的优势是大模型推理跑得更快,我们跟几个合作伙伴测过,在英伟达的A卡系列上,70B模型比他们评测的好几家都要快。尤其针对A800这样国内“阉割版”的卡,我们也针对它裁减掉的这一部分资源做了全局统一优化,所以可以把这一部分通过软件补回来一些。
另外一个最重要的特征,就是我们的大模型推理工具是跨平台的。
我们对底下的芯片厂商可以接Transformer这样的接口,对上是一套,包括模型怎么划分、通讯怎么做等,这些都是用到机器模型技术,把所有机器相关的抽象出来,这样我们整个优化可以跨不同的芯片来做,可以跟云、跟服务器的厂商有一个很好的集成,可以对用户屏蔽底下的环境。
这个产品在今年12月份就会发布,到时候更多的技术细节和评测指标也会发布出来。
我们公司的定位是,国产芯片确实现在有一定生态困境,我们可以以编译为核心,来帮助芯片厂商破解这个困境,我们也很欢迎上下游所有生态伙伴的合作。谢谢大家!
以上是崔慧敏演讲内容的完整整理。
继续阅读
阅读原文