开发者
资源
[object Object]

下文将以Matmul+LeakyRelu融合算子的实现为例,介绍Mix融合算子的设计和实现流程。

算子的设计过程分为算子分析、数据流分析、Tiling策略设计三部分。

[object Object]

算子分析是指明确算子的数学表达式、输入、输出,核函数的名称等信息。

  1. 明确算子的数学表达式及计算逻辑。该算子的计算逻辑为,先进行一个矩阵乘操作,然后将矩阵乘的结果与一个alpha参数进行LeakyRelu操作。数学表达式如下:

    [object Object]
  2. 明确输入和输出。

    • Matmul+LeakyRelu算子输入为a、b、bias,输出为c。alpha作为激活函数LeakyRelu的系数,为固定值,可以在算子实现中直接使用常数值参与计算。
    • 本样例中算子输入a、b支持的数据类型为half(float16),算子输入bias支持的数据类型为float32,算子输出c的数据类型为float32。
    • 输入矩阵a的形状为[M,K],输入矩阵b的形状为[K, N],输出矩阵c的形状为[M,N],输入bias的形状为[1, N]。
    • 算子输入输出支持的数据格式为:ND。
  3. 确定核函数名称和参数。

    • 您可以自定义核函数名称,本样例中核函数命名为matmul_leakyrelu_custom。
    • 根据对算子输入输出的分析,确定核函数的参数a,b,bias,c;a,b, bias为输入在Global Memory上的内存地址,c为输出在Global Memory上的内存地址。

通过以上分析,得到Ascend C Matmul+LeakyRelu算子的设计规格如下:

  • 算子类型(OpType):MATMUL_LEAKYRELU

  • 算子输入输出:

    表 1 MATMUL_LEAKYRELU算子输入输出规格

    [object Object][object Object]

    [object Object]
  • 核函数名称:matmul_leakyrelu_custom

[object Object]

进行算子的数据流分析:数据流向为在Cube核上完成Matmul计算后将数据搬运至Vector核进行LeakyRelu计算。根据上述数据流并结合融合算子的编程范式,规划并行的流水任务。如下图所示:

  1. 将输入数据从Global Memory搬运到Cube核。

  2. 进行Matmul内部的计算,计算公式和计算示意图如下:

    注:bias的shape为[1, N],对A*B结果矩阵的每一行都采用该bias进行偏置。

    图 1 Matmul矩阵乘示意图[object Object][object Object]

  3. 将Matmul的计算结果搬运到Vector核。

  4. 进行Vector矢量计算,该样例中进行LeakyReLU计算。

    Leaky ReLU(带泄露线性整流函数)激活函数,是人工神经网络中一种常用的激活函数,其数学表达式和函数图像如下所示:

  5. 将输出结果搬运到Global Memory。

前三步的内容都封装在Matmul高阶API内,本样例中可以简化为3个stage。如下图所示:

根据上述分析,明确实现过程中会使用到接口,接口。

[object Object]

Tiling策略的设计主要包括多核切分和核内切分策略。

  • 多核切分: 根据当前核数,对输入shape的M, K, N进行多核切分,得到单核内shape大小singleCoreM, singleCoreK, singleCoreN。
  • 核内切分: 根据Local Memory的大小约束,对单核内的shape大小进一步切分,得到A、B、C矩阵参与一次矩阵乘指令的shape大小baseM, baseN, baseK。切分时需要注意:GetTensorC的结果如果放在LocalMemory(UB)上,需要注意,baseM * baseN的大小不能超出UB的限制。

切分策略示意图如下,更多切分策略相关原理请参考

[object Object]

,我们得知Ascend C提供一组Matmul高阶API,封装了常用的切分和数据搬运、计算的算法逻辑,方便用户快速实现Matmul矩阵乘法的运算操作。融合算子中矩阵编程部分的实现与之类似,开发者在host侧通过调用API自动获取Tiling参数,该参数传递到kernel侧后,在初始化操作时传入,通过几个简单的API即可完成矩阵乘操作。再结合上文的融合算子的编程范式,融合算子实现的步骤如下。完整样例请参考

kernel侧实现的代码框架如下,在完成Matmul对象的初始化、左矩阵A、右矩阵B、Bias的设置后,通过单次Iterate叠加while循环的方式完成后续的Matmul计算、LeakyRelu计算、CopyOut流程。

[object Object]

Matmul计算、LeakyRelu计算、CopyOut的具体实现代码如下:

  1. Matmul计算:

    1. 在Cube核上,进行Matmul内部的计算。

    2. 将Matmul的计算结果搬运到Vector核。

      [object Object]
  2. LeakyRelu计算。

    [object Object]
  3. CopyOut,将输出结果搬运到Global Memory。

    [object Object]

host侧实现GenerateTiling函数,在该函数中自动获取Tiling参数,关键步骤介绍如下:

  1. 创建Tiling对象

    [object Object]

    创建对象时需要传入硬件平台信息,硬件平台信息可以通过GetPlatformInfo获取。

  2. 设置A、B、Bias的数据类型和格式。

    设置示例如下,其中TPosition::LCM是Unified Buffer上的逻辑位置,等同于TPosition::VECCALC,关于TPosition的详细内容请参考

    [object Object]
  3. 设置矩阵shape信息。

    [object Object]
  4. 设置可用空间大小信息。

    设置Matmul计算时可用的L1 Buffer/L0C Buffer/Unified Buffer空间大小,-1表示AI处理器对应Buffer的大小。

    [object Object]
  5. 按需设置其他参数,比如设置bias参与计算。

    [object Object]
  6. 获取Tiling参数。

    [object Object]
  7. Tiling参数的序列化保存等其他操作。

[object Object]