做算子优化,最怕遇到什么?不是算子慢,而是不知道它为什么慢。很多时候我们拿到一个模型在征程6平台上跑,第一版性能惨不忍睹,打开profile一看,某个算子吃掉了一半时间。你要是这时候一拍脑袋把那个算子重写一遍,多半是白忙活——因为工具链已经把常用算子调得很不错了,真正的瓶颈往往藏在数据搬运、布局转换、统计过程中重复读取这些"看不见"的位置。这篇文章是给那些已经在征程6工具链上跑通基础部署流程、想再把性能往上拉的开发者看的,我把做车规级模型算子优化踩过的坑、用过的方案,按算子类型整理成一份集锦,希望能让你少走弯路。
1. 算子优化的三层作战地图——先搞清在优化什么
1.1 算法层:让算子"怎么算"变得更省
算法层解决的是"计算量本身能不能更小"的问题。比如把7x7卷积拆成两个串联的3x3卷积、用Winograd减少乘法次数、用多项式近似替代exp函数、把除法改成乘倒数。这一步改的是数学表达,跟硬件无关,收益上限最高,但风险也最大。
我见过不少开发者在卷积性能不好时,第一反应是上Winograd,结果发现3x3小卷积还行,5x5以上省下的乘法和额外的矩阵变换开销一对比,经常是亏的。而且Winograd对量化模型更不友好,误差容易被放大。算法层优化不是越花哨越好,而是要在速度收益和数值稳定性之间找平衡,每次改动都要跟着精度回归走。
1.2 图调度层:让算子之间少搬数据
图调度层的目标很直白:减少DDR和片上SRAM之间的数据往返。最典型的操作是算子融合。Conv+BN+ReLU如果分开执行,每两个算子之间都要把中间结果写一遍DDR再读回来;融合成一个kernel后,中间结果留在SRAM里,搬运量能降40%以上。征程6工具链通常会自动做一部分融合,但自定义算子、复杂分支结构就得靠手动处理。
这一层还有一个重要手段是layout重排。比如某些卷积在NHWC下访存更连续,而矩阵乘在特定layout下向量化更容易。layout转换本身也有成本,要比较"换layout的代价"和"换完之后获得的加速"哪个更划算。很多profile里看起来很高的算子,其实是布局切换造成的假象。
1.3 微架构层:跟向量宽度和存储层次较劲
微架构层是最琐碎、也最接近硬件的一层。核心要理解的是:征程6平台里有DDR、片上SRAM、寄存器三级存储,计算单元有向量核和矩阵核之分。一个算子的瓶颈要是出在带宽,你优化算法公式也没用;瓶颈出在乘累加单元,你优化数据布局收益也有限。
怎么判断瓶颈?很简单,看工具链给出的带宽利用率和指令利用率。如果算子实际访存量接近理论带宽上限,那是memory bound;如果算力没打满但访存量很低,那是compute bound。两类问题的解法完全不同:前者主打压缩数据规模、提高复用率,后者主打tiling和多核并行。后面所有方案本质上都是围绕这张地图展开的。
2. 在动手重写算子前,先用Profile把"慢在哪"钉死
2.1 三步完成基线:编译报告、运行时统计、带宽数据
我的建议是任何优化动手前,先花半天把基线跑清楚。第一步,打开征程6工具链的编译报告,看每个算子被映射成了什么指令、预估周期是多少。第二步,跑一次带profiling的推理,拿到各算子的实际耗时占比。第三步,把理论访存量算出来,和实测带宽对比。三份数据放一起,才敢说某个算子是真的慢,还是被其他算子的等待拖慢的。
这里有个容易被忽略的点:profile里耗时高的算子不一定是需要重写的算子。比如A算子本身很快,但它从一个布局很糟糕的输出里拿数据,结果前面的layout转换算子耗时很高。这时候真正要优化的可能是生成数据的那个算子,让它直接输出A期望的layout。
理论访存量怎么估?一个算子的访存量大致等于输入张量字节数加输出张量字节数再加权重字节数。假设输入是1x64x64x32的fp16张量,体积就是16464322约262KB,加上输出同等大小,单算子一次执行至少要从DDR走上524KB。如果工具链给的实测时间对应下来的带宽占比已经超过80%,就别折腾访存了,赶紧看能不能融合或者压缩数据。
2.2 从Profile结果到可执行的优化清单
做完基线,我会把热点算子整理成一张表。这张表会跟着整个优化过程持续更新:
| 算子 | 耗时占比 | 主要瓶颈 | 理论可压缩空间 | 优化方向 | 优先级 |
|---|---|---|---|---|---|
| conv1 | 38% | 访存不连续,带宽未打满 | 约2.1倍 | layout重排+融合 | P0 |
| softmax | 12% | 两遍扫描,重复读输入 | 约1.6倍 | online算法+融合 | P1 |
| fc | 21% | 算力打满,单核 | 约2.8倍 | tiling+多核 | P0 |
| relu | 3% | 访问延迟 | 约1.2倍 | 融合进前算子 | P2 |
表只是结果,关键是确立"先做P0、再做P1"的节奏。资源永远是有限的,把时间花在占比最高的算子上才是正路。很多团队优化半天耗时反而涨了,就是因为先去打磨了一个边缘算子,回头看主算子的缓存已经跟不上了。
3. 高频算子的六类优化方案,照着改就行
3.1 卷积:layout重排、隐式GEMM,以及Winograd的取舍
卷积是所有优化里最容易见效的。首先看数据布局:通道方向是否在内存里连续。如果当前是NCHW,试着切到NHWC或工具链推荐的变体,访存命中率经常直接上一个台阶。第二个套路是把卷积转成隐式GEMM——不要像我第一次那样显式构造im2col矩阵:那会让输入数据膨胀kh*kw倍,DDR很快被塞爆。隐式GEMM不真正构建col矩阵,而是在乘累加循环里直接按卷积滑窗偏移取数,内存占用回到正常水平。
Winograd不是不能用,但要挑场景:3x3、步长1的小卷积收益最明显,5x5以上省的那点乘法和额外的矩阵变换开销一对比,经常是亏的。另外Winograd对量化模型更敏感,误差容易放大,务必做精度回归。遇到大kernel卷积(7x7或更大)时,先别急着上Winograd,优先考虑能否拆成小卷积串联,或者直接用隐式GEMM配合通道分块。
3.2 Softmax与LayerNorm:把两遍扫描改成一遍
softmax常规实现是三步:先找最大值、再算指数和、最后归一化。这意味着输入要从DDR被读两次甚至三次。优化办法是用online softmax:一遍遍历中同时维护running_max和running_sum,遇到更大的max时对已经累加的sum做rescale,最后统一归一化。核心代码长这样:
float running_max = -FLT_MAX, running_sum = 0.0f; for (int i = 0; i < N; ++i) { float old_max = running_max; running_max = max(running_max, x[i]); running_sum = running_sum * expf(old_max - running_max) + expf(x[i] - running_max); } for (int i = 0; i < N; ++i) y[i] = expf(x[i] - running_max) / running_sum;LayerNorm同理,可以在一遍里同时累计sum和sum_of_squares,避免第二个循环。更进一步的方案是把归一化操作直接融合到全连接或卷积的输出阶段,让中间结果不出SRAM。虽然online版本比朴素版本要多做一次乘法更新,但少一遍完整的DDR读取,实测里对大序列收益非常明显。
3.3 大矩阵乘:Tiling、Packing和双缓冲
矩阵乘在征程6上的优化节奏一般是这样:先按片上SRAM大小把MxNxK切成小块,比如32x32x64量级的tile;然后对每个tile做Packing,把A、B的小块重排成内存连续、向量化友好的排列;最后用双缓冲方式预取下一块数据,让DMA搬运和向量计算重叠执行。这样处理之后,矩阵乘的耗时能比朴素实现下降3到5倍。
Packing的代价是一次额外的数据重排,但如果重排后每个小块的加载都能连续命中,整体收益远大于损失。多核拆分优先按输出维度N去切,让每个核独立计算不重叠的输出区间,最后不需要复杂的同步逻辑。tile大小需要实测:太大SRAM放不下,太小搬运逻辑开销比计算还明显。
3.4 逐点算子:向量化、循环展开,不如直接融合
ReLU、sigmoid、各种逐点算子几乎都是memory bound,它们在"DDR读进来、计算一下、写回去"这条路上花费了大部分时间。这种算子的优化空间其实很小,最有效的办法不是自己重写,而是把它复合进前一个算子或者后一个算子的kernel里。如果必须独立执行,那就用工具链提供的向量接口一次读取多个元素,并手动做一轮循环展开隐藏访存延迟。
多说一句,我在实践中发现很多人喜欢在逐点算子上抠性能,其实这块壳本身就薄,收益不大。真正的收益在把它融合掉之后省下的那次数千字节的kernel启动开销和DDR来回,这比你写一百行手写循环都有用。
3.5 大kernel与depthwise卷积:把SRAM当搬运用
depthwise卷积的特点是每个通道单独计算,卷积核又大,直接实现时每个通道都要把整个输入通道读一遍,DDR带宽容易被打满。优化思路是按输入通道分块,把一个小块的整池数据预存到SRAM里,然后在这个块上循环使用,算完再换下一块。同时给SRAM分配双份buffer,一份在算、一份由DMA往里面填,这样搬运等待基本被隐藏了。
具体分块多大,主要看两组数:kernel需要的输入通道数,以及片上的SRAM预算。我一般的做法是先算算子理论的访存量,再反推分块大小,保证每个块至少能覆盖一次完整的数据复用周期。你可以这样估算:假设片上可用SRAM为1MB,双缓冲要占用一半,剩下512KB用于数据缓存,再除以输入张量的单通道字节数,就能得出每批最多能缓存多少个通道。
3.6 Attention:Flash式分块是现阶段最稳的解
现在的模型大头基本是Attention系列算子。最直观的实现是生成完整的NxN注意力分数矩阵,然后再Softmax、再乘V。这么大的中间矩阵一旦超过SRAM容量,就得反复读写DDR,性能断崖式下跌。Flash风格的分块思路是把Q、K、V各切成小块,只在SRAM里计算某一局部的分数;Softmax的归一化统计量则用online方式在分块过程中动态更新。这一套组合拳下来,大序列的Attention能比朴素实现对折提速。
如果你的模型不需要处理特别长的序列,也可以先试试简单的饱和策略:把输入长度按固定窗口切块,分别在窗口内做attention再拼接。但注意,这会丢失跨窗口的上下文信息,不是所有任务都能接受,需要根据实际效果来定。
4. 量化与混合精度:优化加速的真正放大器
4.1 用int8/int16替换浮点时的算子改造点
把算子的核心乘累加从float换成int8之后,运算速度通常能翻一倍以上,但代价是要处理量化参数的来回搬运。工程上我会把"反量化-计算-重量化"压缩进同一个kernel,不要把量化参数做成单独的小算子满世界跑。矩阵乘的累加器一般保留int32,最后一次性requantize回int8;偏置直接加在int32阶段,省得再读一遍。
这一步经常出现的细节问题:scale和zero point的精度处理。如果激活值是非对称量化,zero point要小心地融合进偏移量计算里;权重采用对称量化时,计算还能再省一次加法。不要把所有算子都套同一套量化配置,按"权重稳定用per-channel、激活用per-tensor"起步,遇到波动大的层再单独调整。
4.2 敏感层定位与混合精度策略
不是说全模型int8就行。像某些检测头、归一化层、最终的softmax这些层,对数值变化特别敏感,一旦量化掉,召回率或坐标偏移立刻失控。我的做法是把模型按层切成几段,逐段改成int8并对比输出张量的最大绝对误差,误差突然放大的那段就是敏感层。之后对敏感层保留fp16或bf16,其他层继续走int8,性能和精度基本能同时保住。
举个例子:有一个视觉检测模型,前面backbone全量化后误差都很小,唯独到了输出header第一个卷积,最大绝对误差从0.01级别跳到了0.5。单独看那层的权重分布,发现数值跨度特别大,per-tensor量化扛不住,改成per-channel后误差又回到0.05。这个案例后来成了我在团队里讲量化必提的素材。
4.3 校准集和量化参数值得花心思
校准集的质量往往比量化算法本身更影响结果。我见过有人用几十张训练集图片做校准,跑出来精度崩了还怪int8不好用。正确做法是覆盖模型部署后会遇到的全部工况,包括白天、夜间、逆光、雨雾这类变化。如果激活值分布存在明显离群点,静态范围会被单个异常值撑爆,这时要么用动静态混合,要么对离群通道单独做per-channel量化。
校准集还有一个隐性要求:必须经过和部署时完全相同的前处理流水线。有些团队校准用了一套图像缩放参数,推理用另一套,算出来的量化范围压根不对应实际输入分布,精度自然崩。这类问题排查起来非常隐蔽,但记录好数据流就能避免。
5. 从"能算"到"上板稳":实测中拧出来的几个坑
5.1 地址对齐和bank冲突
征程6平台的向量计算单元对地址对齐有硬性要求,比如16字节或32字节对齐。自定义算子时如果指针没对齐,轻则性能下降,重则工具链直接报错。另一个坑是bank冲突:两个不同地址如果映射到同一个bank,向量单元读取时会被迫串行化。之前有同事实现的算子比理论慢三倍,最后发现就是每个通道间没加padding,连续的channel访问全部撞到同一个bank,加了一行padding后耗时立刻掉下来。
排查bank冲突有个笨办法:把算子拆开,分别测单通道访问和多通道访问的时间,如果多通道访问时间远大于单通道时间的简单累加,大概率就是bank冲突。这时候不要猜,直接在代码里对通道维度加padding重新对齐,性能往往自己就回来了。
5.2 手写向量代码被编译器优化掉的怪事
在编写高性能算子时,编译器有时会认为自己比你更懂调度,把你辛辛苦苦安排的指令顺序重新组合,甚至把某些"看似无用"的读操作直接删掉。遇到这种情况,先不要怀疑是自己写错,打开工具链生成的汇编对照一下。解决方法是使用工具链提供的intrinsic接口,或者给关键变量加volatile修饰,再不行就在编译选项里显式关闭某段代码的自动调度。
我个人的倾向是:尽量不在征程6上直接手写汇编,先从intrinsic开始。intrinsic能提供接近汇编的控制力,同时保住可读性。真要用内联汇编,也务必逐行加注释,并保存好编译产物,否则三个月后回来看代码,谁也不记得当初为什么要那么写。
5.3 多核拆分后的结果一致性
多核并行后最头疼的是浮点累加顺序变了,同一个日期的推理结果前后差了零点几个数值,明明算法没变。尤其是对于有多个输出的算子,每个核在自己的局部缓冲里算好了再拼到一起,拼的顺序不同,结果就有微小偏差。这不是bug,但对严格验收来说可能是个问题。我的建议是:如果精度余量紧张,就固定一份累加顺序;如果余量充足,就设置好cmodel或硬件在环测试的容差,别为这点差异反复返工。
更隐蔽的是不同核间因为负载不均衡导致的时序抖动。某个核先算完去忙别的,晚到的核把数据写回来,顺序自然漂移。解决思路是先确认任务切分已经均匀了,再用固定的拼装顺序,不要相信"反正都是求和"这种事。
6. 优化真正拼的不是技巧,是流程
6.1 一份可复用的优化迭代记录模板
把优化过程文档化是我吃过亏之后才养成的习惯。每轮优化只改一个变量,记录版本号、改动内容、对比基线的耗时、精度影响、是否回退原因。这张表看起来简单,但能救你于水火:很多方案单独调都是有效的,合在一起就互相打架,没有记录你根本不知道是哪个动作导致精度崩了。
| 版本 | 改动摘要 | 耗时变化 | 精度变化 | 结论 |
|---|---|---|---|---|
| v1 | baseline | - | - | 基准 |
| v2 | conv1改NHWC | -32% | +0.1% | 保留 |
| v3 | softmax改online | -5% | 0 | 保留 |
| v4 | 全模型int8 | -18% | -1.2% | 回退,敏感层保留fp16 |
记录时不要只写结论,把关键的中间数值也留一份,比如某个算子的实测带宽、SRAM用量。这样过段时间再拿到类似性能问题时,你甚至可以直接翻旧账,不用重新跑一轮profile。
6.2 我的几点操作体会
最后说几句大实话。算子优化这个活,大部分时间不是在写算子,而是在判断"该不该优化"。拿到一个性能问题,先看瓶颈在哪一层,再决定动手;动手前想清楚收益多大,优化完必须做精度回归,三者缺一不可。工具链提供的基础算子和自动融合已经很强了,很多时候你把业务代码从Python侧的逻辑里解放出来、减少不必要的张量转换,比手写一个汇编kernel带来的收益更大。
有一点我想特别提醒:不要为了追求极限性能,把算子代码改到只有你自己能看懂。后维护它的人大概率是三个月后的自己。保留一份朴素的参考实现,把优化版和参考版放在一起对比测试,性能出问题随时能回滚。征程6这套工具链我用下来,最大的感受是它把硬件细节暴露得恰到好处——足够让你做精细控制,又不会让你陷入底层指令的汪洋大海。希望这份集锦能给你省下几个调研的晚上。