大家好啊,我是董董灿。
好久没写文章了,最近收到了一些朋友的微信,问还写不写文章。
当然会继续写,只不过最近工作和生活忙的焦头烂额,一直没抽出时间来好好整理下自己,重整旗鼓继续写文。
总的来说,被两件事牵住了——
第一是房子在搞装修,从年前到现在几乎每个周末都要去逛建材城家具城,看家具家电,周末基本连轴转。
第二是项目到了紧要关头,几乎每人都在加班,项目中接到的神经网络一个接一个,功能调试和性能优化看不到头。
两件事,一边是工作,一边是生活。
这一篇主要想说一个最近项目的事——
神经网络性能优化之——算子融合(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