AI芯片软硬件协同设计：从算子需求到硬件映射的工程实践-开发者社区

这些年一直在做AI芯片相关的软硬件协同设计工作，从最初的算子开发到后来的全栈工具链，踩过的坑不少，也总结了一些自己的方法论。这系列文章到这里已经是第6篇了，前面聊过架构选型、指令集设计、存储层次优化，这篇我想换个角度，专门聊聊从“算子需求”到“硬件映射”这一整条链路到底该怎么落地，以及在工程实践中那些文档里不会写、但能决定项目成败的细节。无论你是刚接触AI芯片设计的学生、准备转行做加速器软件的工程师，还是已经在做ISP、NPU、GPU相关项目的开发者，这篇文章应该都能给你一些参考。

1. 软硬件协同设计的第一性思考

1.1 需求侧的边界到底该由谁定义

做AI芯片的软硬件协同设计，第一个绕不开的问题是“边界”。很多项目一开始就搞错了顺序——先定硬件架构，再回头适配算子，结果就是软件侧天天给硬件擦屁股。我在实际项目里吃过这个亏，所以我现在的原则很明确：先梳理算子集，再定义硬件抽象层，最后才谈微架构。

为什么要这个顺序？因为AI芯片的本质是一个“面向领域特定计算”的加速引擎，它和通用CPU有本质区别。CPU追求指令集通用性，而NPU、GPU这类芯片追求的是在有限面积和功耗下，把特定计算模式的效率拉到极致。如果你不从算子需求出发，你根本不知道到底该给矩阵乘分配多少MAC阵列、给激活函数分配多少特殊函数单元、给数据搬运分配多少带宽。

以常见的卷积神经网络为例，一个典型的主干网络里，卷积算子占比约为90%以上的计算量，但不同类型卷积（标准卷积、深度可分离卷积、空洞卷积、转置卷积）对硬件的要求差异极大。深度可分离卷积里那个 1x1 逐点卷积，本质是个矩阵乘法，它对MAC阵列的利用率要求很高，而3x3的深度卷积则更依赖数据复用策略和缓存带宽。如果一开始只盯着一个算子设计，后边切换模型时性能会崩得很难看。

1.2 从算子到底层原语的解构路径

那么，怎么把一个高层的算子集优雅地映射到硬件上？我的做法是拆成“计算图算子 → 计算原语 → 微架构操作”三层。

计算图算子就是我们说的 Conv2d、MatMul、Softmax 这些。计算原语是硬件真正执行的最小计算单位，例如矩阵乘法外积原语、向量点积原语、规约求和原语。微架构操作则是这些原语在硬件上对应的实际数据流动、计算、写回动作。

举个实际的例子，一个 Conv2d 算子，软件侧会经过如下解析：

把卷积展开成矩阵乘（im2col）或者隐式GEMM（Implicit GEMM）模式；
确定 NHWC 还是 NCHW 的数据布局，这直接影响后续维度的访问连续性；
把整个计算图按 batch、通道、空间位置切分成多个 tile，便于在SRAM里放得下；
映射到计算原语，比如“按 K 维做 16x16 外积累加”；
最后生成硬件能理解的指令流，包括DMA搬运、矩阵计算、向量激活和写回。

在整个过程中，最容易翻车的其实是数据布局的选择。很多人觉得 NCHW 还是 NHWC 只是软件层的习惯区别，但在芯片上这件事直接决定你的DMA能不能 burst 传输、能不能命中片上缓存。NCHW 在训练框架里友好，但 NHWC 在推理芯片上往往存储更连续，特别是做逐通道量化、通道剪枝或者硬件支持向量运算时，NHWC 能减少大量的数据重排开销。我看到过不少项目，硬件设计得很先进，结果因为软件侧坚持用 NCHW，导致片内带宽利用率只有40%左右。

2. 算子映射的工程落地

2.1 计算图优化与融合策略

算子映射的第一步不是直接生成指令，而是对计算图做一遍“编译期优化”。计算图融合是这中间收益最大的一步，典型的是把“卷积 + 批归一化 + ReLU + 量化”融合成一个算子。

