ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

模型编译器如何反向定义AI芯片架构:软硬协同设计实战

模型编译器如何反向定义AI芯片架构:软硬协同设计实战

1. 从一次流片失败说起:模型编译器凭什么能“反向定义”芯片

三年前我参与过一个边缘推理芯片的早期架构评审,当时团队里做硬件的同事普遍有个朴素认知:芯片是地基,编译器是装修队,地基打好了装修随便搞。结果第一版流片回来,跑自家模型还行,一跑客户的主流检测网络,实测吞吐只有理论峰值的18%。复盘会上做编译的同事甩出一张图,说你们这个片上缓存的分块策略根本没考虑卷积算子在实际调度里的复用距离,硬件参数定死了,我在上面怎么排都排不出好流水。那次之后我彻底信了一件事:模型编译器不是芯片的附属品,它和芯片架构之间是互相塑造的关系。

这篇文章想聊的就是这个事——模型编译器为什么能影响AI芯片的架构设计。我会从编译器到底在做什么、它和硬件之间的接口长什么样、哪些编译决策会倒逼硬件改架构、以及实际做架构评估时怎么把编译器拉进来一起迭代这几个角度展开。适合正在做AI芯片架构定义、算子库开发、推理框架适配的同行看,也适合想理解“软硬协同”到底协同在哪一层的人。全文基于我参与过的几个推理芯片项目经验,涉及具体参数的地方我会说明计算逻辑,但不会涉及任何具体公司的产品细节。

先把核心结论摆出来:编译器是模型和硬件之间唯一的翻译层,它决定了模型的算力需求以什么形态落到硬件上。硬件架构里每一个可配置项——缓存层级、数据位宽、计算阵列形状、片上网络拓扑——最终都要通过编译器的调度策略来兑现价值。如果编译器无法把模型映射好,再豪华的硬件参数也是纸面峰值。这就是为什么现在做AI芯片,架构团队和编译器团队必须坐在一起定规格,而不是先定硬件再让编译器去适配。

2. 模型编译器到底在做什么:从计算图到硬件指令的完整链路

2.1 编译器的三层翻译职责

很多人对模型编译器的理解停留在“把PyTorch模型转成芯片能跑的格式”,这个理解太浅了。一个完整的模型编译器至少承担三层翻译职责,每一层都会对硬件提出不同的要求。

第一层是图级优化。模型进来是一张计算图,节点是算子,边是张量。编译器要做的第一件事是图变换:算子融合(把Conv+BN+ReLU合成一个算子)、常量折叠、死代码消除、布局转换。这一层决定了最终有多少个算子要落到硬件上,以及它们之间的依赖关系。举个例子,一个典型的残差块如果不做融合,是Conv、BN、ReLU、Add四个独立算子,融合之后可能变成一个复合算子。算子数量直接决定了硬件调度器的复杂度和指令发射开销。

第二层是算子级调度。每个算子怎么在硬件上切分、怎么排流水、数据怎么在存储层级之间搬运,都在这一层决定。比如一个卷积算子,编译器要决定输出通道怎么分块、输入特征图怎么复用、权重是常驻还是流式加载。这一层的决策直接对应硬件的缓存大小、DMA通道数、计算阵列的并行度。

第三层是指令级生成。把调度方案翻译成具体的硬件指令序列,包括计算指令、访存指令、同步指令。这一层决定了指令集的形态——是VLIW还是RISC,有没有专门的矩阵乘指令,同步是用屏障还是用信号量。

注意:这三层不是严格串行的,现代编译器会在层与层之间反复迭代。图级优化要考虑算子级调度的可行性,算子级调度要预估指令级生成的代价。这种跨层迭代正是编译器复杂度的来源,也是它能影响架构的原因——因为它在多个抽象层级上同时做决策。

2.2 编译器的“成本模型”才是真正的话语权

编译器内部有一个东西叫成本模型(Cost Model),它决定了编译器在面临多个调度方案时选哪个。成本模型里通常包含这几类代价:计算代价(这个方案需要多少个计算周期)、访存代价(需要搬多少数据、搬多远)、同步代价(需要多少次等待)、以及片上存储占用。

