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

mindspore/ops_generate/api/functional_map_cpp_generator.py

@@ -159,14 +159,6 @@ class FunctionalMapCppGenerator(BaseGenerator):
         sig_str += '}\n},'
         return sig_str
 
-    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_single_signature_str(self, func_api_name, tensor_proto, is_tensor_method) -> str:
         """
         Generates a single function signature string for the given operation prototype.
@@ -186,7 +178,7 @@ class FunctionalMapCppGenerator(BaseGenerator):
         arg_valid_types = []
         for _, arg in enumerate(op_proto.op_args):
             arg_name = arg.arg_name
-            if is_tensor_method and self._is_input_arg(arg_name, op_name):
+            if is_tensor_method and _is_input_arg(arg_name, op_name):
                 continue
 
             arg_valid_types = self._handle_arg_valid_types(arg, arg_name, arg_valid_types, func_api_name)
@@ -502,3 +494,12 @@ class FunctionalMapCppGenerator(BaseGenerator):
                 self._get_and_append_single_op_varargs_list(func_protos,
                                                             mint_varargs_list)
         return mint_varargs_list
+
+
+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

mindspore/ops_generate/api/functions_cc_generator.py

@@ -37,6 +37,13 @@ class FunctionsHeaderGenerator(BaseGenerator):
         """
         self.FUNCTIONS_H_TEMPLATE = template.FUNCTIONS_H_TEMPLATE
         self.function_interface_template = Template("${return_type} BACKEND_EXPORT ${op_name}(${input_args});")
+        self.function_interface_template_comm = Template(
+            "${return_type} BACKEND_EXPORT ${op_name}_inner(${input_args}," \
+            "CommHandlePtr comm_handle, const std::string& target);"
+        )
+        self.function_interface_template_comm_return_handle = Template(
+            "${return_type_with_handle} BACKEND_EXPORT ${op_name}(${input_args});"
+        )
 
     def generate(self, work_path, op_protos):
         """
@@ -48,14 +55,30 @@ class FunctionsHeaderGenerator(BaseGenerator):
         """
         functions_list = []
         for op_proto in op_protos:
-            if op_proto.op_dispatch is None or op_proto.op_dispatch.is_comm_op:
+            if op_proto.op_dispatch is None:
                 continue
             input_args_with_type_str = self._get_input_args(op_proto)
             return_type_str = _get_return_type_str(op_proto)
-            functions = self.function_interface_template.replace(op_name=op_proto.op_name,
-                                                                 input_args=input_args_with_type_str,
-                                                                 return_type=return_type_str)
+            function_template = (
+                self.function_interface_template
+                if not op_proto.op_dispatch.is_comm_op
+                else self.function_interface_template_comm
+            )
+            functions = function_template.replace(op_name=op_proto.op_name,
+                                                  input_args=input_args_with_type_str,
+                                                  return_type=return_type_str)
             functions_list.append(functions)
+            if op_proto.op_dispatch.is_comm_op:
+                return_type_with_handle = _get_return_type_with_handle_str(return_type_str)
+                functions_with_handle = (
+                    self.function_interface_template_comm_return_handle.replace(
+                        op_name=op_proto.op_name,
+                        input_args=input_args_with_type_str,
+                        return_type=return_type_str,
+                        return_type_with_handle=return_type_with_handle,
+                    )
+                )
+                functions_list.append(functions_with_handle)
         pyboost_func_h_str = self.FUNCTIONS_H_TEMPLATE.replace(op_call_with_grad=functions_list)
         save_path = os.path.join(work_path, K.MS_PYBOOST_FUNCTIONS_AUTO_GEN_PATH)
         file_name = "functions.h"
@@ -73,7 +96,7 @@ class FunctionsHeaderGenerator(BaseGenerator):
         """
         args_list = []
         for op_arg in op_proto.op_args:
-            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg))
+            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg), op_proto.op_view)
             args_list.append("const " + input_dtype + " &" + op_arg.arg_name)
         return args_list
 
@@ -89,12 +112,19 @@ class FunctionsGenerator(BaseGenerator):
         """
         self.FUNCTIONS_CC_TEMPLATE = template.FUNCTIONS_CC_TEMPLATE
         self.FUNCTION_BODY_TEMPLATE = template.FUNCTION_BODY_TEMPLATE
