做AI芯片编译优化时,有一个直觉很容易踩坑:算子越大,融合得越彻底,性能就越好。于是有人把一串算子融成一个大内核,或者反过来,把一个大算子拆成几个小算子指望并行和流水。两种做法都可能翻车,因为算子切分的真正账本不在算子个数,而在访存:中间结果写出去再读回来,带宽和延迟都要花钱。

先说为什么要拆。大算子的问题是资源胃口太大:一个融合内核可能需要同时装下多个张量的中间结果,寄存器和片上缓存装不下,就会溢出到显存或者全局内存,溢出本身就是访存开销。拆成小算子后,每个小内核的资源需求变小,更容易塞进片上存储,调度也更灵活,还能把不同阶段流水起来,让计算和访存重叠。这就是拆分的收益,看起来合情合理。

但拆分的第一笔账是中间结果的往返。原来在片上寄存器里递过去的数据,拆开后要先写回全局内存,下一个算子再读回来。写一次、读一次,两次访存都要走带宽,还要付启动延迟和同步开销。如果中间张量很大、而计算本身很轻,比如逐元素操作,这笔往返开销可能比计算本身贵一个数量级。算子越小、切得越碎,这种往返就越频繁,最后带宽被中间结果吃光,算力利用率反而掉下来。

第二笔账是并行度的损失。大算子内部往往有天然的并行维度,拆开后每个小算子的并行规模变小,填充不满计算阵列时,硬件利用率就下来了。更隐蔽的是同步点变多:每个小算子结束都要同步、都要重新调度,调度开销和尾部效应累积起来,在小规模张量上尤其明显。切分不是免费的午餐,每切一刀都是一次同步和一次访存。

第三笔账是数据布局的反复横跳。不同算子对数据排布的要求可能不同:一个算子喜欢按通道分块,另一个喜欢按空间分块。拆开后,中间可能要多一次转置或者重排,这又是访存。融合大算子时,编译器可以在内部直接按最合适的布局流转,拆开后这个自由度就没了。所以切分决策不能只看算子图,要看数据在切分点前后的布局是不是一致。

那么切分点该怎么选?实用的办法是算访存强度:对每个候选切分点,估算中间结果的字节数除以它带来的计算量,再对比硬件的带宽与算力比值。如果中间结果很小、计算很重,拆开的代价低,可以拆;如果中间结果很大、计算很轻,宁可忍受大算子的资源压力,也别切。还有一个经验法则:切在“数据自然缩小”的地方,比如池化、下采样之后,中间结果已经变小了,往返成本低,切分就划算;切在数据膨胀的地方,比如上采样、升维之后,代价最大。

最后提醒一点:切分策略要跟着硬件走。片上缓存大的芯片容得下大算子,切分可以保守;片上缓存小的芯片,大算子一溢出就全是访存,还不如早点切。编译器的自动切分、算子融合、内存规划其实是同一盘棋:目标从来不是算子最少或者最多,而是让数据在片上待得越久越好,让每次访存都物有所值。记住这个原则,切分时就不会只盯着算子个数做文章。