CANN/asc-devkit矩阵批量计算接口

IterateNBatch产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品xAtlas 推理系列产品 AI CorexAtlas 推理系列产品 Vector CorexAtlas 训练系列产品xKirin X90xKirin 9030x功能说明调用一次IterateNBatch会进行N次IterateBatch计算计算出N个多Batch的singleCoreM * singleCoreN大小的C矩阵。在调用该接口前需将MatmulConfig中的isNBatch参数设为true使能多Batch输入多Batch输出功能并调用SetWorkspace接口申请临时空间用于缓存计算结果即IterateNBatch的结果输出至SetWorkspace指定的Global Memory内存中。对于BSNGD、SBNGD、BNGS1S2的Layout格式调用该接口之前需要在tiling中使用SetALayout/SetBLayout/SetCLayout/SetBatchNum设置A/B/C的Layout轴信息和最大BatchNum数对于Normal数据格式则需使用SetBatchInfoForNormal设置A/B/C的M/N/K轴信息和A/B矩阵的BatchNum数。实例化Matmul时通过MatmulType设置Layout类型当前支持3种Layout类型BSNGD、SBNGD、BNGS1S2。函数原型template bool sync true, bool waitIterateBatch false __aicore__ inline void IterateNBatch(const uint32_t batchLoop, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA 0, const uint32_t matrixStrideB 0, const uint32_t matrixStrideC 0)参数说明表 1模板参数说明参数名描述sync获取C矩阵过程分为同步和异步两种模式同步需要同步等待IterateNBatch执行结束后续由开发者自行获取输出到Global Memory上的计算结果。异步不需要同步等待IterateNBatch执行结束。通过该参数设置同步或者异步模式同步模式设置为true异步模式设置为false。默认为同步模式。waitIterateBatch是否需要通过WaitIterateBatch接口等待IterateNBatch执行结束仅在异步场景下使用。默认为false。true需要通过WaitIterateBatch接口等待IterateNBatch执行结束然后由开发者自行获取输出到Global Memory上的计算结果。false不需要通过WaitIterateBatch接口等待IterateNBatch执行结束。调用本接口后需要调用GetBatchTensorC接口获取C矩阵或者由开发者自行处理等待IterateNBatch执行结束的过程。参数名输入/输出描述batchLoop输入当前计算的BMM个数。batchA输入当前单次BMM调用计算左矩阵的batch数。batchB输入当前单次BMM调用计算右矩阵的batch数brc场景batchA/B不相同。enSequentialWrite输入输出是否连续存放数据。matrixStrideA输入A矩阵源操作数相邻nd矩阵起始地址间的偏移默认值是0。matrixStrideB输入B矩阵源操作数相邻nd矩阵起始地址间的偏移默认值是0。matrixStrideC输入该参数预留保持默认值0即可。返回值说明无约束说明单BMM内计算遵循之前的约束条件。对于BSNGD、SBNGD、BNGS1S2 Layout格式输入A、B矩阵多Batch数据总和应小于L1 Buffer的大小。当使能MixDualMaster双主模式场景时即模板参数enableMixDualMaster设置为true不支持使用该接口。调用示例实例功能完成aGM、bGM矩阵乘结果保存到cGm上其中aGM数据的layout格式为BSNGDbGM数据的layout格式为BSNGDcGM的layout格式为BNGS1S2左矩阵每次计算batchA个SD数据右矩阵每次计算batchB个SD数据。// 创建Matmul实例 AscendC::MatmulaType, bType, cType, biasType mm1; AscendC::TPipe pipe; g_cubeTPipePtr pipe; REGIST_MATMUL_OBJ(pipe, GetSysWorkSpacePtr(), mm1); mm1.Init(tiling); int g_lay tiling.ALayoutInfoG tiling.BLayoutInfoG ? tiling.ALayoutInfoG : tiling.BLayoutInfoG; int for_extent tiling.ALayoutInfoB * tiling.ALayoutInfoN * g_lay / tiling.BatchNum; mm1.SetTensorA(gm_a[0], isTransposeAIn); mm1.SetTensorB(gm_b[0], isTransposeBIn); mm1.SetWorkspace(workspaceGM, 0); if (tiling.isBias) { mm1.SetBias(gm_bias[0]); } // 多batch Matmul计算 mm1.IterateNBatch(for_extent, batchA, batchB, false);创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考