mindspore 2.4.10__cp310-cp310-win_amd64.whl → 2.6.0rc1__cp310-cp310-win_amd64.whl

mindspore/ops_generate/pyboost/pyboost_grad_function_cpp_generator.py

@@ -0,0 +1,155 @@
+# Copyright 2024 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.
+# ============================================================================
+"""
+This module defines the PyboostGradFunctionsGenerator class, which is responsible for generating PyBoost gradient
+functions and saving the corresponding C++ source files. The generator parses operator prototypes and constructs
+function definitions, includes necessary headers, and generates the registration code for the PyBoost functions.
+"""
+
+
+import os
+
+import common.template as template
+from common.template import Template
+from common.gen_utils import save_file
+import common.gen_constants as K
+from common.op_proto import OpProto
+from common.base_generator import BaseGenerator
+from pyboost import pyboost_utils
+
+from .op_template_parser import OpTemplateParser
+
+
+class PyboostGradFunctionsGenerator(BaseGenerator):
+    """
+    PyboostGradFunctionsGenerator generates PyBoost gradient functions based on operator prototypes (instances of
+    OpProto).
+
+    This class processes operator prototypes (`op_protos`) to generate PyBoost functions. It defines the function
+    body, handles value conversion, creates contiguous tensor values, and generates the necessary header includes and
+    registration code. The generated content is then saved to a specified location in the file system.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.pyboost_func_include_header_template = Template(
+            f'#include "{K.MS_PYBOOST_BASE_PATH}/auto_generate/${{operator_name}}.h"\n')
+        self.GEN_OPS_DEF_HEADER_TEMPLATE = template.GEN_OPS_DEF_HEADER_TEMPLATE
+        self.contiguous_template = Template(
+            "convert_$arg_name = runtime::ValueConverter::ContiguousTensorValue($device_target, convert_$arg_name);\n")
+
+    def generate(self, work_path, op_protos):
+        """
+        Generates the PyBoost gradient functions and writes them to the appropriate files.
+
+        This method processes a list of operator prototypes (`op_protos`), extracting necessary information such as
+        operator names, arguments, and conversion types. It uses this data to construct function bodies, includes,
+        and registration code. The generated content is saved to a specified path as a C++ source file.
+
+        Args:
+            work_path (str): The file path where the generated files will be saved.
+            op_protos (list): A list of operator prototypes containing information about the operators to be processed.
+
+        Returns:
+            None
+        """
+        pyboost_func_str = ''
+        pyboost_func_reg_def = ''
+        pyboost_func_include_headers_str = ''
+        for op_proto in op_protos:
+            if (op_proto.op_dispatch is None) or (not op_proto.op_dispatch.enable):
+                continue
+            op_parser = OpTemplateParser(op_proto)
+            op_pyboost_func_name = op_parser.get_pyboost_func_name()
+            operator_name = op_proto.op_name
+            op_name_str = op_proto.op_class.name
+            op_args_str = [op_arg.arg_name for op_arg in op_proto.op_args]
+            convert_value_type_str = self._convert_value_type(op_proto)
+            device_target = "op_runner_info->device_target"
+            convert_value_type_str += self._contiguous_tensor_value(op_proto, device_target)
+
+            call_args_str = []
+            for op_arg in op_proto.op_args:
+                call_arg = 'convert_' + op_arg.arg_name
+                call_args_str.append(call_arg)
+            pyboost_func_str += template.PYBOOST_GRAD_FUNCTION_TEMPLATE.replace(
+                func_name=op_pyboost_func_name,
+                op_name=op_name_str,
+                op_args=op_args_str,
+                convert_body=convert_value_type_str,
+                call_args=call_args_str)
+            pyboost_func_str = pyboost_func_str + template.NEW_LINE
+            pyboost_func_reg_def += template.REGISTER_PYBOOST_GRAD_DEFINE_TEMPLATE.replace(
+                pyboost_op_name=op_proto.op_class.name,
+                pyboost_cfunc_name=op_pyboost_func_name)
+            pyboost_func_include_headers_str += self.pyboost_func_include_header_template.replace(
+                operator_name=operator_name)
+
+        register_func_str = template.REGISTER_PYBOOST_GRAD_TEMPLATE.replace(register_func=pyboost_func_reg_def)
+        pyboost_func_file = \
+            template.PYBOOST_GRAD_HEADER_TEMPLATE.replace(include_op_header=pyboost_func_include_headers_str,
+                                                          function_body=pyboost_func_str,
+                                                          register_function_body=register_func_str)
+        save_path = os.path.join(work_path, K.PYBOOST_GRAD_FUNC_GEN_PATH)
+        file_name = "pyboost_grad_functions.cc"
+        save_file(save_path, file_name, pyboost_func_file)
+
+    def _convert_value_type(self, op_proto: OpProto) -> str:
+        """
+        Generates the code for converting the operator's input values to the required types.
+
+        This method iterates over the operator's arguments, checks if they are optional, and generates the appropriate
+        conversion code based on the argument's data type.
+
+        Args:
+            op_proto (OpProto): The operator prototype containing information about the operator's arguments.
+
+        Returns:
+            str: A string containing the code for converting the input values to the required types.
+        """
+        convert_template = Template(
+            "auto convert_$arg_name = ValueConverter::${convert_func}(op_runner_info->inputs[$arg_index]);\n")
+        parser_func_str = ''
+        for index, arg in enumerate(op_proto.op_args):
+            is_optional = pyboost_utils.is_optional_param(arg)
+            convert_type_str = pyboost_utils.get_value_convert_type_str(arg.arg_dtype, is_optional)
+            parser_func_str += convert_template.replace(arg_name=arg.arg_name, convert_func=convert_type_str,
+                                                        arg_index=pyboost_utils.get_index(index))
+        return parser_func_str
+
+    def _contiguous_tensor_value(self, op_proto: OpProto, device_target: str) -> str:
+        """
+        Generates the code for converting tensors to contiguous format if required.
+
+        This method checks the data type of the operator's arguments and generates code for converting tensors
+        to contiguous format, which is necessary for certain types of tensors.
+
+        Args:
+            op_proto (OpProto): The operator prototype containing information about the operator's arguments.
+            device_target (str): The device target string used in the conversion code.
+
+        Returns:
+            str: A string containing the code for converting tensors to contiguous format.
+            If the operator is a view operation, an empty string is returned.
+        """
+        if op_proto.op_view:
+            return ''
+        contiguous_func_str = ''
+        need_contiguous_dtype = {'tensor', 'tuple[tensor]'}
+        for arg in op_proto.op_args:
+            if arg.arg_dtype not in need_contiguous_dtype:
+                continue
+            contiguous_func_str += self.contiguous_template.replace(arg_name=arg.arg_name, device_target=device_target)
+        return contiguous_func_str

mindspore/ops_generate/pyboost/pyboost_inner_prim_generator.py

@@ -0,0 +1,132 @@
+# Copyright 2024 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.
+# ============================================================================
+"""
+This module defines the PyboostInnerPrimGenerator class, which is responsible for generating Python primitive
+wrappers for Pyboost operations. The generator constructs Python function definitions based on operator prototypes,
+generates necessary import statements, and writes the generated content into Python source files.
+
+The primary functionality is to take operator prototypes, extract relevant fields, and create Python function wrappers
+that can be used to call the Pyboost primitive implementations.
+"""
+
+import os
+import common.template as template
+import common.gen_constants as K
+from common.gen_utils import save_file
+
+from common.op_proto import OpProto
+from common.base_generator import BaseGenerator
+
+from .op_template_parser import OpTemplateParser
+
+
+class PyboostInnerPrimGenerator(BaseGenerator):
+    """
+    PyboostInnerPrimGenerator is responsible for generating Python primitive wrappers for Pyboost operators.
+
+    This class processes operator prototypes (`op_protos`) to generate Python function implementations. It handles the
+    inclusion of necessary headers, processes operator arguments, and creates Python functions that wrap Pyboost
+    primitives.
+
+    Attributes:
+        IMPORT_PYBOOST_PRIM_HEADER (Template): Template for importing Pyboost primitive headers.
+        PYBOOST_PY_FUNC_IMPORT_HEADER (Template): Template for importing Python functions related to Pyboost.
+        PYTHON_PRIM_TEMPLATE (Template): Template for generating Python primitive functions.
+    """
+
+    def __init__(self):
+        """
+        Initializes the PyboostInnerPrimGenerator class.
+
+        This constructor sets up the required templates for generating import headers, Python function imports,
+        and Python primitive function wrappers.
+        """
+        self.IMPORT_PYBOOST_PRIM_HEADER = template.IMPORT_PYBOOST_PRIM_HEADER
+        self.PYBOOST_PY_FUNC_IMPORT_HEADER = template.PYBOOST_PY_FUNC_IMPORT_HEADEAR
+        self.PYTHON_PRIM_TEMPLATE = template.PYTHON_PRIM_TEMPLATE
+
+    def generate(self, work_path, op_protos):
+        """
+        Generates Python wrappers for Pyboost primitives and writes them to a Python source file.
+
+        This method processes a list of operator prototypes (`op_protos`), extracting necessary information such as
+        class names, arguments, and handlers. It constructs Python function wrappers for the Pyboost primitives and
+        generates the required import statements. The generated Python code is saved to a specified path.
+
+        Args:
+            work_path (str): The file path where the generated Python file will be saved.
+            op_protos (list): A list of operator prototypes containing information about the operators to be processed.
+
+        Returns:
+            None
+        """
+        gen_py = ''
+        gen_header = template.PY_LICENSE_STR + self.IMPORT_PYBOOST_PRIM_HEADER
+        for op_proto in op_protos:
+            # only process pyboost enabled scenario
+            if op_proto.op_dispatch is None:
+                continue
+            if not op_proto.op_dispatch.enable:
+                continue
+            op_parser = OpTemplateParser(op_proto)
+            if not op_parser.has_prim_init():
+                continue
+
+            gen_header += self.PYBOOST_PY_FUNC_IMPORT_HEADER.replace(class_name=op_proto.op_class.name)
+            input_args, process_func, processed_args = self._get_fields_for_prim_tpl(op_proto)
+            gen_py += self.PYTHON_PRIM_TEMPLATE.replace(class_name=op_proto.op_class.name,
+                                                        input_args=input_args,
+                                                        process_func=process_func,
+                                                        func_impl_name=op_proto.op_name,
+                                                        processed_args=processed_args)
+
+        save_file(os.path.join(work_path, K.PY_AUTO_GEN_PATH), "pyboost_inner_prim.py", gen_header + gen_py)
+
+    def _get_fields_for_prim_tpl(self, op_proto: OpProto):
+        """
+        Extracts the necessary fields for the primitive template from the operator prototype.
+
+        This method processes the arguments of the operator prototype and generates the input arguments, the function
+        that handles argument processing, and the processed arguments list, which will be used in the final Python
+        function definition.
+
+        Args:
+            op_proto (OpProto): The operator prototype from which the argument data will be extracted.
+
+        Returns:
+            tuple: A tuple containing three return values required for the primitive template to be generated:
+                - input_args (list): List of input argument names for the Python function.
+                - process_func (str): String representing the argument processing logic for the function.
+                - processed_args (list): List of processed argument names used in the function call.
+        """
+        args = op_proto.op_args
+        operator_name = op_proto.op_name
+
+        input_args = []
+        process_func = ''
+        processed_args = []
+
+        for arg in args:
+            arg_name = arg.arg_name
+            arg_handler = arg.arg_handler
+            processed_arg = arg_name
+            if arg_handler != '' and arg_handler != 'dtype_to_type_id':
+                process_func += \
+                    f"""converted_{arg_name} = {arg_handler}('{operator_name}', '{arg_name}', {arg_name})\n"""
+                processed_arg = 'converted_' + arg_name
+            input_args.append(arg_name)
+            processed_args.append(processed_arg)
+
+        return input_args, process_func, processed_args

mindspore/ops_generate/pyboost/pyboost_native_grad_functions_generator.py

@@ -0,0 +1,272 @@
+# Copyright 2024 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.
+# ============================================================================
+"""
+This module defines the PyboostGradFunctionsCppGenerator and PyboostGradFunctionsHeaderGenerator classes,
+which are responsible for generating C++ gradient function implementations and headers for PyBoost operations.
+
+The PyboostGradFunctionsCppGenerator generates the actual function definitions for the gradient functions,
+while the PyboostGradFunctionsHeaderGenerator creates the corresponding function declarations in header files.
+"""
+
+import os
+
+from pyboost import pyboost_utils
+from pyboost.pyboost_utils import is_optional_param
+import common.template as template
+from common.template import Template
+import common.gen_constants as K
+from common.gen_utils import save_file
+from common.op_proto import OpProto
+from common.base_generator import BaseGenerator
+
+
+class PyboostGradFunctionsCppGenerator(BaseGenerator):
+    """
+    PyboostGradFunctionsCppGenerator generates C++ implementations for PyBoost gradient functions.
+
+    This class processes operator prototypes (`op_protos`) to create C++ function definitions that
+    wrap PyBoost gradient functionality. It constructs function bodies, handles value conversion,
+    and generates necessary include statements for each operator.
+
+    Attributes:
+       PYBOOST_NATIVE_GRAD_FUNCTION_TEMPLATE (Template): Template for generating native gradient function definitions.
+       PYBOOST_NATIVE_GRAD_FUNCTIONS_TEMPLATE (Template): Template for generating the overall gradient functions file.
+       native_function_multi_output_template (Template): Template for handling multiple output functions.
+       native_function_single_output_template (str): Template for handling single output functions.
+       native_include_header_template (Template): Template for generating include headers for each operator.
+       convert_template (Template): Template for converting argument values to native types.
+    """
+
+    def __init__(self):
+        self.PYBOOST_NATIVE_GRAD_FUNCTION_TEMPLATE = template.PYBOOST_NATIVE_GRAD_FUNCTION_TEMPLATE
+        self.PYBOOST_NATIVE_GRAD_FUNCTIONS_TEMPLATE = template.PYBOOST_NATIVE_GRAD_FUNCTIONS_TEMPLATE
+        self.native_function_multi_output_template = template.MULTI_OUTPUT_TEMPLATE
+        self.native_function_single_output_template = "const auto &output_value = op->outputs()[0];\n"
+        self.native_include_header_template = Template(
+            f'#include "{K.MS_PYBOOST_BASE_PATH}/auto_generate/${{operator_name}}.h"\n')
+        self.convert_template = Template(
+            "auto convert_$arg_name = runtime::ValueConverter::${convert_func}(ConvertNode2Value($arg_name));\n")
+
+    def generate(self, work_path, op_protos):
+        """
+        Generates C++ gradient function implementations and writes them to a source file.
+
+        This method processes a list of operator prototypes (`op_protos`), extracting necessary information such as
+        operator names, arguments, and conversion types. It constructs C++ function bodies and generates the required
+        include statements. The generated content is saved to a specified path.
+
+        Args:
+            work_path (str): The file path where the generated C++ file will be saved.
+            op_protos (list): A list of operator prototypes containing information about the operators to be processed.
+
+        Returns:
+            None
+        """
+        pyboost_func_str = ''
+        pyboost_func_include_headers_str = ''
+        ops_inc_head_set = set()
+        for op_proto in op_protos:
+            if op_proto.op_dispatch is None:
+                continue
+            if not op_proto.op_dispatch.enable:
+                continue
+
+            op_args_str = [op_arg.arg_name for op_arg in op_proto.op_args]
+            convert_value_type_str = self._convert_native_value_type(op_proto)
+            convert_value_type_str += self._contiguous_tensor_value(op_proto, "device_target_")
+            call_args_str = self._get_call_args(op_proto)
+            call_args_with_type = self._get_call_args_with_type(op_proto)
+
+            first_var_name = op_proto.op_args[0].arg_name
+            output_expr = self._get_output_expr(op_proto)
+
+            pyboost_func_str += \
+                template.PYBOOST_NATIVE_GRAD_FUNCTION_TEMPLATE.replace(func_name=op_proto.op_class.name,
+                                                                       op_name=op_proto.op_class.name,
+                                                                       op_args=op_args_str,
+                                                                       convert_body=convert_value_type_str,
+                                                                       call_args=call_args_str,
+                                                                       call_args_with_type=call_args_with_type,
+                                                                       first_var_name=first_var_name,
+                                                                       output_expr=output_expr)
+            pyboost_func_str = pyboost_func_str + template.NEW_LINE
+            pyboost_func_include_headers_str += (
+                self.native_include_header_template.replace(operator_name=op_proto.op_name))
+            ops_inc_head_set.add(
+                template.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_proto.op_class.name[0].lower()))
+        native_grad_func_file = \
+            self.PYBOOST_NATIVE_GRAD_FUNCTIONS_TEMPLATE.replace(include_op_header=pyboost_func_include_headers_str,
+                                                                function_body=pyboost_func_str,
+                                                                ops_inc=list(sorted(ops_inc_head_set)))
+        save_file(os.path.join(work_path, K.PYBOOST_NATIVE_GRAD_FUNC_GEN_PATH),
+                  "pyboost_native_grad_functions.cc", native_grad_func_file)
+
+    def _convert_native_value_type(self, op_proto: OpProto) -> str:
+        """
+        Generates native value conversion functions for operator arguments.
+
+        This method processes each argument of the operator prototype and constructs conversion statements
+        based on the argument's data type.
+
+        Args:
+            op_proto (OpProto): The operator prototype from which the argument data will be extracted.
+
+        Returns:
+            str: A string containing the conversion statements for the operator's arguments.
+        """
+        parser_func_str = ''
+        for op_arg in op_proto.op_args:
+            is_optional = is_optional_param(op_arg)
+            convert_type_str = pyboost_utils.get_value_convert_type_str(op_arg.arg_dtype, is_optional)
+            parser_func_str += self.convert_template.replace(arg_name=op_arg.arg_name, convert_func=convert_type_str)
+        return parser_func_str
+
+    def _contiguous_tensor_value(self, op_proto: OpProto, device_target: str) -> str:
+        """
+        Generates contiguous tensor value conversion functions if applicable.
+
+        This method constructs conversion statements for tensors that need to be contiguous. If the operator is a view
+        operation, no conversion is performed.
+
+        Args:
+            op_proto (OpProto): The operator prototype that contains the argument data.
+            device_target (str): The device target to be used in the conversion statements.
+
+        Returns:
+            str: A string containing the contiguous tensor conversion statements.
+        """
+        if op_proto.op_view:
+            return ''
+        contiguous_template = Template(
+            "convert_$arg_name = runtime::ValueConverter::ContiguousTensorValue($device_target, convert_$arg_name);\n")
+        contiguous_func_str = ''
+        need_contiguous_dtype = {'tensor', 'tuple[tensor]'}
+        for op_arg in op_proto.op_args:
+            if op_arg.arg_dtype not in need_contiguous_dtype:
+                continue
+            contiguous_func_str += contiguous_template.replace(arg_name=op_arg.arg_name, device_target=device_target)
+        return contiguous_func_str
+
+    def _get_output_expr(self, op_proto: OpProto):
+        """
+        Determines the output expression based on the operator prototype.
+
+        This method checks if the operator produces multiple outputs and returns the corresponding output expression.
+
+        Args:
+            op_proto (OpProto): The operator prototype to evaluate.
+
+        Returns:
+            str: The output expression used in the function implementation.
+        """
+        output_expr = self.native_function_single_output_template
+        if pyboost_utils.is_op_multi_output(op_proto.op_returns):
+            output_expr = self.native_function_multi_output_template
+        return output_expr
+
+    def _get_call_args(self, op_proto: OpProto):
+        """
+        Generates the list of call arguments for the operator function.
+
+        This method constructs a list of argument names prefixed with 'convert_' for use in the function call.
+
+        Args:
+            op_proto (OpProto): The operator prototype containing the argument information.
+
+        Returns:
+            list: A list of formatted argument names to be used in the function call.
+        """
+        call_args_str = []
+        for op_arg in op_proto.op_args:
+            call_arg = 'convert_' + op_arg.arg_name
+            call_args_str.append(call_arg)
+        return call_args_str
+
+    def _get_call_args_with_type(self, op_proto: OpProto):
+        """
+        Generates the list of call arguments with type information.
+
+        This method constructs a list of argument declarations with the appropriate type for the function definition.
+
+        Args:
+            op_proto (OpProto): The operator prototype containing the argument information.
+
+        Returns:
+            list: A list of argument declarations with types for the function definition.
+        """
+        call_args_with_type = []
+        for op_arg in op_proto.op_args:
+            call_args_with_type.append('const NodePtr &' + op_arg.arg_name)
+        return call_args_with_type
+
+
+class PyboostGradFunctionsHeaderGenerator(BaseGenerator):
+    """
+    PyboostGradFunctionsHeaderGenerator generates C++ header declarations for PyBoost gradient functions.
+
+    This class processes operator prototypes to create function declarations for the grad functions in a header file.
+    """
+
+    def __init__(self):
+        self.native_function_header_template = Template("static NodePtr $func_name(${call_args_with_type});\n")
+
+    def generate(self, work_path, op_protos):
+        """
+        Generates C++ header declarations for gradient functions and writes them to a header file.
+
+        This method processes a list of operator prototypes (`op_protos`), extracting necessary information
+        and constructing the function declarations. The generated content is saved to a specified path.
+
+        Args:
+            work_path (str): The file path where the generated header file will be saved.
+            op_protos (list): A list of operator prototypes containing information about the operators to be processed.
+
+        Returns:
+            None
+        """
+        native_function_headers_str = ''
+        for op_proto in op_protos:
+            if op_proto.op_dispatch is None:
+                continue
+            if not op_proto.op_dispatch.enable:
+                continue
+
+            call_args_with_type = self._get_call_args_with_type(op_proto)
+
+            func_header = self.native_function_header_template.replace(func_name=op_proto.op_class.name,
+                                                                       call_args_with_type=call_args_with_type)
+            native_function_headers_str += func_header
+        native_grad_func_header_file = template.PYBOOST_NATIVE_GRAD_FUNCTIONS_HEADER_TEMPLATE.replace(
+            native_grad_func_def=native_function_headers_str)
+
+        save_file(os.path.join(work_path, K.PYBOOST_NATIVE_GRAD_FUNC_GEN_PATH),
+                  "pyboost_native_grad_functions.h", native_grad_func_header_file)
+
+    def _get_call_args_with_type(self, op_proto: OpProto):
+        """
+        Generates the list of call arguments with type information for the header declarations.
+
+        This method constructs a list of argument declarations with the appropriate type for the function declaration.
+
+        Args:
+            op_proto (OpProto): The operator prototype containing the argument information.
+
+        Returns:
+            list: A list of argument declarations with types for the function declaration.
+        """
+        call_args_with_type = []
+        for op_arg in op_proto.op_args:
+            call_args_with_type.append('const NodePtr &' + op_arg.arg_name)
+        return call_args_with_type