SetOrgShape

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

Function

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.
orgN	Input	Sets the size of the original complete shape N. 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.
orgKa	Input	Sets the size of the original complete shape Ka of matrix A. The unit is element.
orgKb	Input	Sets the size of the original complete shape Kb of matrix B. 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.

Returns

None

Restrictions

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

Examples

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);

Parent topic: Matmul Kernel APIs