以推理场景为例，BatchNorm 在推理阶段可以折算成对卷积权重和偏置的缩放：假设卷积的权重为 W，BN层的缩放因子为 gamma，均值为 mean，方差为 var，那么融合后的等效权重就是 W * gamma / sqrt(var + eps)，等效偏置是 (b - mean) * gamma / sqrt(var + eps) + beta。这一步做完之后，前向计算里就不需要单独做 BN 的方差归一化了，省掉的是一次完整的数据读改写。

ReLU 融合更是常规操作。如果你的NPU的计算单元在输出阶段支持直接饱和截断，那么ReLU就是零成本的。但要注意，不是所有情况下融合都是好事。比如一些对数值范围敏感的网络，如果融合了量化操作，中间结果的位宽不够，精度会掉得很厉害。所以我通常建议：推理场景默认全融合，训练场景只做算子级别的融合，不做量化融合。

计算图优化还有一个核心点是数据格式转换的插入。硬件往往只能高效处理某一种或几种数据格式，比如推理芯片喜欢 INT8 和 INT4，训练芯片喜欢 FP16 和 BF16。从 FP32 转成 INT8 需要量化节点，这个节点的插入位置决定了精度损失。我的经验是，量化参数的校准数据要从训练集中抽一部分，覆盖不同的数值分布区间，不能只拿一两张图去统计。

2.2 访存密集型与计算密集型算子的处理差异

算子根据其瓶颈类型，映射策略完全不同。计算密集型算子（如大规模 MatMul、大通道数 Conv）应该靠加大MAC利用率来提升性能，访存密集型算子（如 Elementwise、Gather、Softmax）则应该靠减少片外访问次数来提升性能。

我做过的某款芯片上，Elementwise 算子的性能瓶颈完全不在计算，而在 DMA 搬运。如果每做一次 y = x1 + x2 就要从 DRAM 读两个数据再写一个数据，那么哪怕计算单元跑到了 1GHz，整体吞吐还是被内存带宽锁死。解决办法是：把多个连续的 Elementwise 操作合并成一次流水线处理，或者利用片上SRAM做中转，减少对 DRAM 的重复读写。

这里有一个链式融合的经典做法：把“加偏置 → 乘缩放 → 截断到int8 → 写回”四步融合成一个向量流水，中间不需要任何数据落片外。实测下来，在跑一些轻量级检测模型时，整个模型的访存时间能缩短 30% 到 40%。这个数字对于终端推理芯片来说非常可观。

Softmax 这种算子也是有讲究的。常规实现需要先找全局最大值，再算 exp 和 sum，最后做除法。在芯片上实现时，正确的做法是做一个两遍扫描：第一遍算 max，第二遍算 exp 和 sum，前提是数据要能放进SRAM。如果输入序列特别长，SRAM装不下，就得做分块Softmax，也就是 online softmax。这里有个容易踩的坑：分块softmax的 rescale 系数很容易写成浮点，但硬件上要尽量避免浮点除法，否则周期数飙升。我用的办法是把 rescale 系数定点化，用查表或者向左移位来逼近，精度损失微乎其微，但速度提升明显。

2.3 tile切分和调度顺序的定量分析

接下来聊聊算子的硬件映射中最核心的环节——tile切分。写AI芯片编译器的人都知道一句话：调度即性能。同样的计算图，tile大小不一样、循环顺序不一样，性能差距可能是一个数量级。

假设一个卷积层：输入是 1x64x56x56，权重是 128x64x3x3，输出是 1x128x56x56。我们做一个简化分析，忽略batch维。SRAM大小为 512KB，MAC阵列是 128x128。

首先判断，整个输出能不能直接放进SRAM？1x128x56x56x4字节 = 1.6MB，显然放不下。输入呢？1x64x56x56x4 = 0.8MB，也放不下。所以必须切分。

切分策略的选择上，我倾向先按输出通道切，再按空间位置切。为什么？因为权重数据是不能被共享的，每个输出通道对应一组独立的权重，如果先切空间，意味着每一块空间tile都要读取全部128组权重，权重的重复读取率太高。反过来按输出通道切，每个输出通道组只需要读取对应的那组权重，输入数据可以在多个输出通道组之间共享。

举个量化例子：把输出通道切分成4份，每份32个输出通道，对应权重大小为 32x64x3x3x4 = 288KB，加上输入块 SRAM 放得下。空间位置再切分成4x4的tile，每个tile是 1x64x14x14x4 = 200KB。这样叠加起来刚好在512KB以内。

