SetOrgShape

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 original complete shapes M, N, and K during Matmul computation. The unit is the number of elements. It is used to modify the shape during running. For example, reuse the same Matmul object and obtain data from different matrix blocks for computation.

Prototype

__aicore__ inline void SetOrgShape(int orgM, int orgN, int orgK)

__aicore__ inline void SetOrgShape(int orgM, int orgN, int orgKa, int orgKb, int orgKc = 0)

Parameters

**Table 1** Parameters
Parameter	Input/Output	Description
orgM	Input	Sets the size of the original complete shape M. The unit is element. For the Atlas 350 Accelerator Card, when this API is called by a Matmul object created via the MDL template, this parameter is used to set the size of the complete shape M on GM or L1. The unit is element.
orgN	Input	Sets the size of the original complete shape N. The unit is element. For the Atlas 350 Accelerator Card, when this API is called by a Matmul object created via the MDL template, this parameter is used to set the size of the complete shape N on GM or L1. The unit is element.
orgK	Input	Sets the size of the original complete shape K. The unit is element. This parameter can be set if the original complete shape meets the condition of Ka = Kb. For the Atlas 350 Accelerator Card, when this API is called by a Matmul object created via the MDL template, this parameter is used to set the size of the complete shape K on GM or L1. The unit is element. This parameter can be set if the original complete shape meets the condition of Ka = Kb.
orgKa	Input	Sets the size of the original complete shape Ka of matrix A. The unit is element. For the Atlas 350 Accelerator Card, when this API is called by a Matmul object created via the MDL template, this parameter is used to set the size of the complete shape Ka on GM or L1. The unit is element.
orgKb	Input	Sets the size of the original complete shape Kb of matrix B. The unit is element. For the Atlas 350 Accelerator Card, when this API is called by a Matmul object created via the MDL template, this parameter is used to set the size of the complete shape Kb on GM or L1. The unit is element.
orgKc	Input	Sets N of the output matrix C. The unit is element. This parameter can be set when N of the input matrix B is different from that of the output matrix C. The default value is 0, indicating that N of matrix B is used.

For a Matmul object created via the MDL template on the Atlas 350 Accelerator Card, this API does not need to be called if the shape of data on L1 is consistent with that described by orgMIn/orgNIn/orgKIn/orgKaIn/orgKbIn in the tiling API SetOrgShape.
For a Matmul object created via the MDL template on the Atlas 350 Accelerator Card, this API must be called to specify the value of orgM/orgN/orgK/orgKa/orgKb on GM or L1 if the shape of data on L1 is inconsistent with that described by orgMIn/orgNIn/orgKIn/orgKaIn/orgKbIn in the tiling API SetOrgShape.
For example, when the MDL template is used, input matrix A is on L1, input matrix B is on GM, and the shape and size of matrix A on L1 are inconsistent with those described by the original orgMIn/orgKIn/orgKaIn on the tiling side, call the SetOrgShape(orgM, orgN, orgK) or SetOrgShape(orgM, orgN, orgKa, orgKb) API to specify the orgM/orgK/orgKa parameters that are related to matrix A on L1.

Returns

None

Restrictions

This API must be called before SetTensorA, SetTensorB, SetBias, and SetSingleShape.

Examples

Set the original complete shape of a matrix.

REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling);
mm.SetTensorA(gm_a);
mm.SetTensorB(gm_b);
mm.SetBias(gm_bias);
mm.IterateAll(gm_c);
// Reuse the mm object.
mm.SetOrgShape(orgM, orgN, orgK);
mm.SetTensorA(gm_a1);
mm.SetTensorB(gm_b1);
mm.SetBias(gm_bias1);
mm.IterateAll(gm_c1);

For a Matmul object created via the MDL template on the Atlas 350 Accelerator Card, set the complete shape on GM or L1.

For details about a complete example, see matmul_tscm_mdl_setorgshape sample.

REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling);
for (int m = 0; m < mIter_; m++) {
     for (int n = 0; n < nIter_; n++) {
          for (int k = 0; k < kIter_; k++) {
               // Reuse mm and specify the shape of A in L1 and the shape of B in GM.
               mm.SetOrgShape(alignedSingleM, tiling.N, alignedSingleK, tiling.Kb, tiling.N);
               mm.SetSingleShape(curBaseM, curBaseN, curBaseK);
               mm.SetTensorA(tscm_a[offset_a]); // Set aMatrix tscm input
               mm.SetTensorB(gm_b[offset_b]);
               mm.SetBias(gm_bias[offset_bias]);
               mm.Iterate(k != 0);
            }
            matmulObj.GetTensorC(gm_c[offset_c]);
     }
}

Parent topic: Matmul Kernel APIs