SetTensorB

功能说明

设置矩阵乘的右矩阵B。

函数原型

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

参数说明

表1 参数说明
参数名	输入/输出	描述
gm	输入	B矩阵在Global Memory上的首地址。 Atlas推理系列产品AI Core，支持的数据类型为：half/float/int8_t Atlas A2训练系列产品/Atlas 800I A2推理产品，支持的数据类型为：half/float/bfloat16_t/int8_t/int4b_t Atlas 200/500 A2推理产品，支持的数据类型为：half/float/bfloat16_t/int8_t
rightMatrix	输入	B矩阵在TSCM上的首地址或者在VECOUT上的首地址。 Atlas推理系列产品AI Core，支持的数据类型为：half/float/int8_t Atlas A2训练系列产品/Atlas 800I A2推理产品，支持的数据类型为：half/float/bfloat16_t/int8_t/int4b_t Atlas 200/500 A2推理产品，支持的数据类型为：half/float/bfloat16_t/int8_t 若设置TSCM首地址，默认矩阵可全载，已经位于TSCM，Iterate接口无需再进行GM->A1/B1搬运。
bScalar	输入	B矩阵中设置的值。支持传入标量数据，标量数据会被扩展为一个形状为[1, K]的tensor参与矩阵乘计算，tensor的数值均为该标量值。例如，开发者可以通过将bScalar设置为1来实现矩阵A在K方向的reduce sum操作。 Atlas A2训练系列产品/Atlas 800I A2推理产品，支持的数据类型为：half/float Atlas推理系列产品AI Core不支持该参数。 Atlas 200/500 A2推理产品不支持该参数。
isTransposeB	输入	B矩阵是否需要转置。注意：若B矩阵MatmulType ISTRANS参数设置为true，此参数可以为true也可以为false，即运行时可以转置和非转置交替使用；若B矩阵MatmulType ISTRANS参数设置为false，此参数只能设置为false，若强行设置为true，精度会有异常。

返回值

无

支持的型号

Atlas A2训练系列产品/Atlas 800I A2推理产品

Atlas推理系列产品AI Core

Atlas 200/500 A2推理产品

注意事项

传入的TensorB地址空间大小需要保证不小于singleK * singleN。

调用示例

REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling);
mm.SetTensorA(gm_a);
mm.SetTensorB(gm_b);    // 设置右矩阵B
mm.SetBias(gm_bias);
mm.IterateAll(gm_c);

父主题： Matmul