AI芯片软硬件协同设计：从指令集到数据流的真实性能之道

上个月我对比了两款标称 AI 算力差距接近两倍的芯片跑同一份 YOLOv5s 量化模型开满 batch size 做端到端吞吐测试。结果非常打脸理论算力高的那颗只快了不到 20%。翻了一晚上 profiler 数据发现大量时间花在了数据搬运和算子调度上两个 MAC 阵列几乎都在等饭吃。这个场景我遇到过很多次也是 AI 芯片软硬件设计这个系列写到第六期我最想拿出来认真聊的东西——芯片的真实水平从来不由算力单元单独决定而是由软件栈和硬件架构共同咬合出来的。前五期我们分别聊过 AI 芯片的宏观分类、MAC 阵列的组织方式、多核缓存一致性、网络互联和工具链的基本盘。这期把视角收回到芯片内部专门讲软件和硬件之间的接缝怎么设计。搞清楚这一层你才能理解为什么同样制程、同样算力的两颗芯片实际表现能差出去一大截也才知道一颗芯片在设计阶段就该把哪些软件问题提前扛在自己肩上。1. 一个反直觉的现象算力翻倍了端到端性能却纹丝不动先说这个现象的根源。所有 AI 芯片厂商都会在 PPT 上写一个漂亮的 TOPS 数字那是 MAC 阵列跑满、数据全部驻留在最内层 buffer 时的理想值。但真实模型推理时片上内存放不下整个网络权重和激活值必须不断从外部 DRAM 搬进来搬出去。这个搬运动作一旦成为瓶颈MAC 阵列就只能空转。1.1 从理论峰值到实际吞吐的三级衰减我习惯把 AI 芯片的真实性能理解为三层衰减之后的结果。第一层是算力利用率。任何一个真实算子都不可能是 MAC 阵列满负荷运转。卷积核尺寸、stride、通道对齐方式、tile 边界处的补零操作都会让一部分 MAC 单元在某个周期闲置。这一层通常能拿到 60%~80% 的利用率已经是工程上做得不错的水平。第二层是数据供给率。算力单元每个周期要吃多少 byte 数据片上 buffer 能不能喂得饱喂不饱就得等DRAM 带宽再高也架不住计算单元消耗数据的速度是带宽的好几倍。这一层往往是最大的水分来源。以 128 MAC/cycle 的小阵列为例跑 FP16 累加 FP32一个周期要消耗 128 2 2 512 byte 的输入数据。假设芯片主频 1GHz纯算力是 128 GOPS对应的瞬时数据需求是 512GB/s。而一颗中端芯片的 DRAM 带宽可能只有 50GB/s。差了十倍。唯一的解决办法就是让数据在片上流动起来反复复用这时的实际性能取决于你能把数据局部性利用到什么程度。第三层是调度开销。算子之间的切换、同步、同步 barrier、流水线排空、驱动和 runtime 的调用损耗这些都在啃噬端到端的时间。小模型尤其吃亏因为计算本身只有几百微秒调度却可能吃掉同样量级的时间。1.2 软硬件接口设计直接决定三层衰减的加乘效果这三层衰减不是独立的。指令集定得不好编译器就没法把循环变换成数据复用率高的形态——第一层和第三层会同时劣化。片上 buffer 的容量和 bank 数目设计得抠门底层软件再会调也喂不饱 MAC——第二层就死死卡住。所以软硬件接口设计的本质是在芯片物理实现之前先把这三层衰减的上限估算清楚然后把整个架构往那个上限上顶。我评估一颗 AI 芯片时第一件事从来不是看 TOPS而是找它的 roofline 模型文档。看完算力、片上 SRAM 总容量、DRAM 带宽这三个数心里基本有数了。后面所有评测本质上都在验证这张 roofline 是不是真的兑现了。2. 指令集设计软硬件之间签下的第一份合同指令集是软件能看到的那部分硬件。AI 芯片的指令集设计和通用 CPU 完全不同CPU 指令追求通用、规整、便于流水化而 AI 芯片的指令集必须围绕张量计算这个核心场景做大量裁剪和定制。这个裁剪的分寸感是整个软硬件设计里最见功力的一件事。2.1 从标量、向量到矩阵指令指令粒度三连跳早期的 DSP 和 GPU 走的是标量向量指令路线一条指令操作一组数据。到了 AI 专用芯片时代大家不约而同地走向了矩阵指令一条指令直接触发一个二维 tile 的乘累加操作。这三者的差别我用量化数据说明。完成一个 16x16x16 的矩阵乘标量指令单条 MAC 指令计算一次乘累加需要 4096 条指令每条指令都要取指、译码、发射——光指令开销就把计算时间淹没。向量指令一条指令操作 16 个数据需要 256 条指令效率提升明显但仍不够。矩阵指令一条指令触发整个 16x16 的 MAC 阵列工作一个周期一次完成 4096 次乘累加指令数直接降到 1。VLIW 和 SIMD 处理器之所以在 AI 大模型面前显得吃力核心原因就在这里它们也有并行度但并行度体现在数据位宽上而矩阵指令的并行度体现在操作数矩阵的二维结构上。2.2 指令粒度越粗越好代价正在别处如果把矩阵指令的 tile 做得特别大比如 64x64一次触发 4096 个 MAC 干活指令发射次数确实大幅减少功耗也更低。但这会带来两个新的麻烦。第一个麻烦是灵活性下降。实际模型里的矩阵形状千奇百怪Attention 里的矩阵乘经常是 (B, H, N, D) 这种四维视角下的变形如果指令只接受规整的 64x64 tile碰到 17x33 这种非对齐形状就非常痛苦要么补零浪费算力要么靠软件做切碎重组。第二个麻烦是寄存器压力。Tile 指令的输入来自片上 buffertile 越大需要的 buffer 空间越大。一个 64x64x16 的输入就是 64 64 2 byte 8KBFP16两个输入加一个输出20KB 起步。如果片上总共就几百 KB SRAM这种大手大脚的指令格式会把 buffer 吃光多核场景下更是养不起。所以我见过的比较成熟的商业架构普遍把矩阵指令的 tile 定在 16x16 到 32x32 之间配合多级 buffer 和可配置的 data path。这个粒度既给了编译器足够的表达空间又不会让硬件在非对齐形状面前手足无措。2.3 指令集设计里容易被忽略的三个细节第一地址生成器。AI 芯片的访存模式高度规律卷积的滑动窗口、通道维度的 stride、矩阵转置这些如果全交给通用寻址模式指令里就得带着一堆地址偏移量白白浪费编码空间。成熟的方案是让指令携带地址生成模式字段硬件内部有专门的地址生成单元按固定模式批量产生访存地址一条指令管一整轮数据搬运软件只需提供基地址和步长。第二显式 buffer vs 隐式缓存。GPU 走隐式缓存路线软件只管访问地址硬件负责缓存局部性。AI 芯片多走显式 buffer 路线片上 SRAM 被划分为多个 buffer 区域软件编译器显式管理数据在各 buffer 间的流转。前者的优势是软件简单劣势是缓存未命中不可预测后者的优势是性能确定性极高编译器能精确安排每个数据的生命周期劣势是软件复杂度暴增。做软硬件协同设计时这个问题必须是最先拍板的决策之一因为它决定了编译器团队后面的全部工作方式。第三同步机制。多核 AI 芯片里一个算子可能被切分到多个核上并行执行核之间要么有数据依赖要么需要共享中间结果。指令集层面提供的同步原语有没有barrier、sync tile这类语义直接决定了编译器做多核切分时敢不敢激进。我个人的经验是设计指令集时最好有编译器工程师坐在旁边随时问一句话这个字段你生成器真的生成得出来吗很多硬件团队把指令集设计成自认为完备的状态结果编译器后端一写发现某个模式根本无法高效表达到头来要么支持要么让硬件吞下一条废指令。这种返工成本极高必须在指令集定稿前用真实 workload 的 trace 去验证。3. 数据流设计真正的性能瓶颈全在搬数据如果只能从一个维度判断 AI 芯片软硬件设计的好坏我会看它对数据流的理解深度。算力是死的一颗芯片有多少 MAC 摆在那设计定稿就固定了。但数据怎么流动在哪个环节被复用了几次决定了这些 MAC 能不能持续吃饱。3.1 卷积算子的滑动窗口和片上 buffer 的配合拿一个经典的 3x3 卷积来拆。输入 feature map 是 HxWxC输出同样尺寸卷积核在 H 和 W 方向上滑动。朴素实现是每算一个输出点都要把输入窗口里的 3x3xC 数据从 DRAM 读一遍。而相邻两个输出点的窗口有重叠2/3 的输入数据被重复读取了。这张局部性就是块设计的核心依据。软件侧编译器要做的是把整张 feature map 切成一个个 tile让一个 tile 的输入数据尽可能地留在片上 buffer 里反复复用只有 tile 边界的数据需要重新从 DRAM 加载。硬件侧要做的是在 buffer 到 MAC 阵列之间设计一个能按滑动窗口规律循环送数的数据路径。两者必须严丝合缝。我见过一个案例编译器输出一个 tile 的循环顺序是 W 方向优先但硬件的数据路径是按 H 方向优先流水设计的结果每次流水线都要重新建立MAC 利用率掉了一半。最后是软件改了循环嵌套顺序才解决。3.2 Transformer 的注意力矩阵为什么是硬件的噩梦Transformer 里的 Attention 计算有一个让硬件很头疼的特征全局依赖。Q 和 K 的每个位置都要互算点积得到一个 NxN 的 attention 矩阵N 是序列长度然后还要对每一行做 SoftmaxSoftmax 需要整行的最大值和归一化分母。这意味着 Softmax 操作不能像卷积那样分块流水必须等一整行的点积结果都齐了才能开始。这在软硬件设计上带来两个难题一是 attention 矩阵本身占空间。序列长度 2048 时单头的 attention 分数矩阵是 2048x2048x4 byte 16MB片上 SRAM 根本放不下必须分块。二是 Softmax 的两次遍历特性要么把中间结果存下来再读一遍要么用在线 Softmax 算法只存统计量但这要求数据路径支持额外的临时寄存器沟通。设计时如果没考虑到这种全局依赖算子数据流就是线性的计算-搬运-计算遇到 Attention 就会频繁打断。3.3 循环变换是软件侧性价比最高的优化手段软件侧能做的数据流优化本质上是排列组合循环顺序。核心方法论就是两个循环分块tiling和循环重排permutation。以矩阵乘 C A x B 为例如果不做 tiling内层循环每算一个输出元素都要把 A 的一行和 B 的一列搬进片上访问量是 O(N^3) 的数据量。做了 tiling 之后让 A 的一个 tile 和 B 的一个 tile 驻留在片上反复复用DRAM 访问量可以降到 O(N^3 / tile_size)。tile 取 32x32 时理论 DRAM 流量能降两个数量级。编译器干的核心活之一就是自动找出最优 tiling 参数和循环顺序。我做算子优化时给团队的硬性要求是任何算子的第一版优化先别碰汇编先把循环变换和 tiling 做对再看 profiler 数据。因为大半的带宽问题都是循环写得太粗糙、数据复用率太低导致的根本不是汇编技巧层面的事。4. 算子库与编译器自动化与手写之间的阵地战AI 芯片软件栈的终极目标是让上层框架随便扔一个模型下来编译器就能产出一份接近手写质量的算子实现。这个目标到今天也没完全实现。现实是所有量产芯片的软件栈都是编译器 手写算子库 图优化三者的拼盘。4.1 静态图、动态图与 JIT三种模式对应三种硬件压力PyTorch 的动态图模式对 AI 芯片软件栈是最不友好的。算子边界在运行时才确定芯片没法提前做全局的算子融合和内存规划。所以现在的通行做法是引入torch.compile或FX Graph这类机制先把模型变成一张静态图再做硬件相关的优化。静态图对硬件的意义非常直接编译器能看到全局数据流能够提前规划所有中间 tensor 的放置位置可以决定哪些算子在硬件上融合成一个大算子哪些 buffer 可以复用。这就是为什么做 AI 芯片的公司必养一个图优化团队的原因。以我接触过的芯片软件栈来说图优化阶段对端到端性能的贡献通常比算子本身还大尤其是那种算子小而多的网络融合几个算子就能省掉好几轮 DRAM 往返。4.2 算子融合为什么是必选项一个 Conv-BN-ReLU 的实例Conv、BatchNorm、ReLU 这三个算子在推理模式下是链式关系Conv 输出 - BN 缩放平移 - ReLU 截断。如果分开实现每个算子的输出都得落到 DRAM下一个算子再读回来两段流量浪费掉了。融合的做法是把 BN 的缩放系数融合进 Conv 的卷积核里ReLU 作为激活函数挂在 Conv 算子的输出端。硬件实现时MAC 阵列计算完一个输出点直接经过一个激活单元再写回 buffer完全不用出片。一次 DRAM 访问变成了零次。对于 CNN 类模型三轮这样的融合端到端延迟能降 30% 以上。这个优化的执行者可以是图优化器也可以是手写算子库。区别在于图优化器做的是模式匹配碰到 Conv-BN-ReLU 就匹配替换手写算子库做的是把这个融合后的算子实现调优到极致。一个真正可用的软件栈两者都要。4.3 算子库的维护成本被严重低估很多芯片团队在写软件栈规划时把算子库的工程量估成照着 CUDA 抄一遍就完事。这是最常见、也最致命的误判。实际问题有这些每个算子的 shape 组合爆炸。卷积有 1x1/3x3/5x5/7x7有 depthwise/group/strided/dilated有各种 NCHW/NHWC 布排每个组合都要验证性能达标。硬件特性会变。芯片做了一次小改版buffer bank 数从 4 改到 8所有手工调过 bank 冲突的算子都得重新调一遍。新模型引入新算子。大模型时代 FlashAttention、RoPE、GQA 这类新算子层出不穷算子库必须持续投入人力去实现和调优。我见过的比较务实的方案是算子库只覆盖高频热点算子剩下的紧急算子靠编译器 JIT 生成兜底后续再人工迭代优化。这样既保证上线速度又可持续打磨性能。把算子库团队和编译器团队放在同一个项目组里而不是分属不同部门能省掉大量沟通和推诿的成本。5. 软硬件协同设计的落地流程从需求到流片前的产线分工前面的部分都是在讲理念和原理最后这部分聊聊我实际推进一颗 AI 芯片软硬件协同设计时用的那套可复现的流程。这个过程有明确的分工边界也有大量需要两边摊开来说的灰色地带。5.1 第一步定义 workload 矩阵任何协同设计都要从业务场景开始。我的做法是把目标场景拆成一张 workload 表模型类型CNN 为主还是 Transformer 为主混合占比多少。算子分布哪些算子占了端到端时间的 80%这些算子的特征访存密集还是计算密集。精度要求是否必须 FP16 全精度INT8 的量化误差能否接受。部署形态端侧低功耗还是云端高吞吐batch size 多大。这张表直接决定硬件取舍的方向。比如以 Transformer 为主那么对 tiling 灵活性和全局 buffer 的需求就远高于对大 MAC 阵列的需求以 CNN 为主3x3 卷积的滑动窗口数据路径就是最优先实现的部分。5.2 第二步架构探索阶段的软硬件联合建模在写 RTL 之前先建一个 cycle approximate 的性能模型同时让编译器团队基于虚拟指令集做 prototype。性能模型能估算每个算子在给定 buffer 大小、带宽、MAC 数量下的耗时编译器 prototype 能暴露指令集的表达死角。这一步跑几周比流片后发现问题再改版便宜一整个量级。我在上一个项目里吃过亏硬件团队拍板把 L2 buffer 砍掉一半省面积性能模型跑 MatMul 显示只慢了 8%大家都觉得能接受。结果真实模型跑起来Attention 矩阵的中间结果根本放不下被迫拆成小块性能直接掉了 40%。原因就是性能模型里没建模中间结果缓冲不足导致的额外 DRAM 流量。从那以后我的性能模型必定包含所有中间 tensor 的 buffer 容量检查绝不给这种侥幸留空间。5.3 第三步软硬件接口的冻结仪式软件栈团队和硬件团队必须定期做接口冻结审查。审查内容不是口头对齐而是把交付物明确下来指令集手册的最终版本每条指令的语义、延迟、限制条件。地址映射和内存模型包括 buffer 的地址空间划分和权限。中断和同步语义多核任务的启动方式。硬件的可配置寄存器清单软件初始化和低功耗流程的配置方式。我坚持做接口冻结的时间点是 RTL 冻结前至少两个月。晚于这个时间硬件任何小改动都可能导致编译器后端代码重写而编译器后端代码往往是软硬件整个链条里最复杂、最晚开始写的部分。5.4 踩坑记录三条写进团队文档的教训这些教训我在前几期提过一些这一期再集中补三个典型的。一是不要相信硬件支持软件后面再说。对 AI 芯片来说软件栈的开发周期和硬件一样长。硬件支持一个特性只要有文档和 RTL软件支持这个特性需要编译器、算子库、runtime 三层的改动和验证。凡是被推迟到后面再说的功能最后基本都会被砍掉因为没有人有时间补上这个窟窿。二是不要用单算子的 best case 衡量整体。MatMul 跑到了 90% 利用率不代表整体性能好。我要求所有性能报告必须附带端到端模型数据以及对应 profiler 的算子级耗时占比。只有这样才能看出每个算子的优化是否值得投入。三是程序中尽量降低同步粒度。多核协同执行一个大算子时如果每个核都在每次算子边界做全局同步会引入大量等待。成熟的做法是基于数据依赖做细粒度的 producer-consumer 同步机制让核与核之间直接通过 buffer 状态标志传递数据而不是每次都走全局 barrier。这个方案的实现复杂度高但收益非常直接。我第一次把一个大卷积切分成 4 核执行时全局 barrier 版本效率只有 65%改成细粒度同步后直接拉到了 85%。6. 写在最后的个人体会这个系列写到这里我已经把 AI 芯片从算力单元、缓存、互联、工具链到软硬件接口的全链条都过了至少一遍。每次写类似指令集粒度多大的性价比最高这种问题我给出的答案都是取决于你的软件栈有多强。这句话不是和稀泥而是这个行业的真实规律。软硬件设计做得好的芯片往往是那颗从一开始就让硬件工程师和编译器工程师坐在同一张桌子前、用同一套 benchmark 数据去迭代的芯片。如果你正在做一个 AI 芯片相关的项目哪怕只是定义一个小的加速模块我的建议是从第一天就把谁负责指令集语义定义谁负责编译器后端谁负责算子库性能验收这三件事定下来。然后在每个里程碑节点问一句现在的硬件能让编译器生成干净的代码吗如果不能那就是设计出了毛病趁早改。最后补一个工具层面的小技巧架构探索阶段在性能模型和 RTL 之间做一个对比校准的脚本每个月跑一遍真实 workload 的 trace把性能模型预测的耗时和 RTL 仿真的实际耗时放在同一张表里对比。误差超过 10% 的点必须有人跟进解释。这颗胶水环节看起来不起眼却能让架构团队和验证团队之间少吵无数架这是我做过的所有项目里投入产出比最高的一项基建工作。