def CusMatMulCube(input_x1, input_x2, bias=None, output_y={}, trans_a=False, trans_b=False, kernel_name="matmulcube"):
"""
calculating matrix multiplication with bias, C = A*B + bias, support input
data with fractal format.
Parameters:
shape_a: list or tuple
Shape of the first tensor a with rank > 1
shape_b: list or tuple
Shape of the second tensor b with the same type with a,
and shape_a, shape_b must be 2 dims
src_dtype: str
The data type of input, support "float32", "float16"
dst_dtype: str
The data type of output, support "float32", "float16"
trans_a: bool
If True, shape_a == transposed before multiplication
trans_b: bool
If True, shape_b == transposed before multiplication
is_fractal: bool
If True, the input data format of a and b must be fractal format
shape_bias: list or tuple
Shape of bias, only support the input data format with ND
Returns
-------
None
"""
shape_a = input_x1.get("ori_shape")
shape_b = input_x2.get("ori_shape")
if shape_a is not None:
if len(shape_a) < 2:
shape_a = input_x1.get("shape")
if shape_b is not None:
if len(shape_b) < 2:
shape_b = input_x2.get("shape")
shape_a = list(shape_a)
shape_b = list(shape_b)
if input_x1.get("format") == "FRACTAL_NZ":
shape_a = _get_input_shape(shape_a)
shape_b = _get_input_shape(shape_b)
util.check_kernel_name(kernel_name)
util.check_shape_rule(shape_a)
util.check_shape_rule(shape_b)
util.check_shape_size(shape_a, SHAPE_SIZE_LIMIT)
util.check_shape_size(shape_b, SHAPE_SIZE_LIMIT)
if input_x1.get("format") == "FRACTAL_NZ":
shape_a = [shape_a[1], shape_a[0]]
trans_a = bool(1 - trans_a)
if input_x2.get("format") == "FRACTAL_NZ":
shape_b = [shape_b[1], shape_b[0]]
trans_b = bool(1 - trans_b)
shape_bias = ()
if bias is not None and bool(bias):
shape_bias = bias.get("shape")
shape_bias = list(shape_bias)
shape_bias = _get_bias(shape_bias)
src_dtype = input_x1.get("dtype").lower()
dst_dtype = output_y.get("dtype").lower()
if src_dtype in ("float32", "int32"):
matmul_vector_cce(shape_a, shape_b, src_dtype, trans_a, trans_b, shape_bias, kernel_name)
return
_shape_check(shape_a, shape_b, shape_bias, src_dtype, trans_a, trans_b)
m_shape = shape_a[len(shape_a) - 2]
km_shape = shape_a[len(shape_a) - 1]
kn_shape = shape_b[len(shape_a) - 2]
n_shape = shape_b[len(shape_a) - 1]
if src_dtype == "float16":
block_reduce = cce.BLOCK_REDUCE
block_in = cce.BLOCK_IN
block_out = cce.BLOCK_OUT
if trans_a and km_shape == 1:
block_in = cce.BLOCK_VECTOR
if not trans_a and m_shape == 1:
block_in = cce.BLOCK_VECTOR
if trans_b and kn_shape == 1:
block_out = cce.BLOCK_VECTOR
if not trans_b and n_shape == 1:
block_out = cce.BLOCK_VECTOR
if trans_a:
shape_a_temp = (m_shape // block_reduce, km_shape // block_in, block_reduce, block_in)
else:
shape_a_temp = (m_shape // block_in, km_shape // block_reduce, block_in, block_reduce)
if trans_b:
shape_b_temp = (kn_shape // block_out, n_shape // block_reduce, block_reduce, block_out)
else:
shape_b_temp = (kn_shape // block_reduce, n_shape // block_out, block_out, block_reduce)
if input_x1.get("format") == "FORMAT_FRACTAL_Z":
shape_a_temp = (shape_a_temp[0], shape_a_temp[1], shape_a_temp[2], shape_a_temp[3])
format_a = "fractal"
elif input_x1.get("format") == "FRACTAL_NZ":
shape_a_temp = (shape_a_temp[0], shape_a_temp[1], shape_a_temp[2], shape_a_temp[3])
format_a = "FRACTAL_NZ"
else:
shape_a_temp = (shape_a[len(shape_a) - 2], shape_a[len(shape_a) - 1])
format_a = "ND"
if input_x2.get("format") == "FORMAT_FRACTAL_Z":
shape_b_temp = (shape_b_temp[0], shape_b_temp[1], shape_b_temp[2], shape_b_temp[3])
format_b = "fractal"
elif input_x2.get("format") == "FRACTAL_NZ":
shape_b_temp = (shape_b_temp[0], shape_b_temp[1], shape_b_temp[2], shape_b_temp[3])
format_b = "FRACTAL_NZ"
else:
shape_b_temp = (shape_b[len(shape_b) - 2], shape_b[len(shape_b) - 1])
format_b = "ND"
tensor_bias = None
tensor_a = tvm.placeholder(shape_a_temp, name='tensor_a',
dtype=src_dtype)
tensor_b = tvm.placeholder(shape_b_temp, name='tensor_b',
dtype=src_dtype)
if shape_bias:
tensor_bias = tvm.placeholder(shape_bias, name='tensor_bias',
dtype=dst_dtype)
result = te.lang.cce.matmul(tensor_a, tensor_b, trans_a, trans_b, format_a=format_a,
format_b=format_b, dst_dtype=dst_dtype, tensor_bias=tensor_bias)
with tvm.target.cce():
schedule = generic.auto_schedule(result)
tensor_list = [tensor_a, tensor_b, result]
if shape_bias:
tensor_list = [tensor_a, tensor_b, tensor_bias, result]
config = {"print_ir": False,
"name": kernel_name,
"tensor_list": tensor_list}
te.lang.cce.cce_build_code(schedule, config)
def mat_mul(input_x1,
input_x2,
bias,
offset_w={},
output_y={},
trans_a=False,
trans_b=False,
offset_x=0,
kernel_name="matmul"):
"""
calculating matrix multiplication with bias, C = A*B + bias, support input
data with fractal format.
Parameters:
input_x1: dict
A dict object, contains a matrix(2D Tensor) 's type and
shape and format, the type can be float16,
float32, int32, the shape must be 2-dimensional,
the format can be [ND, NHWC, FRACTAL_NZ]
input_x2: dict
A dict object, contains a matrix(2D Tensor) 's type and
shape and format, the type can be float16,
float32, int32, the shape must be 2-dimensional,
the format can be [ND, NHWC, FRACTAL_NZ]
bias: dict
A dict object, contanis a 1-dimensional tensor's info:
the shape and type and format, the type can be float16,
float32, int32, the shape must be 1-dimensional,
the format can be [ND, NHWC]
output_y: dict
A dict object, contains a matrix(2D Tensor) 's type and
shape and format, the type can be float16,
float32, int32, the shape must be 2-dimensional,
the format can be [ND, NHWC, FRACTAL_NZ]
trans_a: bool
If True, shape_a == transposed before multiplication
trans_b: str
If true, the shape in input_x2 must be transposed before multiplication
kernel_name: str
cce kernel name, default value is "matmul"
Returns
-------
None
"""
shape_a = input_x1.get("ori_shape")
shape_b = input_x2.get("ori_shape")
shape_a_length = len(shape_a)
shape_b_length = len(shape_b)
if shape_a is not None:
if shape_a_length < 2:
shape_a = input_x1.get("shape")
if shape_b is not None:
if shape_b_length < 2:
shape_b = input_x2.get("shape")
shape_a = list(shape_a)
shape_b = list(shape_b)
if input_x1.get("format") == "FRACTAL_NZ":
shape_a = _get_input_shape(shape_a, trans_a)
shape_b = _get_input_shape_b(shape_b, trans_b)
check_shape(shape_a, param_name="input_x1")
check_shape(shape_b, param_name="input_x2")
shape_bias = ()
if bias is not None and bool(bias):
shape_bias = bias.get("shape")
shape_bias = list(shape_bias)
shape_bias = _get_bias(shape_bias)
src_dtype = input_x1.get("dtype").lower()
dst_dtype = output_y.get("dtype").lower()
target_type = ["float32", "int32"]
if src_dtype in target_type:
matmul_vector_cce(shape_a, shape_b, src_dtype, trans_a, trans_b,
shape_bias, kernel_name)
return
if src_dtype != "int8":
_shape_check(shape_a, shape_b, shape_bias, src_dtype, trans_a, trans_b)
else:
_shape_check_quantification(shape_a, shape_b, trans_a, trans_b,
input_x1.get("format"))
shape_a_temp = input_x1.get("shape")
shape_b_temp = input_x2.get("shape")
tensor_bias = None
if src_dtype == "int8":
format_a = "FRACTAL_NZ"
format_b = "FRACTAL_Z"
else:
format_a = "FRACTAL_NZ"
format_b = "FRACTAL_NZ"
tensor_a = tvm.placeholder(shape_a_temp,
name='tensor_a',
attrs={'format': format_a},
dtype=src_dtype)
tensor_b = tvm.placeholder(shape_b_temp,
name='tensor_b',
attrs={'format': format_b},
dtype=src_dtype)
shape_bias_length = len(shape_bias)
if shape_bias_length > 0:
tensor_bias = tvm.placeholder(shape_bias,
name='tensor_bias',
dtype=dst_dtype)
result = mat_mul_compute_self(tensor_a, tensor_b, tensor_bias, offset_w,
output_y, trans_a, trans_b, offset_x,
kernel_name)
with tvm.target.cce():
schedule = generic.auto_schedule(result)
tensor_list = [tensor_a, tensor_b, result]
if shape_bias_length > 0:
tensor_list = [tensor_a, tensor_b, tensor_bias, result]
config = {
"print_ir": False,
"name": kernel_name,
"tensor_list": tensor_list
}
te.lang.cce.cce_build_code(schedule, config)