文档HOT
下载
中文
注册

Mmad

功能说明

完成矩阵乘加（C += A * B）操作。矩阵ABC分别为A2/B2/CO1中的数据。

ABC矩阵的数据排布格式分别为ZZ，NZ，ZN。

下图中每个小方格代表一个分形矩阵，Z字形的黑色线条代表数据的排列顺序，起始点是左上角，终点是右下角。

矩阵A：每个分形矩阵内部是行主序，分形矩阵之间是行主序。简称小Z大Z格式。分形shape为16 x (32B/sizeof(AType))，大小为512Byte。

矩阵B：每个分形矩阵内部是列主序，分形矩阵之间是行主序。简称小N大Z格式。分形shape为 (32B/sizeof(BType)) x 16，大小为512Byte。

矩阵C：每个分形矩阵内部是行主序，分形矩阵之间是列主序。简称小Z大N格式。分形shape为16 x 16，大小为256个元素。

以下是一个简单的例子，假设分形矩阵的大小是2x2（并不符合真实情况，仅作为示例），矩阵ABC的大小都是4x4。

0

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

矩阵A的排列顺序：0，1，4，5，2，3，6，7，8，9，12，13，10，11，14，15。

矩阵B的排列顺序：0，4，1，5，2，6，3，7，8，12，9，13，10，14，11，15。

矩阵C的排列顺序：0，1，4，5，8，9，12，13，2，3，6，7，10，11，14，15。

函数原型

  • 不传入biasLocal
    1
2
    template <typename DstT, typename Src0T, typename Src1T>
__aicore__ inline void Mmad(const LocalTensor<DstT>& dstLocal, const LocalTensor<Src0T>& fmLocal, const LocalTensor<Src1T>& filterLocal, const MmadParams& mmadParams)
  • 传入biasLocal
    1
2
    template <typename DstT, typename Src0T, typename Src1T, typename BiasT>
__aicore__ inline void Mmad(const LocalTensor<DstT>& dstLocal, const LocalTensor<Src0T>& fmLocal, const LocalTensor<Src1T>& filterLocal, const LocalTensor<BiasT>& biasLocal, const MmadParams& mmadParams)

参数说明

表1 参数说明

参数名称

输入/输出

含义

dstLocal

输出

目的操作数，结果矩阵，类型为LocalTensor，支持的TPosition为CO1。

LocalTensor的起始地址需要256个元素对齐。

fmLocal

输入

源操作数，左矩阵a，类型为LocalTensor，支持的TPosition为A2。

LocalTensor的起始地址需要512字节对齐。

filterLocal

输入

源操作数，右矩阵b，类型为LocalTensor，支持的TPosition为B2。

LocalTensor的起始地址需要512字节对齐。

biasLocal

输入

源操作数，bias矩阵，类型为LocalTensor，支持的TPosition为C2、CO1。

LocalTensor的起始地址需要128字节对齐。

mmadParams

输入

矩阵乘相关参数，类型为MmadParams，结构体具体定义为：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
struct MmadParams
{
    uint16_t m;
    uint16_t n;
    uint16_t k;
    bool isBias;
    int32_t fmOffset;
    bool enSsparse;
    bool enWinogradA;
    bool enWinogradB;
    uint8_t unitFlag;
    bool cmatrixInitVal;
    bool cmatrixSource;
};

参数说明请参考表2。
表2 MmadParams结构体内参数说明

参数名称

含义

m

左矩阵Height，取值范围：m∈[0, 4095] 。默认值为0。

n

右矩阵Width，取值范围：n∈[0, 4095] 。默认值为0。

k

左矩阵Width、右矩阵Height，取值范围：k∈[0, 4095] 。默认值为0。

fmOffset

预留参数，用户无需关心，使用默认值0即可。

enSsparse

预留参数，用户无需关心，使用默认值false即可。

enWinogradA

预留参数，用户无需关心，使用默认值false即可。

enWinogradB

预留参数，用户无需关心，使用默认值false即可。

unitFlag

预留参数，用户无需关心，使用默认值0即可。

cmatrixInitVal

配置C矩阵初始值是否为0。默认值true。

  • true：C矩阵初始值为0；
  • false：C矩阵初始值通过cmatrixSource参数进行配置。

cmatrixSource

