mindspore 2.5.0__cp310-cp310-win_amd64.whl → 2.6.0__cp310-cp310-win_amd64.whl

mindspore/ops_generate/{ops_name_h_generator.py → op_def/ops_name_h_generator.py}

@@ -21,14 +21,23 @@ constants for operator names based on given prototypes.
 
 import os
 
-import gen_constants as K
-import gen_utils
-import pyboost_utils
+import common.gen_constants as K
+import common.gen_utils as gen_utils
+import common.template as template
+from common.base_generator import BaseGenerator
+from pyboost import pyboost_utils
 
-# refactored
-import template
 
-from base_generator import BaseGenerator
+OP_NAME_OP_DEF = """
+#ifndef MINDSPORE_CORE_OP_NAME_${suffix}_H_
+#define MINDSPORE_CORE_OP_NAME_${suffix}_H_
+
+namespace mindspore::ops {
+$ops_namespace_body
+}  // namespace mindspore::ops
+
+#endif  // MINDSPORE_CORE_OP_NAME_${suffix}_H_
+"""
 
 
 class OpsNameHGenerator(BaseGenerator):
@@ -40,7 +49,7 @@ class OpsNameHGenerator(BaseGenerator):
         """
         Initializes the OpsNameHGenerator instance.
         """
-        self.op_name_op_def_template = template.Template(K.OP_NAME_OP_DEF)
+        self.op_name_op_def_template = template.Template(OP_NAME_OP_DEF)
         self.op_def_body_template = template.Template("""constexpr auto kName${k_name_op} = "${k_name_op}";\n""")
 
     def generate(self, work_path, op_protos):
