Help Center/ Atlas 300 Application (Model 3000)/ TE API Reference/ Compute APIs/ te.lang.cce.matmul(tensor_a, tensor_b, trans_a=False, trans_b=False, alpha_num=1.0, beta_num=0.0, tensor_c=None)
Updated on 2022-03-13 GMT+08:00
View PDF
Share

te.lang.cce.matmul(tensor_a, tensor_b, trans_a=False, trans_b=False, alpha_num=1.0, beta_num=0.0, tensor_c=None)

Performs matrix multiplication as follows: tensor_c=alpha_num * trans_a(tensor_a) * trans_b(tensor_b) + beta_num * tensor_c.

For tensor_a and tensor_b, the last two dimensions of shape (after transposition) must meet the following condition: Matrix multiplication (M, K) * (K, N) = (M, N). Only one batch is supported. The tensor_a data layout must comply with the fractal structure of L0A, and the tensor_b data layout must comply with the fractal structure of L0B. In mini mode, only the float16 data type is supported.

This API is defined in mmad_compute.py.

Parameter Description

  • tensor_a: matrix A, tvm.tensor type
  • tensor_b: matrix B, tvm.tensor type
  • trans_a: whether matrix A is transposed, Boolean type
  • trans_b: whether matrix B is transposed, Boolean type
  • alpha_num: matrix A*B coefficient. Only 1.0 is supported.
  • beta_num: matrix C coefficient. Only 0.0 is supported.
  • tensor_c: matrix C, tvm.tensor type. Since beta_num supports only 0.0, this parameter is used as a reserved interface.

Return Value

tensor_c: tensor calculated based on the relational calculation, tvm.tensor type

Calling Example

import tvm
import te.lang.cce
a_shape = (1024, 256)
b_shape = (256, 512)
a_fractal_shape = (a_shape[0] // 16, a_shape[1] // 16, 16, 16)
b_fractal_shape = (b_shape[0] // 16, b_shape[1] // 16, 16, 16)
in_dtype = "float16"
tensor_a = tvm.placeholder(a_fractal_shape, name='tensor_a', dtype=in_dtype)
tensor_b = tvm.placeholder(a_fractal_shape, name='tensor_b', dtype=in_dtype)
res = te.lang.cce.matmul(tensor_a, tensor_b, False, False)
Parent topic: Compute APIs