适配部分算子动态shape场景下infer min max value

mindspore/python/mindspore/ops/_op_impl/tbe/__init__.py

@@ -319,6 +319,7 @@ from .resize_nearest_neighbor_grad_ds import _resize_nearest_neighbor_grad_ds_tb
 from .pad_d import _pad_d_tbe
 from .pad_d_ds import _pad_d_ds_tbe
 from .arg_max_with_value import _arg_max_with_value_tbe
+from .arg_max_with_value_ds import _arg_max_with_value_ds_tbe
 from .arg_min_with_value import _arg_min_with_value_tbe
 from .smooth_l1_loss import _smooth_l1_loss_tbe
 from .smooth_l1_loss_ds import _smooth_l1_loss_ds_tbe

mindspore/python/mindspore/ops/_op_impl/tbe/arg_max_with_value_ds.py

@@ -0,0 +1,39 @@
+# Copyright 2022 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.
+# ============================================================================
+
+"""ArgMaxWithValue op"""
+from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType
+
+arg_max_with_value_ds_op_info = TBERegOp("ArgMaxWithValue") \
+    .fusion_type("ELEMWISE") \
+    .async_flag(False) \
+    .binfile_name("arg_max_with_value.so") \
+    .compute_cost(10) \
+    .kernel_name("arg_max_with_value") \
+    .partial_flag(True) \
+    .dynamic_shape(True) \
+    .attr("axis", "required", "int", "all") \
+    .input(0, "x", False, "required", "all") \
+    .output(0, "indice", False, "required", "all") \
+    .output(1, "values", False, "required", "all") \
+    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default) \
+    .get_op_info()
+
+
+@op_info_register(arg_max_with_value_ds_op_info)
+def _arg_max_with_value_ds_tbe():
+    """ArgMaxWithValue TBE register"""
+    return

mindspore/python/mindspore/ops/operations/array_ops.py

@@ -355,6 +355,20 @@ class Cast(PrimitiveWithInfer):
         if 'min_shape' in x and 'max_shape' in x:
             out['min_shape'] = x['min_shape']
             out['max_shape'] = x['max_shape']
+        if 'min_value' in x and 'max_value' in x:
+            np_dst_type = mstype.dtype_to_nptype(dst_type)
+            if isinstance(x['min_value'], (int, float, tuple, list)):
+                min_value = Tensor(np.array(x['min_value']).astype(np_dst_type))
+            else:
+                min_value = Tensor(x['min_value'].asnumpy().astype(np_dst_type))
+            min_value = tuple(min_value.asnumpy())
+            if isinstance(x['max_value'], (int, float, tuple, list)):
+                max_value = Tensor(np.array(x['max_value']).astype(np_dst_type))
+            else:
+                max_value = Tensor(x['max_value'].asnumpy().astype(np_dst_type))
+            max_value = tuple(max_value.asnumpy())
+            out['min_value'] = min_value
+            out['max_value'] = max_value
         return out
 
 
@@ -892,7 +906,8 @@ class Gather(Primitive):
         >>> output = ops.Gather()(input_params, input_indices, axis)
         >>> print(output)
         [1. 3. 5. 3. 7.]
-        >>> # case2: input_indices is a Tensor with shape (2, 2). When the input_params has one dimension, the output shape is equal to the input_indices shape.
+        >>> # case2: input_indices is a Tensor with shape (2, 2). When the input_params has one dimension,
+        the output shape is equal to the input_indices shape.
         >>> input_indices = Tensor(np.array([[0, 2], [2, 6]]), mindspore.int32)
         >>> axis = 0
         >>> output = ops.Gather()(input_params, input_indices, axis)
@@ -2043,29 +2058,19 @@ class Tile(PrimitiveWithInfer):
             return (True, base_tensor)
         return (False, None)
 
-    def __infer__(self, x, multiples):
-        multiples_v = multiples['value']
-        if multiples_v is None:
-            if len(multiples['shape']) != 1:
-                raise ValueError(f'For \'{self.name}\' the dim of multiples must be 1.')
-            rank = max(len(x['shape']), multiples['shape'][0])
-            out_shape = [-1] * rank
-            # tile can't infer min/max shape if multiples_v is None
-            return {'shape': out_shape,
-                    'dtype': x['dtype'],
-                    'value': None,
-                    'min_shape': [1] * rank,
-                    'max_shape': [1] * rank
-                    }
-
+    def _get_shape_and_range(self, x, multiples):
+        """calculate tile shape and value"""
         x_shp = x['shape']
-        validator.check_value_type(
-            "multiples", multiples_v, [tuple], self.name)
-        for i, multiple in enumerate(multiples_v):
-            validator.check_positive_int(
-                multiple, "multiples[%d]" % i, self.name)
-        validator.check_value_type(
-            "x[\'dtype\']", x["dtype"], mstype.tensor_type, self.name)
+        multiples_v = multiples['value']
+        value = None
+        if multiples_v is None:
+            multiples_v = multiples['min_value']
+        if 'max_shape' in x and 'min_shape' in x:
+            max_shape = x['max_shape']
+            min_shape = x['min_shape']
+        else:
+            max_shape = list(x_shp)
+            min_shape = list(x_shp)
         len_sub = len(multiples_v) - len(x_shp)
         multiples_w = None
         if len_sub == 0:
@@ -2073,19 +2078,85 @@ class Tile(PrimitiveWithInfer):
         if len_sub > 0:
             for i in range(0, len_sub):
                 x_shp.insert(0, 1)
+                min_shape.insert(0, 1)
+                max_shape.insert(0, 1)
             multiples_w = multiples_v
         elif len_sub < 0:
             raise ValueError(f"For '{self.name}', the length of 'multiples' can not be smaller than "
                              f"the dimension of 'input_x', but got length of 'multiples': {len(multiples_v)} "
                              f"and dimension of 'input_x': {len(x_shp)}.")
-        for i, a in enumerate(multiples_w):
-            x_shp[i] *= a
-        value = None
-        if x['value'] is not None:
-            value = Tensor(np.tile(x['value'].asnumpy(), multiples_w))
-        return {'shape': x_shp,
+        if 'max_value' in multiples and 'min_value' in multiples:
+            multiples_v_max = multiples['max_value']
+            multiples_v_min = multiples['min_value']
+            i = 0
+            for a, b in zip(multiples_v_min, multiples_v_max):
+                if isinstance(a, (Tensor_, Tensor)):
+                    a = a.asnumpy()
+                if isinstance(b, (Tensor_, Tensor)):
+                    b = b.asnumpy()
+                x_shp[i] *= a
+                if a != b:
+                    x_shp[i] = -1
+                min_shape[i] *= a
+                max_shape[i] *= b
+                i += 1
+        else:
+            for i, a in enumerate(multiples_w):
+                x_shp[i] *= a
+                max_shape[i] *= a
+                min_shape[i] *= a
+            if x['value'] is not None:
+                value = Tensor(np.tile(x['value'].asnumpy(), multiples_w))
+        out_shape = {
+            'shape': x_shp,
+            'max_shape': max_shape,
+            'min_shape': min_shape
+        }
+        return out_shape, value
+
+    def __infer__(self, x, multiples):
+        multiples_v = multiples['value']
+        if multiples_v is None:
+            if 'max_value' not in multiples or 'min_value' not in multiples:
+                if len(multiples['shape']) != 1:
+                    raise ValueError(f'For \'{self.name}\', the dim of multiples must be 1.')
+                rank = max(len(x['shape']), multiples['shape'][0])
+                out_shape = [-1] * rank
+                return {
+                    'shape': out_shape,
+                    'dtype': x['dtype'],
+                    'value': None,
+                    'max_shape': [1] * rank,
+                    'min_shape': [1] * rank
+                }
+            out_shape, value = self._get_shape_and_range(x, multiples)
+            max_shape = out_shape.get('max_shape', None)
+            min_shape = out_shape.get('min_shape', None)
+            shape = out_shape.get('shape', None)
+            return {
+                'shape': shape,
                 'dtype': x['dtype'],
-                'value': value}
+                'value': value,
+                'max_shape': max_shape,
+                'min_shape': min_shape
+            }
+
+        validator.check_value_type(
+            "multiples", multiples_v, [tuple], self.name)
+        for i, multiple in enumerate(multiples_v):
+            validator.check_positive_int(
+                multiple, "multiples[%d]" % i, self.name)
+        validator.check_value_type(
+            "x[\'dtype\']", x["dtype"], mstype.tensor_type, self.name)
+        out_shp, value = self._get_shape_and_range(x, multiples)
+        shp = out_shp.get('shape', None)
+        out = {'shape': shp,
+               'dtype': x['dtype'],
+               'value': value}
+        if 'max_shape' in x and 'min_shape' in x:
+            out['max_shape'] = out_shp.get('max_shape', None)
+            out['min_shape'] = out_shp.get('min_shape', None)
+        return out
 
 
 class UnsortedSegmentSum(PrimitiveWithInfer):
@@ -4099,6 +4170,28 @@ class TensorScatterUpdate(PrimitiveWithInfer):
     def __init__(self):
         self.init_prim_io_names(inputs=['input_x', 'indices', 'updates'], outputs=['y'])
 