@@ -54,15 +63,21 @@ class OpsNameHGenerator(BaseGenerator):
         Returns:
             None
         """
-        op_name_gen_list = []
+        import os
+        import collections
+
+        op_name_gen_dict = collections.defaultdict(list)
+
         for op_proto in op_protos:
             k_name_op = pyboost_utils.get_op_name(op_proto.op_name, op_proto.op_class.name)
-            op_name_gen_list.append(self.op_def_body_template.replace(k_name_op=k_name_op))
-
-        op_name_code = self.op_name_op_def_template.replace(ops_namespace_body=op_name_gen_list)
+            first_char = k_name_op[0].lower()
+            op_name_gen_dict[first_char].append(self.op_def_body_template.replace(k_name_op=k_name_op))
 
-        op_name_code = template.CC_LICENSE_STR + op_name_code
+        for first_char, op_name_gen_list in op_name_gen_dict.items():
+            op_name_code = self.op_name_op_def_template.replace(ops_namespace_body=op_name_gen_list,
+                                                                suffix=first_char.upper())
+            op_name_code = template.CC_LICENSE_STR + op_name_code
 
-        save_path = os.path.join(work_path, K.MS_OP_DEF_AUTO_GENERATE_PATH)
-        file_name = "gen_ops_name.h"
-        gen_utils.save_file(save_path, file_name, op_name_code)
+            save_path = os.path.join(work_path, K.MS_OP_DEF_AUTO_GENERATE_PATH)
+            file_name = f"gen_ops_name_{first_char}.h"
+            gen_utils.save_file(save_path, file_name, op_name_code)

mindspore/ops_generate/op_def/ops_primitive_h_generator.py

@@ -0,0 +1,125 @@
+# 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.
+# ============================================================================
+"""
+Module for generating C++ header files for operator primitives.
+
+This module defines the `OpsPrimitiveHGenerator` class, which creates C++ header files
+containing definitions for operator primitives based on provided operator prototypes.
+"""
+
+import os
+
+import common.gen_constants as K
+import common.gen_utils as gen_utils
+import common.template as template
+from common.base_generator import BaseGenerator
+from pyboost import pyboost_utils
+
+
+OP_PRIM_OP_DEF_H = """
+#ifndef MINDSPORE_CORE_OPS_GEN_OPS_PRIMITIVE_${suffix}_H_
+#define MINDSPORE_CORE_OPS_GEN_OPS_PRIMITIVE_${suffix}_H_
+
+#include "ir/primitive.h"
+#include "mindapi/base/macros.h"
+#include "$auto_gen_path/gen_ops_name_${suffix}.h"
+
+namespace mindspore::prim {
+$ops_prim_gen
+}  // namespace mindspore::prim
+#endif  // MINDSPORE_CORE_OPS_GEN_OPS_PRIMITIVE_${suffix}_H_
+"""
+
+
+OP_PRIM_OP_DEF_CC = """
+
+#include "$auto_gen_path/gen_ops_primitive_${suffix}.h"
+
+#include <memory>
+
+namespace mindspore::prim {
+$ops_prim_gen
+}  // namespace mindspore::prim
+
+"""
+
+
+class OpsPrimitiveHGenerator(BaseGenerator):
+    """
+    This class generates the header file for operator primitives.
+    """
+
+    def __init__(self):
+        """
+        Initializes the generator with templates for operator primitive definitions.
+        """
+        self.op_prim_op_def_template = template.Template(OP_PRIM_OP_DEF_H)
+        self.op_prim_op_def_cc_template = template.Template(OP_PRIM_OP_DEF_CC)
+        self.op_def_h_template = template.Template(
+            "OPS_API extern const PrimitivePtr kPrim${k_name_op};\n")
+        self.op_def_template = template.Template(
+            "const PrimitivePtr kPrim${k_name_op} = std::make_shared<Primitive>(ops::kName${k_name_op});\n")
+        self.op_def_rw_template = template.Template(
+            "const PrimitivePtr kPrim${k_name_op} = std::make_shared<Primitive>(ops::kName${k_name_op}, "
+            "true, kPrimTypeBuiltIn, true);\n")
+
+    def generate(self, work_path, op_protos):
+        """
+        Generates the header file content for operator primitives and saves it.
+
+        Args:
+            work_path (str): The directory to save the generated files.
+            op_protos (list): A list of operator prototypes.
+
+        Returns:
+            None
+
+        The method generates the content of the header file for each operator primitive
+        defined in the 'op_protos' list and saves it to the specified work path.
+        """
+        import os
+        import collections
+        ops_prim_gen_dict = collections.defaultdict(list)
+        ops_prim_cc_gen_dict = collections.defaultdict(list)
+
+        for op_proto in op_protos:
+            k_name_op = pyboost_utils.get_op_name(op_proto.op_name, op_proto.op_class.name)
+            first_char = k_name_op[0].lower()
+            if op_proto.op_args_signature:
+                if op_proto.op_args_signature.rw_write:
+                    ops_prim_gen_dict[first_char].append(self.op_def_h_template.replace(k_name_op=k_name_op))
+                    ops_prim_cc_gen_dict[first_char].append(self.op_def_rw_template.replace(k_name_op=k_name_op))
+                    continue
+
+            ops_prim_gen_dict[first_char].append(self.op_def_h_template.replace(k_name_op=k_name_op))
+            ops_prim_cc_gen_dict[first_char].append(self.op_def_template.replace(k_name_op=k_name_op))
+
+        for first_char, ops_prim_gen_list in ops_prim_gen_dict.items():
+            op_prim_op_def = self.op_prim_op_def_template.replace(auto_gen_path=K.MS_OP_DEF_AUTO_GENERATE_PATH,
+                                                                  ops_prim_gen=ops_prim_gen_list, suffix=first_char)
+            res_str = template.CC_LICENSE_STR + op_prim_op_def
+
+            save_path = os.path.join(work_path, K.MS_OP_DEF_AUTO_GENERATE_PATH)
+            file_name = f"gen_ops_primitive_{first_char}.h"
+            gen_utils.save_file(save_path, file_name, res_str)
+
+        for first_char, ops_prim_gen_list in ops_prim_cc_gen_dict.items():
+            op_prim_op_def = self.op_prim_op_def_cc_template.replace(auto_gen_path=K.MS_OP_DEF_AUTO_GENERATE_PATH,
+                                                                     ops_prim_gen=ops_prim_gen_list, suffix=first_char)
+            res_str = template.CC_LICENSE_STR + op_prim_op_def
+
+            save_path = os.path.join(work_path, K.MS_OP_DEF_AUTO_GENERATE_PATH)
+            file_name = f"gen_ops_primitive_{first_char}.cc"
+            gen_utils.save_file(save_path, file_name, res_str)

mindspore/ops_generate/op_def_py/gen_op_def_py.py

@@ -0,0 +1,47 @@
+# Copyright 2025 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 operator definition python files.
+"""
+import os
+import shutil
+
+from op_def_py.op_prim_py_generator import OpPrimPyGenerator
+from op_def_py.op_def_py_generator import OpDefPyGenerator
+from resources.resource_list import ResourceType
+from common import gen_constants as K
+
+
+def generate_ops_prim_file(work_path, op_protos, doc_dict, file_pre):
+    generator = OpPrimPyGenerator()
+    generator.generate(work_path, op_protos, doc_dict, file_pre)
+
+
+def generate_ops_def_file(work_path, os_protos, doc_dict, file_pre):
+    generator = OpDefPyGenerator()
+    generator.generate(work_path, os_protos, doc_dict, file_pre)
+
+
+def generate_ops_py_files(resource_mgr, file_pre='gen'):
+    """
+    Generate ops python file from yaml.
+    """
+    op_protos = resource_mgr.get_resource(ResourceType.OP_PROTO)
+    op_protos += resource_mgr.get_resource(ResourceType.FUNC_OP_PROTO)
+    doc_dict = resource_mgr.get_resource(ResourceType.OP_DOC_YAML)
+    generate_ops_prim_file(K.WORK_DIR, op_protos, doc_dict, file_pre)
+    generate_ops_def_file(K.WORK_DIR, op_protos, doc_dict, file_pre)
+    shutil.copy(os.path.join(K.WORK_DIR, K.PY_OPS_GEN_PATH, 'ops_auto_generate_init.txt'),
+                os.path.join(K.WORK_DIR, K.PY_AUTO_GEN_PATH, "__init__.py"))