关键点在于:成本模型的权重设置,直接反映了编译器对硬件瓶颈的判断。如果成本模型里访存代价权重很高,编译器就会倾向于选择数据复用率高的调度方案,这会倒逼硬件增加缓存或者优化片上网络。如果计算代价权重高,编译器就会倾向于把计算阵列填满,这会倒逼硬件增加计算单元或者优化指令发射。

我见过一个案例,某团队的成本模型最初把计算代价和访存代价设成1:1,结果编译器总是选择把大算子整个加载到片上再算,导致片上缓存需求极大。后来把访存代价权重调高到3:1,编译器立刻转向分块流式的调度策略,片上缓存需求降了一半,但要求DMA有更高的并发能力。你看,改一个成本模型的权重,硬件规格就得跟着动。这就是编译器影响架构的最直接路径。

2.3 为什么说编译器是“架构的镜子”

有个比喻我觉得很贴切:编译器是架构的镜子。你设计了一个什么样的硬件,编译器就会照出一个什么样的调度方案。反过来,如果你对编译器提出某种调度要求,它就会照出硬件需要补什么能力。

比如你希望编译器能自动做算子融合,那硬件就得支持融合后算子的高效执行——可能需要更大的寄存器文件来暂存中间结果,可能需要更灵活的指令流水。如果你希望编译器能做动态形状推理,那硬件就得支持运行时的形状计算和动态调度,不能全是静态编译。

所以做架构定义的时候,我习惯先问编译器团队三个问题:你们现在最头疼的调度约束是什么?你们希望硬件提供什么原语来简化调度?如果给你们一个理想硬件,你们会怎么设计?这三个问题的答案,往往比架构师自己拍脑袋定的规格更接近真实需求。

3. 编译器决策如何倒逼硬件架构:四个真实的影响路径

3.1 算子融合策略决定片上存储层级设计

算子融合是编译器最常用的优化手段,但它对硬件的影响极其深远。融合意味着中间结果不写回主存,直接在片上传递。这就要求硬件有足够的片上存储来暂存这些中间结果,而且这些存储的访问模式要匹配融合算子的数据流。

举个具体的例子。假设编译器要把一个1x1卷积和一个3x3深度卷积融合。1x1卷积的输出是3x3深度卷积的输入,如果融合执行,这个中间张量必须留在片上。假设中间张量是64通道、56x56分辨率、INT8精度,那就是64×56×56=200704字节,约200KB。如果硬件片上缓存只有128KB,这个融合就做不了,编译器只能放弃融合,把中间结果写回主存再读回来,带宽开销直接翻倍。

所以编译器团队会向架构团队提需求:片上缓存至少要能容纳典型融合模式下的最大中间张量。架构团队再去权衡:是增加SRAM面积,还是接受融合率下降带来的带宽开销。这个权衡的输入参数,就是编译器给出的融合模式统计。

更细一点,融合还影响存储的bank划分。如果融合算子的数据流是逐通道处理的,那片上缓存按通道分bank就能并行访问;如果数据流是逐像素处理的,那按像素分bank更合适。编译器在调度时会根据硬件的bank结构来调整数据排布,但如果硬件bank结构设计得不好,编译器再怎么调也调不出高并行度。

3.2 数据复用模式决定计算阵列的形状

计算阵列是AI芯片的核心,它的形状(比如128x128的脉动阵列,或者16x16的乘加树)直接决定了矩阵乘的吞吐。但阵列形状不是拍脑袋定的,它要匹配编译器能实现的最高数据复用模式。

矩阵乘的数据复用有三种基本模式:输出 stationary(输出结果留在计算单元里,输入和权重流过)、权重 stationary(权重留在计算单元里,输入和输出流过)、输入 stationary(输入留在计算单元里,权重和输出流过)。不同的复用模式对阵列形状的要求不同。

如果编译器主要采用权重 stationary 模式,那阵列的每一列可以共享同一个权重,阵列形状可以做得“瘦高”,比如32x256,这样权重广播的代价低。如果编译器主要采用输出 stationary 模式,那阵列的每个计算单元要独立累加,阵列形状可以做得“矮胖”,比如256x32,这样输出累加的并行度高。

