mindspore 2.6.0__cp311-cp311-win_amd64.whl → 2.7.0rc1__cp311-cp311-win_amd64.whl

mindspore/ops_generate/pyboost/op_template_parser.py

@@ -29,7 +29,7 @@ from common.op_proto import OpProto
 from . import pyboost_utils
 from .pyboost_utils import get_input_dtype, tuple_input_to_cpp_type, get_return_type, \
     number_input_to_cpp_type, get_const_number_convert, get_tuple_input_convert, is_optional_param, \
-    get_input_args_type_str
+    get_input_args_type_str, basic_type_convert_str, input_dtype_to_cpp_type
 
 
 class OpTemplateParser:
@@ -45,8 +45,42 @@ class OpTemplateParser:
 
     def __init__(self, op_proto: OpProto):
         self.op_proto = op_proto
-
-    def _parse_call_args_types(self, op_args):
+        self.tensor_arg_handler_prt_template = Template(
+            "parse_args.arg_list_[${idx}] = "
+            "py::cast((*pynative::${func_str}(\"${func_name}\", \"${op_arg_name}\", "
+            "parse_args.arg_list_[${idx}]))->value());\n"
+            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
+            "parse_args.dst_types_[${idx}] = ${new_type};\n"
+        )
+        self.function_arg_handler_prt_template = Template(
+            "parse_args.arg_list_[${idx}] = "
+            "py::cast((*${func_str}(\"${func_name}\", \"${op_arg_name}\", parse_args.arg_list_[${idx}]))->value());\n"
+            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
+            "parse_args.dst_types_[${idx}] = ${new_type};\n"
+        )
+        self.arg_handler_template = Template(
+            "parse_args.arg_list_[${idx}] = "
+            "py::cast(pynative::${func_str}(\"${func_name}\", \"${op_arg_name}\", parse_args.arg_list_[${idx}]));\n"
+            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
+            "parse_args.dst_types_[${idx}] = ${new_type};\n"
+        )
+        self.arg_handler_optional_template = Template(
+            'if (!py::isinstance<py::none>(parse_args.arg_list_[${idx}])) {\n'
+            '  ${arg_handler_str}\n'
+            '}\n'
+        )
+        self.arg_handler_type_map = {"to_2d_paddings": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                     "dtype_to_type_id": "ops::OP_DTYPE::DT_INT",
+                                     "to_kernel_size": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                     "to_strides": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                     "str_to_enum": "ops::OP_DTYPE::DT_INT",
+                                     "to_pair": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                     "to_dilations": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                     "to_output_padding": "ops::OP_DTYPE::DT_TUPLE_INT",
+                                     "to_rates": "ops::OP_DTYPE::DT_TUPLE_INT"}
+
+    @staticmethod
+    def _parse_call_args_types(op_proto, basic_type=False, is_convert=False):
         """
         Parses the data types of the call arguments for the operator.
 
@@ -57,22 +91,37 @@ class OpTemplateParser:
             list: A list of data types for the call arguments.
         """
         call_args_types = []
-        for op_arg in op_args:
+        if basic_type and not op_proto.op_view:
+            basic_type = False
+            raise Exception("Only view op support basic type now, please check.")
+        for op_arg in op_proto.op_args:
             is_optional = is_optional_param(op_arg)
-            call_args_types.append(get_input_dtype(
-                op_arg.arg_dtype, is_optional))
+            if is_convert:
+                if op_arg.is_type_id:
+                    call_args_types.append('TypeId')
+                    continue
+                call_args_types.append(input_dtype_to_cpp_type(
+                    op_arg.arg_dtype, is_optional))
+            else:
+                call_args_types.append(get_input_dtype(
+                    op_arg.arg_dtype, is_optional, basic_type))
         return call_args_types
 
-    def parse_call_args_with_types(self):
+    def parse_call_args_with_types(self, basic_type=False, is_convert=False):
         """
         Parses the original call arguments and their types for the operator.
 
         Returns:
             list: A list of formatted strings representing the call arguments with their types.
         """
