SetTensorB

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 right matrix B for matrix multiplication.

Prototype

__aicore__ inline void SetTensorB(const GlobalTensor<SrcBT>& gm, bool isTransposeB = false)

__aicore__ inline void SetTensorB(const LocalTensor<SrcBT>& rightMatrix, bool isTransposeB = false)

__aicore__ inline void SetTensorB(SrcBT bScalar)

The Atlas inference product AI Core does not support the SetTensorB(SrcBT bScalar) API prototype.

The Atlas 200I/500 A2 inference product do not support the SetTensorB(SrcBT bScalar) API prototype.

Parameters

**Table 1** Template parameters
Parameter	Description
SrcBT	Operand data type.

**Table 2** Parameters
Parameter	Input/Output	Description
gm	Input	Matrix B. The type is GlobalTensor. SrcBT indicates the data type of matrix B. For the Atlas 350 Accelerator Card, the supported data types are half, float, bfloat16_t, int8_t, fp8_e4m3fn_t, fp8_e5m2_t, and hifloat8_t. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half, float, bfloat16_t, int8_t, and int4b_t. For the Atlas inference product AI Core, the supported data type is half, float, and int8_t. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half, float, bfloat16_t, int8_t, and int4b_t. For the Atlas 200I/500 A2 inference product, the supported data types are half, float, bfloat16_t, and int8_t.
rightMatrix	Input	Matrix B. The type is LocalTensor, and TPosition can be TSCM or VECOUT. SrcBT indicates the data type of matrix B. For the Atlas 350 Accelerator Card, the supported data types are half, float, bfloat16_t, int8_t, fp8_e4m3fn_t, fp8_e5m2_t, and hifloat8_t. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half, float, bfloat16_t, int8_t, and int4b_t. For the Atlas inference product AI Core, the supported data types are half, float, and int8_t. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half, float, bfloat16_t, int8_t, and int4b_t. For the Atlas 200I/500 A2 inference product, the supported data types are half, float, bfloat16_t, and int8_t. If the initial address in the TSCM is set, the matrix can be fully loaded by default. In this case, the Iterate API does not need to transfer data from GM to A1/B1.
bScalar	Input	Value set in matrix B. Scalar data can be passed. The passed scalar data is expanded into a tensor with the shape of [1, K] to participate in matrix multiplication. The tensor values are all scalar values. For example, a developer can implement a reduce sum operation on matrix B in the K direction by setting bScalar to 1. SrcBT indicates the data type of matrix B. For the Atlas 350 Accelerator Card, the supported data types are half and float. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half and float. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half and float. This parameter is not supported by the Atlas inference product AI Core. This parameter is not supported by the Atlas 200I/500 A2 inference product.
isTransposeB	Input	Whether matrix B should be transposed. Notes: If MatmulType ISTRANS of matrix B is set to true, this parameter can be set to either true or false at runtime, allowing for alternating use of transposed and non-transposed matrix B. If MatmulType ISTRANS of matrix B is set to false, this parameter can only be set to false. If it is forcibly set to true, the precision will be abnormal. For input types other than half and bfloat16_t, to ensure that the L1 Buffer space calculation size is consistent between the tiling side and the kernel side, and to guarantee correct result precision, the value of this parameter must be the same as the value of ISTRANS defined for matrix B's MatmulType on the kernel side and the value of isTrans of the SetBType() API on the tiling side. In other words, all three parameters must be set to true or false simultaneously.

Returns

None

Restrictions

Ensure that the size of the input TensorB address space is greater than or equal to singleK x singleN.

Example

REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling);
mm.SetTensorA(gm_a);
// Example 1: The right matrix is in the global memory.
mm.SetTensorB(gm_b, isTransposeB);
if (tiling.isBias) {
    mm.SetBias(gmBias);
}
mm.IterateAll(gm_c);
mm.End();
// Example 2: The right matrix is in the local memory.
mm.SetTensorB(local_a, isTransposeB);
// Example 3: Set scalar data.
mm.SetTensorB(scalar_a, isTransposeB);

Parent topic: Matmul Kernel APIs