tile调度顺序上，我的习惯是：最内层循环保持空间tile按行连续，这样DMA可以从DRAM连续搬数；中间层循环切换输出通道，充分利用输入数据的片上复用；最外层循环切换空间位置。这和CPU上做循环分块优化的思路类似，只是约束条件换成了SRAM容量和DMA burst长度。

3. 数据通路与片上存储的静态规划

3.1 存储层次与带宽的匹配关系

AI芯片的性能公式并不复杂：算力乘上数据复用率要小于等于存储带宽，否则芯片就饿死了。很多芯片标称算力很高，实际跑模型跑不到标称的一半，原因就是数据通路跟不上。

典型的AI芯片存储层次从外到内是：DRAM（以GB计，带宽通常几十GB/s到上百GB/s）、片上SRAM（以MB计，带宽可达数百GB/s到数TB/s）、寄存器文件（每周期可读多份数据，带宽最高）。一个设计良好的数据通路，关键是把最热的数据往最内层放。

举个例子，一个 DDR4-3200 双通道的带宽理论值约为 51.2GB/s。一个 128x128 MAC 阵列在 1GHz 下，以INT8计算，算力是 128x128x2x1GHz = 32TOPS。如果数据没有复用，每个MAC每周期都需要从外部读入两个新操作数，需要的带宽是 128x128x2x1B = 32GB/s。看起来刚好卡在带宽范围内，但如果你再加上权重读取、偏置读取、激活写回，带宽立刻爆炸。所以，没有片上复用方案的AI芯片，设计上就是不合格的。

空间复用是解决这个问题的核心手段。在卷积计算中，输入特征图在一个输出通道上是共享的，权重在batch维上是共享的。通过引入数据复用，可以把外部带宽需求降低一个甚至两个数量级。实际工程里，我通常会做一次理论分析：计算每个算子的算术强度（Arithmetic Intensity，单位是FLOP/Byte），低于某个阈值就判定为访存密集型，优先优化数据复用，而不是堆算力。

3.2 双缓冲机制和DMA时序规划

DMA和计算单元之间的协调，经典方案是双缓冲（Double Buffering）。原理很简单：当计算单元在处理 tile N 时，DMA在预取 tile N+1 的数据，计算结束后，两个角色对调。这样做可以把访存时间和计算时间重叠起来，理想情况下芯片的总执行时间就是纯计算时间。

但是，真正做起来有几个坑。

第一个坑是DMA描述符的开销。如果你每个tile都重新配置一次DMA，中断和同步的开销会吃掉所有流水带来的收益。我的做法是：在启动计算之前，一次性生成整条链式DMA描述符，让DMA硬件自己按照链表顺序搬运，CPU/控制器只负责启动链首。

第二个坑是对齐问题。DMA搬运时，起始地址和数据长度必须满足对齐要求，比如64字节对齐。如果你的tile大小算出来不对齐，就要重新调整，否则DMA效率骤降甚至报错。这个在切分tile阶段就要考虑到，不要等代码跑挂了才去查。

第三个坑是同步粒度。有的设计是每个tile都做一次同步，这样安全性高，但性能被同步开销拖累。我的实践经验是：只在“写回同一个SRAM地址”的边界上同步，其他的全部异步。换句话说，如果下一个tile的数据会覆盖上一个tile还没被计算单元读完的缓冲，才需要插入同步点，除此以外不加任何同步。这个优化做下来，流水线利用率能从70%提升到90%以上。

3.3 权重驻留与在线重加载的策略平衡

在实际部署场景中，模型权重的大小差异非常大。以MobileNet这类轻量模型为例，权重往往只有几MB到十几MB，如果SRAM足够大，可以把整个模型权重常驻在片上，彻底消除权重搬运。但像一些大模型或者超分辨率模型，权重动辄几十MB上百MB，SRAM塞不下，这时就要考虑分块加载。

权重驻留和在线重加载的选择，背后是一次性能与灵活性的权衡。权重常驻的优点是执行时零额外等待，缺点是模型适配性差，换成一个大模型就废了。在线重加载的优点是灵活，缺点是频繁搬运会占用DMA带宽，甚至成为瓶颈。

