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

mindspore/ops_generate/common/template_utils.py

@@ -0,0 +1,212 @@
+# 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.
+# ============================================================================
+"""Templates for code auto generation."""
+import re
+
+
+class Template:
+    """
+    template for generate c++/python code
+    """
+    regular_str = r"(^[^\n\S]*)?\$([^\d\W]\w*|\{,?[^\d\W]\w*\,?})"
+    regular_match = re.compile(regular_str, re.MULTILINE)
+
+    def __init__(self, code_pattern):
+        self.code_pattern = code_pattern
+
+    @staticmethod
+    def load_from_file(file_path):
+        """load template from file"""
+        with open(file_path, "r") as f:
+            return Template(f.read())
+
+    def replace(self, **kwargs):
+        """
+        replace param.
+        :param kwargs:
+        :return:
+        """
+
+        def find(key: str):
+            if key in kwargs:
+                return kwargs[key]
+            raise TypeError(f"{key} should be in kwargs!")
+
+        def add_indent(indent, var):
+            return "".join([indent + line + "\n" for data in var for line in str(data).splitlines()]).rstrip()
+
+        def extract_variable(key):
+            start = ""
+            end = ""
+            if key[0] == "{":
+                key = key[1:-1]
+                if key[0] == ",":
+                    start = ","
+                    key = key[1:]
+                if key[-1] == ",":
+                    end = ", "
+                    key = key[:-1]
+            return find(key), start, end
+
+        def match_rule(match):
+            indent = match.group(1)
+            key = match.group(2)
+            var, start, end = extract_variable(key)
+            if indent is not None:
+                if not isinstance(var, list):
+                    return add_indent(indent, [var])
+                return add_indent(indent, var)
+            if isinstance(var, list):
+                code = ", ".join(str(x) for x in var)
+                if not var:
+                    return code
+                return start + code + end
+            return str(var)
+
+        return self.regular_match.sub(match_rule, self.code_pattern)
+
+
+NEW_LINE = "\n"
+
+PYTHON_PRIM_TEMPLATE = Template("""
+
+class _Pyboost${class_name}Prim(${class_name}Prim_):
+    def __call__(self, ${input_args}):
+        ${process_func}
+        return super().__call__([${processed_args}])
+
+
+${func_impl_name}_impl = _Pyboost${class_name}Prim()
+""")
+
+IMPORT_PYBOOST_PRIM_HEADER = f"""
+from mindspore.ops._utils.arg_handler import *
+"""
+
+IMPORT_PYBOOST_FUNC_HEADER = f"""
+from mindspore.common import dtype as mstype
+from mindspore.ops.auto_generate.pyboost_inner_prim import *
+
+"""
+
+PY_LICENSE_STR = f"""# Copyright 2023 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.
+# ============================================================================
+"""
+
+OPS_PY_PRIM_HEADER = f"""
+\"\"\"Operators definition generated by gen_ops.py, includes primitive classes.\"\"\"
+
+from mindspore.ops.primitive import Primitive, prim_arg_register
+from mindspore.ops import signature as sig
+from mindspore.common import dtype as mstype
+from mindspore.common._decorator import deprecated
+from mindspore.ops._primitive_cache import _get_cache_prim
+from mindspore.ops._utils.arg_dtype_cast import type_it
+from mindspore.ops._utils.arg_handler import *
+from mindspore._c_expression import OpDtype
+from mindspore.common.jit_context import jit_context
+from mindspore._checkparam import is_stub_tensor
+"""
+
+OPS_PY_DEF_HEADER = f"""
+\"\"\"Operators definition generated by gen_ops.py, includes functions.\"\"\"
+
+from .gen_ops_prim import *
+from .pyboost_inner_prim import *
+from mindspore.ops.operations.manually_defined.ops_def import *
+from mindspore.ops._primitive_cache import _get_cache_prim
+"""
+
+CUSTOM_OPS_PY_DEF_HEADER = f"""
+\"\"\"Operators definition generated by gen_ops.py, includes functions.\"\"\"
+
+from .gen_ops_prim import *
+"""
+
+PRIMITIVE_CLASS_DESC = """    r\"\"\"
+    .. code-block::
+
+        prim = ops.$class_name($init_args_str)
+        out = prim($input_args_str)
+
+    is equivalent to
+
+    .. code-block::
+
+        ops.$func_name($args_str)
+
+    Refer to :func:`mindspore.ops.$func_name` for more details.
+    \"\"\"
+"""
+
+CC_LICENSE_STR = f"""/**
+ * Copyright 2023-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.
+ */"""
+
+OP_PROTO_TEMPLATE = Template("""
+${func_impl_declaration}
+OpDef g${class_name} = {
+  /*.name_=*/"${class_name}",
+  /*.args_=*/ {
+    ${input_args}
+  },
+  /* .returns_ = */ {
+    ${return_args}
+  },
+  /*.signatures_ =*/ {
+    ${signatures}
+  },
+  /*.indexes_ =*/ {
+    ${indexes}
+  },
+  /*.func_impl_=*/${func_impl_define},
+  /*.enable_dispatch_ =*/${enable_dispatch},
+  /*.is_view_ =*/${is_view},
+  /*.is_graph_view_ =*/${is_graph_view},
+};
+REGISTER_PRIMITIVE_OP_DEF(${class_name}, &g${class_name});
+""")
+
+OP_PRIM_CLASS_DEFINE_TEMPLATE = Template("""
+
+class ${class_name}(Primitive):
+${class_desc}${signature_code}${deprecated_code}
+${init_method}
+
+${call_method}""")

