ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

AI芯片软硬件协同设计:从数据搬运到编译器栈的实战经验

AI芯片软硬件协同设计:从数据搬运到编译器栈的实战经验

做AI芯片这行,圈内常讲一句话:“硬件决定天花板,软件决定地板。”过去几年我见过不少团队,硬件流片了才发现编译器跟不上,驱动写不出来,模型跑起来性能只有理论峰值的零头。反过来也有团队把软件栈堆得无比豪华,结果芯片架构在数据搬运上吃了大亏,再好的算子库也救不回来。AI芯片的软硬件设计,本质上是在算力、带宽、存储、编译、运行时之间寻找一个能落地的平衡点,它不是硬件工程师的独角戏,也不是软件工程师的闭门造车。

这几年国产AI芯片、边缘NPU、端侧推理加速器层出不穷,随便翻一个项目的技术文档,里面都会写“支持主流深度学习框架”“提供高效编译器”。但真做起来你就会发现,最难的往往不是某个算子怎么实现,而是从算法需求到芯片微架构再到工具链,整个链条能不能对得上。这篇文章不打算讲教科书式的理论,我想用我做加速器项目的实际经验,把这摊事拆开揉碎讲一遍:芯片怎么设计、软件栈怎么搭、软硬件怎么协同迭代,以及那些文档里不会写的坑。适合正在做AI芯片相关项目的人,也适合想转行做异构计算、编译器、体系结构的同学参考。

1. 先搞清楚AI芯片到底在解决什么问题

1.1 传统芯片为什么跑不动神经网络

很多人第一次接触AI芯片时会有一个困惑:CPU和GPU不是都能跑神经网络吗,为什么还要专门做一颗芯片?答案藏在两个词里:数据搬运和并行度。

神经网络的核心操作是卷积、矩阵乘法、归一化、激活,其中计算量最集中的是矩阵乘法和卷积。这类运算有一个共同特点:大量数据被反复使用。一个卷积层,输入特征图的一个像素会被多个输出通道的卷积核读取,权重也会被所有空间位置的输入复用。传统CPU为了灵活性,把数据和指令都存在片外内存里,每算一次都要从内存取指令、取数据,算完再写回。这个过程里真正用于计算的晶体管可能只有一小部分,大部分时间都耗在了“搬数据”上。

GPU比CPU聪明,它用大量线程掩盖延迟,带宽也比CPU高得多,但GPU毕竟是通用并行计算设备,它在处理不规则逻辑、动态分支时很强,遇到固定模式的矩阵乘也算得很快。问题是它的功耗和面积代价太高,对于推理场景来说,能效比远远不够。AI芯片的思路很直接:把神经网络里最频繁出现的计算模式固化成专用硬件单元,让数据在片上流动,尽量减少对片外内存的访问。方向对了,但这也带来了新的麻烦——硬件一旦固化,软件就要想尽办法去适配它。这就是AI芯片软硬件设计一体性的根源。

1.2 软硬件协同到底在协同什么

我见过最离谱的一个项目,硬件团队花了一年半时间做了一个性能峰值很高的加速器,流片回来后才发现编译器团队连指令集都还没定稿。后来逼不得已,软件团队在驱动层里绕了好多弯,牺牲了一半性能才算跑通。这事给我留下的教训是:AI芯片从第一天起,软硬件就是一起设计的,不是两个项目。

协同的核心可以拆成三个层面。第一是计算模式的协同,你要明确芯片主要跑哪些模型、哪些算子,卷积和矩阵乘的比例是多少,激活函数用哪种,这些直接决定硬件单元怎么排布。第二是数据流的协同,权重放哪、激活放哪、中间结果怎么缓存,这个决策既要看硬件上SRAM够不够,也要看编译器能不能把数据分块调度好。第三是编程接口的协同,软件开发者用一个什么样的抽象去看待芯片,是一堆寄存器,还是一种类似kernel的编程模型,会直接决定编译器的实现难度和驱动层的复杂度。

换句话说,硬件设计不能只考虑峰值算力,软件设计也不能只考虑框架兼容。真正决定一颗AI芯片好不好用的,是一条完整的链路:模型在框架里走完图优化,算子被编译器映射到硬件指令,指令在运行时被调度到计算单元,数据在DMA和片上缓存之间流转。这条链路上任何一环出现断层,最终都会表现为“项目延期”和“性能达不到预期”。

2. 硬件设计:从计算单元到数据流