-        call_args = self.parse_original_call_args(self.op_proto.op_args)
-        call_args_types = self._parse_call_args_types(self.op_proto.op_args)
+        call_args = OpTemplateParser.parse_original_call_args(self.op_proto.op_args)
+        call_args_after_convert, _, _ = self.op_args_converter()
+        call_args_types = self._parse_call_args_types(self.op_proto, basic_type, is_convert)
         call_args_with_types = []
+        if is_convert:
+            for type_name, arg_name in zip(call_args_types, call_args_after_convert):
+                call_args_with_types.append("const " + type_name + " &" + arg_name)
+            return call_args_with_types
         for type_name, arg_name in zip(call_args_types, call_args):
             call_args_with_types.append("const " + type_name + " &" + arg_name)
         return call_args_with_types
@@ -90,7 +139,7 @@ class OpTemplateParser:
         need_malloc_tensors = []
         tensor_list_convert = []
         call_args_with_tensor = []
-        call_args = self.parse_original_call_args(self.op_proto.op_args)
+        call_args = OpTemplateParser.parse_original_call_args(self.op_proto.op_args)
         for op_arg, call_arg in zip(self.op_proto.op_args, call_args):
             if pyboost_utils.is_tensor(op_arg):
                 call_arg = op_arg.arg_name + "_tensor"
@@ -105,7 +154,8 @@ class OpTemplateParser:
                 call_args_with_tensor.append(call_arg)
         return need_malloc_tensors, tensor_list_convert, call_args_with_tensor
 
-    def parse_original_call_args(self, op_args):
+    @staticmethod
+    def parse_original_call_args(op_args):
         """
         Parses the original call arguments from the operator prototype.
 
@@ -139,7 +189,7 @@ class OpTemplateParser:
         call_args_after_convert = []
         value_tuple_convert = []
         const_number_convert = []
-        call_args = self.parse_original_call_args(self.op_proto.op_args)
+        call_args = OpTemplateParser.parse_original_call_args(self.op_proto.op_args)
         for op_arg, call_arg in zip(self.op_proto.op_args, call_args):
             if number_input_to_cpp_type(op_arg.arg_dtype):
                 call_args_after_convert.append(call_arg + "_imm")
@@ -199,8 +249,7 @@ class OpTemplateParser:
         args_signature = self.op_proto.op_args_signature
         if args_signature is not None:
             dtype_group = args_signature.dtype_group
-            indexes = {arg.arg_name: index for index,
-                       arg in enumerate(self.op_proto.op_args)}
+            indexes = {arg.arg_name: index for index, arg in enumerate(self.op_proto.op_args)}
             if dtype_group is not None:
                 match = re.findall(r'\((.*?)\)', dtype_group)
                 for item in match:
@@ -223,10 +272,10 @@ class OpTemplateParser:
             list: A list of call arguments that are tensors.
         """
         call_args_tensor = []
-        call_args_types = self._parse_call_args_types(self.op_proto.op_args)
-        call_args = self.parse_original_call_args(self.op_proto.op_args)
+        call_args_types = self._parse_call_args_types(self.op_proto)
+        call_args = OpTemplateParser.parse_original_call_args(self.op_proto.op_args)
         for _type, arg_name in zip(call_args_types, call_args):
-            if _type in ("mindspore::tensor::BaseTensorPtr", "std::optional<mindspore::tensor::BaseTensorPtr>"):
+            if _type in ("mindspore::tensor::TensorPtr", "std::optional<mindspore::tensor::TensorPtr>"):
                 call_args_tensor.append(arg_name)
         return call_args_tensor
 
