AI芯片软硬件协同设计:从架构到编译优化的核心命题
1. 从能跑就行到跑得高效AI芯片软硬件协同设计的核心命题聊到AI芯片的软硬件设计很多人的第一反应是这不就是硬件堆算力、软件调驱动吗。我刚开始接触这个方向时也这么想直到真正把一个模型从训练到部署完整跑过一遍才发现事情远没有这么简单。AI芯片和通用CPU最大的区别在于它的价值几乎完全取决于软硬件配合的紧密度——一颗峰值算力很高的芯片如果软件调度跟不上、数据搬运成了瓶颈实际利用率可能连三成都不到。这就是为什么业内常说AI芯片是三分硬件、七分软件虽然这个比例各家有各家的说法但软件栈的重要性是公认的。这篇文章想聊的是我在实际做AI芯片软硬件协同设计过程中积累的一些理解和经验。核心围绕几个问题展开AI芯片的硬件架构到底特殊在哪里、软件栈需要解决哪些关键问题、编译器和运行时是怎么把模型翻译成硬件能执行的指令、以及在实际调优中哪些环节最容易成为性能瓶颈。内容会涉及计算单元设计、片上存储层次、数据流调度、算子融合、量化映射这些具体话题适合对AI加速器感兴趣、想搞清楚芯片到底怎么跑模型的读者。不管你是偏硬件的数字设计方向还是偏软件的编译优化方向相信都能从中找到对自己有用的部分。需要提前说明的是下面涉及的架构方案和参数都是基于公开的通用设计思路做的合理推演具体到不同厂商、不同产品会有很大差异但底层的设计逻辑是相通的。2. AI芯片硬件架构的特殊性为什么不能照搬CPU的设计思路2.1 计算单元从通用ALU到脉动阵列的取舍逻辑通用CPU的算力核心是少量功能强大的ALU配合复杂的乱序执行、分支预测、多级缓存来提升单线程性能。这套设计对逻辑控制密集的任务非常高效但放到深度学习场景里就很不划算。神经网络的计算主体是矩阵乘加特点是计算模式规整、数据复用率高、对单次计算延迟不敏感但对吞吐量极度敏感。用CPU跑矩阵乘大量的晶体管面积花在了指令译码、乱序调度、分支预测这些和矩阵计算本身无关的逻辑上能效比自然上不去。AI芯片的主流做法是换一条路用大量简单的乘加单元MAC组成阵列通过规整的数据流让每个MAC每个时钟周期都在干活。最典型的就是脉动阵列结构。它的核心思想是让数据像心跳一样在阵列中有节奏地流动每个计算单元从相邻单元接收数据、算完再传给下一个权重预先加载到阵列中保持不动激活值沿着一个方向流动部分和沿着另一个方向累积。这种设计的精妙之处在于它把数据搬运的次数降到了最低——每个数据进入阵列后会被多个计算单元复用而不是算一次就从内存重新读一次。但脉动阵列也不是没有代价。它的效率高度依赖于矩阵维度是否和阵列尺寸匹配。如果矩阵是128×128而阵列是256×256利用率直接掉到25%。所以实际芯片设计时阵列尺寸的选择要结合目标模型的主流维度来定同时软件侧需要通过矩阵分块tiling来尽量填满阵列。我见过一些设计为了追求峰值算力把阵列做得很大结果实际模型跑上去利用率很低反而不如小阵列高频来得实在。2.2 存储层次AI芯片真正的瓶颈往往不在算力很多人评估AI芯片时只看TOPS每秒万亿次操作这个指标但实际部署过模型的人都知道算力从来不是唯一瓶颈甚至很多时候不是主要瓶颈。真正卡脖子的是存储带宽和片上存储容量。有一个经典的经验数据做一次MAC运算需要的能量大约是做一次32位浮点乘法的几分之一但从DRAM读一个数据进来的能量可能是这个计算能量的几十倍甚至上百倍。也就是说数据搬运的代价远大于计算本身。这就引出了AI芯片存储层次设计的核心原则尽可能让数据在离计算单元近的地方被复用。典型的层次是寄存器文件→片上SRAM也叫便笺存储器→L2缓存→DRAM。越靠近计算单元带宽越高、延迟越低、功耗越小但容量也越小、成本越高。设计的艺术就在于在有限的芯片面积和功耗预算下合理分配每一层的大小并且让数据在层间的搬运次数最少。这里有个容易被忽略的点片上SRAM的容量直接决定了能支持多大的算子融合。比如一个卷积层后面接一个激活层再接一个池化层如果SRAM够大可以把中间结果留在片上不用写回DRAM再读出来省下的带宽非常可观。但如果SRAM不够就只能分块计算中间结果反复进出DRAM性能立刻打折。所以我在看一颗AI芯片的设计时会特别关注它的片上存储容量和带宽这往往比峰值算力更能说明实际性能。2.3 数据流架构权重 stationary、输出 stationary 还是行 stationary数据流的选择是AI芯片硬件设计的另一个核心决策。简单说就是在矩阵计算过程中让哪一类数据保持在计算单元附近不动以减少搬运。常见的有三种权重固定weight stationary权重预加载到计算阵列中保持不动激活值流过阵列。这是脉动阵列的典型做法适合权重复用率高的场景比如卷积层。输出固定output stationary每个计算单元负责一个输出元素累加结果留在本地激活值和权重都流动进来。这种方式对全连接层比较友好。行固定row stationary一种更细粒度的混合策略把卷积运算拆成行级别的复用在Eyeriss等学术架构中被提出能同时利用权重复用和激活复用。实际芯片往往不会只用一种数据流而是根据算子类型动态切换或者用可重构的互连来支持多种模式。软件侧在映射算子时需要知道硬件支持哪种数据流才能生成高效的调度方案。这也是软硬件协同设计的一个典型接口硬件提供数据流模式软件负责选择和执行映射。3. 软件栈要解决的核心问题把模型翻译成硬件指令3.1 图编译与算子融合减少访存才是第一优先级深度学习框架训练出来的模型是一张计算图节点是算子卷积、矩阵乘、激活、归一化等边是张量。软件栈的第一项工作就是把这张图编译成硬件能执行的指令序列。这个过程里最重要的优化就是算子融合。为什么融合这么关键举个具体例子。假设有一个卷积层后面接一个ReLU激活如果不融合执行流程是卷积计算→结果写入DRAM→从DRAM读出→ReLU计算→结果再写入DRAM。中间那次写和读完全是浪费因为数据本来可以留在片上。融合之后卷积算完的结果直接留在SRAM里ReLU从SRAM读、算完再写回SRAM省掉了两次DRAM访问。对于大模型来说这种访存节省累积起来非常可观。但融合不是想融就能融的。它受限于片上SRAM容量和算子的数据依赖关系。比如两个算子之间有复杂的跨通道依赖或者中间张量太大放不下就融不了。所以编译器需要有一套代价模型评估每种融合方案的收益和代价选择最优的组合。我实际调优时发现融合策略对性能的影响有时候比换一颗芯片还大同一个模型在不同融合方案下性能差两三倍是常有的事。3.2 内存分配与复用张量生命周期管理模型推理过程中会同时存在多个张量每个张量有各自的生存周期。如果每个张量都分配独立的内存空间内存需求会迅速膨胀。软件栈需要做内存复用分析当一个张量不再被后续算子使用时它占用的空间可以分配给新的张量。这个问题的本质是图着色问题的一个变种属于NP难问题实际中通常用贪心策略加启发式规则来求解。常见的做法是按张量大小降序排列依次分配能复用就复用。但这里有个坑如果只按大小贪心可能会把一个大张量放在低地址导致后续小张量无法复用它的空间反而浪费。更好的做法是结合张量的生命周期区间来做区间调度类似寄存器分配里的线性扫描算法。另外片上SRAM和DRAM之间的分配策略也很讲究。频繁访问的中间结果尽量放SRAM权重这种只读且复用率高的数据可以放DRAM但要做好预取。我见过一些实现把所有中间张量都往SRAM里塞结果SRAM不够就频繁换入换出性能反而比老老实实放DRAM还差。3.3 量化与精度映射INT8不是简单地把浮点砍掉几位量化是AI芯片软件栈里另一个重头戏。把FP32模型转成INT8推理理论上能省四分之三的存储和带宽计算单元也能做得更小更快。但量化不是简单地把浮点数截断它涉及到缩放因子scale和零点zero point的选取以及每一层量化误差的累积控制。实际做量化时常见的方法有训练后量化PTQ和量化感知训练QAT。PTQ简单快捷但对一些对精度敏感的模型比如检测、分割可能会掉点明显。QAT在训练过程中模拟量化误差让模型自己去适应精度保持得更好但需要重新训练成本高。我个人的经验是分类模型用PTQ通常够用检测和分割模型最好上QAT尤其是小目标检测量化误差很容易把微弱信号淹没。还有一个容易被忽略的点不同层对量化的敏感度差异很大。第一层和最后一层通常最敏感因为第一层直接处理输入数据最后一层直接决定输出。很多芯片支持混合精度就是让这些敏感层保持FP16其余层用INT8。软件栈需要能自动分析每层的敏感度并做出精度分配决策这背后又是一套搜索和评估的逻辑。4. 编译器与运行时的协同从计算图到指令流的完整链路4.1 算子库与代码生成手写汇编还是自动生成AI芯片的算子实现有两条路一是手写高度优化的算子库二是用编译器自动生成代码。手写库的优点是性能极致针对特定硬件做了精细调优比如循环展开、指令重排、双缓冲预取这些。缺点是开发成本高每换一代硬件就要重写而且覆盖的算子种类有限遇到新算子就得重新开发。自动代码生成比如基于TVM、Halide这类框架的思路的优点是通用性强新算子只要用DSL描述就能生成代码而且能针对不同硬件后端做自动调优。缺点是生成的代码在极端情况下可能不如手写的尤其是涉及复杂数据流和特殊指令时。实际产品里通常是混合策略核心高频算子比如3×3卷积、GEMM用手写库保证性能长尾算子用自动生成保证覆盖率。编译器需要有一套调度机制决定每个算子走哪条路径。这里的关键是接口要统一手写库和生成代码要能无缝拼接在同一个执行图里否则调度开销会把性能吃掉。4.2 运行时调度静态编排与动态调度的平衡模型编译完之后运行时负责把指令序列真正推到硬件上执行。这里有两种风格静态编排和动态调度。静态编排是在编译期就把所有指令的顺序、依赖、资源分配确定好运行时只管按顺序发射。好处是开销小、可预测性强适合推理场景。坏处是灵活性差遇到动态shape或者控制流就不好处理。动态调度是运行时根据实际情况决定下一步执行什么灵活但开销大。对于AI推理来说大部分模型的shape是固定的静态编排通常更优。但有些场景比如NLP里的变长序列shape动态变化就需要运行时做动态分块和调度。我实际用下来比较好的方案是静态为主、动态为辅主体计算图静态编排在几个关键决策点比如分块大小选择、精度切换留出运行时接口让运行时根据实际输入做微调。这样既保证了效率又不至于完全僵化。4.3 多核与多芯片扩展数据并行、模型并行还是流水并行单颗AI芯片的算力总有上限大模型往往需要多核甚至多芯片协同。这时候并行策略的选择就很重要。数据并行是最简单的每个核跑完整的模型处理不同的batch数据。优点是实现简单缺点是每个核都要存一份完整权重对存储压力大而且batch小的时候效率低。模型并行是把模型切开放到不同核上每个核算一部分。缺点是核间通信频繁切分点选不好通信开销会很大。流水并行是把模型按层切成多个阶段不同阶段在不同核上数据像流水线一样流过。适合层数多、层间依赖线性的模型。实际大模型推理里通常是混合并行注意力层用模型并行因为参数量大前馈层用数据并行整体用流水线串起来。软件栈需要能自动搜索最优的并行策略这又是一个组合优化问题。我见过一些框架用强化学习来搜并行策略效果不错但搜索成本高实际中更多还是用启发式规则加人工调优。5. 实际调优中的性能瓶颈定位与解决思路5.1 用roofline模型快速判断是算力瓶颈还是带宽瓶颈调优第一步是搞清楚瓶颈在哪。我常用的工具是roofline模型横轴是计算强度每字节内存访问对应的计算量纵轴是性能。每个算子在图上是一个点如果点在斜线下方说明是带宽瓶颈如果在水平线下方说明是算力瓶颈。举个例子一个逐元素加法算子计算强度极低读两个数、写一个数、做一次加法几乎必然落在带宽瓶颈区。这种算子优化方向就是减少访存比如融合到前一个算子里。而一个大的矩阵乘计算强度高通常落在算力瓶颈区优化方向就是提升计算单元利用率比如调整分块大小让数据复用更充分。实际调优时我会先把模型里所有算子的roofline点画出来看看哪些落在带宽瓶颈区且耗时占比高这些就是优先优化对象。很多时候把几个带宽瓶颈的逐元素算子融合掉整体性能就能提升一大截。5.2 分块大小与双缓冲让数据搬运和计算重叠起来分块tiling是AI芯片上最基础也最重要的优化手段。核心思想是把大矩阵切成小块让每块能放进片上SRAM计算时从DRAM预取下一块实现计算和搬运的重叠双缓冲。分块大小的选择是个权衡块太小数据复用率低DRAM访问次数多块太大放不进SRAM或者留给双缓冲的空间不够重叠效果差。理论最优分块大小可以通过分析SRAM容量、DRAM带宽、计算吞吐三者的关系来算但实际中还要考虑bank冲突、对齐要求等细节通常需要在理论值附近做实验微调。我踩过的一个坑是只考虑了单层的最优分块没考虑层与层之间的衔接。结果单层看性能很好但层间切换时SRAM要清空重填开销累积起来很可观。后来改成全局视角做分块让相邻层能共享部分片上数据整体性能才上来。5.3 指令级优化循环展开、软件流水与指令重排到了指令级别优化空间依然很大。循环展开可以减少循环控制开销提高指令级并行软件流水可以让不同迭代的计算和访存重叠指令重排可以填补流水线气泡。这些优化手写起来很繁琐但编译器如果能做好收益很直接。我见过一个卷积算子经过循环展开和软件流水优化后性能提升了40%以上。关键是要让编译器知道硬件的具体参数流水线深度、指令发射宽度、访存延迟等。这些信息通常通过后端的调度模型传给编译器模型越准确生成的代码越优。不过也要注意过度优化可能导致代码膨胀指令缓存命中率下降反而拖累性能。实际中要在优化收益和代码大小之间找平衡点。6. 软硬件接口设计那些容易扯皮的边界问题6.1 指令集架构CISC还是RISC定制指令怎么加AI芯片的指令集设计是个容易引发争论的话题。一派主张用RISC风格的简洁指令集靠编译器做优化另一派主张加定制指令把常见算子模式固化到硬件里。我的观察是纯RISC对AI场景效率偏低因为矩阵计算里很多模式是固定的用通用指令描述会有大量冗余。但定制指令加太多又会导致编译器复杂、灵活性差。比较务实的做法是基础指令集保持简洁针对最高频的几种计算模式比如矩阵乘加、卷积、激活加定制指令其余用基础指令组合实现。定制指令的设计要特别注意和编译器的配合。指令的语义要清晰编译器能准确建模它的延迟和吞吐否则调度会出问题。我见过一些定制指令因为延迟不透明编译器不敢调度结果硬件能力发挥不出来。6.2 同步与通信多核之间怎么高效交换数据多核AI芯片里核间通信的开销往往被低估。如果同步机制设计得不好核间等待会吃掉大量性能。常见的同步方式有屏障同步、信号量、消息传递等。屏障同步简单但粒度粗所有核必须等到最慢的那个信号量灵活但容易死锁消息传递适合松耦合但延迟高。实际芯片通常提供多种同步原语软件根据场景选择。一个实用的技巧是尽量用数据流驱动代替显式同步核A算完把结果写到约定地址核B轮询这个地址有数据就算。这样避免了全局屏障的开销但需要硬件提供高效的内存一致性或至少是释放一致性模型。如果一致性模型太弱软件要加很多内存屏障指令开销也不小。6.3 调试与性能分析接口没有可观测性就没有优化最后但同样重要的是调试和性能分析接口。AI芯片如果只提供最终性能数字不提供中间过程的计数器、trace、profiling信息优化就无从下手。好的性能分析接口应该能回答这些问题每个算子的执行时间是多少、DRAM访问次数和带宽利用率是多少、计算单元利用率是多少、SRAM bank冲突率是多少、指令缓存命中率是多少。这些计数器在硬件设计时就要预留不能等流片完了再补。软件侧要有对应的工具链能把这些计数器数据可视化和计算图对应起来。我实际调优时一半以上的时间花在看profiling数据上没有这些数据就是盲人摸象。所以我现在看一颗AI芯片的成熟度会特别关注它的调试和profiling工具是否完善这往往比纸面算力更能反映实际可用性。7. 一些踩坑之后的个人体会做AI芯片软硬件设计这些年最大的体会是这个领域没有银弹每个决策都是权衡。阵列做大做小、SRAM给多给少、指令集加不加定制、并行策略怎么选都没有绝对正确的答案只有适合当前目标场景的答案。而且这些决策之间高度耦合硬件改了软件要跟着改软件发现瓶颈了硬件可能也要调所以协同设计不是一句口号是必须落到流程里的。另一个体会是可观测性再怎么强调都不为过。我早期参与的一个项目硬件性能其实不错但因为没有足够的性能计数器软件团队只能靠猜来优化走了很多弯路。后来补上了profiling接口很快就定位到了几个关键瓶颈性能提升比之前几个月的盲调都大。最后一个建议是做AI芯片不要只盯着峰值算力。峰值算力是给外行看的实际性能取决于软硬件配合的每一个细节。把软件栈做扎实、把调试工具做完善、把常见模型的端到端性能调好比在纸面上多堆几个TOPS有意义得多。