大家好啊，我是董董灿。

好久没写文章了，最近收到了一些朋友的微信，问还写不写文章。

当然会继续写，只不过最近工作和生活忙的焦头烂额，一直没抽出时间来好好整理下自己，重整旗鼓继续写文。

总的来说，被两件事牵住了——

第一是房子在搞装修，从年前到现在几乎每个周末都要去逛建材城家具城，看家具家电，周末基本连轴转。

第二是项目到了紧要关头，几乎每人都在加班，项目中接到的神经网络一个接一个，功能调试和性能优化看不到头。

两件事，一边是工作，一边是生活。

这一篇主要想说一个最近项目的事——

神经网络性能优化之——算子融合（operator fusion）。

最近接了很多奇奇怪怪的神经网络需要实现和优化。

里面有不少常见、好实现并且好优化的算子，如 element-wise 类的 add 算子、cast 算子等。

一般遇到这类算子，我都是直接将他们从神经网络中忽略掉的。

因为它们属于维度无关的算子，不论你在内存上的摆布（layout）是 NHWC，还是 NCHW，还是其他奇奇怪怪的摆布。

又或者是将数据在单核运算，还是拆开放在多核上运算，都无所谓。只要对位相加得到正确结果就行了。

稍微难一点的算子比如 reduce_max、reduce_sum 类的算子，它们维度相关，尤其是在进行多核拆分时，需要考虑拆的维度最好不要是 reduce 维度，因为一旦在核间拆了 reduce 维度，那么势必还要做核间的 reduce。

不过这类算子也不是很难，毕竟如果做多核运算，不拆 reduce 维度就好很多。

再稍微难一点的算子像是多维转置（transpose，也有叫 permute ) 算子，维度更加相关。如果做多核运算，一旦在核间拆到转置的那一维，就会出现核间的数据搬运操作。

除非你找个地方将数据暂存起来，或者直接放到DDR上。总之，多核的transpose比较难。

更难的，像是 scatter 类的算子，它需要根据其中一个 tensor(indice) 中的数据作为坐标索引来完成数据（data）的更新(update)，多核拆分场景下需要确保更新的数据维度不要被拆散。

scatter类的算子难点在于，取 indice 中的坐标几乎无法向量化，只能标量的取用，然后向量化的完成数据更新。

这类算子，一直是性能优化的重灾区。

这次遇到的网络，很不幸，上面的几类算子都有。

做AI推理优化的同学都知道，如果神经网络层与层之间的数据是放在外存（ddr）上的话，一个好处是可以连续访存。但同时带来一个坏处就是带宽可能会成为性能瓶颈。

毕竟ddr上的数据需要不断的与片上存储进行 IO 交换，带宽不够的话，很影响推理性能。

因此，大部分做AI推理加速的同学，都会考虑将层与层之间的数据放在离计算核心更近的位置，比如片内的SRAM上，前提是SRAM足够大，能放的下这些数据。

那么问题就来了，如果将数据都放在SRAM上，为了更好的推理性能，更小的推理延时，网络优化或算子优化的同学几乎都要做多核间的数据拆分（AI加速芯片几乎都是多核或众核架构），一旦做了拆分，就引出了问题，不同算子对于拆分有不同的友好度。

像是上面说的， element-wise 类的算子属于维度无关，怎么拆都行。我们暂时不考虑此类算子。

但如果一个网络片段中存在如 reduce_sum -> transpose 的层。

reduce_sum 一个友好的拆分方法是在core间拆非累加维。

如 NHWC 在 channel 维度做累加的话，我们可以在core间拆非channel维，比如可以拆 W 维（对于NHWC而言，W维属于第2维）。

但如果下一层 transpose 做的是 NHWC 转置变成 NCHW，一旦拆了W维，每个核内只有一部分的W，转置之后的W又处于最低维（对于NCHW而言，W维属于第3维）。

不同算子的拆分规则不同，这就乱了。

为了让最终数据不乱，同时遵守一个拆分规则，需要做多余的数据搬运操作。

最终，在我们花了大量的时间将这个奇怪的网络所有算子调试完后，一个自然而然的想法一下子在脑海中涌现出来：

这个网络这么不好优化，为什么不在一开始将所有算子融合成一个大算子来做呢？

融合是有好处的——

• 没有数据对ddr的访存，杜绝了带宽瓶颈的存在

• 层与层之间不用遵守固定的layout定义，想怎么存就怎么存，只要确保把数据算对就行

• 指令条数明显会少，通过融合算子内部的流水排布，更容易做图优化

于是，融合就这么开始了。

其实，融合的想法，在最开始适配网络的时候也想过，但限于开发周期和交付压力，还是选择了用小算子来拼网络的思路。

不过融合作为一个神经网络的优化大杀器，还是很值得做的。

简单说了下最近遇到的一个小项目。一些算法和拆分看不懂没关系，后面我会逐步写一写神经网络推理优化方面的文章。

AI推理和训练优化是一个很难啃的骨头，大火的chatGPT光训练一次，就要花费几千万美元的成本。如果将网络优化的足够好，可以大幅度降低训练成本，而这，都是白花花的银子啊。

---

v v v v v v

**本文为作者原创，请勿转载，转载请联系作者。**
**点击下方卡片，关注我的公众号，有最新的文章和项目动态。**

v v v v v v