识别融合算子

在分析算子耗时已经统计了NPU算子个数，与GPU算子个数进行对比，如果算子个数有较明显的差距，那么极大可能是由融合算子导致。融合算子包括两种，一种是自动融合的算子，一种是手写的融合算子。

• 自动融合算子识别。

  在PyTorch2.0版本之前，成图优化通常使用torch.jit.script的装饰器，对一段小算子进行融合优化，后端会使用nvfuser生成融合算子。

  在脚本中搜索“torch.jit.script”，就可以找到使用此种方式生成的融合算子。

  图1 通过torch.jit.script检索融合算子
  

  在GPU的kernel统计表中，通常是CudaCodeGen::kernelxx，也可能是kernelxx。

  图2 通过kernel统计表查看融合算子
  
• 手写融合算子识别。

  手写融合算子需要对profiling进行分析，根据以往经验，目前发现几个典型的GPU融合算子：

  1. aten::baddbmm形式。

     

     在BatchMatMul的基础上增加了2个乘和1个加，请注意，当β为0的时候可以减少1个乘和1个加，这也是当前较常用的优化点。

  2. ScaledUpperTriangMaskedSoftmax和ScaledUpperTriangMaskedSoftmaxBackward：

     这个算子是maskdfill，softmax及前后cast，dropout等算子的融合算子，NPU上存在类似的融合算子，叫ScaledMaskedSoftmax，NPU中该融合算子的入参比GPU的多一个，主要是因为NPU的dropout的genmask是从外面传入的，而GPU中是内部生成的。

     以上2个融合算子的组合构成了self-attention部分，这部分也可以直接使用flashattention算子替换。

  3. For_each xxx形式：

     另外一个比较典型的融合算子是优化器部分的for_each算子，这类算子可以看作是一组并行执行算子的融合，如aten::_amp_foreach_non_finite_check_and_unscale_，这部分另外的融合算子就是for_each的adam算子，如void multi_tensor_apply_kernel<TensorListMetadata<4>。