forked from huawei/mindspore2022
123 lines
4.8 KiB
Python
123 lines
4.8 KiB
Python
# Copyright 2019 Huawei Technologies Co., Ltd
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""dsl create helping function"""
|
|
import _akg
|
|
from _akg.utils import format_transform as ft_util
|
|
|
|
class TensorUtils:
|
|
"""Class for creating tensor."""
|
|
CREATE_SCH_ONLY = 'create_sch_only'
|
|
|
|
@classmethod
|
|
def get_tensor_attrs(cls, tensor):
|
|
"""get tensor attrs."""
|
|
tensor_attrs = dict()
|
|
if "attrs" in dir(tensor.op):
|
|
tensor_attrs = dict(tensor.op.attrs.items())
|
|
return tensor_attrs
|
|
|
|
@classmethod
|
|
def update_tensor_attrs(cls, tensor, attrs):
|
|
"""update tensor attrs."""
|
|
tensor_attrs = cls.get_tensor_attrs(tensor)
|
|
tensor_attrs.update(attrs)
|
|
tensor = _akg.tvm.compute(tensor.shape,
|
|
lambda *indice: tensor[indice],
|
|
name=tensor.op.name,
|
|
tag=tensor.op.tag,
|
|
attrs=tensor_attrs)
|
|
return tensor
|
|
|
|
@classmethod
|
|
def is_create_sch_only(cls, tensor):
|
|
tensor_attrs = cls.get_tensor_attrs(tensor)
|
|
if cls.CREATE_SCH_ONLY in tensor_attrs.keys():
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def is_output_value(cls, tensor):
|
|
"""check output value."""
|
|
return not cls.is_create_sch_only(tensor)
|
|
|
|
@classmethod
|
|
def inplace_set(cls, input_tensor, output_tensor, buffer_name="data_buf"):
|
|
"""inplace set."""
|
|
input_tensor_shape = ft_util.get_shape(input_tensor)
|
|
output_tensor_shape = ft_util.get_shape(output_tensor)
|
|
if not input_tensor_shape == output_tensor_shape:
|
|
raise RuntimeError("Shape of the input_tensor and the output_tensor should be equal, "
|
|
"but got %s and %s"%(input_tensor_shape, output_tensor_shape))
|
|
output_tensor = cls.update_tensor_attrs(output_tensor, {cls.CREATE_SCH_ONLY: 1})
|
|
data_buf = _akg.tvm.decl_buffer(input_tensor.shape, input_tensor.dtype, name=buffer_name)
|
|
binds_info = {input_tensor: data_buf, output_tensor: data_buf}
|
|
return output_tensor, binds_info
|
|
|
|
@classmethod
|
|
def inplace_set_tensors(cls, input_tensors, output_tensors, buffer_names=None):
|
|
"""
|
|
inplace set for tensors
|
|
|
|
Args:
|
|
in_tensors (Union[list, tuple]): Origin input tensors.
|
|
out_tensors (Union[list, tuple]): Origin output tensors.
|
|
buffer_names (Union[list, tuple] or None): Buffer names used to bind.
|
|
|
|
Return:
|
|
inplace_tensors (list): Output tensors with the inplace info.
|
|
binds_infos (dict): Dictionary that maps the input tensor and the output
|
|
tensor to buffer.
|
|
"""
|
|
if not buffer_names:
|
|
buffer_names = ["data_buf_%s" % i for i in range(len(input_tensors))]
|
|
for arg in (input_tensors, output_tensors, buffer_names):
|
|
if not isinstance(arg, (tuple, list)):
|
|
raise RuntimeError("arg must be tuple or list!")
|
|
if len(input_tensors) != len(output_tensors) or len(input_tensors) != len(buffer_names):
|
|
raise RuntimeError("length of the input_tensors, output_tensors and buffer_names must be equal!")
|
|
|
|
inplace_tensors = []
|
|
binds_infos = dict()
|
|
for input_tensor, output_tensor, buffer_name in zip(input_tensors, output_tensors, buffer_names):
|
|
inplace_tensor, binds_info = cls.inplace_set(input_tensor, output_tensor, buffer_name)
|
|
inplace_tensors.append(inplace_tensor)
|
|
binds_infos.update(binds_info)
|
|
return inplace_tensors, binds_infos
|
|
|
|
def produce_shapes(shape1, shape2):
|
|
"""two input shapes produce three output shape."""
|
|
shape1 = list(shape1)
|
|
shape2 = list(shape2)
|
|
flag = 0
|
|
if len(shape1) < len(shape2):
|
|
shape1, shape2 = shape2, shape1
|
|
flag = 1
|
|
|
|
output_shape_len = len(shape1)
|
|
dec = output_shape_len - len(shape2)
|
|
for i in range(dec):
|
|
shape2 = [1] + shape2
|
|
|
|
out_shape = []
|
|
for i in range(output_shape_len):
|
|
if (shape1[i] != shape2[i]) and (shape1[i] != 1) and (shape2[i] != 1):
|
|
raise RuntimeError("input shapes not match!")
|
|
out_shape.append(shape1[i] if shape1[i] > shape2[i] else shape2[i])
|
|
|
|
if flag == 1:
|
|
shape1, shape2 = shape2, shape1
|
|
|
|
return shape1, shape2, out_shape
|