首页 / 资讯中心 / 文章详情

AI芯片软硬件协同设计:从计算模式到编译映射的联合优化

AI芯片软硬件协同设计:从计算模式到编译映射的联合优化 ★ FEATURED ARTICLE
1. 从一颗芯片的诞生说起为什么软硬件必须一起设计很多人第一次接触AI芯片脑子里浮现的画面是一块焊在板子上的黑色方块上面印着某个厂商的logo旁边配着散热片和风扇。但如果你真的拆开一台推理服务器或者拿到一颗刚流片回来的加速卡你会发现真正决定它好不好用的从来不是那块硅片本身而是围绕它的一整套软硬件协同体系。我见过太多团队硬件指标漂亮得不行峰值算力标到几百TOPS结果模型跑上去利用率不到百分之二十最后项目延期、预算超支问题就出在软硬件设计脱节上。所谓AI芯片的软硬件设计说白了就是两件事一是设计一颗适合特定计算模式的芯片架构二是设计一套能让算法工程师无痛把模型跑上去的软件栈。这两件事必须同步推进不能先做完硬件再补软件也不能先写软件再倒逼硬件。原因很简单AI计算负载变化太快今天流行的算子结构半年后可能就被新的注意力机制取代。如果硬件设计时没有预留足够的可编程性和数据通路灵活性软件再怎么优化也救不回来。这篇文章适合谁看如果你是刚入行的芯片设计工程师想了解软件侧到底在关心什么如果你是算法工程师好奇自己的模型在芯片上到底经历了什么或者你是项目负责人需要判断一个AI芯片方案是否靠谱那这篇内容应该能给你一些实在的参考。我会从计算模式、架构选型、数据流设计、软件栈分层、编译映射、实测调优几个角度把软硬件协同设计的核心逻辑拆开讲清楚中间会穿插一些实际项目中踩过的坑和总结出来的经验参数。先给一个最核心的判断AI芯片的软硬件设计本质上是在做“计算模式与数据流动的联合优化”。谁能把数据搬运的代价降到最低谁能让计算单元尽量不空转谁就能在同样的工艺节点下拿到更好的实际性能。下面我们一层一层往下拆。2. 计算模式决定架构走向从矩阵乘到注意力机制的硬件映射2.1 为什么AI芯片偏爱矩阵乘加单元要理解AI芯片的硬件设计得先回到AI计算本身。无论是卷积神经网络还是Transformer核心计算都可以归结为大量的矩阵乘法与累加操作。一个典型的全连接层输入向量乘以权重矩阵再加上偏置本质上就是一堆乘加运算的集合。卷积层稍微复杂一点但通过im2col展开后同样可以转化为矩阵乘。所以硬件设计的第一反应就是既然乘加这么多那就堆乘加单元。但堆多少、怎么堆、数据怎么喂这才是真正考验设计的地方。我参与过一个边缘推理芯片的项目最初架构师拍脑袋定了256个乘加单元觉得算力够了。结果仿真发现权重和激活值从片外内存搬到计算阵列的带宽根本跟不上计算单元大部分时间在等数据。后来把乘加单元砍到128个省下来的面积用来做更大的片上缓存和更宽的数据总线实际吞吐反而提升了近四成。这个教训很直接算力不是堆出来的是喂出来的。2.2 数据复用模式与片上存储层次AI计算有一个天然优势就是权重和激活值存在大量复用机会。卷积核在输入特征图上滑动时同一个权重会被反复使用同一张输入特征图也会和多个卷积核做运算。硬件设计要做的就是设计合理的片上存储层次让这些复用发生在离计算单元最近的地方。常见的做法是三级存储寄存器文件、片上缓存、片外内存。寄存器文件离乘加单元最近延迟最低但容量最小片上缓存容量大一些可以存放当前计算块需要的权重和激活值片外内存容量最大但带宽有限、延迟高。设计的关键是确定每一级存什么、存多少、什么时候换入换出。我习惯用一个简单的经验公式来估算片上缓存需求假设计算阵列有N个乘加单元每个周期需要M个权重和K个激活值那么片上缓存至少要能容纳一个计算块完整的数据量通常取计算阵列规模的4到8倍。比如128个乘加单元的阵列片上缓存做到512到1024个数据单元比较合理。太小了频繁换入换出太大了面积和功耗吃不消。2.3 从卷积到注意力可编程性的取舍前几年很多AI芯片是专用卷积加速器硬件只认卷积和池化其他算子靠CPU补。这在当时没问题因为主流模型就是卷积网络。但现在Transformer大行其道注意力机制里的矩阵乘、softmax、层归一化和卷积的计算模式差别很大。如果硬件只针对卷积优化跑Transformer时要么效率暴跌要么根本跑不了。所以现在的AI芯片设计普遍强调可编程性。但可编程性不是免费的通用性越强单位面积的能效通常越低。这里有一个很实际的取舍是做一个能跑所有算子的通用加速器还是做几个专用单元分别处理卷积、矩阵乘和向量运算再用一个调度器把它们串起来。我个人的经验是对于训练芯片可编程性优先级更高因为训练涉及的算子多且变化快对于推理芯片可以适当专用化但至少要支持矩阵乘和向量运算这两类核心操作。具体到硬件设计上可以设计一个可配置的乘加阵列通过改变数据通路和累加模式让它既能算卷积也能算矩阵乘。这样面积增加不多但灵活性提升明显。3. 软件栈的分层逻辑从驱动到框架的完整链路3.1 驱动层芯片对外的第一道门软件栈最底下是驱动层负责和硬件直接打交道。驱动要做的事情包括初始化芯片、分配内存、下发计算任务、处理中断、上报状态。听起来简单但实际写起来坑很多。比如内存分配AI计算需要大块连续内存但操作系统给的是虚拟地址驱动得做地址映射和页表管理。再比如任务下发硬件队列深度有限驱动要负责排队和流控不能让硬件溢出。我见过一个项目驱动里内存分配用了标准的页分配接口结果每次分配出来的物理页不连续硬件DMA传输时频繁跨页带宽直接打对折。后来改成预留大块连续内存池自己管理分配释放问题才解决。这个经验说明驱动层不能只求功能正确还要考虑硬件的物理特性。3.2 运行时与编译器把模型翻译成硬件指令驱动之上是运行时和编译器。运行时负责资源管理、任务调度、多流并发编译器负责把上层框架的模型图翻译成硬件能执行的指令序列。这是软硬件协同设计中最核心也最复杂的部分。编译器要做的事情包括算子融合、内存分配、指令调度、数据搬运插入。举个例子一个卷积层后面跟着激活函数编译器可以把它们融合成一个硬件指令激活函数直接在计算阵列的输出端完成省去一次写回和读取。再比如编译器要决定哪些数据放在片上缓存、哪些放片外内存什么时候预取下一块数据。这些决策直接影响实际性能。我通常会用“编译器的质量决定芯片实际性能的下限”这句话来提醒团队。硬件峰值算力再高编译器映射不好实际跑出来可能只有峰值的百分之二三十。所以芯片设计早期就要让编译器团队介入根据编译器的需求来调整硬件指令集和数据通路。3.3 框架适配层让算法工程师无感迁移最上面是框架适配层对接主流深度学习框架。算法工程师用熟悉的接口写模型底层自动调用芯片的运行时和驱动。这一层的目标是让迁移成本尽可能低最好改一行代码就能切换设备。但现实往往没那么美好。不同框架的算子定义有差异有些算子芯片不支持需要回退到CPU或者用多个算子组合实现。适配层要处理这些差异还要做精度对齐、内存格式转换、动态shape支持。我参与过一个项目适配层里光算子映射表就维护了上千条记录每条都要测试验证。这活儿很枯燥但省不得因为算法工程师不会关心底层细节他们只关心模型能不能跑通、精度对不对、速度快不快。4. 数据流设计决定实际性能的隐形战场4.1 片内数据搬运的代价估算很多人评估AI芯片性能时只看峰值算力但实际跑模型时瓶颈往往在数据搬运上。片内数据搬运的功耗和延迟通常比计算本身高一个数量级。有研究数据表明在同等工艺下从片外内存读一个字节的功耗大约是进行一次乘加运算的几十倍甚至上百倍。所以硬件设计的一个核心原则就是能片上复用就片上复用能近端计算就近端计算。具体怎么估算搬运代价我习惯用一个简化模型假设计算阵列规模为N每个周期需要从片上缓存读取W个权重和A个激活值那么片上缓存的带宽至少要达到N乘以(WA)字节每周期。如果片上缓存带宽不够计算阵列就会空转。同样片外内存带宽要满足最坏情况下的数据换入换出需求通常取片上缓存带宽的十分之一到五分之一作为设计目标。4.2 双缓冲与流水线设计为了让计算和数据搬运重叠起来双缓冲是常用手段。把片上缓存分成两块一块供当前计算使用另一块预取下一块数据。计算和搬运并行进行理论上可以把搬运时间隐藏掉。但双缓冲要处理好同步问题预取完成才能切换计算完成才能覆盖。硬件上可以用乒乓缓冲加信号量来实现软件上编译器要负责插入同步指令。流水线设计也是类似思路。把计算过程拆成多个阶段每个阶段处理不同的数据块阶段之间用缓冲区连接。这样即使某个阶段偶尔变慢整体吞吐也不会断崖式下跌。但流水线级数不是越多越好级数多了同步开销和缓冲区面积都会增加。我一般建议边缘推理芯片用两到三级流水训练芯片可以用四到五级具体要看计算阵列的规模和任务粒度。4.3 稀疏化与量化对数据流的改变稀疏化和量化是降低数据搬运压力的两个有效手段。稀疏化利用权重或激活值中的零元素跳过不必要的计算和存储量化把浮点数压缩成低比特整数直接减少数据量。但这两项技术对硬件设计提出了新要求。稀疏化需要硬件支持稀疏编码和索引计算时根据索引跳过零元素。这增加了控制逻辑的复杂度但能显著降低有效计算量和存储需求。量化则需要硬件支持低比特乘加比如INT8或INT4同时要处理好累加过程中的精度问题。我实测过一个INT8量化的模型相比FP16版本片上缓存需求减少一半计算吞吐提升近一倍精度损失控制在百分之一以内。但量化不是万能的有些对精度敏感的层还得保留高精度硬件要支持混合精度计算。5. 编译映射与指令调度把模型图变成硬件时序5.1 算子融合的边界与收益算子融合是编译器最常用的优化手段。把多个连续的小算子合并成一个大的硬件任务减少中间结果的写回和读取。比如卷积加偏置加激活可以融合成一个指令矩阵乘加层归一化也可以融合。但融合不是越多越好融合太多会导致单个任务过大片上缓存放不下反而要频繁换入换出。我通常会用“融合后的任务数据量不超过片上缓存容量的百分之七十”作为一条经验边界。留出百分之三十的余量给预取和临时变量。另外融合还要考虑算子的依赖关系有数据依赖的算子才能融合并行分支不能强行合并。编译器在做融合决策时要结合硬件资源约束和任务图结构不能只追求融合数量。5.2 内存分配与生命周期分析编译器另一个重要任务是内存分配。模型中的张量有各自的生存周期从产生到最后一次使用。编译器要分析这些生命周期把不再使用的内存回收分配给新的张量。理想情况下内存占用可以压缩到峰值需求而不是所有张量之和。但AI模型的内存分配有个难点动态shape。很多模型支持变长输入张量大小在编译时不确定。编译器要么做动态内存分配要么预留最大可能尺寸。前者灵活但有运行时开销后者简单但浪费内存。我一般建议对性能敏感的模型用静态shape加最大尺寸预留对灵活性要求高的场景用动态分配加内存池。5.3 指令调度的时序约束指令调度决定了硬件任务的执行顺序和并行度。编译器要分析任务之间的依赖关系把没有依赖的任务并行执行有依赖的按顺序执行。同时要考虑硬件资源的限制比如计算阵列同时只能执行一个任务DMA通道有限片上缓存容量有限。调度做得好硬件利用率能提升百分之三四十调度做得差硬件大部分时间在等。我见过一个项目编译器调度时没有考虑DMA和计算的并行所有数据搬运都串在计算前面结果实际性能只有理论值的三分之一。后来改成搬运和计算重叠性能直接翻倍。这个例子说明指令调度不是简单的排序而是要在时间轴上把计算、搬运、同步这些操作合理安排让硬件资源尽量不空闲。6. 实测调优从仿真到流片的性能验证6.1 仿真阶段的性能建模芯片流片之前性能验证主要靠仿真和建模。仿真可以跑真实指令序列但速度慢适合验证功能正确性建模速度快适合做设计空间探索和性能预估。我通常会用两种手段结合先用分析模型快速评估不同架构参数下的理论性能选出几个候选方案再用周期级仿真验证候选方案的实际表现。分析模型不用太复杂抓住主要矛盾就行。比如计算受限还是带宽受限片上缓存命中率大概多少流水线效率如何。这些参数可以用简单的公式估算不需要精确到每个周期。关键是快速迭代在早期把明显不合理的方案排除掉。6.2 流片后的实测与调优流片回来之后实测才是真正的考验。实测会发现很多仿真阶段没暴露的问题工艺偏差导致频率上不去电压波动影响稳定性散热不足触发降频驱动和固件的bug导致任务失败。这时候软硬件团队要紧密配合硬件问题改版解决软件问题打补丁优化。我印象最深的一次芯片实测性能只有仿真值的六成。排查了很久最后发现是驱动里一个内存屏障指令用错了位置导致DMA传输和计算任务之间出现了不必要的等待。改掉之后性能恢复到九成以上。这个经历告诉我软硬件协同设计不只是架构和编译的事驱动和固件里的细节同样能决定成败。6.3 持续迭代与生态建设AI芯片不是一锤子买卖流片成功只是开始。后续要持续迭代软件栈支持新算子、新模型、新框架要建设开发者生态提供文档、示例、工具链要收集用户反馈优化实际使用体验。硬件可能两三年才更新一代但软件可以每月每周迭代。软硬件协同设计的最终目标是让用户觉得这颗芯片好用、易用、愿意用。我个人的体会是做AI芯片软硬件设计最怕的就是闭门造车。硬件团队觉得软件团队不懂硬件软件团队觉得硬件团队不懂算法最后做出来的东西谁都不满意。打破这个局面的办法只有一个让软硬件团队坐在一起从需求分析到架构设计从编译映射到实测调优全程共同决策。只有这样才能把芯片的潜力真正释放出来。
阅读完成 · 觉得有帮助?
咨询建站