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

mindspore/ops_generate/pyboost/auto_grad_reg_cc_generator.py

@@ -0,0 +1,93 @@
+# 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 provides a generator class for creating C++ header files for AutoGrad registration functionality.
+"""
+
+import os
+
+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.base_generator import BaseGenerator
+from pyboost.pyboost_utils import is_optional_param, get_input_dtype
+
+
+class AutoGradRegHeaderGenerator(BaseGenerator):
+    """
+    Generates C++ header files for the AutoGrad registration functionality based on operator prototypes.
+    """
+
+    def __init__(self):
+        """
+        Initialize the AutoGrad registration header generator with templates for code generation.
+        """
+        self.AUTO_GRAD_REG_H_TEMPLATE = template.AUTO_GRAD_REG_H_TEMPLATE
+        self.op_type_enum_template = Template("k${class_name} = ${enum_val},\n")
+        self.op_grad_func_template = Template("using ${class_name}GradFunc = std::function<void(${grad_func_args})>;")
+        self.op_grad_func_obj_template = Template("${class_name}GradFunc ${class_name}GradFuncObj;")
+        self.op_grad_func_args_template = Template(
+            "const kernel::pyboost::OpPtr &, ${input_tensor_prt_args}"
+        )
+
+    def generate(self, work_path, op_protos):
+        """
+        Generate the AutoGrad registration header file.
+
+        Args:
+            work_path (str): The directory where the generated file should be saved.
+            op_protos (list): A list of operator prototypes used to generate the header.
+        """
+        op_type_enum_list = []
+        op_grad_func_list = []
+        op_grad_func_obj_list = []
+        index = 0
+        for op_proto in op_protos:
+            if op_proto.op_dispatch is None or op_proto.op_dispatch.is_comm_op:
+                continue
+            op_type_enum_list.append(self.op_type_enum_template.replace(class_name=op_proto.op_class.name,
+                                                                        enum_val=index))
+            grad_func_args_with_type_str = self._get_grad_func_args_with_type_str(op_proto)
+            op_grad_func_list.append(self.op_grad_func_template.replace(class_name=op_proto.op_class.name,
+                                                                        grad_func_args=grad_func_args_with_type_str))
+            op_grad_func_obj_list.append(self.op_grad_func_obj_template.replace(class_name=op_proto.op_class.name))
+            index += 1
+
+        pyboost_func_h_str = self.AUTO_GRAD_REG_H_TEMPLATE.replace(op_enum=op_type_enum_list,
+                                                                   op_grad_func=op_grad_func_list,
+                                                                   op_grad_func_obj=op_grad_func_obj_list)
+
+        save_path = os.path.join(work_path, K.MS_PYBOOST_FUNCTIONS_AUTO_GEN_PATH)
+        file_name = "auto_grad_op_reg.h"
+        save_file(save_path, file_name, pyboost_func_h_str)
+
+    def _get_grad_func_args_with_type_str(self, op_proto):
+        """
+        Get the gradient function arguments with type information.
+
+        Args:
+            op_proto: The operator prototype.
+
+        Returns:
+            str: A string of input tensor pointer arguments with types.
+        """
+        input_tensor_prt_args_str = ""
+        for op_arg in op_proto.op_args:
+            is_optional = is_optional_param(op_arg)
+            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional)
+            input_tensor_prt_args_str += f"const {input_dtype} &, "
+
+        return self.op_grad_func_args_template.replace(input_tensor_prt_args=input_tensor_prt_args_str.rstrip(', '))

mindspore/ops_generate/pyboost/gen_pyboost_func.py

@@ -0,0 +1,175 @@
+# 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.
+# ============================================================================
+"""
+Generate pyboost function from pyboost_op.yaml
+"""
+
+import os
+import shutil
+import logging
+from resources.resource_list import ResourceType
+from common import gen_constants as K
+from api.functions_cc_generator import FunctionsGenerator, FunctionsHeaderGenerator
+from .pyboost_inner_prim_generator import PyboostInnerPrimGenerator
+from .pyboost_functions_py_generator import PyboostFunctionsPyGenerator
+from .pyboost_functions_h_generator import PyboostFunctionsHeaderGenerator
+from .pyboost_functions_cpp_generator import PyboostFunctionsGenerator
+from .pyboost_grad_function_cpp_generator import PyboostGradFunctionsGenerator
+from .pyboost_native_grad_functions_generator import (
+    PyboostGradFunctionsHeaderGenerator,
+    PyboostGradFunctionsCppGenerator,
+)
+from .pyboost_op_cpp_code_generator import (
+    PyboostCommonOpHeaderGenerator,
+    PyboostOpFunctionGenerator,
+    PyboostOpHeaderGenerator,
+    delete_residual_files,
+    PyboostOpRegisterCppCodeGenerator,
+)
+from .pyboost_overload_functions_cpp_generator import PyboostOverloadFunctionsGenerator
+from .auto_grad_impl_cc_generator import AutoGradImplGenerator
+from .auto_grad_reg_cc_generator import AutoGradRegHeaderGenerator
+
+
+def clear_old_generated_code(work_path):
+    """ delete old generated files to prevent compilation failure """
+    files_to_clear = ['mindspore/ops/kernel/common/pyboost',
+                      'mindspore/ops/kernel/functions/auto_generate',
+                      'mindspore/ccsrc/runtime/pynative/op_function']
+    for f in files_to_clear:
+        real_path = os.path.join(work_path, f)
+        if os.path.exists(real_path):
+            shutil.rmtree(real_path)
+            logging.warning("rm file %s", real_path)
+
+
+def gen_pyboost_code(resource_mgr):
+    """ gen_pyboost_code """
+    work_path = K.WORK_DIR
+    op_protos = resource_mgr.get_resource(ResourceType.OP_PROTO)
+    doc_yaml_data = resource_mgr.get_resource(ResourceType.OP_DOC_YAML)
+    mint_func_protos = resource_mgr.get_resource(ResourceType.MINT_FUNC_PROTOS)
+    alias_func_mapping = resource_mgr.get_resource(ResourceType.ALIAS_API_MAPPING)
+    clear_old_generated_code(work_path)
+    call_pyboost_inner_prim_generator(work_path, op_protos)
+    call_pyboost_functions_py_generator(work_path, op_protos, doc_yaml_data)
+    call_pyboost_functions_h_generator(work_path, op_protos)
+    call_pyboost_functions_cpp_generator(work_path, op_protos)
+    call_pyboost_overload_functions_cpp_generator(work_path, op_protos, mint_func_protos, alias_func_mapping)
+    call_pyboost_grad_functions_cpp_generator(work_path, op_protos)
+    call_pyboost_native_grad_functions_generator(work_path, op_protos)
+    call_pyboost_op_cpp_code_generator(work_path, op_protos)
+    # op splice
+    call_pyboost_auto_grad_cpp_code_generator(work_path, op_protos)
+
+
+def call_pyboost_auto_grad_cpp_code_generator(work_path, op_protos):
+    call_auto_grad_impl_cc_generator(work_path, op_protos)
+    call_auto_grad_reg_header_generator(work_path, op_protos)
+    call_functions_header_generator(work_path, op_protos)
+    call_functions_cc_generator(work_path, op_protos)
+
+
+def call_auto_grad_impl_cc_generator(work_path, op_protos):
+    generator = AutoGradImplGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_auto_grad_reg_header_generator(work_path, op_protos):
+    generator = AutoGradRegHeaderGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_functions_header_generator(work_path, op_protos):
+    generator = FunctionsHeaderGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_functions_cc_generator(work_path, op_protos):
+    generator = FunctionsGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_pyboost_inner_prim_generator(work_path, op_protos):
+    generator = PyboostInnerPrimGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_pyboost_functions_py_generator(work_path, op_protos, doc_yaml_data):
+    generator = PyboostFunctionsPyGenerator()
+    generator.generate(work_path, op_protos, doc_yaml_data)
+
+
+def call_pyboost_functions_h_generator(work_path, op_protos):
+    generator = PyboostFunctionsHeaderGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_pyboost_functions_cpp_generator(work_path, op_protos):
+    generator = PyboostFunctionsGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_pyboost_overload_functions_cpp_generator(work_path, op_protos, mint_func_protos, alias_func_mapping):
+    generator = PyboostOverloadFunctionsGenerator()
+    generator.generate(work_path, op_protos,
+                       mint_func_protos, alias_func_mapping)
+
+
+def call_pyboost_grad_functions_cpp_generator(work_path, op_protos):
+    generator = PyboostGradFunctionsGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_pyboost_native_grad_functions_generator(work_path, op_protos):
+    h_generator = PyboostGradFunctionsHeaderGenerator()
+    h_generator.generate(work_path, op_protos)
+
+    cc_generator = PyboostGradFunctionsCppGenerator()
+    cc_generator.generate(work_path, op_protos)
+
+
+def call_pyboost_op_cpp_code_generator(work_path, op_protos):
+    delete_residual_files(work_path, op_protos)
+    call_PyboostCommonOpCppCodeGenerator(work_path, op_protos)
+    call_PyboostOpHeaderGenerator(work_path, op_protos)
+    call_merge_pyboost_op_cpp_code_generator(work_path, op_protos)
+    call_PyboostOpRegisterCppCodeGenerator(work_path, op_protos)
+
+
+def call_merge_pyboost_op_cpp_code_generator(work_path, op_protos):
+    generator = PyboostOpFunctionGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_PyboostCommonOpCppCodeGenerator(work_path, op_protos):
+    generator = PyboostCommonOpHeaderGenerator()
+    generator.generate(work_path, op_protos)
+
+
+def call_PyboostOpHeaderGenerator(work_path, op_protos):
+    generator = PyboostOpHeaderGenerator('ascend')
+    generator.generate(work_path, op_protos)
+
+    generator = PyboostOpHeaderGenerator('gpu')
+    generator.generate(work_path, op_protos)
+
+    generator = PyboostOpHeaderGenerator('cpu')
+    generator.generate(work_path, op_protos)
+
+
+def call_PyboostOpRegisterCppCodeGenerator(work_path, op_protos):
+    op_register_cpp_generator = PyboostOpRegisterCppCodeGenerator()
+    op_register_cpp_generator.generate(work_path, op_protos)