2.1 MAC阵列为什么成为绝对主流

几乎所有AI芯片的计算核心都离不开MAC阵列,也就是乘累加单元阵列。一个MAC单元执行一次a*b+c的操作,对于定点数来说,这几乎就是AI芯片里最基础的肌肉记忆。

以最简单的卷积为例,一个输出通道的某一个位置的结果,是把输入特征图对应区域的所有通道、所有空间点,与卷积核权重做乘加后累加得到的。如果硬件里有一排MAC单元,它们可以同时计算不同的输出位置,或者同时处理不同的输入通道,这就实现了并行。把MAC单元排成二维阵列,就形成了我们常说的脉动阵列或运算阵列。设计阵列大小时,很多团队会直接参考目标模型的算子形状。比如一个边缘场景的神经网络,输入是224x224的图片,第一层卷积输出通道可能是64或128,那阵列宽度设计成64、128就比较自然,太小浪费带宽,太大面积又受不了。

MAC阵列的工作方式分很多流派。有的设计让权重固定在片上,输入数据依次流过阵列,这叫权重固定型数据流;有的设计让输入特征图固定不动,权重在整个计算过程中循环,这叫输入固定型数据流。名字听起来玄乎,本质上就是决定“到底谁在片上来回搬运,谁在片外慢慢取”。我在实际项目中通常会让编译器直接感知数据流类型,因为编译器在排调度顺序时,必须知道哪些数据是常驻的、哪些是流动的,才能决定分块大小。

2.2 存储层次和带宽,才是真正的性能瓶颈

如果说MAC阵列是AI芯片的心脏,那存储层次和互连就是血管。芯片不吃饭不行,数据喂不进去,计算单元再猛也是空转。

一个很简单的估算公式:峰值算力是P(单位TOPS),每次计算需要的操作数如果是2字节(比如INT8的一个乘加需要1字节激活和1字节权重),那么理论上每秒需要从某个地方读到2P字节的数据,才能让计算单元满负荷。假设一颗芯片算力是32TOPS,那在完全没有数据复用的情况下,需要64TB/s的带宽——这个数字比现在最先进的高带宽内存还夸张得多。所以芯片设计者必须用另一招:把数据缓存到离计算单元最近的SRAM里,让数据在片上被反复使用,片外只承担“分块”级别的搬运。

片上存储怎么分,是个大学问。很多NPU会把SRAM分成权重缓冲、输入缓冲和输出缓冲三块,分别对应计算过程中三种不同的数据流。为了让软件效率最大化,存储大小和编译器能做的分块策略直接相关。比如权重缓冲只有256KB,那编译器一次就只能装载一定大小的权重块,算子维度就要按这个尺寸切分。如果硬件团队和编译器团队没有商量好SRAM容量和端口数量,后面几乎一定会出现“编译器生成的代码在硬件上跑出大量空泡”的悲剧。

另外,用DMA(直接内存访问)做异步搬运也是标准操作。计算单元在算上一块数据的时候,DMA可以把下一块数据从DDR搬到片上SRAM。这种双缓冲机制,能大幅隐藏访存延迟。但这里有一个软件必须配合的点:编译器要把搬运指令和计算指令交错安排好,否则硬件设计了这个流水能力,软件不调度,照样白搭。

2.3 硬件设计流程和工具选型

硬件设计本身有一套成熟流程:规格定义、微架构设计、RTL编码、仿真验证、FPGA原型、逻辑综合、布局布线、流片。RTL编码阶段,国内多数团队现在还在用SystemVerilog,但我个人非常推荐小规模项目尝试Chisel或SpinalHDL这类基于高级语言描述硬件的方案。它们最大的优势是可以参数化地生成运算阵列,改一个并行度参数,整个RTL配套的流水逻辑和接口也跟着变。我们之前做过一个模拟项目X,用Chisel描述了一个可配置的脉动阵列,从32x32改成64x64的算力配置,只需要改几个参数,省掉了大量复制粘贴RTL的时间。

但要注意,高级语言生成RTL并不能替代验证。仿真验证仍然要花掉硬件团队60%以上的时间。AI芯片验证最常用的方法是用一个参考模型(通常是用Python或C写的算法模型)和RTL仿真做对齐。喂同样的数据,跑同样的算子,比对每次输出的结果。这个参考模型不只是硬件验证要用,软件模拟器也要用,后面编译器生成的指令跑出来的结果也要跟它比对。所以这个参考模型一定要维护好,它是整个软硬件协同的中枢。