mindspore/ops_generate/gen_custom_ops.py

@@ -0,0 +1,69 @@
+# 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.
+# ============================================================================
+"""
+Auto generate custom ops files.
+"""
+
+import logging
+import argparse
+from resources.resource_manager import ResourceManager
+from resources.resource_list import ResourceType
+from resources.yaml_loader import CustomOpDocYamlLoader
+from op_def.ops_def_cc_generator import CustomOpsDefCcGenerator
+from op_def_py.custom_op_prim_py_generator import CustomOpPrimPyGenerator
+from op_def_py.op_def_py_generator import CustomOpDefPyGenerator
+from common.op_proto import CustomOpProtoLoader
+
+
+def get_config():
+    """get config from user"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-m", "--module_name", type=str, required=True)
+    parser.add_argument("-i", "--input_path", type=str, required=True)
+    parser.add_argument("-o", "--output_path", type=str, required=True)
+    parser.add_argument("-d", "--doc_path", type=str, required=True)
+    return parser.parse_args()
+
+
+def generate_custom_op_def(module_name, input_path, doc_path, output_path):
+    """Automatically generate all necessary files for custom operators."""
+    resource_mgr = ResourceManager()
+    resource_mgr.register_resource(CustomOpProtoLoader(input_path))
+    op_protos = resource_mgr.get_resource(ResourceType.OP_PROTO)
+    doc_dict = dict()
+    if doc_path != "":
+        resource_mgr.register_resource(CustomOpDocYamlLoader(doc_path))
+        doc_dict = resource_mgr.get_resource(ResourceType.OP_DOC_YAML)
+
+    generator = CustomOpsDefCcGenerator()
+    generator.generate(output_path, op_protos)
+    generator = CustomOpPrimPyGenerator()
+    generator.generate(output_path, module_name, op_protos, doc_dict, "gen")
+    generator = CustomOpDefPyGenerator()
+    generator.generate(output_path, op_protos, doc_dict, "gen")
+
+
+def main():
+    """main function"""
+    args = get_config()
+    generate_custom_op_def(args.module_name, args.input_path, args.doc_path, args.output_path)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as e:
+        logging.critical("Auto generate failed, err info: %s", e)
+        raise e

mindspore/ops_generate/op_def/ops_def_cc_generator.py

@@ -24,11 +24,10 @@ import common.gen_utils as gen_utils
 
 # refactored
 from common.op_proto import OpProto
-import common.template as template
+import common.template_utils as template
 
 from common.base_generator import BaseGenerator
 
-
 CC_OPS_DEF = """
 
 #include "$auto_generate_path/gen_ops_def.h"
@@ -72,8 +71,12 @@ class OpsDefCcGenerator(BaseGenerator):
             operator_name = op_proto.op_name
             class_name = op_proto.op_class.name
             if not op_proto.func_op:
-                gen_include_list.append(self.include_template.replace(path=K.MS_OPS_FUNC_IMPL_PATH,
-                                                                      operator_name=operator_name))
+                if op_proto.op_dispatch and op_proto.op_dispatch.is_comm_op:
+                    gen_include_list.append(self.include_template.replace(path=K.MS_OPS_COMM_FUNC_IMPL_PATH,
+                                                                          operator_name=operator_name))
+                else:
+                    gen_include_list.append(self.include_template.replace(path=K.MS_OPS_FUNC_IMPL_PATH,
+                                                                          operator_name=operator_name))
                 func_impl_declaration_str = self.func_impl_declaration_template.replace(class_name=class_name)
             else:
                 func_impl_declaration_str = self.empty_func_impl_declaration_template.replace(class_name=class_name)
@@ -85,7 +88,6 @@ class OpsDefCcGenerator(BaseGenerator):
             # Process outputs.
             return_args_str = get_cc_op_def_return(args_dict, op_proto)
 
