gemm

功能说明

通用矩阵乘，计算表达式如下：

A' = transpose(A) if transA else A

B' = transpose(B) if transB else B

Y = alpha * A' * B' + beta * C。

Tensor A'与Tensor B'的shape后两维(经过对应转置)需要满足矩阵乘(M, K) * (K, N) = (M, N)，支持多维度。

当alpha=1.0且beta=1.0时，计算结果与矩阵乘Matmul计算接口结果相同。

函数原型

gemm(tensor_a, tensor_b, para_dict)

参数说明

• tensor_a：A矩阵，tvm.tensor类型。
• tensor_b：B矩阵，tvm.tensor类型。
• para_dict：扩展参数字典，dict类型，包含如下参数：
  • tensor_alpha：矩阵乘系数，tvm.tensor类型或float类型。
  • tensor_beta：C矩阵系数，tvm.tensor类型或float类型。
  • trans_a：A矩阵是否需要转置，bool类型。
  • trans_b：B矩阵是否需要转置，bool类型。
  • format_a：输入矩阵A的数据排布格式，string类型，取值为：“ND”与“fractal”。
  • format_b：输入矩阵B的数据排布格式，string类型，取值为：“ND”与“fractal”。
  • dst_dtype：输出数据类型，string类型，支持："float16"与"float32"。
  • tensor_c：输入矩阵C，默认值为None，如果取值不为空，矩阵A和矩阵B相乘后的计算结果加上矩阵C，矩阵C的shape支持broadcast，其数据类型要和dst_dtype保持一致。
  • format_out：输出tensor的format，string类型。
  • compress_index：压缩权重矩阵的索引，tvm.tensor类型，仅在alpha=1.0且beta=1.0时使用。
  • offset_a: tensor_a的偏移，仅在alpha=1.0且beta=1.0时使用。
  • offset_b: tensor_b的偏移，仅在alpha=1.0且beta=1.0时使用
  • kernel_name：算子在内核中的名称（即生成的二进制文件与算子描述文件的名称）
  • quantize_params：量化相关参数的。
    量化相关入参，采用字典格式。quantize_params为None表示去使能量化，非None表示使能量化，包含以下参数：

    • quantize_alg：量化模式，取值范围包含“NON_OFFSET” 、“HALF_OFFSET_A”，默认值为“NON_OFFSET”。
    • scale_mode_a：预留参数，暂不使用。
    • scale_mode_b：预留参数，暂不使用。
    • scale_mode_out：出口反量化，量化参数的取值类型，取值范围包含“SCALAR”和“VECTOR”，默认值为“SCALAR”。
    • sqrt_mode_a：预留参数，暂不使用。
    • sqrt_mode_b：预留参数，暂不使用。
    • sqrt_mode_out：表示scale_drq是否开方，取值范围包含“NON_SQRT”和“SQRT”，默认值为“NON_SQRT”。
    • scale_q_a：预留参数，暂不使用。
    • offset_q_a：预留参数，暂不使用。
    • scale_q_b：预留参数，暂不使用。
    • offset_q_b：预留参数，暂不使用。
    • scale_drq：出口反量化/重量化的量化参数的placeholder，默认值为None。

输入tensor支持的数据类型：

Atlas 200/300/500 推理产品：支持的数据类型有float16，float32，int8，uint8，int32。其中int8，uint8，int32会被转换为float16。

Atlas 训练系列产品：支持的数据类型有float16，float32，int8，uint8，int32。其中int8，uint8，int32会被转换为float16。

返回值

y：根据关系运算计算后得到的tensor，tvm.tensor类型。

约束说明

无。

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

调用示例

from tbe import tvm
from tbe import dsl
a_shape = (1024, 256)
b_shape = (256, 512)
bias_shape = (512, )
in_dtype = "float16"
dst_dtype = "float32"
tensor_a = tvm.placeholder(a_shape, name='tensor_a', dtype=in_dtype)
tensor_b = tvm.placeholder(b_shape, name='tensor_b', dtype=in_dtype)
tensor_bias = tvm.placeholder(bias_shape, name='tensor_bias', dtype=dst_dtype)
para_dict = {
    "tensor_bias":tensor_bias,
    "dst_dtype": dst_dtype
}
res = dsl.gemm(tensor_a, tensor_b, para_dict)