等待IterateBatch异步接口返回,支持连续输出到Global Memory。
__aicore__ inline void WaitIterateBatch()
参数名 |
输入/输出 |
描述 |
|---|---|---|
无 |
无 |
NA |
无
Atlas A2训练系列产品/Atlas 800I A2推理产品
Matmul<aType, bType, cType, biasType> mm;
mm.SetTensorA(queryGm[tensorACoreOffset]);
mm.SetTensorB(keyGm[tensorBCoreOffset + sInnerStart * singleProcessSInnerSize *
tilingData->attentionScoreOffestStrideParams.matmulHead], true);
mm.SetTail(singleProcessSOuterSize, mmNNum);
mm.template IterateBatch<false>(workspaceGm[tmp_block_idx * mmResUbSize * sInnerLoopTimes],batchA, batchB, false);
// do some others compute
mm.WaitIterateBatch(); // 等待IterateBatch完成
DataCopy(dstUB, GM); // 进行GM到UB的拷贝