我参与过一个项目,硬件团队最初定了一个64x64的方形阵列,觉得对称好看。编译器团队评估后说,我们的主要模型是深度可分离卷积,深度卷积部分的权重复用率极低,方形阵列会导致大量计算单元空转。后来改成16x256的长条形阵列,深度卷积的利用率从35%提升到72%。阵列形状不是美学问题,是编译器调度可行性的问题。

3.3 指令调度粒度决定控制通路的复杂度

编译器生成的指令序列,其调度粒度直接影响硬件控制通路的复杂度。如果编译器做的是静态调度(所有指令的执行顺序在编译时确定),那硬件控制通路可以做得非常简单,不需要复杂的动态调度逻辑,省面积省功耗。但静态调度的前提是编译器能准确预估每条指令的执行时间,这要求硬件的行为高度确定性——没有缓存缺失导致的随机延迟,没有资源冲突导致的动态等待。

如果编译器做的是动态调度(硬件在运行时决定指令发射顺序),那硬件就需要一个复杂的调度器,面积和功耗都会上去,但能更好地处理不确定性。这两种选择对应完全不同的架构风格。

实际项目中,我倾向于混合调度:计算指令用静态调度,保证流水线效率;访存指令用动态调度,容忍主存延迟的不确定性。这就要求硬件在计算通路和访存通路之间有一个解耦的接口,比如用FIFO或者令牌环来传递数据。这个接口的设计,就是编译器调度策略和硬件架构之间的契约。

提示:如果你在做架构评估,一定要让编译器团队给出指令调度的粒度需求。他们说要“细粒度动态调度”和说要“粗粒度静态调度”,对应的硬件复杂度可能差三倍。

3.4 动态形状支持决定运行时的设计

现在的模型越来越倾向于动态形状——输入分辨率可变、序列长度可变、batch size可变。编译器要支持动态形状,就不能把所有东西都静态编译死,必须留一部分决策到运行时。这就要求硬件支持运行时的形状计算、动态的算子选择、以及可变的资源分配。

具体来说,动态形状支持对硬件的影响包括:需要可配置的计算阵列(比如可以关闭部分计算单元来匹配小形状)、需要可变的片上缓存分配(不同形状需要不同大小的缓存)、需要运行时的指令生成或者指令patch机制。

我见过一个设计,为了支持动态形状,在芯片上放了一个小型的微控制器来跑运行时的调度逻辑。这个微控制器的存在,直接改变了芯片的架构——它需要自己的指令存储器、自己的数据通路、以及和主计算通路的同步机制。如果编译器不需要支持动态形状,这个微控制器完全可以省掉,芯片面积能小15%。

4. 实操:怎么在架构定义阶段就把编译器拉进来

4.1 建立编译器-架构联合评估流程

很多团队的流程是:架构团队定规格→硬件团队实现→编译器团队适配→发现性能不达标→回头改架构。这个流程的周期太长,一次迭代至少六个月。我推荐的做法是联合评估,在架构定义阶段就让编译器团队参与,用编译器的成本模型来评估架构候选方案。

具体流程是这样的:架构团队提出几个候选架构(比如不同缓存大小、不同阵列形状的组合),编译器团队针对每个候选架构跑一遍典型模型的编译,输出预估的吞吐、带宽需求、片上存储占用。然后两个团队一起看结果,找出瓶颈在哪儿,再调整架构参数。

这个流程的关键是编译器要能快速评估一个架构,不能等硬件RTL写好再评估。所以编译器团队需要有一个架构模拟器,能根据架构参数快速估算编译后的性能。这个模拟器的精度不需要很高,但趋势要准——能区分出A架构比B架构好20%还是差20%。

4.2 用“编译器友好度”作为架构评估指标

除了传统的PPA(功耗、性能、面积)指标,我建议加一个“编译器友好度”指标。这个指标怎么量化?可以从这几个维度打分:

维度评估方法权重建议
算子覆盖率编译器能高效映射的算子比例25%
调度灵活度编译器可选的调度方案数量20%
存储适配度片上存储与典型融合模式的匹配程度20%
指令效率生成指令的填充率(有效指令/总指令)20%
动态支持对动态形状和动态调度的支持程度15%