+        self.FUNCTION_COMM_BODY_TEMPLATE = template.FUNCTION_COMM_BODY_TEMPLATE
         self.pyboost_func_include_header_template = Template(
             f'#include "{K.MS_PYBOOST_BASE_PATH}/auto_generate/${{operator_name}}.h"\n'
         )
         self.clone_inplace_input_template = Template(
             'GetCloneFunc()(op, prim::kPrim${class_name}, device_target, {${grad_args}});'
         )
+        self.create_aclnn_op_template = Template(
+            'auto op = CREATE_PYBOOST_OP(${class_name}, device_target);'
+        )
+        self.create_internal_op_template = Template(
+            'auto op = CREATE_PYBOOST_SELECTED_OP(${class_name}, device_target);'
+        )
 
     def generate(self, work_path, op_protos):
         """
@@ -108,7 +138,7 @@ class FunctionsGenerator(BaseGenerator):
         op_call_with_grad_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:
+            if op_proto.op_dispatch is None:
                 continue
             func_include_headers_list.append(
                 self.pyboost_func_include_header_template.replace(operator_name=op_proto.op_name))
@@ -137,8 +167,22 @@ class FunctionsGenerator(BaseGenerator):
         inplace_clone_args = self._get_clone_input_args(op_proto, False, False)
         clone_func_str = self._get_clone_inplace_str(op_proto.op_inplace, op_proto.op_class.name, inplace_clone_args)
         return_type_str = _get_return_type_str(op_proto)
+        if op_proto.op_dispatch.is_comm_op:
+            return_type_with_handle = _get_return_type_with_handle_str(return_type_str)
+            comm_body = self.FUNCTION_COMM_BODY_TEMPLATE.replace(op_name=op_proto.op_name,
+                                                                 class_name=op_proto.op_class.name,
+                                                                 input_args=input_args,
+                                                                 clone_func=clone_func_str,
+                                                                 input_args_with_type=input_args_with_type,
+                                                                 return_type=return_type_str,
+                                                                 return_type_with_handle=return_type_with_handle)
+            return comm_body
+        create_op_str = self.create_aclnn_op_template.replace(class_name=op_proto.op_class.name)
+        if getattr(op_proto.op_dispatch, 'internal_op_ascend') != 'None':
+            create_op_str = self.create_internal_op_template.replace(class_name=op_proto.op_class.name)
         return self.FUNCTION_BODY_TEMPLATE.replace(op_name=op_proto.op_name,
                                                    class_name=op_proto.op_class.name,
+                                                   create_op=create_op_str,
                                                    input_args=input_args,
                                                    clone_func=clone_func_str,
                                                    input_args_with_type=input_args_with_type,
@@ -157,7 +201,7 @@ class FunctionsGenerator(BaseGenerator):
         """
         args_list = []
         for op_arg in op_proto.op_args:
-            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg))
+            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg), op_proto.op_view)
             if has_type:
                 args_list.append("const " + input_dtype + " &" + op_arg.arg_name)
             else:
@@ -194,7 +238,7 @@ class FunctionsGenerator(BaseGenerator):
         """
         args_list = []
         for op_arg in op_proto.op_args:
-            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg))
+            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg), op_proto.op_view)
             if has_type:
                 args_list.append(f"const {input_dtype} &{op_arg.arg_name}")
             else:
@@ -205,6 +249,10 @@ class FunctionsGenerator(BaseGenerator):
         return args_list
 
 
+def _get_return_type_with_handle_str(return_type_str):
+    return f"std::tuple<{return_type_str}, CommHandlePtr>"
+
+
 def _get_return_type_str(op_proto):
     """
     Get the return type string for the function.
