更新时间:2024-04-30 GMT+08:00
分享

优化算子执行

优化算子执行有两个思路:

  • 减少不必要的算子执行。比如减少不必要的格式转换算子和存储转连续算子。
  • 加速慢算子的执行速度。遇到此类问题,尝试基于AOE调优(详见性能调优五板斧)或者联系华为工程师分析处理。
图1 优化思路

减少不必要的算子执行

  • 减少不必要的存储转连续算子

    PyTorch的tensor对象由表示层和存储层(Storage)构成,表示层主要包含tensor的形状、步长、类型和是否连续等信息,存储层为连续内存的一维数组。对tensor的维度进行转换的操作(如transpose、permute等)后,转换前和转换后的两个tensor的表示层不同,但存储层相同,如下图所示。

    图2 表示层和存储层

    所以,会存在数学上相邻的元素在内存上不再连续排布的tensor。NPU的硬件工作机制要求所有的NPU算子只处理内存连续的tensor,为满足这一硬件限制,下发NPU算子前,PyTorch Adaptor的API会调用format_contiguousV2 API,format_contiguousV2会检测输入的tensor是否内存连续,如果不连续,会下发相关NPU算子完成新内存申请与数据copy保证输入tensor内存连续。开发者也可主动调用tensor的contiguous()方法提前完成这一操作,这样返回的就是内存连续的tensor。转存储连续操作涉及到内存申请以及写入,会导致NPU等待内存数据完成,而相对NPU执行速度来说,内存操作是比较慢的,因此应该尽量减少转存储连续操作。

    图3 format_contiguous

    首先,统计一个step内的CPU profiling信息,查看是否有大量contiguous字样的API调用(如format_contiguous、aten::contiguous等。特别需要强调:不需要关注format_contiguousV2,format_contiguousV2只有调用了aten::contiguous才代表对应tensor做了内存转连续操作)。

    图4 CPU profiling信息

    然后,切换到Operator视图,基于Operator分组,搜索contiguous关键字,进而从每一个调用栈信息定位到代码,逐个排查是否可以减少存储转连续算子的下发。

    图5 Operator视图

    可以通过如下几种优化策略尝试减少存储转连续算子的下发:

    • 如果存在对同一tensor的多次非连续操作,则可通过优先将其转连续以避免多次转换。如下图,因为box1和box2进行转置后存储不连续,后续每一行代码都会触发4次转存储操作(4行代码总共触发16次),修改后,总共触发两次转存储操作。

      修改前:

      box1 = box1.transpose(0, 1)
      box2 = box2.transpose(0, 1)
      b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
      b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
      b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
      b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2

      修改后:

      box1 = box1.transpose(0, 1).contiguous()
      box2 = box2.transpose(0, 1).contiguous()
      b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
      b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
      b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
      b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
    • 使用计算类算子代替维度变换的View类算子。如下所示使用“index_select”代替“torch.transpose(x, 1, 2).contiguous”

      原始channel_shuffle操作:

      def channel_shuffle(x, groups):
          # type: (torch.Tensor, int) -> torch.Tensor
          batchsize, num_channels, height, width = x.data.size()
          channels_per_group = num_channels // groups    # reshape
          x = x.view(batchsize, groups,
                     channels_per_group, height, width)
          x = torch.transpose(x, 1, 2).contiguous()
          # flatten
          x = x.view(batchsize, -1, height, width)
          return x
      同等语义修改:
      def channel_shuffle_index_select(x, groups=2):
          N, C, H, W = x.shape
          inp = C
          # channel_shuffle操作是对C维按一定规则的重排的工作,可以被表达为一次简单的重排
          group_len = inp // groups
          index = torch.from_numpy(np.array(list(range(inp))).reshape(groups, group_len).transpose(1, 0).flatten()).long()
          x = x.index_select(1, index)
          return x
    • 涉及到非0轴的slice操作,尝试先做permute操作将对应轴改变到0轴后调用。contiguous()方法后,再做slice操作。背后的原理是:tensor的存储机制使得存储连续的tensor对象的0维slice的结果也是存储连续的。

      举例如下,修改前需要做两次存储转连续操作,修改后只需一次存储转连续操作。

      修改前:

      position_ids, block_position_ids = position_ids[:, 0, :].transpose(0, 1).contiguous(), \
      position_ids[:, 1, :].transpose(0, 1).contiguous()

      修改后:

      position_ids_t = position_ids.permute(1, 2, 0).contiguous()
      position_ids, block_position_ids = position_ids_t[0, :, :], position_ids_t[1, :, :]
    • Transpose+reshape操作使用自定义融合算子torch_npu.npu_confusion_transpose替换,使用例子如下。
      >>> x = torch.arange(24).reshape(2,3,4).npu()
      >>> x
      tensor([[[ 0,  1,  2,  3],
      [ 4,  5,  6,  7],
      [ 8,  9, 10, 11]],
      [[12, 13, 14, 15],
      [16, 17, 18, 19],
      [20, 21, 22, 23]]], device='npu:0')
      >>> y = torch_npu.npu_confusion_transpose(x, (0,2,1), (2,2,6), True)
      >>> y
      tensor([[[ 0,  4,  8,  1,  5,  9],
      [ 2,  6, 10,  3,  7, 11]],
      [[12, 16, 20, 13, 17, 21],
      [14, 18, 22, 15, 19, 23]]], device='npu:0')
      >>> x.shape
      torch.Size([2, 3, 4])
      >>> y.shape
      torch.Size([2, 2, 6])
  • 减少不必要的格式转换算子

    关于数据排布格式,NPU在NCHW基础格式上,定义了众多与硬件强相关的私有格式,用于加速硬件计算,NPU对于不同的计算单元(cube、vector)所使用的默认私有格式不一致,导致在数据在不同的计算单元下流通时需要进行格式转化。此外,基础格式到NPU私有格式也需要进行格式转换,ACL编译期间根据算子的参数格式定义和当前输入的数据格式来决定是否向NPU侧下发格式转换的算子TransData。

    图6 算子TransData

    更多信息详见此处,算子的定义见此处,可以查看CANN算子参数的数据格式定义。可以通过tensor的“npu_format_cast(format_type)”方法进行tensor的格式转换以适配算子。