+    def _infer_min_max_value(self, input_x_value, indices_value, updates_value):
+        """TensorScatterUpdate infer min max value"""
+        if isinstance(input_x_value, tuple):
+            input_x_value = list(input_x_value)
+        if isinstance(input_x_value, (Tensor, Tensor_)):
+            input_x_value = input_x_value.asnumpy()
+        indices = indices_value.asnumpy()
+        input_x = np.array(input_x_value)
+        updates = np.array(updates_value)
+        for i, indice in enumerate(indices):
+            input_x[indice] = updates[i]
+        output = tuple(input_x.tolist())
+        return output
+
+    def infer_min_value(self, input_x_value, indices_value, updates_value):
+        """TensorScatterUpdate infer min value"""
+        return self._infer_min_max_value(input_x_value, indices_value, updates_value)
+
+    def infer_max_value(self, input_x_value, indices_value, updates_value):
+        """TensorScatterUpdate infer max value"""
+        return self._infer_min_max_value(input_x_value, indices_value, updates_value)
+
     def infer_shape(self, input_x_shape, indices_shape, updates_shape):
         if len(indices_shape) < 2:
             raise ValueError(f"For '{self.name}', the dimension of 'indices' cannot be less than 2,"

mindspore/python/mindspore/ops/operations/math_ops.py

@@ -25,6 +25,7 @@ from ...common.tensor import Tensor
 from ...common._decorator import deprecated
 from .._utils import get_broadcast_shape
 from ..primitive import Primitive, PrimitiveWithInfer, PrimitiveWithCheck, prim_attr_register, _run_op
+from ..._c_expression import Tensor as Tensor_
 
 
 def _infer_shape_reduce(x, axis, keep_dims, prim_name):
@@ -234,6 +235,27 @@ class Add(_MathBinaryOp):
         >>> print(output.dtype)
         Float32
     """
+    def _add_infer_special_value(self, a, b):
+        """Add infer min max value"""
+        if a is not None and b is not None:
+            if isinstance(a, (Tensor, Tensor_)):
+                a = a.asnumpy()
+            if isinstance(b, (Tensor, Tensor_)):
+                b = b.asnumpy()
+            a = np.array(a)
+            b = np.array(b)
+            out = a + b
+            out = tuple(out.tolist())
+            return out
+        return None
+
+    def infer_min_value(self, x, y):
+        """Add infer min value"""
+        return _add_infer_special_value(x, y)
+
+    def infer_max_value(self, x, y):
+        """Add infer min value"""
+        return _add_infer_special_value(x, y)
 
     def infer_value(self, x, y):
         if x is not None and y is not None:
@@ -1915,6 +1937,27 @@ class Mul(_MathBinaryOp):
         >>> print(output)
         [ 4. 10. 18.]
     """
+    def _infer_min_max_value(self, x, y):
+        """Mul infer min max value"""
+        if x is not None and y is not None:
+            if isinstance(x, (Tensor, Tensor_)):
+                x = x.asnumpy()
+            if isinstance(y, (Tensor, Tensor_)):
+                y = y.asnumpy()
+            x = np.array(x)
+            y = np.array(y)
+            out = x * y
+            out = tuple(out.tolist())
+            return out
+        return None
+
+    def infer_min_value(self, x, y):
+        """Mul infer min value"""
+        return self._infer_min_max_value(x, y)
+
+    def infer_max_value(self, x, y):
+        """Mul infer max value"""
+        return self._infer_min_max_value(x, y)
 
     def infer_dtype(self, x_dtype, y_dtype):
         mul_valid_type = mstype.number_type + (mstype.bool_,)
@@ -2829,6 +2872,27 @@ class Div(_MathBinaryOp):
         >>> print(output.dtype)
         Float32
     """
+    def _div_infer_special_value(self, x, y):
+        """Div infer min max value"""
+        if x is not None and y is not None:
+            if isinstance(x, (Tensor, Tensor_)):
+                x = x.asnumpy()
+            if isinstance(y, (Tensor, Tensor_)):
+                y = y.asnumpy()
+            x = np.array(x)
+            y = np.array(y)
+            out = x / y
+            out = tuple(out.tolist())
+            return out
+        return None
+
+    def infer_min_value(self, x, y):
+        """Div infer min value"""
+        return _div_infer_special_value(x, y)
+
+    def infer_max_value(self, x, y):
+        """Div infer max value"""
+        return _div_infer_special_value(x, y)
 
     def infer_value(self, x, y):
         if x is not None and y is not None:

mindspore/python/mindspore/ops/primitive.py

@@ -576,7 +576,6 @@ class PrimitiveWithInfer(Primitive):
         # in non-graph_mode, it is not necessary to infer min/max shape
         if not is_graph_mode:
             return out
-
         # output does not contain dynamic shape, no need to calculate min/max shape
         def has_dynamic_shape(shp):
             if isinstance(shp, int):
@@ -585,6 +584,27 @@ class PrimitiveWithInfer(Primitive):
                 return any(has_dynamic_shape(e) for e in shp)
             return False
 
+        # calculate min/max value for output
+        def get_specified_value(elems, attr):
+            has_specified_value = False
+            ret_vals = []
+            for elem in elems:
+                if attr in elem:
+                    has_specified_value = True
+                    ret_vals.append(elem[attr])
+                else:
+                    ret_vals.append(elem['value'])
+            return has_specified_value, tuple(ret_vals)
+
+        has_min_value, min_values = get_specified_value(args, 'min_value')
+        has_max_value, max_values = get_specified_value(args, 'max_value')
+        if has_min_value and has_max_value:
+            if hasattr(self, 'infer_min_value'):
+                fn_infer_min_value = getattr(self, 'infer_min_value')
+                out['min_value'] = fn_infer_min_value(*min_values)
+            if hasattr(self, 'infer_max_value'):
+                fn_infer_max_value = getattr(self, 'infer_max_value')
+                out['max_value'] = fn_infer_max_value(*max_values)
         if not has_dynamic_shape(out['shape']):
             return out