配置C矩阵初始值是否来源于C2（存放Bias的硬件缓存区）。默认值为false。

  • false：来源于CO1；
  • true：来源于C2。

Atlas 训练系列产品，仅支持配置为false。

Atlas 推理系列产品AI Core，仅支持配置为false。

Atlas A2 训练系列产品/Atlas 800I A2 推理产品，支持配置为true/false。

Atlas 200I/500 A2 推理产品, 支持配置为true/false。

注意：带biasLocal输入的接口配置该参数无效，会根据biasLocal输入的位置来判断C矩阵初始值是否来源于CO1还是C2。

isBias

该参数废弃，新开发内容不要使用该参数。如果需要累加初始矩阵，请使用带biasLocal的接口来实现；也可以通过cmatrixInitVal和cmatrixSource参数配置C矩阵的初始值来源来实现。推荐使用带biasLocal的接口，相比于配置cmatrixInitVal和cmatrixSource参数更加简单方便。

配置是否需要累加初始矩阵，默认值为false，取值说明如下：

  • false：矩阵乘，无需累加初始矩阵，C = A * B。
  • true：矩阵乘加，需要累加初始矩阵, C += A * B。
表3 dstLocal、fmLocal、filterLocal支持的精度类型组合（Atlas 训练系列产品）

左矩阵fmLocal type

右矩阵filterLocal type

结果矩阵dstLocal type

uint8_t

uint8_t

uint32_t

int8_t

int8_t

int32_t

uint8_t

int8_t

int32_t

half

half

half

说明：

该精度类型组合，精度无法达到双千分之一，且后续处理器版本不支持该类型转换，建议直接使用half输入float输出。

half

half

float
表4 dstLocal、fmLocal、filterLocal支持的精度类型组合（Atlas 推理系列产品AI Core ）

左矩阵fmLocal type

右矩阵filterLocal type

结果矩阵dstLocal type

int8_t

int8_t

int32_t

uint8_t

int8_t

int32_t

uint8_t

uint8_t

int32_t

half

half

half

说明：

该精度类型组合，精度无法达到双千分之一，且后续处理器版本不支持该类型转换，建议直接使用half输入float输出。

half

half

float
表5 dstLocal、fmLocal、filterLocal支持的精度类型组合（Atlas 200I/500 A2 推理产品）（Atlas A2 训练系列产品/Atlas 800I A2 推理产品）

左矩阵fmLocal type

右矩阵filterLocal type

结果矩阵dstLocal type

int8_t

int8_t

int32_t

half

half

float

float

float

float

bf16

bf16

float

int4b_t

int4b_t

int32_t
表6 dstLocal、fmLocal、filterLocal、biasLocal支持的精度类型组合

左矩阵fmLocal type

右矩阵filterLocal type

biasLocal type

结果矩阵dstLocal type

int8_t

int8_t

int32_t

int32_t

half

half

float

float

float

float

float

float

bf16

bf16

float

float

注意事项

  • dstLocal只支持位于CO1，fmLocal只支持位于A2，filterLocal只支持位于B2。
  • 操作数地址偏移对齐要求请参见通用约束。
  • 无效数据与有效数据排布方式如下：以下是一个具体的例子，数据为half类型。

    如图，当M=30，K=70，N=40的时候，A2中 2x5个16x16矩阵，B2中有5x3个16x16矩阵，CO1中有2x3个16x16矩阵。在这种场景下M、K和N都不是 16 的倍数，A2中右下角的矩阵实际有效的数据只有14x6个，但是也需要占一个16x16矩阵的空间，其他无效数据在计算中会被忽略。一个16x16分形的数据块中，无效数据与有效数据排布的方式示意如下：

支持的型号

  • 不传入biasLocal的接口

    Atlas 训练系列产品

    Atlas 推理系列产品AI Core

    Atlas A2 训练系列产品/Atlas 800I A2 推理产品

    Atlas 200I/500 A2 推理产品

  • 传入biasLocal的接口

    Atlas A2 训练系列产品/Atlas 800I A2 推理产品

    Atlas 200I/500 A2 推理产品

调用示例

不含矩阵乘偏置的样例请参考Mmad样例。

包含矩阵乘偏置的样例请参考包含矩阵乘偏置的Mmad样例。

父主题： 矩阵计算