@@ -217,8 +265,8 @@ def _get_return_type_str(op_proto):
     """
     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 op_proto.op_returns:
         temp_return = get_return_type(return_obj.arg_dtype)

mindspore/ops_generate/api/tensor_func_reg_cpp_generator.py

@@ -81,7 +81,7 @@ class TensorFuncRegCppGenerator(BaseGenerator):
             'MS_LOG(INFO) << "Call Tensor${class_name}";\n'
             'auto res = mindspore::pynative::'
             '${pyboost_function}(mindspore::prim::kPrim${class_name}, parse_args.src_types_, ${convert_args});\n'
-            'trace::Capture(parse_args.arg_list_, "${class_name}", &res);\n'
+            'trace::Capture(parse_args.arg_list_, mindspore::prim::kPrim${class_name}, &res);\n'
             'return res;\n'
         )
         self.callback_python_template = Template(

mindspore/ops_generate/common/base_generator.py

@@ -1,3 +1,17 @@
+# 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.
+# ============================================================================
 """
 Base generator interface for other pyboost file generator classes
 """

mindspore/ops_generate/common/gen_constants.py

@@ -55,13 +55,13 @@ TENSOR_PY_CC_PATH = "mindspore/ccsrc/pybind_api/ir/tensor_register/auto_generate
 
 # yaml keys def
 OP_KEYS = {'args', 'args_signature', 'returns', 'function', 'class', 'view', 'graph_view', 'dispatch', 'labels',
-           'bprop_expander'}
+           'bprop_expander', 'non-differentiable'}
 ARG_KEYS = {'dtype', 'default', 'prim_init', 'type_cast', 'arg_handler'}
 RETURN_KEYS = {'dtype', 'inplace', 'type_cast'}
 ARG_SIGNATURE_KEYS = {'rw_write', 'rw_read', 'rw_ref', 'dtype_group'}
 CLASS_KEYS = {'name', 'disable'}
 FUNCTION_KEYS = {'name', 'disable'}
-DISPATCH_KEYS = {'enable', 'is_comm_op', 'Ascend', 'GPU', 'CPU'}
+DISPATCH_KEYS = {'enable', 'is_comm_op', 'Ascend', 'InternalOpAscend', 'GPU', 'CPU'}
 TENSOR_FUNC_KEYS = {'op_yaml', 'py_method', 'kwonlyargs',
                     'varargs', 'alias', 'Ascend', 'GPU', 'CPU', 'interface'}
 
@@ -89,3 +89,8 @@ MS_OPS_KERNEL_PATH = "mindspore/ops/kernel"
 MS_PYBOOST_FUNCTIONS_AUTO_GEN_PATH = "mindspore/ccsrc/pyboost/functions/auto_generate"
 MS_COMMON_PYBOOST_KERNEL_PATH = os.path.join(MS_OPS_KERNEL_PATH, "common/pyboost")
 MS_PYBOOST_BASE_PATH = "mindspore/ccsrc/pyboost"
+MS_PYBOOST_INTERNAL_FUNCTIONS_AUTO_GEN_PATH = os.path.join(MS_OPS_KERNEL_PATH,
+                                                           "ascend/pyboost/internal/functions")
+MS_INTERNAL_PYBOOST_GEN_PATH = "mindspore/ccsrc/plugin/device/ascend/kernel/internal/pyboost/auto_gen"
+MS_PLUGIN_INTERNAL_PATH = "mindspore/ccsrc/plugin/device/ascend/kernel/internal"
+MS_OPS_PYBOOST_INTERNAL = "mindspore/ops/kernel/ascend/pyboost/internal"

mindspore/ops_generate/common/gen_utils.py

@@ -121,25 +121,6 @@ def merge_files(origin_dir, merged_file_path, file_format):
     merge_files_to_one_file(op_yaml_file_names, merged_file_path)
 
 
-def merge_files_append(origin_dir, merged_file_path, file_format):
-    """
-    Merge multiple files into one file.
-    origin_dir: indicates the origin file directory.
-    merged_file_path: indicates the merged file path.
-    file_format: indicates the format of regular matching.
-    Files whose names meet the regular matching in 'origin_dir' directory will be merged into one file.
-    """
-    file_paths = glob.glob(os.path.join(origin_dir, file_format))
-    merged_content = ''
-    file_paths.sort()
-    for file_path in file_paths:
-        with open(file_path, 'r') as file:
-            merged_content += file.read()
-            merged_content += '\n'
-    with open(merged_file_path, 'a') as file:
-        file.write(merged_content)
-
-
 def safe_load_yaml(yaml_file_path):
     """
     Load yaml dictionary from file.

mindspore/ops_generate/common/op_proto.py

