大家好啊，我是董董灿。
最近接了很多奇奇怪怪的神经网络需要实现和优化。里面有不少常见、好实现并且好优化的算子，如 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光训练一次，就要花费几千万美元的成本。如果将网络优化的足够好，可以大幅度降低训练成本，而这，都是白花花的银子啊。
本文作者原创，请勿转载，转载请联系作者