3. 软件栈设计:让硬件真正跑起来

3.1 编程接口与运行时:软件如何“看到”芯片

硬件做出来了,软件团队要开始面对一个更现实的问题:写程序的开发者用什么样的接口来操作这块芯片。上一代很多加速器的做法是直接给寄存器,软件工程师写一堆底层驱动去设置寄存器开启搬运、触发计算。这种做法对固件开发是没问题的,但对上层应用开发者来说太痛苦。现在主流做法是给一层抽象,类似“命令队列”或者“Kernel任务”的接口。

我印象很深的一次经历是,我们设计了一个简单的软件抽象:应用层把一个算子的输入输出和参数打包成一个任务描述,扔进一个异步执行队列,硬件跑完任务后通过中断通知完成。就这么一个看似简单的设计,把驱动的工作量立刻分成两块。同步版本比较简单,提交任务后原地等待,但计算单元空转严重;异步版本复杂一些,需要维护多个任务在队列里并行执行,DMA搬运和计算单元的重叠度也更高。我们最终在接口上直接限制了“驱动层必须支持至少四个任务并发”,因为测试下来,小于这个并发数,DMA的搬运延迟根本藏不住。

运行时还有一个绕不开的点:动态Shape。很多端侧模型输入分辨率不定,或者用到了动态batch。硬件设计如果只支持固定Shape,那你软件上就得偷偷用padding补到某个对齐的规格,浪费计算资源。我们一开始把硬件设计成了支持范围检查、不支持动态地址跳跃的模式,后来发现编译器要额外生成很多边界判断代码,性能损失很大。如果你做边缘推理芯片,我建议从第一天起就想清楚是否要支持动态Shape,如果要,地址计算单元和相关寄存器得提前留好。

3.2 编译器与算子映射:从框架落到指令

软件栈里最难啃的是编译器。AI编译器做的事情通俗讲三步:把深度学习框架(比如PyTorch的模型)解析成计算图,在图上做优化(算子融合、常量折叠、布局转换),然后把优化后的算子映射到芯片的指令集上。

算子融合是性价比最高的一步。卷积后面通常跟着偏置加、激活函数,甚至可以跟归一化层合并。卷积如果单独计算,要把输出写回DDR,再读出来做归一化,再写回去做激活。在芯片上这些“中间结果”完全可以留在SRAM里传递。硬件设计时,通常会把卷积单元、激活运算单元设计成可以直接拼接的流水。编译器要做的,就是识别出这种融合模式,生成一条能一气呵成的指令序列。我们之前测试过一个典型的MobileNet类模型,做了算子融合之后,DDR访问量能减少差不多一半,这不是小数目。

再往下是算子映射。编译器要把一个卷积算子切分成多个子任务,依次搬运输入块、加载权重、执行计算、写回输出。这个切分的逻辑叫“Tiling”。Tiling策略依赖于硬件参数:SRAM多大、MAC阵列多大、权重缓冲多大、DMA搬运粒度多大。所以编译器要有一个硬件描述文件,把芯片的存储容量、对齐要求、指令编码都写清楚。这里特别要强调的是,硬件团队给编译器的参数必须和实际RTL一致。我们有一次改了SRAM容量上限但没有同步更新编译器配置,结果生成的代码在仿真环境里频频出现存储越界,排查了整整三天,最后才发现是参数没同步。

3.3 低精度量化:性能翻倍的隐形利器

芯片能不能用好,还有很大程度取决于是否支持低精度计算。INT8已经把推理性能翻过一倍的,INT4和INT8混用更是近年来的热门方向。硬件支持低精度的代价,不是换个计算单元那么简单,而是整个数据通路——包括缓存位宽、DMA搬运宽度、累加器精度——都要跟着调整。最常见的做法是累加器用INT32,防止多次乘加累加时溢出,然后写回时再根据量化参数缩放到INT8。

软件这边,量化算法的选择也很讲究。最简单的方案是预先算好每层权重和激活的缩放因子,推理时直接按固定点计算。但这个方案对激活分布的鲁棒性差,需要在校准集上统计真实分布,一般是跑几百张有代表性的图片,看每层激活的最大值、最小值或分布的百分位点。更激进的做法是训练感知量化,在训练阶段就模拟低精度的舍入误差,让模型权重去适应量化噪声。很多团队只把量化当成“训练完以后加一个转换脚本”,这是不够的。深度优化的模型通常要用量化感知训练才能保证精度不掉。

