Ascend C提供一组Matmul高阶API,方便用户快速实现Matmul矩阵乘法的运算操作。
Matmul的计算公式为:C = A * B + Bias,其示意图如下。
下文中提及的M轴方向,即为A矩阵纵向;K轴方向,即为A矩阵横向或B矩阵纵向;N轴方向,即为B矩阵横向。
实现Matmul矩阵乘运算的具体步骤如下:
创建Matmul对象的示例如下:
// 纯cube模式(只有矩阵计算)场景下,需要设置该代码宏,并且必须在#include "lib/matmul_intf.h"之前设置 // #define ASCENDC_CUBE_ONLY #include "lib/matmul_intf.h" typedef matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, half> aType; typedef matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, half> bType; typedef matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float> cType; typedef matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float> biasType; matmul::Matmul<aType, bType, cType, biasType> mm;
创建对象时需要传入A、B、C、Bias的参数类型信息, 类型信息通过MatmulType来定义,包括:内存逻辑位置、数据格式、数据类型。
template <AscendC::TPosition POSITION, CubeFormat FORMAT, typename TYPE, bool ISTRANS = false, LayoutMode LAYOUT = LayoutMode::NONE, bool IBSHARE = false> struct MatmulType { constexpr static AscendC::TPosition pos = POSITION; constexpr static CubeFormat format = FORMAT; using T = TYPE; constexpr static bool isTrans = ISTRANS; constexpr static LayoutMode layout = LAYOUT; constexpr static bool ibShare = IBSHARE; };
参数 |
说明 |
---|---|
POSITION |
内存逻辑位置 针对Atlas A2训练系列产品/Atlas 800I A2推理产品:
针对Atlas推理系列产品AI Core:
针对Atlas 200I/500 A2推理产品:
|
CubeFormat |
针对Atlas A2训练系列产品/Atlas 800I A2推理产品:
针对Atlas推理系列产品AI Core:
|
TYPE |
针对Atlas A2训练系列产品/Atlas 800I A2推理产品:
针对Atlas推理系列产品AI Core:
针对Atlas 200I/500 A2推理产品:
注意:A矩阵和B矩阵数据类型需要一致,具体数据类型组合关系请参考表2。 |
ISTRANS |
是否开启使能矩阵转置的功能。
默认为false不使能转置。 |
LAYOUT |
表征数据的排布 NONE:默认值,表示不使用BatchMatmul;其他选项表示使用BatchMatmul。 NORMAL:BMNK的数据排布格式。 BSNGD:原始BSH shape做reshape后的数据排布,具体可参考IterateBatch中对该数据排布的介绍。 SBNGD:原始SBH shape做reshape后的数据排布,具体可参考IterateBatch中对该数据排布的介绍。 BNGS1S2:一般为前两种数据排布进行矩阵乘的输出,S1S2数据连续存放,一个S1S2为一个batch的计算数据。 |
IBSHARE |
是否使能IBShare。IBShare的功能是能够复用L1上相同的A矩阵或B矩阵数据。当A矩阵和B矩阵同时使能IBShare时,表示L1上的A矩阵和B矩阵同时复用,此时只支持Norm模板。 除A、B矩阵同时复用的场景外,与Matmul模板参数中的IBShare模板配合使用,具体参数设置详见表2。 |
A矩阵 |
B矩阵 |
Bias |
C矩阵 |
支持平台 |
---|---|---|---|---|
float |
float |
float/half |
float |
|
half |
half |
float |
float |
|
half |
half |
half |
float |
|
int8_t |
int8_t |
int32_t |
int32_t/half |
|
int4b_t |
int4b_t |
int32_t |
int32_t/half |
|
bfloat16_t |
bfloat16_t |
float |
float |
|
bfloat16_t |
bfloat16_t |
half |
float |
|
half |
half |
float |
int8_t |
|
bfloat16_t |
bfloat16_t |
float |
int8_t |
|
int8_t |
int8_t |
int32_t |
int8_t |
|
half |
half |
float |
half |
|
half |
half |
half |
half |
|
bfloat16_t |
bfloat16_t |
float |
bfloat16_t |
|
REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling);
mm.SetTensorA(gm_a); // 设置左矩阵A mm.SetTensorB(gm_b); // 设置右矩阵B mm.SetBias(gm_bias); // 设置Bias // Atlas推理系列产品AI Core上需要额外调用SetLocalWorkspace接口设置计算所需的UB空间 mm.SetLocalWorkspace(ubBuf);
while (mm.Iterate()) { mm.GetTensorC(gm_c); }
mm.IterateAll(gm_c);
mm.End();