@@ -258,8 +307,8 @@ class OpTemplateParser:
         """
         returns_type = []
         type_convert_to_base = {
-            'std::vector<mindspore::tensor::TensorPtr>': 'std::vector<mindspore::tensor::BaseTensorPtr>',
-            'mindspore::tensor::TensorPtr': 'mindspore::tensor::BaseTensorPtr'
+            'std::vector<mindspore::tensor::TensorPtr>': 'std::vector<mindspore::tensor::TensorPtr>',
+            'mindspore::tensor::TensorPtr': 'mindspore::tensor::TensorPtr'
         }
         for return_obj in self.op_proto.op_returns:
             temp_return = get_return_type(return_obj.arg_dtype)
@@ -316,14 +365,6 @@ class OpTemplateParser:
 
         return op_outputs, call_outputs
 
-    def _is_input_arg(self, arg_name, op_name):
-        res = False
-        if op_name in K.INPUT_NAME_MAP and arg_name == K.INPUT_NAME_MAP[op_name]:
-            res = True
-        elif op_name not in K.INPUT_NAME_MAP and arg_name in K.INPUT_ARGS_NAME:
-            res = True
-        return res
-
     def generate_signature_str(self, kw_only_args=None, varargs=None, *, is_tensor_api: bool) -> str:
         """
         Generates a single function signature string for the given operation prototype.
@@ -348,7 +389,7 @@ class OpTemplateParser:
         for arg in self.op_proto.op_args:
             arg_name = arg.arg_name
 
-            if is_tensor_api and self._is_input_arg(arg_name, op_name):
+            if is_tensor_api and _is_input_arg(arg_name, op_name):
                 continue
 
             single_arg = ''
@@ -398,42 +439,9 @@ class OpTemplateParser:
         Returns:
             str: Generated argument handler processing code.
         """
-        tensor_arg_handler_prt_template = Template(
-            "parse_args.arg_list_[${idx}] = "
-            "py::cast((*pynative::${func_str}(\"${func_name}\", \"${op_arg_name}\", \
-parse_args.arg_list_[${idx}]))->value());\n"
-            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
-            "parse_args.dst_types_[${idx}] = ${new_type};\n"
-        )
-        function_arg_handler_prt_template = Template(
-            "parse_args.arg_list_[${idx}] = "
-            "py::cast((*${func_str}(\"${func_name}\", \"${op_arg_name}\", parse_args.arg_list_[${idx}]))->value());\n"
-            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
-            "parse_args.dst_types_[${idx}] = ${new_type};\n"
-        )
         arg_handler_prt_template = (
-            tensor_arg_handler_prt_template) if is_tensor_api else function_arg_handler_prt_template
+            self.tensor_arg_handler_prt_template) if is_tensor_api else self.function_arg_handler_prt_template
 
-        arg_handler_template = Template(
-            "parse_args.arg_list_[${idx}] = "
-            "py::cast(pynative::${func_str}(\"${func_name}\", \"${op_arg_name}\", parse_args.arg_list_[${idx}]));\n"
-            "parse_args.src_types_[${idx}] = ops::OP_DTYPE::DT_BEGIN;\n"
-            "parse_args.dst_types_[${idx}] = ${new_type};\n"
-        )
-        arg_handler_optional_template = Template(
-            'if (!py::isinstance<py::none>(parse_args.arg_list_[${idx}])) {\n'
-            '  ${arg_handler_str}\n'
-            '}\n'
-        )
-        arg_handler_type_map = {"to_2d_paddings": "ops::OP_DTYPE::DT_TUPLE_INT",
-                                "dtype_to_type_id": "ops::OP_DTYPE::DT_INT",
-                                "to_kernel_size": "ops::OP_DTYPE::DT_TUPLE_INT",
-                                "to_strides": "ops::OP_DTYPE::DT_TUPLE_INT",
-                                "str_to_enum": "ops::OP_DTYPE::DT_INT",
-                                "to_pair": "ops::OP_DTYPE::DT_TUPLE_INT",
-                                "to_dilations": "ops::OP_DTYPE::DT_TUPLE_INT",
-                                "to_output_padding": "ops::OP_DTYPE::DT_TUPLE_INT",
-                                "to_rates": "ops::OP_DTYPE::DT_TUPLE_INT"}
         arg_handler_processor = []
         op_args = op_proto.op_args
         for idx, op_arg in enumerate(op_args):
@@ -442,7 +450,7 @@ parse_args.arg_list_[${idx}]))->value());\n"
                                for word in arg_handler.split('_'))
             if arg_handler:
                 op_arg_name = op_arg.arg_name
-                new_type = arg_handler_type_map[arg_handler]
+                new_type = self.arg_handler_type_map.get(arg_handler, "Not exist")
                 if func_str in ("StrToEnum", "DtypeToTypeId"):
                     arg_handler_str = arg_handler_prt_template.replace(func_str=func_str,
                                                                        func_name=func_name,
@@ -450,20 +458,21 @@ parse_args.arg_list_[${idx}]))->value());\n"
                                                                        idx=idx,
                                                                        new_type=new_type)
                 else:
-                    arg_handler_str = arg_handler_template.replace(func_str=func_str,
-                                                                   func_name=func_name,
-                                                                   op_arg_name=op_arg_name,
-                                                                   idx=idx,
-                                                                   new_type=new_type)
+                    arg_handler_str = self.arg_handler_template.replace(func_str=func_str,
+                                                                        func_name=func_name,
+                                                                        op_arg_name=op_arg_name,
+                                                                        idx=idx,
+                                                                        new_type=new_type)
 
                 if op_arg.default == "None":
-                    arg_handler_str = arg_handler_optional_template.replace(idx=idx,
-                                                                            arg_handler_str=arg_handler_str)
+                    arg_handler_str = self.arg_handler_optional_template.replace(idx=idx,
+                                                                                 arg_handler_str=arg_handler_str)
                 arg_handler_processor.append(arg_handler_str)
 
         return arg_handler_processor
 
-    def get_input_tensor_index(self, op_proto):
+    @staticmethod
+    def get_input_tensor_index(op_proto):
         """
         Get index of input.
 