4. 软硬件协同设计的关键流程

4.1 规格定义:先从目标模型反推需求

做AI芯片最忌讳的就是直接从“我要做一个128TOPS的芯片”开始。正确的做法是从目标模型和应用场景反推。比如说你的芯片要跑一个车载场景的检测网络,你先把这套网络的FLOPs算出来,再把片外存储带宽算出来,再看在目标帧率下每个算子允许消耗多少时间。

这里我给一个简化的换算方法。假设一个模型总共需要50GMACs,目标帧率是30FPS,那每秒要完成1.5TMACs。如果芯片的MAC阵列利用率为60%,那理论峰值至少要做到1.5/0.6=2.5TMACs,换算成INT8大概5TOPS。这个数字决定了MAC阵列大小、工作频率和对应的SRAM容量范围。片外带宽则根据每个MAC操作平均要搬多少字节来估。有了初版算力、带宽、存储三个参数,才能进入微架构设计。否则直接拍脑袋定规格,很容易出现算力足够但带宽不够,或者存储小到编译器完全没发挥空间的情况。

我们内部习惯把这套推理过程固化成一张“性能预算表”,每一层算子列出计算量、访存量、允许时间、片上缓存需求,然后对着这张表去调整硬件参数。这张表也是软件团队做性能建模的基础,两个团队必须对同一组数字负责。

4.2 性能建模与模拟器:先模拟再流片

软件团队最怕的是硬件还没回来,没有地方可以联调。解决办法是做一个周期级或近似周期级的软件模拟器。模拟器的本质是把硬件的计算单元、DMA搬运、缓冲区、命令队列抽象出来,按照同样的时序规则计算每个操作要多少个cycle。芯片架构师用这个模拟器验证峰值性能,编译器开发者用这个模拟器看指令调度的效果,驱动团队甚至可以把它当“虚拟硬件”来跑完整的软件栈。

我在模拟器上踩过一个坑:模拟器性能模型做得太乐观,没有考虑不同DMA通道间的仲裁开销,结果实际流片后性能比模拟器预测值低了20%。这个教训后来变成了一个规定:模拟器里必须包含资源的竞争模型,至少要把总线带宽上限、SRAM端口冲突、中断开销这些因素加进去。宁可模拟器悲观一点,也比流片后被现实打脸强。

做性能建模时还有一个很实用的指标:Roof-line模型。它把芯片性能画成两条边的屋顶,斜边代表计算密集区域的“算术强度”和带宽决定性能,平台部分代表算力峰值。把目标模型的算子放到这张图上,一眼就能看出每个算子是“计算受限”还是“访存受限”。对应地,优化手段完全不同:计算受限的算子要优化MAC阵列利用率和指令调度,访存受限的算子要优化数据复用、调整分块大小、减少DDR访问。这个模型听上去简单,但几乎是我做每颗AI芯片最常用的决策工具。

4.3 硬件和软件如何在流片前共同推进

既然硬件还没回来,软件团队能不能提前干活?当然能。关键是搭建一个“带着走”的软硬件联调环境。我在参与某跨平台系统设计时,就把环境分成了三层。最底层是RTL仿真,速度慢但准确,适合验证关键指令的时序。中间层是模拟器,速度较快,适合跑完整算子库的测试。最顶层是FPGA原型,速度最快,适合把整个编译器和运行时框架跑起来,甚至可以跑真实模型。

FPGA原型并不是把ASIC设计原封不动烧进去就行。由于FPGA频率低、片上存储容量小,很多ASIC设计跑不起来或者跑得太慢。我们通常会把FPGA原型里的大阵列缩小、存储深度降级,但保持指令接口和数据流向不变。这样编译器和驱动看到的还是一个“缩小的芯片”,能把整套软件流程打通。等ASIC回来了,只需要调整性能配置文件,大多数代码可以直接复用。这个方法让我避免了“硬件回来才开始写软件”的死亡节奏。

4.4 典型四步迭代流程

我总结了一个比较通用的软硬件协同迭代流程,适用于从零设计一颗AI芯片或者给已有芯片做定义阶段优化。

