IterateNBatch

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	x
Atlas inference product AI Core	x
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

When IterateNBatch is called once, IterateBatch computation is performed for N times to compute N multi-batch matrix C with a size of singleCoreM × singleCoreN. Before calling this API, set isNBatch in MatmulConfig to true to enable the multi-batch input and output function, and call SetWorkspace to allocate temporary space for caching computation results. That is, the result of IterateNBatch is output to the global memory specified by SetWorkspace.

For the BSNGD, SBNGD, and BNGS1S2 layouts, before calling this API, you need to use SetALayout, SetBLayout, SetCLayout, and SetBatchNum in tiling to set the layout axis information and maximum number of batches for matrices A, B, and C. For the NORMAL layout, use SetBatchInfoForNormal to set the M, N, and K axes of matrices A, B, and C, and the number of batches of matrices A and B. During Matmul instantiation, MatmulType is used to set the layout type. Currently, three layout types are supported: BSNGD, SBNGD, and BNGS1S2.

Prototype

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)

Parameters

**Table 1** Template parameters
Parameter	Description
sync	Matrix C can be obtained in synchronous or asynchronous mode. Synchronous: Wait until IterateNBatch is executed. Developers can obtain the computation result output to the global memory. Asynchronous: Do not need to wait until IterateNBatch is executed. This parameter specifies the two modes: true for the synchronous mode and false for the asynchronous mode. The synchronous mode is used by default.
waitIterateBatch	Used only in asynchronous scenarios, indicating whether to use WaitIterateBatch to wait for the completion of IterateNBatch execution. The default value is false. true: Use the WaitIterateBatch API to wait until IterateNBatch is executed. Then, developers can obtain the computation result output to the global memory. false: Do not need to use the WaitIterateBatch API to wait until IterateNBatch is executed. After calling this API, you need to call the GetBatchTensorC API to obtain matrix C or wait until IterateNBatch is executed.

Parameter	Input/Output	Description
batchLoop	Input	Number of currently computed BMMs.
batchA	Input	Number of batches of the left matrix in a single BMM call.
batchB	Input	Number of batches of the right matrix in a single BMM call. batchA and batchB are different in the brc scenario.
enSequentialWrite	Input	Whether the output data is stored continuously.
matrixStrideA	Input	Offset between the start addresses of adjacent nd matrices of the matrix A's source operand. The default value is 0.
matrixStrideB	Input	Offset between the start addresses of adjacent nd matrices of the matrix B's source operand. The default value is 0.
matrixStrideC	Input	This parameter is reserved. Retain the default value 0.

Returns

None

Restrictions

The computation within a single BMM complies with the previous constraints.
For the BSNGD, SBNGD, and BNGS1S2 layout formats, the total batch data input to matrices A and B must be less than the size of L1 Buffer.
This API is not supported when enableMixDualMaster (dual-master mode) is set to true.

Example

In this example, the aGM and bGM matrices are multiplied and the result is saved to cGM. The layout format of the aGM, bGM, and cGM data is BSNGD, BSNGD, and BNGS1S2, respectively. The left matrix computes batchA SD data each time, and the right matrix computes batchB SD data each time. For details about the complete operator sample, see IterateNBatch sample.

// Create a Matmul instance.
AscendC::Matmul<aType, 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]);
}
// Execute multi-batch Matmul computation.
mm1.IterateNBatch(for_extent, batchA, batchB, false);

Parent topic: Matmul Kernel APIs