AscendC矩阵批量计算接口

IterateBatch【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit产品支持情况Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品AI Core支持Atlas 推理系列产品Vector Core不支持Atlas 训练系列产品不支持Kirin X90支持功能说明单次Matmul计算处理的shape比较小时由于每次计算均涉及到内部的通信可能会影响性能该接口提供批量处理Matmul的功能调用一次IterateBatch可以计算出多个singleCoreM * singleCoreN大小的C矩阵。在使用该接口前需要了解一些必备的数据排布格式通用数据格式NORMALBMNK的数据排布格式BBatch批处理的大小M、N、K为矩阵乘[M, K]*[K, N]的矩阵维度其数据排布格式如下BSH/SBHBBatch批处理的大小 Ssequence length序列长度H N * D其中N为head的数量D为head的大小。Layout格式如下图所示BSNGD为原始BSH shape做reshape后的shapeS和D为单Batch的矩阵乘的M轴或N轴和K轴一个SD为一个batch的计算数据Layout格式如下图所示SBNGD为原始SBH shape做reshape后shapeS和D为矩阵乘的M轴或N轴和K轴一个SD为一个Batch的计算数据Layout格式如下图所示BNGS1S2一般为前两种Layout进行矩阵乘的输出S1S2数据连续存放一个S1S2为一个Batch的计算数据Layout格式如下图所示实例化Matmul时需要通过MatmulType设置输入输出的Layout格式当前支持4种Layout类型BSNGD、SBNGD、BNGS1S2、NORMALBMNK的数据排布格式使用NORMAL表示。对于BSNGD、SBNGD、BNGS1S2 Layout格式调用该接口之前需要在host Tiling实现中使用SetALayout、SetBLayout、SetCLayout、SetBatchNum设置A/B/C的Layout轴信息和最大BatchNum数对于NORMAL Layout格式则需使用SetBatchInfoForNormal设置A/B/C的M/N/K轴信息和A/B矩阵的BatchNum数。使用ND格式输入时是否开启MatmulConfig中的intrinsicsCheck取决于单个BMM内A或B矩阵的源矩阵内轴长度。该内轴可理解为源ND数据中一行的连续跨度即一次ND2NZ搬运所使用的源行宽。不同Layout下的取值规则如下NORMAL Layout下按普通矩阵内轴判断BNGS1S2 Layout下按源ND矩阵一行的实际跨度判断BSNGD Layout下A和B矩阵源矩阵内轴分别为ALayoutInfoN * ALayoutInfoG * ALayoutInfoD、BLayoutInfoN * BLayoutInfoG * BLayoutInfoDSBNGD Layout下A和B矩阵源矩阵内轴分别为ALayoutInfoB * ALayoutInfoN * ALayoutInfoG * ALayoutInfoD、BLayoutInfoB * BLayoutInfoN * BLayoutInfoG * BLayoutInfoD。单个矩阵乘迭代顺序可通过tiling参数iterateOrder调整。函数原型mix模式输出至GMtemplate bool sync true, bool waitIterateBatch false __aicore__ inline void IterateBatch(const GlobalTensorDstT gm, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0, const bool enPartialSum false, const uint8_t enAtomic 0)输出至UBVECINtemplate bool sync true __aicore__ inline void IterateBatch(const LocalTensorDstT ubCmatrix, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0, const bool enPartialSum false, const uint8_t enAtomic 0)纯cube模式使用前需先调用SetBatchNum接口设置batchA和batchB的大小。输出至GM__aicore__ inline void IterateBatch(const GlobalTensorDstT gm, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0)输出至UBVECIN__aicore__ inline void IterateBatch(const LocalTensorDstT ubCmatrix, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0)参数说明表1模板参数说明参数名描述sync获取C矩阵过程分为同步和异步两种模式同步需要同步等待IterateBatch执行结束。异步不需要同步等待IterateBatch执行结束。通过该参数设置同步或者异步模式同步模式设置为true异步模式设置为false。默认为同步模式。异步场景需要配合WaitIterateBatch接口使用。waitIterateBatch是否需要通过WaitIterateBatch接口等待IterateBatch执行结束仅在异步场景下使用。默认为false。true需要通过WaitIterateBatch接口等待IterateBatch执行结束。false不需要通过WaitIterateBatch接口等待IterateBatch执行结束开发者自行处理等待IterateBatch执行结束的过程。表2接口参数说明参数名输入/输出描述gm输出C矩阵。类型为GlobalTensor。不同型号支持的数据类型请参考支持的数据类型。ubCmatrix输出C矩阵。类型为LocalTensor。不同型号支持的数据类型请参考支持的数据类型。batchA输入左矩阵的batch数。batchB输入右矩阵的batch数。在batchA/batchB不相同的情况下默认做broadcast操作。多batch计算支持在G轴上做输入broadcast和输出reduce左矩阵、右矩阵G轴维度必须是整数倍的关系。enSequentialWrite输入输出是否连续存放数据即是否开启连续写模式连续写写入[baseM, baseN]非连续写写入[singleCoreM, singleCoreN]中对应的位置。左右矩阵和输出矩阵的存储位置为Unified Buffer则enSequentialWrite参数应配置为true输出矩阵的存储位置为GM则enSequentialWrite参数应配置为false。matrixStrideA输入A矩阵源操作数相邻nd矩阵起始地址间的偏移单位是元素默认值是0。matrixStrideB输入B矩阵源操作数相邻nd矩阵起始地址间的偏移单位是元素默认值是0。matrixStrideC输入该参数预留保持默认值0即可。enPartialSum输入是否将矩阵乘的结果累加于现有的CO1数据默认值为false。在L0C累加时只支持A矩阵和B矩阵相乘的输出C矩阵规格为singleMbaseM singleNbaseN。enAtomic输入是否开启Atomic操作默认值为0。参数取值0不开启Atomic操作1开启AtomicAdd累加操作2开启AtomicMax求最大值操作3开启AtomicMin求最小值操作返回值说明无约束说明该接口只支持Norm模板即BatchMatmul只支持Norm模板。使用该接口时A、B矩阵的Layout格式必须相同。对于BSNGD、SBNGD、BNGS1S2 Layout格式输入A、B矩阵按分形对齐后的多Batch数据总和应小于L1 Buffer的大小对于NORMAL Layout格式没有这种限制但需通过MatmulConfig配置输入A、B矩阵多Batch数据大小与L1 Buffer的大小关系对于BSNGD、SBNGD、BNGS1S2 Layout格式称左矩阵、右矩阵的G轴分别为ALayoutInfoG、BLayoutInfoG则ALayoutInfoG / batchA BLayoutInfoG / batchB对于NORMAL Layout格式batchA、batchB必须满足倍数关系。如果接口输出到Unified Buffer上输出C矩阵大小BaseM*BaseN应小于分配的Unified Buffer内存大小。如果接口输出到Unified Buffer上且单核计算的N方向大小singleCoreN非32字节对齐C矩阵的CubeFormat仅支持ND_ALIGN格式输出C矩阵片时自动将singleCoreN方向上的数据补齐至32字节。对于BSNGD、SBNGD Layout格式输入输出只支持ND格式数据。对于BNGS1S2、NORMAL Layout格式输入支持ND/NZ格式数据。对于BSNGD、SBNGD Layout格式不支持连续写模式。该接口不支持量化模式即不支持SetQuantScalar、SetQuantVector接口。BSNGD场景不支持一次计算多行SD需要算子程序中循环计算即(ALayoutInfoN * ALayoutInfoG) / batchA、(BLayoutInfoN * BLayoutInfoG) / batchB均为整数。异步模式不支持IterateBatch搬运到UB上。当开启MixDualMaster双主模式场景时即模板参数enableMixDualMaster设置为true不支持使用该接口。Ascend 950PR/Ascend 950DT上不支持A、B矩阵内存逻辑位置为TPosition::VECOUT的输入输出至VECIN的函数原型当前只支持BSNGD Layout格式。Atlas 推理系列产品AI Core上只支持NORMAL Layout格式。Atlas 推理系列产品AI Core上不支持A、B矩阵内存逻辑位置为TPosition::TSCM的输入。Ascend 950PR/Ascend 950DT上A、B矩阵内存逻辑位置为TPosition::TSCM时只支持NORMAL Layout格式。Atlas 推理系列产品AI Core上Bias不支持复用Bias的shape大小必须为Batch * N。使用该接口时A矩阵、B矩阵不支持int4b_t类型的输入即BatchMatmul不支持int4b_t类型的矩阵输入。支持的数据类型Ascend 950PR/Ascend 950DT支持的数据类型为half、bfloat16_t、int32_t、float。Atlas A3 训练系列产品/Atlas A3 推理系列产品支持的数据类型为half、bfloat16_t、int32_t、float。Atlas A2 训练系列产品/Atlas A2 推理系列产品支持的数据类型为half、bfloat16_t、int32_t、float。Atlas 推理系列产品AI Core支持的数据类型为half、bfloat16_t、int32_t、float。Kirin X90支持的数据类型为half、int32_t。调用示例纯cube模式的函数调用示例请参考调用示例。该示例完成aGM、bGM矩阵乘结果保存到cGm上其中aGM、bGM、cGM数据的layout格式均为NORMAL左矩阵每次计算batchA个MK数据右矩阵每次计算batchB个KN数据。// 定义matmul type typedef AscendC::MatmulTypeAscendC::TPosition::GM, CubeFormat::ND, half, false, LayoutMode::NORMAL aType; typedef AscendC::MatmulTypeAscendC::TPosition::GM, CubeFormat::ND, half, true, LayoutMode::NORMAL bType; typedef AscendC::MatmulTypeAscendC::TPosition::GM, CubeFormat::ND, float, false, LayoutMode::NORMAL cType; typedef AscendC::MatmulTypeAscendC::TPosition::GM, CubeFormat::ND, float biasType; // 创建Matmul实例 constexpr static MatmulConfig MM_CFG GetNormalConfig(false, false, false, BatchMode::BATCH_LESS_THAN_L1); AscendC::MatmulaType, bType, cType, biasType, MM_CFG mm1; REGIST_MATMUL_OBJ(pipe, GetSysWorkSpacePtr(), mm1); mm1.Init(tiling); mm1.SetTensorA(gm_a, isTransposeAIn); mm1.SetTensorB(gm_b, isTransposeBIn); if (tiling.isBias) { mm1.SetBias(gm_bias); } // 多batch Matmul计算 mm1.IterateBatch(gm_c, batchA, batchB, false);该示例完成aGM、bGM矩阵乘结果保存到cGm上其中aGM数据的layout格式为BSNGDbGM数据的layout格式为BSNGDcGM的layout格式为BNGS1S2左矩阵每次计算batchA个SD数据右矩阵每次计算batchB个SD数据。// 定义matmul type typedef AscendC::MatmulTypeAscendC::TPosition::GM, CubeFormat::ND, half, false, LayoutMode::BSNGD aType; typedef AscendC::MatmulTypeAscendC::TPosition::GM, CubeFormat::ND, half, true, LayoutMode::BSNGD bType; typedef AscendC::MatmulTypeAscendC::TPosition::GM, CubeFormat::ND, float, false, LayoutMode::BNGS1S2 cType; typedef AscendC::MatmulTypeAscendC::TPosition::GM, CubeFormat::ND, float biasType; // 创建Matmul实例 AscendC::MatmulaType, bType, cType, biasType mm1; REGIST_MATMUL_OBJ(pipe, GetSysWorkSpacePtr(), mm1); mm1.Init(tiling); int batchC batchA batchB ? batchA : batchB; int g_lay tiling.ALayoutInfoG tiling.BLayoutInfoG ? tiling.ALayoutInfoG : tiling.BLayoutInfoG; // 计算需要多Batch计算循环次数 int for_exent tiling.ALayoutInfoB * tiling.ALayoutInfoN * g_lay / tiling.BatchNum; for (int i 0; i for_exent; i) { // 计算每次多batch计算A/B矩阵的起始地址 int batchOffsetA i * tiling.ALayoutInfoD * batchA; int batchOffsetB i * tiling.BLayoutInfoD * batchB; mm1.SetTensorA(gm_a[batchOffsetA], isTransposeAIn); mm1.SetTensorB(gm_b[batchOffsetB], isTransposeBIn); int idx_c i * batchC; if (tiling.CLayoutInfoG 1 (tiling.BLayoutInfoG ! 1 || tiling.ALayoutInfoG ! 1)) { idx_c idx_c / (tiling.BLayoutInfoG tiling.ALayoutInfoG ? tiling.BLayoutInfoG : tiling.ALayoutInfoG); } if (tiling.isBias) { int batchOffsetBias idx_c * tiling.CLayoutInfoS2; mm1.SetBias(gm_bias[batchOffsetBias]); } int batchOffsetC idx_c * tiling.CLayoutInfoS2; if (C_TYPE::layout LayoutMode::BNGS1S2) { batchOffsetC idx_c * tiling.CLayoutInfoS2 * tiling.CLayoutInfoS1; } // 多batch Matmul计算 mm1.IterateBatch(gm_c[batchOffsetC], batchA, batchB, false); }【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考