这个打分表不是绝对的,但能强迫架构团队从编译器的视角看问题。我实际用下来,算子覆盖率和存储适配度是最容易出问题的两个维度。很多架构在纸面上算力很强,但编译器一跑发现30%的算子映射不了,只能回退到CPU,整体性能直接崩掉。

4.3 典型模型的编译预演

在架构定稿之前,一定要做一次典型模型的编译预演。选5到10个有代表性的模型——覆盖卷积、Transformer、循环网络等不同结构——让编译器团队做完整的编译,输出每个模型的调度方案和资源需求。

预演的重点不是看绝对性能,而是看资源需求的分布。比如卷积模型可能对片上缓存需求大,Transformer模型可能对矩阵乘阵列的利用率要求高,循环网络可能对指令调度的灵活性要求高。如果某个架构在某一类模型上资源需求爆表,那就要针对性调整。

我踩过的一个坑是:预演时只用了小模型,觉得资源需求很宽松,结果架构定稿后跑大模型发现片上缓存完全不够。后来学乖了,预演必须包含至少一个“压力测试模型”——参数量或者计算量接近目标场景上限的模型。

5. 常见问题与排查技巧实录

5.1 编译器映射失败:算子不支持怎么办

这是最常见的问题。编译器在映射某个算子时发现硬件没有对应的计算原语,或者原语的精度/形状不匹配,只能报错或者回退。

排查思路是这样的:先看这个算子在模型里的占比。如果占比很低(比如不到5%),回退到CPU或者用近似算法都能接受。如果占比很高,那就必须解决。解决路径有两条:一是改硬件,增加对应的计算原语;二是改编译器,用多个已有原语组合出这个算子。

我倾向于优先改编译器,因为改硬件的周期太长。但改编译器也有代价——组合出来的算子可能效率很低。这时候就要算一笔账:组合算子的效率损失,和增加硬件原语的面积代价,哪个更划算。

注意:有些算子看起来简单,但组合起来非常低效。比如某些激活函数,用查表法实现需要额外的存储,用多项式近似需要多个乘加。如果模型里这类算子很多,还是值得加硬件原语的。

5.2 性能不达标:是硬件问题还是编译器问题

性能不达标的时候,最容易出现的情况是硬件团队和编译器团队互相甩锅。硬件说编译器调度不行,编译器说硬件参数不合理。这时候需要一个客观的排查方法。

我的做法是分层定位。先看计算阵列的利用率,如果利用率低于50%,说明调度没把阵列填满,大概率是编译器的问题。如果利用率高于80%但整体性能还是不行,说明阵列本身太小或者频率太低,是硬件的问题。如果利用率在50%到80%之间,那可能是存储带宽或者指令发射的问题,需要进一步看访存通路的利用率。

这个分层定位法能快速缩小问题范围。我实际用下来,大部分性能问题最终都落在存储带宽上——计算阵列利用率很高,但数据供不上,阵列在等数据。这时候要么加带宽,要么改编译器的数据复用策略。

5.3 常见问题速查表

现象可能原因排查方法解决方向
算子映射失败硬件缺少对应原语统计失败算子的模型占比改编译器组合或加硬件原语
阵列利用率低调度方案不佳看阵列空闲周期分布调整成本模型权重
带宽瓶颈数据复用率低看DMA利用率改分块策略或加缓存
指令填充率低指令集设计不合理看指令序列的有效指令比例优化指令集或改调度粒度
动态形状性能骤降运行时开销大对比静态和动态的指令数优化运行时或加硬件支持

5.4 几个独家避坑技巧

第一个技巧:在架构定义阶段就建立编译器回归测试集。选一批典型模型,每次架构参数调整后都跑一遍编译,看性能预估的变化。这样能及时发现架构调整的负面影响。

第二个技巧:成本模型的权重不要拍脑袋定,用实际数据拟合。收集一批已知性能的调度方案,用回归分析反推成本模型的权重。这样成本模型才能真实反映硬件瓶颈。

第三个技巧:给编译器留“后门”。有些调度决策编译器自动做不好,允许开发者手动指定。比如手动指定某个算子的分块大小,或者手动指定某个张量的存储位置。这个后门不需要很优雅,但能解决很多自动调度的死角。