@@ -112,10 +112,12 @@ class OpDispatch:
         gpu (str): The dispatch type for the GPU.
     """
 
-    def __init__(self, enable=False, is_comm_op=False, ascend='default', cpu='default', gpu='default'):
+    def __init__(self, enable=False, is_comm_op=False, ascend='default',
+                 internal_op_ascend='None', cpu='default', gpu='default'):
         self.enable = enable
         self.is_comm_op = is_comm_op
         self.ascend = ascend
+        self.internal_op_ascend = internal_op_ascend
         self.cpu = cpu
         self.gpu = gpu
 
@@ -151,7 +153,8 @@ class OpProto:
                  op_inplace=False,
                  op_labels=None,
                  op_deprecated=None,
-                 bprop_expander=True):
+                 bprop_expander=True,
+                 non_differentiable=False):
         self.op_name = op_name
         self.op_args = op_args
         self.op_function = op_function
@@ -165,6 +168,7 @@ class OpProto:
         self.op_labels = op_labels
         self.op_deprecated = op_deprecated
         self.bprop_expander = bprop_expander
+        self.non_differentiable = non_differentiable
 
     @staticmethod
     def load_from_yaml(op_name, op_data):
@@ -208,10 +212,12 @@ class OpProto:
         # get op deprecated
         op_deprecated = op_data.get('deprecated', None)
         bprop_expander = op_data.get('bprop_expander', True)
+        non_differentiable = op_data.get('non-differentiable', False)
         op_proto = OpProto(op_name=op_name, op_args=op_args, op_returns=op_returns, op_function=op_function,
                            op_class=op_class, op_dispatch=op_dispatch, op_args_signature=op_args_signature,
                            op_view=op_view, op_graph_view=op_graph_view, op_inplace=op_inplace, op_labels=op_labels,
-                           op_deprecated=op_deprecated, bprop_expander=bprop_expander)
+                           op_deprecated=op_deprecated, bprop_expander=bprop_expander,
+                           non_differentiable=non_differentiable)
         return op_proto
 
 
@@ -401,9 +407,10 @@ def get_op_dispatch(op_name, op_data):
             f'The dispatch enable value should be bool, but get {type(enable)}, op name is {op_name}.')
     is_comm_op = op_dispatch.get('is_comm_op', False)
     ascend = op_dispatch.get('Ascend', 'default')
+    internal_op_ascend = op_dispatch.get('InternalOpAscend', 'None')
     cpu = op_dispatch.get('CPU', 'default')
     gpu = op_dispatch.get('GPU', 'default')
-    return OpDispatch(enable, is_comm_op, ascend, cpu, gpu)
+    return OpDispatch(enable, is_comm_op, ascend, internal_op_ascend, cpu, gpu)
 
 
 def get_op_class(op_name, op_data) -> OpClass:

mindspore/ops_generate/common/template.py

@@ -93,7 +93,6 @@ ${func_impl_name}_impl = _Pyboost${class_name}Prim()
 """)
 
 IMPORT_PYBOOST_PRIM_HEADER = f"""
-from mindspore.common._stub_tensor import _convert_stub
 from mindspore.ops._utils.arg_handler import *
 """
 
@@ -118,19 +117,89 @@ REGISTER_PYBOOST_GRAD_DEFINE_TEMPLATE = Template(
 REGISTER_PYBOOST_GRAD_TEMPLATE = Template("${register_func}")
 
 PYBOOST_FUNCTION_HEADER_TEMPLATE = Template.load_from_file(
-    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_functions_h.tpl'))
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_api_h.tpl'))
 
-PYBOOST_FUNCTION_TEMPLATE = Template.load_from_file(
-    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_function.tpl'))
+PYBOOST_CORE_HEADER_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_core_header.tpl'))
+PYBOOST_INTERNAL_OP_HEADER_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/internal/template/pyboost_internal_header_template.tpl'))
+
+PYBOOST_INTERNAL_SINGLE_OP_HEADER_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'./{K.MS_OPS_PYBOOST_INTERNAL}/template/pyboost_internal_single_op_header_template.tpl'))
+
+PYBOOST_INTERNAL_SINGLE_OP_SOURCE_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'./{K.MS_OPS_PYBOOST_INTERNAL}/template/pyboost_internal_single_op_source_template.tpl'))
+
+PYBOOST_INTERNAL_SINGLE_OP_CUSTOMIZE_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'./{K.MS_OPS_PYBOOST_INTERNAL}/template/pyboost_internal_single_op_customize_source_template.tpl'))
+
+PYBOOST_INTERNAL_OP_SOURCE_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'./{K.MS_OPS_PYBOOST_INTERNAL}/template/pyboost_internal_source_template.tpl'))
+
+PYBOOST_INTERNAL_FUNCTION_HEADER_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_OPS_PYBOOST_INTERNAL}/template/pyboost_internal_functions_header_template.tpl'))
+
+PYBOOST_INTERNAL_FUNCTION_SOURCE_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_OPS_PYBOOST_INTERNAL}/template/pyboost_internal_functions_source_template.tpl'))
+
+PYBOOST_INTERNAL_FUNCTION_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_OPS_PYBOOST_INTERNAL}/template/pyboost_internal_function_template.tpl'))
+
+PYBOOST_KERNEL_INFO_ADAPTER_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_PLUGIN_INTERNAL_PATH}/pyboost/template/kernel_info_adapter.tpl'))
+
+PYBOOST_KERNEL_INFO_ADAPTER_H_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_PLUGIN_INTERNAL_PATH}/pyboost/template/kernel_info_adapter_h.tpl'))
+
+PYBOOST_INTERNAL_KERNEL_INFO_ADAPTER_H_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_PLUGIN_INTERNAL_PATH}/pyboost/template/internal_kernel_info_adapter_h.tpl'))
+
+PYBOOST_INTERNAL_KERNEL_INFO_ADAPTER_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_PLUGIN_INTERNAL_PATH}/pyboost/template/internal_kernel_info_adapter.tpl'))
+
+PYBOOST_INTERNAL_KERNEL_INFO_ADAPTER_CPP_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_PLUGIN_INTERNAL_PATH}/pyboost/template/internal_kernel_info_adapter_cpp.tpl'))
+
+PYBOOST_INTERNAL_KERNEL_INFO_ADAPTER_SINGLE_CPP_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'{K.MS_PLUGIN_INTERNAL_PATH}/pyboost/template/internal_kernel_info_adapter_single_cpp.tpl'))
 