-
             inputs_args = self.process_args(op_proto.op_args)
             signature_code = generate_cc_op_signature(op_proto.op_args_signature, inputs_args)
             enable_dispatch = "true" if op_proto.op_dispatch and op_proto.op_dispatch.enable else "false"
@@ -112,9 +114,9 @@ class OpsDefCcGenerator(BaseGenerator):
         save_path = os.path.join(work_path, K.MS_OP_DEF_AUTO_GENERATE_PATH)
         for numbering in range(math.ceil(op_size / max_op_size_in_one_file)):
             gen_include = ''.join(
-                gen_include_list[numbering*max_op_size_in_one_file: (numbering+1)*max_op_size_in_one_file])
+                gen_include_list[numbering * max_op_size_in_one_file: (numbering + 1) * max_op_size_in_one_file])
             gen_cc = ''.join(
-                gen_cc_list[numbering*max_op_size_in_one_file: (numbering+1)*max_op_size_in_one_file])
+                gen_cc_list[numbering * max_op_size_in_one_file: (numbering + 1) * max_op_size_in_one_file])
             cc_ops_def = self.CC_OPS_DEF_TEMPLATE.replace(auto_generate_path=K.MS_OP_DEF_AUTO_GENERATE_PATH,
                                                           gen_include=gen_include,
                                                           gen_cc_code=gen_cc)
@@ -148,6 +150,75 @@ class OpsDefCcGenerator(BaseGenerator):
         return inputs_name
 
 
+class CustomOpsDefCcGenerator(OpsDefCcGenerator):
+    """
+    Generates C++ definition files for operators.
+    """
+
+    def __init__(self):
+        """
+        Initializes templates for generating C++ operator definitions.
+        """
+        super(CustomOpsDefCcGenerator, self).__init__()
+
+        self.include_template = template.Template("""#include "${path}/${operator_name}.h\"\n""")
+        self.func_impl_declaration_template = template.Template("extern OpFuncImpl &g${class_name}FuncImpl;")
+        self.empty_func_impl_declaration_template = template.Template("static OpFuncImpl g${class_name}FuncImpl;")
+        self.func_impl_define_template = template.Template("g${class_name}FuncImpl")
+        self.OP_PROTO_TEMPLATE = template.OP_PROTO_TEMPLATE
+        self.CC_OPS_DEF_TEMPLATE = template.Template(CC_OPS_DEF)
+
+    def generate(self, work_path, op_protos):
+        """
+        Generates C++ code for operator definitions and saves it to a file.
+
+        Args:
+            work_path (str): The directory to save the generated files.
+            op_protos (list): A list of operator prototypes.
+        """
+        gen_cc_list = list()
+        gen_include_list = list()
+
+        for op_proto in op_protos:
+            class_name = "Custom_" + op_proto.op_name
+            func_impl_declaration_str = self.func_impl_declaration_template.replace(class_name=class_name)
+            func_impl_define = self.func_impl_define_template.replace(class_name=class_name)
+
+            # process input
+            args_dict, cc_index_str, input_args_str = process_input_args(op_proto)
+
+            # Process outputs.
+            return_args_str = get_cc_op_def_return(args_dict, op_proto)
+
+            inputs_args = self.process_args(op_proto.op_args)
+            signature_code = generate_cc_op_signature(op_proto.op_args_signature, inputs_args)
+            enable_dispatch = "true" if op_proto.op_dispatch and op_proto.op_dispatch.enable else "false"
+            is_view = "true" if op_proto.op_view else "false"
+            is_graph_view = "true" if op_proto.op_graph_view else "false"
+            op_def_cc = self.OP_PROTO_TEMPLATE.replace(class_name=class_name,
+                                                       input_args=input_args_str,
+                                                       return_args=return_args_str,
+                                                       signatures=signature_code,
+                                                       indexes=cc_index_str,
+                                                       enable_dispatch=enable_dispatch,
+                                                       is_view=is_view,
+                                                       is_graph_view=is_graph_view,
+                                                       func_impl_declaration=func_impl_declaration_str,
+                                                       func_impl_define=func_impl_define)
+
+            gen_cc_list.append(op_def_cc)
+
+        gen_include = ''.join(gen_include_list)
+        gen_cc = ''.join(gen_cc_list)
+        cc_ops_def = self.CC_OPS_DEF_TEMPLATE.replace(auto_generate_path=K.MS_OP_DEF_AUTO_GENERATE_PATH,
+                                                      gen_include=gen_include,
+                                                      gen_cc_code=gen_cc)
+
+        file_name = f"gen_custom_ops_def.cc"
+        ops_def_cc_file_str = template.CC_LICENSE_STR + cc_ops_def
+        gen_utils.save_file(work_path, file_name, ops_def_cc_file_str)
+
+
 def process_input_args(op_proto: OpProto):
     """
     Processes input arguments for C++ code generation.