更新时间:2021-03-18 GMT+08:00
分享

算子分析

开发ScatterNdAdd算子前,我们需要确定算子的功能,输入、输出,算子开发方式,算子类型以及算子实现函数名称等。

  1. 明确算子的功能。

    ScatterNdAdd算子的功能是通过对输入数据中的单个值或切片应用稀疏算法,从而得到输出数据,详细功能示例可参考功能描述

  2. 明确输入和输出。

    • ScatterNdAdd算子有三个输入,一个输出。

      输入:var(需要更新的tensor),indices(指定需要更新的索引位置),updates(更新数据)。

      输出:var(输出更新后的tensor)。

    • 本样例中算子的输入var支持的数据类型为float16、float32、int32、int8、uint8;输入indices支持的数据类型为int32;输入updates支持的数据类型为float16、float32、int32、int8、uint8;输出var支持的数据类型为:float16、float32、int32、int8、uint8。
    • 算子输入支持所有shape,输出shape与输入shape相同。
    • 算子输入支持的format为:ND。

  3. 确定算子开发方式及使用的计算流程设计。

    由于ScatterNdAdd算子涉及对tensor的不同维度上的不同元素同时计算,TBE DSL接口都无法满足此算子的计算要求,所以考虑使用TIK方式进行此算子的实现。

    该算子实现核心的计算流程如下:

    1. 将indices读入到UB Buffer中,然后遍历计算var中需要更新位置的index。
    2. 每算出一个index,就将相应的需要更新的var分片和updates分片搬入到UB Buffer中。
    3. 将两个分片每个对应位置的元素相加后再搬出到GM内存中。
    4. 当遍历完所有的index后,就可以得到最终的计算结果。

    针对这个核心计算流程,我们要设计相应的schedule策略。schedule策略在设计的时候主要考虑两个基本问题:

    1. 首先是shape的泛化。

      由于我们的UB空间有大小限制,所以不是所有的输入shape都能在UB上放下,需要考虑分片搬运入UB buffer进行计算。此时,就需要我们根据输入的shape大小,数据类型,UB buffer空间来计算每次分片搬运的大小和需要搬入的次数。在UB空间划分的时候,要充分合理的利用UB空间来提升性能。相同的输入shape,分10次搬入UB计算完之后再搬回到GM,比分100次搬运和计算性能更优。因此,要满足不同的shape泛化,我们要根据输入的shape来计算和划分UB buffer空间,计算各个指令的参数。

    2. 其次是多核,double buffer等策略。

      当前昇腾AI处理器有多个AI Core可以做并行计算,可以极大的提升计算的性能。对于ScatterNdAdd这个算子,总共需要三层循环,如下所示

      表1 ScatterNdAdd算子循环列表

      空间

      循环

      GM (拆分var)

      多核LOOP

      GM2UB (indices)

      indices LOOP

      GM2UB (var + update)

      Split indices LOOP

      核心是遍历indices,对每个index取出var和update做计算。为了实现多核并行计算,我们将var进行分拆,让不同AI Core进行不同的index位置的计算。例如,假设我们根据前面的公式算出index的取值范围为[0, 6],这时候我们把index=0对应的updates分片放在第一个核处理,index=2对应的updates分片放在第二个核处理,以此类推,最后一个核处理index=6对应的updates分片,这样可以实现多个核的并行处理。

  4. 明确算子实现文件名称、算子实现函数名称以及算子的类型(OpType)。

    • 算子类型需要采用大驼峰的命名方式,即采用大写字符区分不同的语义。
    • 算子文件名称和算子函数名称,可选用以下任意一种命名规则:
      • 用户自定义,此时需要在算子信息定义中配置opFile.valueopInterface.value
      • 不配置算子信息定义中的opFile.valueopInterface.value,FE会将OpType按照如下方式进行转换后进行算子文件名和算子函数名的匹配。
        转换规则如下:
        • 首字符的大写字符转换为小写字符。

          例如:Abc -> abc

        • 小写字符后的大写字符转换为下划线+小写字符。

          例如:AbcDef -> abc_def

        • 紧跟数字以及大写字符后的大写字符,作为同一语义字符串,查找此字符串后的第一个小写字符,并将此小写字符的前一个大写字符转换为下划线+小写字符,其余大写字符转换为小写字符。若此字符串后不存在小写字符,则直接将此字符串中的大写字符转换为小写字符。

          例如:ABCDef -> abc_def;Abc2DEf -> abc2d_ef;Abc2DEF -> abc2def;ABC2dEF -> abc2d_ef。

    因此本例中,算子类型定义为ScatterNdAdd;算子的实现文件名称及实现函数名称定义为scatter_nd_add。

    通过以上分析,得到ScatterNdAdd算子的设计规格如下:

    表2 ScatterNdAdd算子设计规格

    算子类型(OpType)

    ScatterNdAdd

    算子输入

    name:var

    Type:Tensor

    shape:all

    data type:

    float16、float32、int32、int8、uint8

    format:ND

    name:indices

    Type:Tensor

    shape:all

    data type:int32

    format:ND

    name:updates

    Type:Tensor

    shape:all

    data type:

    float16、float32、int32、int8、uint8

    format:ND

    算子输出

    name:var

    Type:Tensor

    shape:all

    data type:

    float16、float32、int32、int8、uint8

    format:ND

    算子实现文件/实现函数名称

    scatter_nd_add

分享:

    相关文档

    相关产品

close