-PYBOOST_COMM_FUNCTION_TEMPLATE = Template.load_from_file(
-    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_comm_function.tpl'))
+PYBOOST_REGISTRY_BODY_CC_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_registry_body_cc.tpl'))
 
-PYBOOST_FUNCTION_DYNAMIC_OUTPUT_TEMPLATE = Template.load_from_file(
-    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_function_dynamic_output.tpl'))
+PYBOOST_CORE_BODY_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_core_body.tpl'))
 
-PYBOOST_FUNCTIONS_CC_TEMPLATE = Template.load_from_file(
-    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_functions_cc.tpl'))
+PYBOOST_CORE_BODY_COMM_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_core_body_comm.tpl'))
+
+PYBOOST_CORE_BODY_SYNC_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_core_body_sync.tpl'))
+
+PYBOOST_REGISTRY_CC_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_registry_cc.tpl'))
+
+PYBOOST_API_CC_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_api_cc.tpl'))
+
+PYBOOST_API_BODY_CC_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_api_body_cc.tpl'))
+
+PYBOOST_CORE_CC_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/pyboost_core_cc.tpl'))
 
 PYBOOST_OVERLOAD_FUNCTIONS_CC_TEMPLATE = Template.load_from_file(
     os.path.join(
@@ -213,6 +282,10 @@ PYBOOST_OP_REGISTER_TEMPLATE = Template.load_from_file(
 PYBOOST_ASCEND_OP_HEADER_TEMPLATE = Template.load_from_file(
     os.path.join(K.WORK_DIR,
                  f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/template/pyboost_aclnn_header_template.tpl'))
+PYBOOST_ASCEND_INTERNAL_OP_HEADER_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR,
+                 f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/internal/template/pyboost_internal_header_template.tpl')
+)
 
 PYBOOST_ASCEND_OP_SOURCE_TEMPLATE = Template.load_from_file(
     os.path.join(K.WORK_DIR,
@@ -226,6 +299,8 @@ PYBOOST_ASCEND_SINGLE_HCLL_OP_HEADER_TEMPLATE = Template.load_from_file(
     os.path.join(K.WORK_DIR,
                  f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/template/pyboost_aclnn_single_hccl_op_header_template.tpl'))
 
+PYBOOST_CALL_FUNC_TEMPLATE = Template('${return_type} Call(${call_args_with_type}) override;')
+
 PYBOOST_ASCEND_SINGLE_OP_SOURCE_TEMPLATE = Template.load_from_file(
     os.path.join(K.WORK_DIR,
                  f'./{K.MS_OPS_KERNEL_PATH}/ascend/pyboost/template/pyboost_aclnn_single_op_source_template.tpl'))
@@ -373,7 +448,6 @@ 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._stub_tensor import _convert_stub
 from mindspore.common.jit_context import jit_context
 from mindspore._checkparam import is_stub_tensor
 """
@@ -509,6 +583,9 @@ FUNCTIONS_CC_TEMPLATE = Template.load_from_file(
 FUNCTION_BODY_TEMPLATE = Template.load_from_file(
     os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/function_body.tpl'))
 
+FUNCTION_COMM_BODY_TEMPLATE = Template.load_from_file(
+    os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/comm_function_body.tpl'))
+
 FUNCTIONS_H_TEMPLATE = Template.load_from_file(
     os.path.join(K.WORK_DIR, './mindspore/ccsrc/pynative/op_function/template/functions_h.tpl'))
 

mindspore/ops_generate/gen_ops.py

@@ -47,6 +47,6 @@ def main():
 if __name__ == "__main__":
     try:
         main()
-    except Exception as e: # pylint: disable=broad-except
+    except Exception as e:
         logging.critical("Auto generate failed, err info: %s", e)
         raise e

mindspore/ops_generate/op_def/lite_ops_cpp_generator.py

@@ -24,7 +24,6 @@ import common.template as template
 from common.base_generator import BaseGenerator
 from pyboost import pyboost_utils
 
-
 LITE_OPS_H = """
 #ifndef MINDSPORE_CORE_OPS_GEN_LITE_OPS_H_
 #define MINDSPORE_CORE_OPS_GEN_LITE_OPS_H_
@@ -122,9 +121,10 @@ class LiteOpsCcGenerator(BaseGenerator):
         self.inc_ops_head_templat = template.Template(INC_OPS_HEAD)
         self.lite_ops_cc_template = template.Template(LITE_OPS_CC)
         self.op_template = template.op_template
-        self.register_primitive_c_template = template.Template("${op_name}::${op_name}():BaseOperator(kName${op_name}) {}\n"
-                                                               "REGISTER_PRIMITIVE_C(kName${op_name}, ${op_name});\n"
-                                                               "MIND_API_OPERATOR_IMPL(${op_name}, BaseOperator);\n\n")
+        self.register_primitive_c_template = template.Template(
+            "${op_name}::${op_name}():BaseOperator(kName${op_name}) {}\n"
+            "REGISTER_PRIMITIVE_C(kName${op_name}, ${op_name});\n"
+            "MIND_API_OPERATOR_IMPL(${op_name}, BaseOperator);\n\n")
 
     def generate(self, work_path, op_protos):
         """

mindspore/ops_generate/op_def/ops_name_h_generator.py

@@ -19,15 +19,12 @@ This module defines the `OpsNameHGenerator` class, which produces C++ code to de
 constants for operator names based on given 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_NAME_OP_DEF = """
 #ifndef MINDSPORE_CORE_OP_NAME_${suffix}_H_
 #define MINDSPORE_CORE_OP_NAME_${suffix}_H_

mindspore/ops_generate/op_def/ops_primitive_h_generator.py

@@ -19,15 +19,12 @@ This module defines the `OpsPrimitiveHGenerator` class, which creates C++ header
 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_
@@ -42,7 +39,6 @@ $ops_prim_gen
 #endif  // MINDSPORE_CORE_OPS_GEN_OPS_PRIMITIVE_${suffix}_H_
 """
 
-
 OP_PRIM_OP_DEF_CC = """
 
 #include "$auto_gen_path/gen_ops_primitive_${suffix}.h"

mindspore/ops_generate/op_def_py/op_prim_py_generator.py

@@ -142,7 +142,10 @@ class OpPrimPyGenerator(BaseGenerator):
                 # step1.3: get args set prim arg expression:
                 assign_str = self._get_assign_str_by_type_it(op_proto.op_class.name, arg)
                 if arg.arg_handler:
-                    assign_str = f"""        self._set_prim_arg_with_handler("{arg.arg_name}", {assign_str}, {arg.arg_handler})"""
+                    assign_str = (
+                        f'        self._set_prim_arg_with_handler('
+                        f'"{arg.arg_name}", {assign_str}, {arg.arg_handler})'
+                    )
                 else:
                     assign_str = f"""        self._set_prim_arg("{arg.arg_name}", {assign_str})"""
                 args_assign.append(assign_str)
@@ -318,7 +321,7 @@ class OpPrimPyGenerator(BaseGenerator):
             call_method_body_str += f"""
         # Add for jit context.
         if jit_context() and jit_context().compiled:
-            return None"""
+            return jit_context().default_output()"""
             pyboost_func_name = pyboost_utils.get_pyboost_name(op_proto.op_name)
             call_method_body_str += f"""
         res = {pyboost_func_name}(self, [{call_args_list_str}])"""