SetBatchNum

Applicability

Product	Supported
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	x
Atlas inference product's AI Core	√
Atlas inference product's Vector Core	x
Atlas training products	x

Function

Resets the number of batches for multi-batch computation without changing tiling.

Prototype

__aicore__ inline void SetBatchNum(int32_t batchA, int32_t batchB)

Parameters

**Table 1** Parameters
Parameter	Input/Output	Description
batchA	Input	Number of batches of matrix A for one computation
batchB	Input	Number of batches of matrix B for one computation

Returns

None

Restrictions

This API is not supported when enableMixDualMaster (dual-master mode) is set to true.
This API can be called only in Cube-only mode (only matrix computation).

Example

// CUBE-ONLY mode
#define ASCENDC_CUBE_ONLY
typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, AType, false, LayoutMode::NORMAL> aType; 
typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BType, false, LayoutMode::NORMAL> bType;
typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, CType, false, LayoutMode::NORMAL> cType;
typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BiasType> biasType;
AscendC::Matmul<aType, bType, cType, biasType> mm1;
mm1.SetTensorA(gm_a, isTransposeAIn);
mm1.SetTensorB(gm_b, isTransposeBIn);
if(tiling.isBias) {
    mm1.SetBias(gm_bias);
}
mm1.SetBatchNum(batchA, batchB);
// Multi-batch Matmul computation
mm1.IterateBatch(gm_c, false, 0, false);

Parent topic: Matmul Kernel APIs