第一步是识别热点算子。从目标应用的代表模型出发,统计每个算子占总运行时间的比例,通常卷积和矩阵乘会占大头,激活函数、归一化等占小头,但也不能忽略,因为访存可能让它们意外地贵。第二步是为热点算子建立参考实现。用Python或者C写一个朴素版本,保证数值正确,这是后续所有验证的基准。第三步是把这些算子映射到硬件方案上,计算存储和带宽需求,然后把结果反馈给微架构设计。第四步是拿性能模型跑一轮仿真,看看瓶颈在哪里,回头调整硬件参数或软件调度策略。这个循环通常要走好几轮,而且越早发现问题代价越低。

注意,这四步里最容易忽视的是第三步和第四步之间的反馈。硬件架构师倾向于认为瓶颈可以通过编译器优化解决,编译器工程师倾向于认为换一个微架构参数就能解决。最后往往是架构师和编译器工程师坐在一起,对着性能模型一条一条抠数据流,才能真正解决问题。

5. 常见问题与排错技巧实录

5.1 算力很高,跑模型就是慢

这个现象背后十有八九是数据搬运问题。我遇到过好几次,一个算子的理论计算时间只有50微秒,但实际跑了500微秒,原因就是权重和输入数据存在DDR里,DMA搬运太碎,或者没有用双缓冲。排查时先在硬件计数器里看DDR空闲和计算单元空闲的比例。如果计算单元大部分时间在等待DMA,优先检查编译器生成的搬运指令和计算指令之间的间隔。另一个常见原因是分块太小,每搬运一小块就触发一次长延迟。解决办法是增大分块尺寸,尽量让数据在片上多次复用。有些时候还得检查命令队列是否串行执行,连续几个任务在等前一个完全结束,这样DMA和计算的重叠就完全没了。

5.2 编译器生成的代码结果正确但性能很差

这种情况通常是编译器没有按硬件微架构的flow来组织指令。比如有些MAC阵列对矩阵的维度有对齐要求,编译器没有把小的矩阵问题padding到硬件对齐尺寸,导致阵列很多行空转。或者循环嵌套的顺序不匹配数据布局,导致每次访问SRAM都发生bank冲突,白白增加延迟。这种问题的排查方式是把一条编译后的指令序列和模拟器的周期计数对齐,看哪一段指令流水停滞,停滞的原因是什么。我们在项目中曾给编译器加了一个“微架构感知的调度后端”,专门根据指令依赖关系和SRAM端口占用情况进行指令重排,性能提升非常明显。

5.3 模拟器性能很好,FPGA上跑不快

FPGA原型和ASIC之间有几个天然差异。第一是频率低,如果你的ASIC目标频率是1GHz,FPGA可能只能跑200MHz,计算能力直接差5倍。第二是片上SRAM的分布和数量不同,FPGA的BRAM配置未必能完整映射你的存储层次。第三是外部DDR接口的行为不同,FPGA开发板的DDR带宽不一定和最终芯片一致。所以用FPGA原型验证性能之前,先要把性能期望值按频率缩放过,而不是直接拿FPGA的帧率除以模型的计算量。如果缩放后仍然明显低于预期,再查数据搬运是否阻塞。

5.4 快速排查速查表

现象可能原因排查方法
计算单元利用率低DMA搬运不及时,同步执行查看流水线和命令队列中的等待事件
跑小模型正常,大模型崩溃分块策略导致SRAM溢出检查编译器内存规划和SRAM容量一致性
数值错误累加器精度不足或量化参数错误比对参考实现,打印每层输出统计
性能模型和实测漂移大总线仲裁和访问冲突未建模给模拟器加入资源竞争模型
指令集改动后驱动失效寄存器描述手工维护不一致用寄存器描述文件自动生成头文件和驱动代码

6. 关于AI芯片软硬件设计,我个人的几点体会

回头再看,我觉得AI芯片的软硬件设计,核心不是选多先进的工艺,也不是把MAC阵列堆得有多大,而是把“模型—算子—指令—数据流”这条链路的每一环都打通。硬件团队和软件团队在项目第一周就要坐在一起,共同定指令集、存储层次、编程接口,然后一起维护一套性能模型和参考实现,这样才不会在流片前后互相抱怨。最后分享一个小技巧:不管项目多紧张,一定保留一份从训练好的模型到芯片上推理结果的端到端测试,哪怕模型很小(比如只用两层卷积的小网络)。只要这条路能每轮迭代都不堵车,整个软硬件系统就永远有一个可以依赖的“生命线”。在芯片还在设计阶段时,这条生命线就是软件模拟器加参考模型;芯片回来后,它就是你最值得信任的验收标尺。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进