我见过一个比较优雅的方案：把SRAM分为权重区和数据区两个部分，权重区用离线分析工具提前算好每个模型的驻留权重集合，数据区用动态分配的方式给激活和中间结果。如果遇到超大模型，则启用“层流式加载”——当前层算完，立即释放它的权重空间，预取下一层的权重。这个方案在跑各种不同规模的模型时都能维持较高的利用率。

还要提一个容易被忽略的细节：权重的重排（Repacking）。硬件读取权重时，往往需要特定的排列格式，比如INT8的M x K矩阵要做16字节对齐切片。如果每次加载都让编译器现场转换，开销巨大。我的建议是在模型转换阶段一次性重排出硬件原生格式，直接存到闪存或者文件系统里，运行时加载后就能直接参与计算，省掉一层软件开销。

4. 编译栈与工具链的结构性设计

4.1 三层IR的设计取舍

AI芯片的编译栈通常分为三层IR：前端IR（如模型转换后的图表示）、中端IR（硬件无关的算子中间表示）、后端IR（硬件相关的指令调度表示）。这个三层结构和通用编译器的逻辑类似，但侧重点完全不同。

前端IR的核心任务是处理“框架异构”。同一个模型可能是从PyTorch导出的、TensorFlow导出的或者ONNX导出的，各个框架的算子语义存在细节差异，前端IR要把它们统一成一个规范的算子集。需要特别注意的是，不能把框架里的兜底算子（比如自定义Python函数）也照单全收，要在前端阶段就识别出来，能映射的映射，不能映射的要么拆解成子图，要么直接报错提示用户换写法。

中端IR重点做的是与硬件无关的优化，例如算子融合、死代码消除、常量折叠、代数简化。很多人认为这些优化没用，因为AI计算图已经是比较规整的了。但实际并非如此，比如推理模型里经常出现多余的Reshape和Transpose，如果能在这个阶段被消除，后端的地址计算会简单很多。

后端IR做的是指令选择、寄存器分配、循环调度、DMA规划。这个层面跟硬件高度耦合，也是不同芯片方案拉开差距的关键。一个好的后端IR设计，应该能在调度阶段显式表达同步依赖、缓冲区和计算资源使用情况，让调度器可以做全局决策。

4.2 调度器如何做全局决策

调度器是我认为整个工具链中算法含量最高的部分。它的输入是一张已经展开好循环、标注好依赖关系的IR图，输出是一段可执行的指令序列。

常见的调度策略包括列表调度（List Scheduling）、模调度（Modulo Scheduling）和整数线性规划（ILP）调度。实际工程中，全用ILP不现实——算子多了求解时间太慢。我常用的方法是分层调度：先在粗粒度上对图级tile做调度，使用启发式策略；再在细粒度上对某个算子内部的循环做模调度，追求极致II（Initiation Interval）。

模调度里经典的模数就是求“资源受限下的最小启动间隔”。比如一个卷积核内部循环有5步依赖链，而可用的MAC阵列只有2份，那么最小II就是max(计算资源限制、依赖链限制)。如果算下来II=3，意味着流水线每3个周期可以启动一个新的迭代。实际编码时，我会先用一个循环分析器算出理论II，再往下做寄存器分配，如果分配失败再逐步增大II。这个流程保证了调度结果不是随手拍脑袋，而是一个有上下界把控的优化过程。

4.3 调试工具在芯片验证中的重要性

很多做硬件出身的人不重视编译器的可调试性，直到芯片样品回来跑不出数据，才发现工具链黑盒得可怕。我在做软硬件协同设计时，坚持把工具链的“可观测性”当成一等公民来设计。

具体做法是在编译器中保留三种调试手段：第一是IR转储，每一层IR优化前后都能导出成可读的文本或者标准图结构，方便自查和回溯；第二是模拟器回放，芯片的指令集模拟器必须和编译器同步更新，任何指令改动都要在模拟器上先行验证；第三是数据比对，模拟器输出的每一层中间结果都要能转成numpy格式，方便跟CPU/GPU的浮点结果做数值对比。

芯片回来后，遇到的第一个难题往往不是性能，而是“功能错误”。有个经验是：要能精确复现到指令级。如果芯片上的输出和模拟器不一致，利用指令级回放工具可以精确锁定是哪一条指令、哪一个cycle开始分歧的。如果这一步做不好，面对几万行的指令日志和上亿次的计算结果，排错基本靠猜，效率极低。