mindspore/ops_generate/{op_def_py_generator.py → op_def_py/op_def_py_generator.py}

@@ -18,13 +18,13 @@ Generate Python operator definitions.
 
 import os
 
-import gen_constants as K
-import gen_utils
+import common.gen_constants as K
+import common.gen_utils as gen_utils
 
 # refactored
-import template
+import common.template as template
 
-from base_generator import BaseGenerator
+from common.base_generator import BaseGenerator
 
 
 class OpDefPyGenerator(BaseGenerator):
@@ -93,7 +93,8 @@ class OpDefPyGenerator(BaseGenerator):
             gen_py += func_code
             gen_py += "\n"
 
-        res_str = template.PY_LICENCE_STR + template.OPS_PY_DEF_HEADER + gen_py[:-len(template.NEW_LINE)]
+        res_str = template.PY_LICENSE_STR + \
+            template.OPS_PY_DEF_HEADER + gen_py[:-len(template.NEW_LINE)]
         save_path = os.path.join(work_path, K.PY_AUTO_GEN_PATH)
         file_name = f"{file_pre}_ops_def.py"
         gen_utils.save_file(save_path, file_name, res_str)

mindspore/ops_generate/{op_prim_py_generator.py → op_def_py/op_prim_py_generator.py}

@@ -18,16 +18,13 @@ Module for generating Python primitive operator definitions from specifications.
 
 import os
 
-import gen_constants as K
-import gen_utils
-import pyboost_utils
-
-# refactored
-from op_proto import OpProto
-import template
-from template import Template
-
-from base_generator import BaseGenerator
+import common.gen_constants as K
+import common.gen_utils as gen_utils
+import common.template as template
+from common.base_generator import BaseGenerator
+from common.op_proto import OpProto
+from common.template import Template
+from pyboost import pyboost_utils
 
 
 class OpPrimPyGenerator(BaseGenerator):
@@ -89,7 +86,8 @@ class OpPrimPyGenerator(BaseGenerator):
                 gen_py += f"\n\n{op_proto.op_name}_op={op_proto.op_class.name}()\n"
 
         pyboost_import_header = self.generate_pyboost_import_header(op_protos)
-        res_str = template.PY_LICENCE_STR + template.OPS_PY_PRIM_HEADER + pyboost_import_header + gen_py
+        res_str = template.PY_LICENSE_STR + \
+            template.OPS_PY_PRIM_HEADER + pyboost_import_header + gen_py
 
         save_path = os.path.join(work_path, K.PY_AUTO_GEN_PATH)
         file_name = f"{file_pre}_ops_prim.py"
@@ -281,8 +279,8 @@ class OpPrimPyGenerator(BaseGenerator):
         call_method_args_str = ", ".join(call_args)
         call_method_body_str = self._get_call_method_body_str(args_handlers, init_args, inputs_args, inputs_default,
                                                               op_proto)
-        call_code_str += f"""    def __call__(self, {call_method_args_str}):\n"""
-        call_code_str += f"""          {call_method_body_str}"""
+        call_code_str += f"""    def __call__(self, {call_method_args_str}):"""
+        call_code_str += f"""{call_method_body_str}"""
         return call_code_str
 
     def _get_call_method_body_str(self, args_handlers, init_args, inputs_args, inputs_default, op_proto: OpProto):
@@ -317,10 +315,21 @@ class OpPrimPyGenerator(BaseGenerator):
         call_method_body_str = ""
         is_pyboost = op_proto.op_dispatch and op_proto.op_dispatch.enable
         if is_pyboost:
+            call_method_body_str += f"""
+        # Add for jit context.
+        if jit_context() and jit_context().compiled:
+            return None"""
             pyboost_func_name = pyboost_utils.get_pyboost_name(op_proto.op_name)
-            call_method_body_str += f"""return _convert_stub({pyboost_func_name}(self, [{call_args_list_str}]))"""
+            call_method_body_str += f"""
+        res = {pyboost_func_name}(self, [{call_args_list_str}])"""
+            call_method_body_str += f"""
+        # Add for jit context.
+        if jit_context():
+            return jit_context().run_op(self, res, {call_args_list_str})
+        return res\n"""
         else:
-            call_method_body_str += f"""return super().__call__({call_args_list_str})\n"""
+            call_method_body_str += f"""
+        return super().__call__({call_args_list_str})\n"""
         return call_method_body_str
 
     def _generate_py_op_signature(self, op_proto: OpProto, args_name, args_default):

mindspore/ops_generate/{auto_grad_impl_cc_generator.py → pyboost/auto_grad_impl_cc_generator.py}

@@ -18,12 +18,12 @@ This module provides a generator class for creating C++ implementation files for
 
 import os
 
-import template
-from template import Template
-import gen_constants as K
-from gen_utils import save_file
-from base_generator import BaseGenerator
-from pyboost_utils import is_optional_param, get_input_dtype, is_op_multi_output
+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, is_op_multi_output
 
 
 class AutoGradImplGenerator(BaseGenerator):
@@ -35,6 +35,7 @@ class AutoGradImplGenerator(BaseGenerator):
         """
         Initialize the AutoGrad implementation generator with templates for code generation.
         """
+        self.OP_DEF_INC_HEAD_TEMPLATE = template.OP_DEF_INC_HEAD_TEMPLATE
         self.AUTO_GRAD_IMPL_CC_TEMPLATE = template.AUTO_GRAD_IMPL_CC_TEMPLATE
         self.DO_GRAD_FUNCTION_BODY_TEMPLATE = template.DO_GRAD_FUNCTION_BODY_TEMPLATE
         self.auto_grad_reg_template = Template("const_cast<kernel::pyboost::${class_name}GradFunc&>(" + \
@@ -55,13 +56,16 @@ class AutoGradImplGenerator(BaseGenerator):
         """
         auto_grad_reg_list = []
         do_grad_op_list = []
+        ops_inc_head_set = set()
         for op_proto in op_protos:
             if op_proto.op_dispatch is None or op_proto.op_dispatch.is_comm_op:
                 continue
             auto_grad_reg_list.append(self.auto_grad_reg_template.replace(class_name=op_proto.op_class.name))
             do_grad_op_list.append(self._get_single_do_grad_op(op_proto))
+            ops_inc_head_set.add(self.OP_DEF_INC_HEAD_TEMPLATE.replace(prefix_char=op_proto.op_class.name[0].lower()))
         pyboost_func_h_str = self.AUTO_GRAD_IMPL_CC_TEMPLATE.replace(do_grad_op=do_grad_op_list,
-                                                                     auto_grad_reg=auto_grad_reg_list)
+                                                                     auto_grad_reg=auto_grad_reg_list,
+                                                                     ops_inc=list(sorted(ops_inc_head_set)))
         save_path = os.path.join(work_path, K.PYBOOST_AUTO_GRAD_FUNC_GEN_PATH)
         file_name = "auto_grad_impl.cc"
         save_file(save_path, file_name, pyboost_func_h_str)

mindspore/ops_generate/{auto_grad_reg_cc_generator.py → pyboost/auto_grad_reg_cc_generator.py}

@@ -18,12 +18,12 @@ This module provides a generator class for creating C++ header files for AutoGra
 
 import os
 
-import template
-from template import Template
-import gen_constants as K
-from gen_utils import save_file
-from base_generator import BaseGenerator
-from pyboost_utils import is_optional_param, get_input_dtype
+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):
@@ -70,7 +70,7 @@ class AutoGradRegHeaderGenerator(BaseGenerator):
                                                                    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_OPS_KERNEL_FUNCTIONS_AUTO_GEN_PATH)
+        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)
 

mindspore/ops_generate/{gen_pyboost_func.py → pyboost/gen_pyboost_func.py}

@@ -16,34 +16,57 @@
 Generate pyboost function from pyboost_op.yaml
 """
 
-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 (
+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 (
+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
-from functions_cc_generator import FunctionsGenerator, FunctionsHeaderGenerator
+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 gen_pyboost_code(work_path, op_protos, doc_yaml_data, tensor_method_protos, mint_func_protos, alias_func_mapping):
+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, tensor_method_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)
@@ -94,14 +117,15 @@ def call_pyboost_functions_h_generator(work_path, op_protos):
     generator.generate(work_path, op_protos)
 
 
-def call_pyboost_functions_cpp_generator(work_path, op_protos, tensor_method_protos):
+def call_pyboost_functions_cpp_generator(work_path, op_protos):
     generator = PyboostFunctionsGenerator()
-    generator.generate(work_path, op_protos, tensor_method_protos)
+    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)
+    generator.generate(work_path, op_protos,
+                       mint_func_protos, alias_func_mapping)
 
 
 def call_pyboost_grad_functions_cpp_generator(work_path, op_protos):