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：量化相关参数。此参数在后续版本会过期废弃，新开发算子请不要使用此参数。

输入tensor支持的数据类型：

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

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

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

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

返回值

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

约束说明

此接口暂不支持与其他TBE DSL计算接口混合使用。

支持的型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

Atlas 推理系列产品

Atlas A2训练系列产品

调用示例

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)