AI芯片软硬件协同设计：从架构到工具链的完整落地

做AI芯片这几年最常被同行和面试的人问的一句话是你们做NPU到底是先写软件还是先画硬件每次听到这个问题我都想笑——正确答案不是先后而是从第一版规格书开始软件团队和硬件团队就坐在同一张桌子前把每一层接口、每一次调度、每一个性能计数器都当成共同产品来设计。这个系列前两篇聊了AI芯片的整体概念和基础架构今天这篇是系列第三篇我打算把重心放到软硬件协同设计的完整落地上架构怎么定、存储数据流怎么设计、编译器和运行时怎么配合、验证工具链怎么搭以及我实打实踩过的一些坑。看这篇内容的人我默认你已经知道AI芯片是干什么的但不一定亲手做过工具链或者微架构。所以后面每一节我都会尽量把为什么这么做讲透不管是做架构、写编译器、写算子库还是做FPGA验证的同学都能找到能直接拿去用的思路。这篇不会讲某个具体芯片的代号和参数而是讲方法、流程和取舍这些都是换一颗芯片、换一个团队也能复用的经验。1. 项目整体设计与软硬件协同思路1.1 为什么AI芯片必须软硬件一起设计很多人容易把AI芯片想象成把很多乘法器堆上去配上大带宽内存就完事。实际完全不是这样。你在纸上画一个计算阵列峰值算力可以标得很高但真正跑模型的时候能不能跑到峰值的十分之一都难说。瓶颈往往不在ALU而在数据搬运、算子调度、指令下发、甚至软件里某个不起眼的同步等待。所以我一直跟团队强调AI芯片本质上是一个软件可定义的硬件平台软件栈和微架构是同一个产品的两个面拆开设计就是给自己挖坑。举个例子同样一个矩阵乘GEMM硬件理论峰值是100TOPS如果软件把数据切块切得不好或者DMA搬运和计算没有重叠实际跑FP16可能只有30到40TOPS。这种差距靠硬件改版很难救但靠软件把数据复用和流水调度理顺经常能直接翻倍。反过来如果硬件不提供足够的控制寄存器、同步原语和性能计数器软件再有想法也施展不开。所以协同设计的第一步不是画电路而是把软硬件契约定清楚。这个契约一般分三层第一层是指令集或任务描述符软件往硬件提交什么第二层是内存模型数据放在哪、怎么对齐、谁负责搬第三层是同步和错误处理硬件怎么通知软件任务干完了。三层里任何一层设计得不顺手后面编译器、运行时都会变得很难写。1.2 架构设计前先回答一份问题清单我每次启动一颗AI芯片的设计第一步不是列微架构特性而是拉着算法、编译器、硬件、软件几个角色一起过一份问题清单。清单大概是这样的目标负载是什么是推理还是训练以卷积为主还是Transformer为主batch大小多少sequence长度多少。精度体系怎么选FP32、FP16、BF16、INT8、INT4哪几档是必选中间累加精度用多宽。算力目标是多少端到端延迟要求多少吞吐要求多少这决定了多核架构的规模。存储怎么规划片上有多少SRAM需要外挂多大的LPDDR/HBM带宽目标是什么。面积和功耗预算是多少这决定了你能塞多少个MAC单元、多少块SRAM。可编程性边界在哪是只支持固定算子还是要能跑任意计算图。这些问题没有标准答案但回答的过程会直接影响架构走向。比如负载主要是Transformer那矩阵乘和attention相关的核函数就要重点优化卷积可能退居其次如果是端侧小模型存储带宽没那么紧张反而要关注启动开销和功耗如果要做训练中间结果的写回和梯度拷出会占掉大量带宽数据流策略就得重新考虑。我见过很多团队一上来就抄成熟产品的框图结果发现自己的编译器后端根本没好物可调或者算法同学提的模型根本塞不进片上存储。所以先回答清单再谈架构。清单里每一条都需要有真实数据支撑比如拿三到五个有代表性的模型跑一遍profile统计出算子耗时占比、张量形状分布、访存流量这些数据才是架构决策的基础。1.3 软硬件接口的契约设计软硬件协同听起来很玄落地其实就落在契约上。我习惯把接口分成两类一类是控制面一类是数据面。控制面回答软件怎么告诉硬件干什么数据面回答数据怎么在内存和计算单元之间流动。控制面最常见的做法是寄存器组加内存映射。硬件把一组控制寄存器映射到一段地址空间软件往里写配置值比如算子的类型、输入地址、输出地址、tile大小、同步方式。这里有一个很容易被低估的设计任务描述符。相比让软件逐个写几百个寄存器更好的做法是在内存里维护一个任务队列每个队列元素是一个描述结构体里面包含了完成一次算子执行所需的全部参数。软件写好描述符之后往一个门铃寄存器写一个值硬件看到门铃就去内存捞描述符自动执行。这个机制能大幅降低指令下发开销多核场景下尤其重要。数据面要设计的是buffer管理规则。哪些内存区域由软件管理哪些由硬件的DMA直接操作地址对齐要求是什么是否支持非连续二维访问。很多时候软件栈卡壳不是因为计算指令不对而是因为DMA描述符里那个二维起点的stride算错了导致搬进来的数据是斜的。所以硬件设计阶段就要把这些规则写得让软件同学能轻松理解并且提供配套的地址计算库而不是丢一句仔细看spec。1.4 用PPA和TCO来衡量设计好坏衡量一颗AI芯片好不好不能只看数据手册上的峰值。工程师更关心的是在真实负载上的有效算力、实际功耗和开发成本。硬件设计团队习惯看PPAPerformance、Power、Area但我建议软件和系统团队还要加一个TCO思维也就是从设计、流片、工具链开发到维护的全周期成本。我用一个简单的表格说明软硬件协同里PPA怎么被重新解读指标传统硬件视角软硬件协同视角Performance峰值TOPS、频率端到端模型耗时、有效算力利用率Power静态功耗、动态功耗跑真实模型时的平均功耗、能耗比Area逻辑面积、SRAM面积编译器能否用更小资源满足需求开发效率硬件验证周期工具链开发周期、算子覆盖速度同样一个乘加阵列面积多10%如果能让编译器少写几千行特殊case这笔账其实是划算的。反过来为了省一点SRAM导致运行时频繁从外存搬数据跑起来功耗翻倍那才是得不偿失。所以我一直建议团队在设计评审时禁止只报峰值指标必须同时报一组真实模型上的端到端数据否则讨论毫无意义。2. 硬件微架构拆解算力、存储与数据流2.1 计算核心怎么选脉动阵列、SIMD还是可重构AI芯片的计算核心没有唯一正确答案但在实际项目中大家选型时其实是在算力密度、灵活性和编译复杂度之间做权衡。我帮你把这几个方案的特点摆出来。脉动阵列是目前最流行的高密度矩阵计算方案。它让数据在阵列里像流水线一样逐周期传递每个PE只做乘加数据复用率极高非常适合GEMM和卷积。缺点是灵活性差对稀疏计算支持不友好遇到形状很怪的张量也会出现PE空转。所以脉动阵列通常不单干旁边会配一组SIMD向量单元专门处理激活、归一化、softmax、逐元素运算这类非矩阵操作。向量单元写起来更像CPU编译器好生成代码硬件也容易验证。可重构阵列听起来很美可以针对不同算子动态改变数据通路但实际上综合工具非常难做后端的布局布线约束也复杂团队没有三年以上的工具链积累我不建议新人碰。我接触过的多数自研AI芯片最终都采用了脉动阵列向量单元通用标量核的组合标量核负责任务调度和复杂控制矩阵计算走阵列杂活走向量单元。这个组合清晰、好验证、好调优也是很多量产芯片的共同选择。2.2 存储层次是成败关键不是计算阵列说句会被架构师们点头的话AI芯片设计真正决定成败的是存储层次和数据流不是那堆乘加器。从能耗角度看一次32位浮点乘加消耗的能量比从DRAM里取一个数低两个数量级。所以软硬件协同的核心目标是让数据在片上SRAM里尽可能多地被复用把访存流量压到最低。具体到微架构片上存储要考虑总容量、逻辑分块、bank数和读写端口。容量决定了一个tile能切多大bank数决定了多个PE同时读写时会不会撞车。比如一个核的SRAM分成16个bankDMA搬运和计算单元同时访问分配不好就会互相拖慢。架构设计阶段就要配合编译器模拟访问模式避免做成看起来容量很大实际有效带宽只有一半的尴尬局面。拿GEMM举例假设片上SRAM可用2MB要计算一个MNK4096的INT8矩阵乘。三个矩阵全放在外存的话总数据量大约是48MB。如果不做分块数据搬运量会严重拖垮执行时间。采用128x128x128的分块策略A、B、C三块分别是128x128字节也就是16KB16KB16KB一共48KB完全能塞进片上SRAM。这个分块下单次块的计算量是2x128x128x128约420万次操作数据总量是48KB计算访存比接近85 FLOP/Byte。这个数值越高说明数据复用越好计算阵列就越不容易被内存带宽卡死。很多性能问题看这个数字就能猜出大概。2.3 DMA、片上互联与任务协同多核AI芯片的瓶颈往往不是计算单元本身而是片上的数据通路和协同机制。核与核之间怎么传数据核与DMA控制器怎么配合任务队列放在哪里这些软硬件协同的细节在项目后期会疯狂决定成败。DMA控制器不是简单的搬数据工具它要能支持多维描述符、链式搬运、中断和事件同步。比如地址不连续时软件可以配置src_stride、dst_stride和次数让硬件自己算步长而不是软件把每个地址都列一遍。这样既能减少描述符数量也能让DMA的一大段搬运和计算重叠起来。运行时和DMA驱动的配合上我建议采用双缓冲或者多级缓冲机制一块buffer在计算另一块buffer正在被DMA搬运下一份数据计算单元永远不会空等。片上互联方面规模小可以用crossbar规模大用NoC网格。NoC灵活性好但要有能力处理多核同时访问同块bank的竞争。硬件里通常会加一个同步单元维护若干事件计数器。软件可以把算子A完成映射成事件另一个核等待该事件后再启动这样同步就不是靠中断和忙等而是靠硬件原语开销小很多。死锁问题也常出在这一层比如核A在等核B的数据核B又在等核A释放buffer所以设计任务依赖图的时候必须强制保证无环并且要有超时看门狗。2.4 面积与功耗低层工程约束如何反推软件设计很多人以为功耗是硬件工程师的事软件不用管。实际上AI芯片上的功耗优化很大一部分要靠软件配合。最典型的例子是电流凸起如果编译器把所有算子都排在同一时刻访问同一个SRAM bank瞬时电流会特别难看供电网络都要多铺一层。为避免这种局面调度器要做功耗感知调度尽量让不同核的访存相位错开。面积约束对软件的反推更直接片上SRAM就这么大编译器必须把算子切块安排得明明白白。别想着硬件把buffer越做越大那个成本太高。软件要做的是在物理容量限制下通过合理的tile大小、生命周期分析和内存复用把峰值内存使用压到最低。编译器的内存分配器和硬件的SRAM分区表其实是同一份规格设计期就要对上。3. 软件栈设计从框架到指令3.1 软件栈的分层与分工AI芯片的软件栈从上到下大致可以分成五层框架接入层、图优化层、中间表示层、算子/代码生成层、运行时层。每一层都要有清晰边界否则团队协作就是灾难。我习惯把分层比喻成开餐厅框架层是客人点菜图优化层是后厨根据客人需求制定做菜顺序IR层是写好的标准菜谱代码生成层是厨师按菜谱颠勺运行时是传菜员和洗碗工负责餐具周转和服务节奏。框架接入层负责对接PyTorch、ONNX这类生态把模型解析成内部图。这一层尽量少做硬件相关的事不然换一个框架就要动一遍后端。图优化层做算子融合、常量折叠、精度转换、布局转换输出一个优化后的计算图。IR层是中间表示设计要稳定不能跟着后端指令集三天两头变。算子/代码生成层把IR里的算子映射到目标指令或者调用算子库。运行时层负责内存管理、任务下发、同步、错误处理。很多团队一上来就想自己写一个全功能编译器我建议反过来先把垂直切面做通也就是从框架到指令直达一个最小模型跑通之后再补水平扩展。软件栈太早铺开你根本不知道哪些层该稳定、哪些层该灵活返工是必然的。3.2 图优化与算子融合编译器里的厨艺图优化是AI编译器最能体现价值的地方算子融合则是其中收益最大的优化。简单说就是把多个算子合并成一个硬件能一次性执行的操作减少数据搬移和中间内存分配。常见的融合例子卷积批归一化融合推理阶段可以把BN的scale和shift折进卷积权重里残差结构的add也可以融合到卷积的累加环节省一次完整的写回和读取Transformer里如果把QK^T、softmax、V乘融合成一个注意力算子整段计算在片上SRAM里完成效果极其明显。实现融合通常要在IR层写pattern匹配规则。我用伪代码表达一下这类逻辑// 伪代码示意匹配 Conv BatchNorm 并改写为一个融合算子 pattern ConvBN { in: Conv(x, w) - BatchNorm(mean, var, gamma, beta) emit: ConvFused(x, w, gamma, beta) }规则看起来简单坑在pattern爆炸。同一段图可能匹配几十种规则组合每种组合的顺序不同结果也不同。所以我建议图优化设计时给pass排序别一把梭全跑否则最后的图可能被改得面目全非数值误差也没法追。3.3 代码生成与算子库硬核的内功算子库是AI芯片软件栈里最硬核的部分。理想情况下通用编译器能把任意算子生成不错的代码但现实是热门算子的性能还得靠人工手写模板或半自动调优才能压榨出来。还是拿GEMM说。代码生成要做的事包括循环切块、循环重排序、向量化、寄存器缓存复用、DMA双缓冲调度。每一步都有参数tile大小、循环展开因子、DMA搬运粒度、累加缓冲存放位置。这些参数互相影响搜起来是个巨大的组合空间。实践中我很少跑全自动搜索而是先用roofline模型估算应该卡计算还是卡访存再基于经验给一组初始参数然后用小步调优去锁范围。这样从小时级收敛到分钟级。算子库还有一个重要原则热门算子手写冷门算子保底。对于已经识别出来的高频算子花大量精力精调是值得的对于冷门或变体算子靠通用代码生成路径跑通就行不要追求极致性能。打造一个全算子全能优化的方案往往是项目悲剧的开端。3.4 运行时与任务调度软硬件的胶水层运行时层是软硬件协同的胶水。它的核心工作不是做计算而是把算子的执行安排得明明白白申请内存、设置任务描述符、下发指令、等待同步、回收buffer。调度设计上我最推荐事件驱动模型。每个算子包一次资源返回值代表该算子执行完毕。运行时维护一个依赖图当一个算子依赖的所有上游事件都就绪就把它插入可执行队列然后下发给硬件。这样多个算子可以在不同核上并行DMA搬运也能和计算重叠。用事件驱动调试时定位慢点非常方便因为你能看到谁在等谁。内存管理同样关键。AI模型的张量生命周期很规则不一定要走复杂的内存池但至少要做流式复用。就是把同一块物理buffer分配给不同生命周期不相交的张量效果可以很惊人。我在一个推理模型上做过一次简单的生命周期分析只靠内存复用就让整体内存占用降了40%以上而这几乎是零成本改动。4. 协同设计落地流程与工具链4.1 用功能模拟器和性能模型并行推进硬件的流片周期很长软件不能等硬件回来才开始写。所以没有芯片的时候软件团队靠什么开发答案是模拟器。我建议至少维护两套模拟模型一套是功能模拟器只要求行为正确速度快用来跑模型验证算法和算子逻辑另一套是性能模型按周期估算执行时间用来做性能调优。功能模拟器的实现可以简单粗暴比如在CPU上模拟硬件指令跑完一遍把结果和浮点参考模型对比。性能模型要有一定精度能够反映计算单元占用、DMA带宽瓶颈、bank冲突、同步等待。性能模型不用对所有模块都周期精确只对瓶颈相关的模块精确就够了否则开发量爆表。这里有个工程技巧给性能模型和真实硬件设计统一的事件输出格式。比如算子开始事件、结束事件、DMA搬运事件、stall事件。这样后来在硬件上抓profile时开发工具几乎可以复用不用重写一套日志系统。4.2 工具链哪些必须自己造哪些可以直接用AI芯片工具链生态还不像CPU那么齐全很多轮子得自己造。我的经验是编译器前端可以直接用开源生态比如对接ONNX或者现有IR但编译器的后端、汇编器、反汇编器、性能剖析器基本都得自己写。不要小看反汇编器和性能剖析器。没有反汇编器你没法确认编译器生成的指令是不是符合预期没有性能剖析器你只能对着跑完的时间和内存流量瞎猜。这两个工具开发成本不高但能极大提升软硬件联调效率。有条件的话把它们做成同一个内部工具框架下的模块复用一套JSON配置会省很多事。还有测试集。我建议每个新项目都准备两套benchmark一套是算子级micro benchmark比如不同形状的GEMM、卷积、softmax、attention另一套是端到端模型集覆盖公司主力业务场景。micro benchmark用来快速定位性能回退端到端模型集用来验证整体收益。每个模型都要记录精度基线、内存占用和延迟任何一次工具链改动都拿它们来做回归。4.3 自动化验证与数值容差RTL仿真环境下跑一次大模型非常慢所以软硬件对齐验证要把重点放在算子级和子图级。我的习惯是维护一个参考实现用高精度浮点在CPU上算一遍然后把同样输入喂给模拟器或者RTL逐层比对输出。数值容差不能拍脑袋。FP16、INT8这些低精度下误差是正常的但你要区分是算法误差还是硬件bug。我经常用的方法是概率分布比对对输出张量统计绝对值误差的分布如果误差集中在最后几个bit那是正常的浮点重排如果误差出现在高位或者大面积异常值基本可以断定是逻辑错误。另一个技巧是让参考实现和待测实现采用相同顺序的浮点运算先排除重排序误差再去查真正的bug。4.4 用roofline模型指导性能调优性能调优不能靠感觉我每次拿到一个性能报告都会先画roofline模型横轴是计算访存比纵轴是可达算力图上画一条计算上限线、一条带宽上限线。任何一个算子的实际性能点只能落在两条线下面。具体调优时先看在哪个区域如果计算访存比低落在带宽受限区那优先优化数据复用和tile大小如果落在计算受限区再去研究指令调度和流水线。我处理过一个attention算子理论计算量不大但因为频繁把Q、K、V从外存读进读出来copy实际性能被访存打到崩溃。后来把整个attention子图融合成一个kernel数据只在SRAM里循环性能立刻从带宽受限跳到计算受限整体端到端快了接近两倍。5. 常见问题与排查实录5.1 高频问题速查表我把团队这几年在软硬件联调中遇到的高频问题整理成一个速查表排查问题的时候可以直接对着查现象可能原因排查方向端到端算力利用率低数据搬运与计算没重叠、同步等待多看profiler里的stall和DMA事件某个算子比预期慢很多tile大小、循环顺序不合适用roofline判断带宽受限还是计算受限数值精度发散低精度表达、重排误差逐层输出对比检查融合和近似函数任务卡死依赖图有环、DMA冲突、同步丢失查任务依赖图加超时看门狗编译时间过长图优化pattern组合爆炸、搜索空间太大压缩pass规则、做缓存、限制搜索深度指令下发开销高每个算子都写大量寄存器换任务描述符门铃机制排查的原则是先看profile数据再看代码最后猜硬件。很多问题不是单靠看代码能发现的一定要让性能计数器说话。5.2 三个让我印象深刻的排坑案例第一个案例是关于端到端性能的。某推理模型在模拟器上怎么看都只能跑到理论峰值的30%左右大家都在纠结是不是计算单元设计浪费了面积。后来我拉了一遍所有算子的时间线发现几乎所有算子的末尾都有一段很长的DMA等待——计算早就做完了但下一份数据还没搬上来。问题根本不在计算阵列而在调度没有做软件流水。我们给运行时加了双缓冲和提前预取把DMA搬运提前到上一算子执行期间最终端到端利用率直接翻倍。这个案例告诉我很多时候性能优化的最大空间不在硬件而在调度策略。第二个案例是数值偏差。某算子在自研芯片上比CPU参考实现差了3%的电量算法同学说是硬件不给力硬件同学说是软件写错了。最后逐层排查发现问题出在softmax的exp实现上硬件里用了一个近似查表而融合优化又把softmax和前面的矩阵乘重排了误差被放大。解决办法是统一定点exp的查表逻辑并且对softmax这种敏感算子保留浮点累加。数值问题真的不能只看最终结果要把中间结果逐层拉出来对比。第三个案例是死锁。多核任务队列上线第一天模型跑半小时后必卡死。查了一周发现两个核在等待对方释放buffer形成了循环依赖。硬件层面其实有事件计数机制但运行时没有做依赖图的环检测。后来我们在下发链路上加了一个合理性检查任何任务如果等待的事件来自自己下游的算子直接报错同时硬件加了一个超时复位寄存器。从此这种问题再也没出现过。5.3 软硬件团队协作中流过的血泪技术问题好查最难的是团队协作问题。软硬件协同项目最常见的地雷是接口冻结不彻底。硬件组觉得寄存器定义改一下影响不大软件组编译器已经按旧定义写了好多模板一改就编译不过。所以项目一开始就要明确接口owner任何改动必须走变更评审同时更新规格文档、模拟器、编译器、RT模型四件套。还有一个协作经验让硬件工程师和软件工程师互相看对方的设计评审文档。硬件评审时软件工程师要能从工具链角度提出这个寄存器能不能这样简化软件评审时硬件工程师要能指出这个调度会导致当前电压不够。互看评审一开始会有点摩擦但坚持一两个月后团队对系统理解会明显提升很多隐患在评审阶段就被消灭了。6. 最后一点个人体会如果这个系列还能继续写下去我其实最想聊的是Chiplet与多芯粒互联和存内计算这两个方向但受限于项目保密和篇幅今天先不展开。最后想分享一个我个人的实操小技巧给新项目做架构spec时先用一个最小示例程序贯穿全流程。这个程序甚至不用是真实模型可以是两个卷积一次add一个softmax但要求从框架解析、图优化、代码生成、运行时下发的每个环节都走通并且在一个可模拟的硬件模型上得到正确结果。这个最小程序就是软硬件协同的一炷香任何接口定义、任何工具链改动都以能跑通它为底线。它能天然地暴露很多早期问题比如寄存器定义缺了某个同步位、DMA描述符字段对不齐、运行时和编译器的buffer生命周期理解不一致。等这个最小程序稳定了再慢慢往里面加算子、加模型、加多核调度架构基本不会走偏。我在这上面吃过亏所以也特别希望你新起的项目能少走这一步弯路。