加速慢算子的执行速度

首先需要寻找执行速度比较慢的NPU算子列表,Kernel视图包含在NPU上执行的所有算子的信息,主要用于确认高耗时算子。

图7 Kernel视图

推荐基于以下思路尝试优化:

  1. 搜索Cast类算子,查看是否Cast类算子最大耗时超过30us或者总耗时占比超过1%,如果超过,需尝试启动混合精度训练,详见此处
    图8 Cast类算子
  2. 基于Accelerator Core排序,统计AI_CPU算子,如果有AI_CPU类算子执行时长超过1000us或者AI_CPU类算子总执行时长占比超过10%,可尝试修改代码替换API_CPU算子。
    需要注意:PyTorch Adaptor针对部分算子,会基于输入类型下发不同运行硬件的算子,所以除了使用同语义算子替换API_CPU算子外,还可以通过修改输入类型使算子下发到API_CORE上(比如torch.topk在参数为一维list使用API_CPU计算,多维参数则基于AI_CORE Vector计算)。
    图9 Accelerator Core排序
  3. 如果遇到算子运行期间NPU的计算单元和存储单元使用率都未达到80%(查看aiv_*_ratio和aic_*_ratio是否达到0.8),或者算子的“Block Dim”小于AI Core/Vector Core,可尝试使用AOE算子调优,提高NPU硬件资源利用率。
    图10 aiv_*_ratio
  4. 针对总耗时最长、平均执行耗时最长以及最大耗时的三种排序的TOP算子,可联系华为工程师获得帮助。
    图11 耗时排序
分享:

    相关文档

    相关产品