SetSingleShape

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	√
Atlas inference product AI Core	√
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

Sets the Matmul single-core computation shape parameters: singleCoreM, singleCoreN, and singleCoreK, with the unit being elements. It is used to modify the shape during runtime, for example, using a Matmul object to process tail blocks. This API has the same function as the SetTail API. You are advised to use this API.

Prototype

__aicore__ inline void SetSingleShape(int singleM, int singleN, int singleK)

Parameters

**Table 1** Parameters
Parameter	Input/Output	Description
singleM	Input	Size of singleCoreM. The unit is element.
singleN	Input	Size of singleCoreN. The unit is element.
singleK	Input	Size of singleCoreK. The unit is element.

Returns

None

Restrictions

None

Example

REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling);

// tailM: number of remaining elements in the M direction; tailN: number of remaining elements in the N direction; tailK: number of remaining elements in the K direction.
// In the tail-core scenario, the number of remaining elements may be less than the number of elements to be computed by a single core. In this case, you need to use SetSingleShape to reset the number of elements for the current computation.
if (tailM < tiling.singleCoreM || tailN < tiling.singleCoreN || tailK < tiling.singleCoreK) {
    matmulObj.SetSingleShape(tailM, tailN, tailK);
}

mm.SetTensorA(gm_a);
mm.SetTensorB(gm_b);
if (tiling.isBias) {
    mm.SetBias(gmBias);
}
mm.IterateAll(gm_c);
mm.End();

Parent topic: Matmul Kernel APIs