@@ -500,14 +509,22 @@ parse_args.arg_list_[${idx}]))->value());\n"
         if is_tensor_api:
             self_index = self.get_input_tensor_index(op_proto)
         convert_args_str = ""
+        arg_basic_convert_template = Template("parse_args.${convert_func}(${index}), ")
         for idx, op_arg in enumerate(op_proto.op_args):
             if is_tensor_api:
                 if self_index == idx:
                     convert_args_str += "input_tensor, "
                     continue
             is_optional = is_optional_param(op_arg)
+            if op_proto.op_view:
+                convert_func = basic_type_convert_str(op_arg.arg_dtype, False)
+                if convert_func != "":
+                    arg_convert_str = arg_basic_convert_template.replace(convert_func=convert_func,
+                                                                         index=idx)
+                    convert_args_str += arg_convert_str
+                    continue
             arg_convert_template = Template("parse_args.ConvertOptional<${des_type}>(${index}), ") if is_optional \
-                                   else Template("parse_args.Convert<${des_type}>(${index}), ")
+                else Template("parse_args.Convert<${des_type}>(${index}), ")
             if op_arg.is_type_id:
                 arg_type_str = get_input_args_type_str('type', False)
             else:
@@ -515,3 +532,12 @@ parse_args.arg_list_[${idx}]))->value());\n"
             convert_args_str += arg_convert_template.replace(index=idx,
                                                              des_type=arg_type_str[:-3])
         return convert_args_str[:-2]
+
+
+def _is_input_arg(arg_name, op_name):
+    res = False
+    if op_name in K.INPUT_NAME_MAP and arg_name == K.INPUT_NAME_MAP[op_name]:
+        res = True
+    elif op_name not in K.INPUT_NAME_MAP and arg_name in K.INPUT_ARGS_NAME:
+        res = True
+    return res