6. 一个具体的架构评估案例:从编译器反馈到规格修改

6.1 初始架构方案与编译器评估结果

假设我们有一个初始架构方案:128x128的脉动阵列,片上缓存256KB,DMA带宽64字节/周期,指令集支持矩阵乘和向量运算。编译器团队拿这个方案跑了三个典型模型:一个卷积网络、一个Transformer、一个混合模型。

评估结果是这样的:卷积网络的阵列利用率78%,Transformer的利用率只有42%,混合模型55%。Transformer利用率低的原因是矩阵乘的K维度(也就是规约维度)比较小,128x128的阵列在K维度上填不满,大量计算单元空转。

6.2 编译器提出的架构修改建议

编译器团队提了两个建议:一是把阵列改成可配置的,支持128x128和64x256两种模式,小K维度时用64x256模式;二是增加一个专门的向量归约单元,处理Transformer里的LayerNorm和Softmax。

架构团队评估后认为,可配置阵列的面积代价约8%,向量归约单元的面积代价约3%,但Transformer的利用率能从42%提升到70%以上,整体收益远大于代价。最终方案采纳了这两个建议。

6.3 修改后的评估结果与经验总结

修改后的架构跑同样的三个模型:卷积网络利用率75%(略降,因为可配置阵列有少量开销),Transformer利用率71%,混合模型68%。整体加权性能提升约35%。

这个案例给我的经验是:编译器的反馈不是“锦上添花”,而是“雪中送炭”。如果没有编译器的评估,这个架构可能就带着Transformer性能的短板流片了,上市后才发现问题,代价就大了。

另外还有一个体会:编译器团队提的建议往往很具体,因为他们是从调度可行性出发的。架构团队要做的不是照单全收,而是理解建议背后的调度需求,然后找到面积、功耗、性能的平衡点。比如可配置阵列不一定非要支持两种模式,也许支持三种模式面积代价只多2%,但能覆盖更多模型。

7. 编译器与架构协同的未来方向

7.1 编译器驱动的架构搜索

现在有一些研究在做编译器驱动的架构搜索——给定一组模型和一组编译器优化策略,自动搜索最优的硬件架构参数。这个方向的思路是:把编译器的成本模型嵌入到架构搜索的循环里,每评估一个架构候选,就用编译器跑一遍映射,用映射结果作为架构的适应度。

这个方向我觉得很有前景,但目前的瓶颈是编译器的评估速度。跑一次完整编译可能要几分钟到几十分钟,架构搜索需要评估成千上万个候选,时间上不现实。所以需要快速评估模型——用解析模型或者机器学习模型来预估编译结果,只在最后几个候选上做完整编译。

7.2 可编程性更强的硬件原语

另一个方向是硬件提供更可编程的原语,让编译器有更大的调度空间。比如提供可配置的SIMD宽度、可配置的缓存层级、可配置的互连拓扑。这样编译器可以根据模型特点动态调整硬件行为,而不是被固定的硬件规格限制死。

但可编程性是有代价的——面积、功耗、验证复杂度都会上升。所以关键是找到“恰到好处”的可编程性:既能覆盖主要模型的调度需求,又不至于过度设计。这个度的把握,还是得靠编译器和架构团队的联合评估。

7.3 对从业者的建议

如果你在做AI芯片架构,我的建议是:把编译器团队当成你的第一用户,而不是下游适配方。在定义规格的时候,多问编译器团队“这个参数你们怎么用”、“这个原语你们需要吗”、“如果改成这样你们会不会更好调度”。他们的回答会帮你避开很多坑。

如果你在做模型编译器,我的建议是:不要只盯着软件优化,要理解硬件的行为。知道硬件的缓存怎么工作、阵列怎么调度、指令怎么发射,你才能写出真正高效的调度方案。而且当你发现硬件设计不合理的时候,要敢于提出来——编译器是离模型最近的人,你最清楚模型需要什么样的硬件。

这个领域还在快速演进,新的模型结构、新的硬件形态、新的编译技术都在不断出现。但核心逻辑不会变:编译器和架构是共生的,谁离了谁都不行。把这两件事放在一起想,才能做出真正好用的AI芯片。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进