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

mindspore/common/api.py

@@ -1,6 +1,6 @@
 # This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
 #
-# Copyright 2020-2024 Huawei Technologies Co., Ltd
+# Copyright 2020-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.
@@ -17,6 +17,8 @@
 """Providing interface methods."""
 from __future__ import absolute_import
 
+__all__ = ['ms_memory_recycle', 'jit', 'jit_class', 'flops_collection']
+
 import gc
 import types
 import sys
@@ -42,23 +44,25 @@ from mindspore.common.sparse_tensor import RowTensor as PythonRowTensor
 from mindspore._c_expression.amp import get_curr_amp_strategy
 from mindspore._c_expression import GraphExecutor_, JitExecutor_, CSRTensor, RowTensor, COOTensor, \
     PyNativeExecutor_, verify_inputs_signature, init_exec_dataset, _set_dataset_mode_config, init_pipeline, \
-    _run_jit_pipeline, _ms_memory_recycle, _bind_device_ctx, StubNode, MSContext, TensorPy as Tensor
+    _run_jit_pipeline, _ms_memory_recycle, _bind_device_ctx, MSContext, TensorPy as Tensor
 from mindspore.parallel._ps_context import _is_role_sched
 from mindspore.parallel._utils import _check_full_batch, _get_parameter_broadcast, _is_in_auto_parallel_mode, \
     _is_parallel_mode
 from mindspore import _checkparam as Validator
 from mindspore._checkparam import is_stub_tensor
-from mindspore.common._utils import is_shape_unknown
+from mindspore.common._utils import is_shape_unknown, get_func
 from mindspore.common.mutable import mutable, _check_element_type
-from mindspore.common.auto_dynamic_shape import get_auto_dynamic_shape_args, update_auto_dynamic_shape_phase, \
-    get_auto_dynamic_shape_args_with_check_input_signature, update_auto_dynamic_shape_phase_with_check_input_signature
+from mindspore.common.dynamic_shape.auto_dynamic_shape import get_auto_dynamic_shape_args, \
+    update_auto_dynamic_shape_phase
+from mindspore.common.dynamic_shape.enable_dynamic import generate_dynamic_tensor_args, ENABLE_DYNAMIC
 from mindspore.common._pijit_context import PIJitCaptureContext
-from mindspore.common.parameter import Parameter, set_parameter_hook_updated, parameter_hook_updated
+from mindspore.common.parameter import Parameter
+from mindspore.common.hook_handle import _hook_version
 from mindspore.common.jit_context import jit_context
 from mindspore.common.jit_trace import _jit_trace
 from mindspore.parallel._utils import _init_auto_parallel_context, _clear_auto_parallel_context
 
-# Store ms_function class compiled pipeline cache.
+# Store jit class compiled pipeline cache.
 ms_compile_cache = set()
 # Store cell compiled pipeline cache.
 cells_compile_cache = {}
@@ -72,6 +76,11 @@ ARG_SPECIFIED = "arg_specified_infos"
 TOTAL_ARG_LEN = "total_arg_length"
 
 
+def _real_phase(phase, obj):
+    real_phase = phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+    return real_phase
+
+
 def _check_recompile_args(compile_args, kwargs):
     """Check recompile of graph"""
 
@@ -134,8 +143,6 @@ def _convert_python_data(data):
     """
     if isinstance(data, PythonTensor):
         return data
-    if isinstance(data, StubNode):
-        return ms.common._stub_tensor._convert_stub(data)
     if data.__class__ is tuple:
         # Handle namedtuple since its type is tuple.
         if hasattr(data, "_fields"):
@@ -278,13 +285,13 @@ def __get_compile_cache_dep_files(file_path, compile_cache_dep_files, pkg):
             module = importlib.util.module_from_spec(module_spec)
             if hasattr(module, '__file__'):
                 dep_file_path = module.__file__
+                # Exclude the installed modules.
+                if not _in_sys_path(dep_file_path) and dep_file_path not in compile_cache_dep_files:
+                    logger.debug(f"dependent file path: {dep_file_path}")
+                    compile_cache_dep_files.append(dep_file_path)
+                    __get_compile_cache_dep_files(dep_file_path, compile_cache_dep_files, module.__package__)
             else:
                 continue
-            # Exclude the installed modules.
-            if not _in_sys_path(dep_file_path) and dep_file_path not in compile_cache_dep_files:
-                logger.debug(f"dependent file path: {dep_file_path}")
-                compile_cache_dep_files.append(dep_file_path)
-                __get_compile_cache_dep_files(dep_file_path, compile_cache_dep_files, module.__package__)
 
 
 def _get_compile_cache_dep_files():
@@ -342,7 +349,7 @@ def _get_parameter_layout():
     return layout
 
 
-def _handle_arg(obj, arg, compile_arg):
+def _handle_arg(obj, arg, has_mutable_arg):
     """Handle arg for runtime .If need handle the arg, return True"""
     from mindspore._extends.parse import compile_config
     if isinstance(arg, PythonTensor):
@@ -352,7 +359,7 @@ def _handle_arg(obj, arg, compile_arg):
             return arg
     elif isinstance(arg, (Tensor, CSRTensor, COOTensor)):
         return arg
-    elif compile_arg is not None and hasattr(compile_arg, "__ms_mutable__") and getattr(compile_arg, "__ms_mutable__"):
+    elif has_mutable_arg:
         # mutable([]) will be eliminated by FuncGraphSpecializer, and empty list is not supported by backend.
         if isinstance(arg, list) and not arg:
             return None
@@ -366,7 +373,7 @@ def _handle_arg(obj, arg, compile_arg):
     return None
 
 
-def _handle_arg_predict(obj, arg, compile_arg):
+def _handle_arg_predict(obj, arg, has_mutable_arg):
     """Handle arg for runtime .If need handle the arg, return True"""
     if arg is None:
         return None
@@ -375,8 +382,7 @@ def _handle_arg_predict(obj, arg, compile_arg):
         return None
 
     if isinstance(arg, (list, tuple)):
-        if compile_arg is not None and hasattr(compile_arg, "__ms_mutable__") and \
-                getattr(compile_arg, "__ms_mutable__"):
+        if has_mutable_arg:
             # mutable([]) will be eliminated by FuncGraphSpecializer, and empty list is not supported by backend.
             if isinstance(arg, list) and not arg:
                 return None
@@ -388,35 +394,30 @@ def _handle_arg_predict(obj, arg, compile_arg):
     return arg
 
 
-def _get_args_for_run(obj, args, kwargs, compile_args):
+def _get_args_for_run(obj, args, kwargs, has_mutable_args_list, is_predict):
     """Get the actual input args and kwargs for runtime."""
     new_args = []
-    for arg, compile_arg in zip(args, compile_args):
-        new_arg = _handle_arg(obj, arg, compile_arg)
+    fn = _handle_arg_predict if is_predict else _handle_arg
+    for arg, has_mutable_arg in zip(args, has_mutable_args_list):
+        new_arg = fn(obj, arg, has_mutable_arg)
         if new_arg is not None:
             new_args.append(new_arg)
 
     for _, value in kwargs.items():
-        new_value = _handle_arg(obj, value, None)
+        new_value = fn(obj, value, None)
         if new_value is not None:
             new_args.append(new_value)
 
     return new_args
 
 
-def _get_args_for_run_predict(obj, args, kwargs, compile_args):
-    """Get the actual input args and kwargs for runtime."""
+def _get_mutable_flags(compile_args):
+    """Get a list of booleans indicating whether each argument is marked as mutable"""
     new_args = []
-    for arg, compile_arg in zip(args, compile_args):
-        new_arg = _handle_arg_predict(obj, arg, compile_arg)
-        if new_arg is not None:
-            new_args.append(new_arg)
-
-    for _, value in kwargs.items():
-        new_value = _handle_arg_predict(obj, value, None)
-        if new_value is not None:
-            new_args.append(new_value)
-
+    for compile_arg in compile_args:
+        has_mutable_arg = compile_arg is not None and hasattr(compile_arg, "__ms_mutable__") and \
+                          getattr(compile_arg, "__ms_mutable__")
+        new_args.append(has_mutable_arg)
     return new_args
 
 
@@ -544,10 +545,12 @@ def _get_parameter_ids(args, kwargs):
             parameter_ids += str(id(value))
     return parameter_ids
 
+
 def _get_tensor_hook_key(tensor):
     """Get the hook key of Tensor/Parameter"""
     return ".".join(map(str, map(id, tensor.hooks())))
 
+
 def _get_hook_key(*args, **kwargs):
     """Get the hook key of Tensors/Parameters"""
     hook_key = ""
@@ -586,13 +589,16 @@ class _JitExecutor:
         The result of pipeline running in graph mode.
     """
 
-    def __init__(self, fn, ms_create_time, input_signature=None, obj=None, jit_config=None, dynamic=0):
+    def __init__(self, fn, ms_create_time, input_signature=None, obj=None, jit_config=None, dynamic=0,
+                 cell_cache_key_extend=''):
         init_pipeline()
         if not isinstance(fn, (types.FunctionType, types.MethodType)):
             raise RuntimeError('fn {} is not function or method'.format(fn))
 
         self.fn = fn
         self.input_signature = input_signature
+        self.dynamic_args_shapes = getattr(get_func(fn), ENABLE_DYNAMIC, None)
+        self.enable_jit_dynamic = self.dynamic_args_shapes is not None
         self.obj = None
         if obj and hasattr(obj, fn.__name__):
             self.obj = obj
@@ -606,9 +612,55 @@ class _JitExecutor:
         self._compile_args = None
         self._enable_auto_dynamic = dynamic == 1
         self.jit_config_dict = jit_config.jit_config_dict if jit_config else None
+        self._cell_cache_key_extend = cell_cache_key_extend
+
+    def _predict(self, *args, **kwargs):
+        """Dedicated routine for predict."""
+        if not hasattr(self.obj, "phase"):
+            return False, None
+
+        predict_vailid_phase = {"prefill", 'increment'}
+        predict_phase = self.obj.phase
+        if predict_phase not in predict_vailid_phase:
+            return False, None
+
+        args_list = args
+        if self.obj is not None:
+            args_list = args_list[1:]
+
+        if predict_phase not in self.obj.phase_cache:
+            try:
+                predict_phase = self.compile(self.fn.__name__, *args_list, **kwargs)
+            except Exception as err:
+                _pynative_executor.clear_res()
+                raise err
+        else:  # get compiled args to generate run args by _generate_run_args
+            compile_args = self._generate_compile_args(args_list)
+            key_id = self._get_key_id()
+            if self.input_signature is None:
+                compile_args = get_auto_dynamic_shape_args(
+                    compile_args, key_id, self._enable_auto_dynamic
+                )
+            self._compile_args = compile_args
+
+        new_inputs = self._generate_run_args(args_list, kwargs)
+        if self.jit_config_dict:
+            jit_config_dict = self.jit_config_dict
+        else:
+            jit_config_dict = JitConfig().jit_config_dict
+        self._graph_executor.set_jit_config(jit_config_dict)
+        output = self._graph_executor(
+            tuple(new_inputs),
+            self.obj.phase_cache[self.obj.phase]
+        )
+        res = _convert_python_data(output)
+        return True, res
 
     @_wrap_func
     def __call__(self, *args, **kwargs):
+        predict, res = self._predict(*args, **kwargs)
+        if predict:
+            return res
         if jit_context() and jit_context().is_nested():
             return jit_context().run_graph("", None, *())
         args_list = args
@@ -616,12 +668,9 @@ class _JitExecutor:
             args_list = args_list[1:]
         phase = ""
         try:
-            if context.get_context("mode") == context.PYNATIVE_MODE:
-                _pynative_executor.set_jit_compile_status(True, phase)
-                phase = self.compile(self.fn.__name__, *args_list, **kwargs)
-                _pynative_executor.set_jit_compile_status(False, phase)
-            else:
-                phase = self.compile(self.fn.__name__, *args_list, **kwargs)
+            _pynative_executor.set_jit_compile_status(True, phase)
+            phase = self.compile(self.fn.__name__, *args_list, **kwargs)
+            _pynative_executor.set_jit_compile_status(False, phase)
         except Exception as err:
             _pynative_executor.clear_res()
             raise err
@@ -630,31 +679,27 @@ class _JitExecutor:
             return None
 
         new_inputs = self._generate_run_args(args_list, kwargs)
-        if context.get_context("mode") == context.PYNATIVE_MODE and not jit_context():
-            output = _pynative_executor.grad_jit(*new_inputs)
+        if self.jit_config_dict:
+            jit_config_dict = self.jit_config_dict
         else:
-            output = self._graph_executor(tuple(new_inputs), phase)
-            if jit_context():
-                if is_stub_tensor(output):
-                    output = output.stub_sync()
-                return jit_context().run_graph(phase, output, *tuple(new_inputs))
-
+            jit_config_dict = JitConfig().jit_config_dict
+        self._graph_executor.set_jit_config(jit_config_dict)
+        output = _pynative_executor.grad_jit(*new_inputs)
+        if jit_context():
+            if is_stub_tensor(output):
+                output = output.stub_sync()
+            return jit_context().run_graph(phase, output, *tuple(new_inputs))
         return output
 
     def compile(self, method_name, *args, **kwargs):
         """Returns pipeline for the given args."""
-        # Check whether hook function registered on Cell object.
-        if self.obj and hasattr(self.obj, "_hook_fn_registered"):
-            if self.obj._hook_fn_registered():
-                logger.warning(f"For 'Cell', it's not support hook function when using 'jit' decorator. "
-                               f"If you want to use hook function, please use context.set_context to set "
-                               f"pynative mode and remove 'jit' decorator.")
         # Chose dynamic shape tensors or actual input tensors as compile args.
         compile_args = self._generate_compile_args(args)
         key_id = self._get_key_id()
-        compile_args = get_auto_dynamic_shape_args_with_check_input_signature(compile_args, key_id,
-                                                                              self.input_signature,
-                                                                              self._enable_auto_dynamic)
+        if self.input_signature is None:
+            compile_args = get_auto_dynamic_shape_args(
+                compile_args, key_id, self._enable_auto_dynamic, self.enable_jit_dynamic
+            )
 
         # Add mutable for compile_args for two scene:
         # 1) Origin args is mutable.
@@ -674,7 +719,7 @@ class _JitExecutor:
                     f'`{self.fn.__module__}`')
             self.obj.__parse_method__ = method_name
             if isinstance(self.obj, ms.nn.Cell):
-                generate_name = generate_name + '.' + str(self.obj.create_time)
+                generate_name = generate_name + '.' + str(self.obj.create_time) + self.obj.phase
                 create_time = str(self.obj.create_time)
             else:
                 generate_name = generate_name + '.' + str(self._create_time)
@@ -694,18 +739,23 @@ class _JitExecutor:
 
         self._graph_executor.set_enable_tuple_broaden(self.enable_tuple_broaden)
         key = self._graph_executor.generate_arguments_key(self.fn, compile_args, kwargs, self.enable_tuple_broaden)
+        key = str(key)
 
         parameter_ids = _get_parameter_ids(args, kwargs)
         if parameter_ids != "":
-            key = str(key) + '.' + parameter_ids
+            key += '.' + parameter_ids
+
+        key += "." + _get_hook_key(*args, **kwargs)
+        key += "." + str(_hook_version())
 
-        key = str(key) + "." + _get_hook_key(*args, **kwargs)
+        phase = generate_name + '.' + key
 
-        phase = generate_name + '.' + str(key)
+        if self.input_signature is None:
+            update_auto_dynamic_shape_phase(compile_args, key_id, phase)
 
-        update_auto_dynamic_shape_phase_with_check_input_signature(compile_args, key_id, phase, self.input_signature)
+        phase = phase + self._cell_cache_key_extend
 
-        if phase in ms_compile_cache and self._graph_executor.has_compiled(phase) and not parameter_hook_updated():
+        if phase in ms_compile_cache and self._graph_executor.has_compiled(phase):
             # Release resource should be released when CompileInner won't be executed, such as cur_convert_input_
             # generated in generate_arguments_key.
             self._graph_executor.clear_compile_arguments_resource()
@@ -716,31 +766,26 @@ class _JitExecutor:
         # If enable compile cache, get the dependency files list and set to graph executor.
         self._set_compile_cache_dep_files()
         if self.jit_config_dict:
-            self._graph_executor.set_jit_config(self.jit_config_dict)
+            jit_config_dict = self.jit_config_dict
         else:
             jit_config_dict = JitConfig().jit_config_dict
-            self._graph_executor.set_jit_config(jit_config_dict)
 
         if self.obj is None:
             # Set an attribute to fn as an identifier.
-            if isinstance(self.fn, types.MethodType):
-                setattr(self.fn.__func__, "__jit_function__", True)
-            else:
-                setattr(self.fn, "__jit_function__", True)
-            is_compile = self._graph_executor.compile(self.fn, compile_args, kwargs, phase)
-            if isinstance(self.fn, types.MethodType):
-                delattr(self.fn.__func__, "__jit_function__")
-            else:
-                delattr(self.fn, "__jit_function__")
+            setattr(get_func(self.fn), "__jit_function__", True)
+            is_compile = self._graph_executor.compile(self.fn, compile_args, kwargs, phase, jit_config_dict)
+            delattr(get_func(self.fn), "__jit_function__")
         else:
             if isinstance(self.obj, ms.nn.Cell):
                 self._graph_executor.set_weights_values(self.obj.parameters_dict())
-            is_compile = self._graph_executor.compile(self.obj, compile_args, kwargs, phase)
+            is_compile = self._graph_executor.compile(
+                self.obj, compile_args, kwargs, phase, jit_config_dict)
 
         if not is_compile:
             raise RuntimeError("Executor compile failed.")
-        set_parameter_hook_updated(False)
         ms_compile_cache.add(phase)
+        if hasattr(self.obj, "phase"):
+            self.obj.phase_cache[self.obj.phase] = phase
 
         return phase
 
@@ -785,41 +830,70 @@ class _JitExecutor:
         if enable_compile_cache is True or enable_compile_cache == "1":
             self._graph_executor.set_compile_cache_dep_files(_get_compile_cache_dep_files())
 
+    def _generate_compile_args_by_enable_dynamic(self, args_list):
+        """Generate compile args by enable_dynamic."""
+        compile_args = generate_dynamic_tensor_args(args_list, self.dynamic_args_shapes)
+        compile_args = _add_mutable_attr(args_list, compile_args, _pynative_executor.requires_grad())
+        if self.obj is not None:
+            _pynative_executor.set_dynamic_input(self.obj, *compile_args)
+        else:
+            _pynative_executor.set_dynamic_input(self.fn, *compile_args)
+        logger.info(f"dynamic shape compile_args: {compile_args}")
+        return compile_args
+
+    def _generate_compile_args_by_set_inputs(self, args_list):
+        """Generate compile args by set_inputs."""
+        compile_args = _generate_dyn_compile_args(args_list, self.obj.get_inputs())
+        if len(compile_args) != len(args_list):
+            raise ValueError(f"The number of actual input tensors: {len(args_list)} is not equal to the number of "
+                             f"dynamic shape tensors: {len(compile_args)}.")
+        self._graph_executor.check_argument_consistency(compile_args, args_list, "set_inputs")
+        Validator.check_symbolic_shape(compile_args, args_list)
+        return compile_args
+
+    def _generate_compile_args_by_input_signature(self, args_list):
+        """Generate compile args by input_signature."""
+        compile_args = list(_generate_dyn_compile_args(args_list, self.input_signature))
+        dyn_shape = any([is_shape_unknown(elem.shape) for elem in compile_args if isinstance(elem, PythonTensor)])
+        Validator.check_symbolic_shape(self.input_signature, args_list)
+        if dyn_shape:
+            # Checkout whether the `sens` has been added to args_list.
+            if len(compile_args) == len(args_list) - 1:
+                logger.warning(f"The number of actual input args '{len(args_list)}' is one more than the number "
+                               f"of input_signature args '{len(compile_args)}'. The last actual args may "
+                               f"be 'sens' and added it to compile args.")
+                compile_args.append(args_list[-1])
+            compile_args = tuple(compile_args)
+            self._graph_executor.check_argument_consistency(compile_args, args_list, "input_signature")
+            if self.obj is not None:
+                _pynative_executor.set_dynamic_input(self.obj, *compile_args)
+            else:
+                _pynative_executor.set_dynamic_input(self.fn, *compile_args)
+        else:
+            if not verify_inputs_signature(compile_args, args_list):
+                raise ValueError("The input args is incompatible with the args in `input_signature`!")
+        return compile_args
+
+    def _check_set_inputs(self):
+        """Check if the `set_inputs()` of Cell object has been set."""
+        return self.fn.__name__ == 'construct' and isinstance(self.obj, ms.nn.Cell) and self.obj.get_inputs()
+
     def _generate_compile_args(self, args_list):
         """Chose dynamic shape tensors or actual input tensors as compile args."""
-        # Case: If the shape of input args is dynamic, get dynamic shape tensor from context and use it to compile.
-        compile_args = _pynative_executor.get_dynamic_input(args_list)
+        # Case: The `enable_dynamic` is provided and `set_inputs()` of Cell object has been set.
+        if self.enable_jit_dynamic and self._check_set_inputs():
+            raise ValueError("When `enable_dynamic` is provided, the `set_inputs()` cannot be set!")
+        # Case: The `enable_dynamic` is provided.
+        if self.enable_jit_dynamic:
+            return self._generate_compile_args_by_enable_dynamic(args_list)
         # Case: The `set_inputs()` of Cell object has been set, using these dynamic shape args as compile args.
-        if self.fn.__name__ == 'construct' and isinstance(self.obj, ms.nn.Cell) and self.obj.get_inputs():
-            compile_args = _generate_dyn_compile_args(args_list, self.obj.get_inputs())
-            if len(compile_args) != len(args_list):
-                raise ValueError(f"The number of actual input tensors: {len(args_list)} is not equal to the number of "
-                                 f"dynamic shape tensors: {len(compile_args)}.")
-            self._graph_executor.check_argument_consistency(compile_args, args_list, "input_signature")
-            Validator.check_symbolic_shape(compile_args, args_list)
-
+        if self._check_set_inputs():
+            return self._generate_compile_args_by_set_inputs(args_list)
         # Case: If dynamic shape tensors have been assigned to `input_signature`, they are preferred as compile args.
         if self.input_signature is not None:
-            compile_args = list(_generate_dyn_compile_args(args_list, self.input_signature))
-            dyn_shape = any([is_shape_unknown(elem.shape) for elem in compile_args if isinstance(elem, PythonTensor)])
-            Validator.check_symbolic_shape(self.input_signature, args_list)
-            if dyn_shape:
-                # Checkout whether the `sens` has been added to args_list.
-                if len(compile_args) == len(args_list) - 1:
-                    logger.warning(f"The number of actual input args '{len(args_list)}' is one more than the number "
-                                   f"of input_signature args '{len(compile_args)}'. The last actual args may "
-                                   f"be 'sens' and added it to compile args.")
-                    compile_args.append(args_list[-1])
-                compile_args = tuple(compile_args)
-                self._graph_executor.check_argument_consistency(compile_args, args_list, "input_signature")
-                if self.obj is not None:
-                    _pynative_executor.set_dynamic_input(self.obj, *compile_args)
-                else:
-                    _pynative_executor.set_dynamic_input(self.fn, *compile_args)
-            else:
-                if not verify_inputs_signature(compile_args, args_list):
-                    raise ValueError("The input args is incompatible with the args in `input_signature`!")
-        return compile_args
+            return self._generate_compile_args_by_input_signature(args_list)
+        # Case: If the shape of input args is dynamic, get dynamic shape tensor from context and use it to compile.
+        return _pynative_executor.get_dynamic_input(args_list)
 
     def _generate_run_args(self, args_list, kwargs):
         """
@@ -832,7 +906,7 @@ class _JitExecutor:
         Returns:
             new_inputs, new input args, which are required for running.
         """
-        return _get_args_for_run(self, args_list, kwargs, self._compile_args)
+        return _get_args_for_run(self, args_list, kwargs, _get_mutable_flags(self._compile_args), False)
 
     def _get_func_graph_proto(self, obj, exec_id, ir_type="onnx_ir", use_prefix=False, incremental=False):
         """Get graph proto from pipeline."""
@@ -993,6 +1067,67 @@ def _check_options(options, backend):
         _check_option_value(option, value)
 
 
+def _jit_ast(hash_obj, dynamic, jit_config, jit_graph_name):
+    """Return the wrapped function for ast mode jit."""
+    def wrap_func(func):
+        nonlocal hash_obj
+        if hasattr(func, "construct"):
+            if isinstance(func, ms.nn.Cell):
+                # Bound the cell object to get the self arg.
+                return types.MethodType(_jit_ast(
+                    hash_obj, dynamic, jit_config, func._jit_graph_name)(func.construct.__func__), func)
+            if isinstance(func, type) and issubclass(func, ms.nn.Cell):
+                func.construct = _jit_ast(
+                    hash_obj, dynamic, jit_config, '')(func.construct)
+            return func
+
+        if isinstance(func, types.MethodType):
+            return types.MethodType(_jit_ast(hash_obj, dynamic, jit_config, '')(func.__func__), func.__self__)
+
+        if not isinstance(func, types.FunctionType):
+            logger.warning(f"The func should be function, method or cell instance/class, but got {func}")
+            return func
+
+        if hasattr(func, "__wrapped_by_jit__"):
+            logger.warning(f"The func {func} should be wrapped by jit only once.")
+
+        if hash_obj is None or not _is_inner_func(func):
+            hash_obj = int(time.time() * 1e9)
+
+        @wraps(func)
+        def staging_specialize(*args, **kwargs):
+            if os.getenv("MS_JIT") == '0':
+                return func(*args, **kwargs)
+
+            args, kwargs = _handle_func_args(func, *args, **kwargs)
+            process_obj = None
+            if args and not isinstance(args[0], PythonTensor) and hasattr(args[0], func.__name__):
+                process_obj = args[0]
+            # Handle auto mixed precision strategy.
+            if not hasattr(func, "amp_strategy"):
+                setattr(get_func(func), "amp_strategy", get_curr_amp_strategy())
+
+            jit_graph_name = ''
+            if hasattr(staging_specialize, "__jit_graph_name__"):
+                jit_graph_name = staging_specialize.__jit_graph_name__
+            jit_executor = _JitExecutor(
+                func, hash_obj, None, process_obj, jit_config, dynamic, jit_graph_name)
+            out = jit_executor(*args, **kwargs)
+            if isinstance(process_obj, ms.nn.Cell):
+                _clear_auto_parallel_context(process_obj)
+            return out
+
+        # `inspect.getfullargspec(func)` will get the specification of the decorated function by default. By set
+        # `__signature__` for the decorated function, `inspect.getfullargspec(func)` will get the specification of
+        # original `func`.
+        staging_specialize.__signature__ = inspect.signature(func)
+        setattr(staging_specialize, "__wrapped_by_jit__", True)
+        setattr(staging_specialize, "__jit_graph_name__", jit_graph_name)
+        return staging_specialize
+
+    return wrap_func
+
+
 def jit(
         function: Optional[Callable] = None,
         *,
@@ -1015,45 +1150,45 @@ def jit(
           and the decoration @jit(capture_mode=“bytecode”) is considered invalid.
 
     Args:
-        function (Function, optional): The Python function that will be run as a graph. Default: ``None``.
+        function (Callable, optional): The Python function or Cell that will be run as a graph. Default: ``None``.
 
     Keyword Args:
         capture_mode (str, optional): The method to create a callable MindSpore graph. The value of capture_mode
-            should be ``ast`` , ``bytecode`` or ``trace`` . Default: ``ast`` .
+            should be ``"ast"`` , ``"bytecode"`` or ``"trace"`` . Default: ``"ast"`` .
 
-            - `ast <https://www.mindspore.cn/tutorials/en/master/compile/static_graph.html>`_ :
-              Parse Python ast to build graph.
-            - `bytecode` :
-              Parse Python bytecode to build graph at runtime. This is an experimental prototype that is subject to
-              change and/or deletion.
-            - `trace` : Trace the execution of Python code to build graph. This is an experimental prototype that is
-              subject to change and/or deletion.
+            - ast: Parse Python ast to build graph.
+            - bytecode: Parse Python bytecode to build graph at runtime. This is an experimental prototype
+              that is subject to change and/or deletion.
+            - trace: Trace the execution of Python code to build graph. This is an experimental prototype
+              that is subject to change and/or deletion.
 
         jit_level (str, optional): Used to control the compilation optimization level. Currently is only effective
-            with default backend. The value of jit_level should be ``O0`` or ``O1`` . Default: ``O0`` .
+            with ms_backend. The value of jit_level should be ``"O0"`` or ``"O1"`` . Default: ``"O0"`` .
 
-            - `O0`: Except for optimizations that may affect functionality, all other optimizations are turned off.
-            - `O1`: Using commonly used optimizations and automatic operator fusion optimizations. This optimization
+            - O0: Except for optimizations that may affect functionality, all other optimizations are turned off.
+            - O1: Using commonly used optimizations and automatic operator fusion optimizations. This optimization
               level is experimental and is being improved.
 
         dynamic (int, optional): Whether dynamic shape compilation should be performed. Default: ``0``. The value range
             is as follows:
 
-            - `0`: Do not perform dynamic shape compilation.
-            - `1`: Enable dynamic shape compilation and automatically detect shape changes.
+            - 0: Do not perform dynamic shape compilation.
+            - 1: Enable dynamic shape compilation and automatically detect shape changes.
 
         fullgraph (bool, optional): Whether to capture the entire function into graph. If False, jit attempts to
             be compatible with all Python syntax in the function as much as possible. If True, we require that the
             entire function can be captured into graph. If this is not possible (that is, if there is Python syntax
-            not supported), then it will raise an exception. This currently only applies when capture_mode is ast.
-            Default: ``False``.
+            not supported), then it will raise an exception. This currently only applies when capture_mode is ``ast``
+            or ``bytecode``. Default: ``False``.
         backend (str, optional): The compilation backend to be used. If this parameter is not set, the framework will
-            use ``GE`` backend for Atlas training series products and ``ms_backend`` backend for others including Atlas
-            A2 training series products by default.
+            use ``"GE"`` backend for Atlas training series products and ``"ms_backend"`` backend for others including
+            Atlas A2 training series products by default.
 
-            - `ms_backend`: Adopt KernelByKernel execution mode.
-            - `GE`: Adopt Sink execution mode. The whole model will be sinked to device to execute, only applicable to
-              the top cell of model. And only can be used in Ascend platform.
+            - ms_backend: Utilizes the built-in backend engine of MindSpore for hardware-related compilation
+              optimization and execution, supporting multiple hardware forms such as Ascend, GPU, and CPU.
+            - GE: Utilizes the GraphEngine, a graph compilation and execution engine within CANN,
+              for Ascend model compilation and execution. Note: This backend takes effect only in static graph mode
+              and can be executed only on Ascend hardware.
 
         **options (dict): A dictionary of options to pass to the compilation backend.
 
@@ -1076,11 +1211,11 @@ def jit(
               `disable_format_transform` can be set to ``True`` to try to improve training performance.
               Default: ``False`` .
             - exec_order (str, optional): Set the sorting method for operator execution, currently only two sorting
-              methods are supported: ``bfs`` and ``dfs`` . Default: ``bfs`` .
+              methods are supported: ``"bfs"`` and ``"dfs"`` . Default: ``"bfs"`` .
 
-              - `bfs`: The default sorting method, breadth priority, good communication masking, relatively good
+              - bfs: The default sorting method, breadth priority, good communication masking, relatively good
                 performance.
-              - `dfs`: An optional sorting method, depth-first sorting. The performance is relatively worse than that
+              - dfs: An optional sorting method, depth-first sorting. The performance is relatively worse than that
                 of bfs execution order, but it occupies less memory. It is recommended to try dfs in scenarios where
                 other execution orders run out of memory (OOM).
 
@@ -1091,11 +1226,11 @@ def jit(
               - global (dict): Set global options.
               - session (dict): Set session options.
 
-            - infer_boost (str, optional): Used to control the inference mode. Default: ``off``, which means
+            - infer_boost (str, optional): Used to control the inference mode. Default: ``"off"``, which means
               the inference mode is disabled. The range is as follows:
 
-              - `on`: Enable inference mode, get better infer performance.
-              - `off`: Disable inference mode, use forward for inference. The performance is poor.
+              - on: Enable inference mode, get better infer performance.
+              - off: Disable inference mode, use forward for inference. The performance is poor.
 
     Returns:
         Function, if `fn` is not None, returns a callable function that will execute the compiled function; If `fn` is
@@ -1114,29 +1249,84 @@ def jit(
         >>> x = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
         >>> y = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
         ...
-        >>> # create a callable MindSpore graph by calling jit
+        >>> # Create a callable MindSpore graph by calling jit.
         >>> def tensor_add(x, y):
         ...     z = x + y
         ...     return z
         ...
         >>> tensor_add_graph = jit(function=tensor_add)
         >>> out = tensor_add_graph(x, y)
+        >>> print(out)
+        Tensor(shape=[1, 1, 3, 3], dtype=Float32, value=
+        [[[[ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00]]]])
         ...
-        >>> # create a callable MindSpore graph through decorator @jit
+        >>> # Create a callable MindSpore graph through decorator @jit.
         >>> @jit
         ... def tensor_add_with_dec(x, y):
         ...     z = x + y
         ...     return z
         ...
         >>> out = tensor_add_with_dec(x, y)
+        >>> print(out)
+        Tensor(shape=[1, 1, 3, 3], dtype=Float32, value=
+        [[[[ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00]]]])
         ...
-        >>> # create a callable MindSpore graph and capture the entire function into the graph
+        >>> # Create a callable MindSpore graph and capture the entire function into the graph.
         >>> @jit(fullgraph=True)
         ... def tensor_add_fullgraph(x, y):
         ...     z = x + y
         ...     return z
         ...
         >>> out = tensor_add_fullgraph(x, y)
+        >>> print(out)
+        Tensor(shape=[1, 1, 3, 3], dtype=Float32, value=
+        [[[[ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00]]]])
+        ...
+        >>> # Create a callable MindSpore graph by trace mode.
+        >>> @jit(capture_mode="trace")
+        ... def tensor_add_by_trace(x, y):
+        ...     z = x + y
+        ...     return z
+        ...
+        >>> out = tensor_add_by_trace(x, y)
+        >>> print(out)
+        Tensor(shape=[1, 1, 3, 3], dtype=Float32, value=
+        [[[[ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00]]]])
+        ...
+        >>> # Create a callable MindSpore graph with ms_backend and jit_level="O1".
+        >>> @jit(backend="ms_backend", jit_level="O1")
+        ... def tensor_add_by_trace(x, y):
+        ...     z = x + y
+        ...     return z
+        ...
+        >>> out = tensor_add_by_trace(x, y)
+        >>> print(out)
+        Tensor(shape=[1, 1, 3, 3], dtype=Float32, value=
+        [[[[ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00]]]])
+        ...
+        >>> # Create a callable MindSpore graph with GE backend and some ge options on Ascend.
+        >>> @jit(backend="GE", ge_options={"global": {"ge.opSelectImplmode": "high_precision"}})
+        ... def tensor_add_by_trace(x, y):
+        ...     z = x + y
+        ...     return z
+        ...
+        >>> out = tensor_add_by_trace(x, y)
+        >>> print(out)
+        Tensor(shape=[1, 1, 3, 3], dtype=Float32, value=
+        [[[[ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00],
+           [ 2.00000000e+00,  2.00000000e+00,  2.00000000e+00]]]])
+        ...
     """
 
     capture_mode = Validator.check_string(capture_mode, ["ast", "bytecode", "trace"], "capture_mode", "jit")
@@ -1155,39 +1345,12 @@ def jit(
     jit_config = JitConfig(jit_level=jit_level, exc_mode=exc_mode, jit_syntax_level=jit_syntax_level,
                            infer_boost=infer_boost, backend=backend, options=options_str)
 
-    def wrap_func(func):
-        nonlocal hash_obj
-        if hash_obj is None or not _is_inner_func(func):
-            hash_obj = int(time.time() * 1e9)
-
-        @wraps(func)
-        def staging_specialize(*args, **kwargs):
-            if os.getenv("MS_JIT") == '0':
-                return func(*args, **kwargs)
-
-            args, kwargs = _handle_func_args(func, *args, **kwargs)
-            process_obj = None
-            if args and not isinstance(args[0], PythonTensor) and hasattr(args[0], func.__name__):
-                process_obj = args[0]
-            # Handle auto mixed precision strategy.
-            if not hasattr(func, "amp_strategy"):
-                if isinstance(func, types.MethodType):
-                    setattr(func.__func__, "amp_strategy", get_curr_amp_strategy())
-                else:
-                    setattr(func, "amp_strategy", get_curr_amp_strategy())
-
-            ms_function_executor = _JitExecutor(func, hash_obj, None, process_obj, jit_config, dynamic)
-            out = ms_function_executor(*args, **kwargs)
-            return out
-
-        return staging_specialize
-
-    if capture_mode == "bytecode":
-        wrap_func = PIJitCaptureContext(jit_config)
-    elif capture_mode == "trace":
-        if function is not None:
-            return _jit_trace(function)
-        return _jit_trace
+    if capture_mode == "ast":
+        wrap_func = _jit_ast(hash_obj, dynamic, jit_config, '')
+    elif capture_mode == "bytecode":
+        wrap_func = PIJitCaptureContext(fullgraph=fullgraph, jit_config=jit_config)
+    else:
+        wrap_func = _jit_trace()
 
     if function is not None:
         return wrap_func(function)
@@ -1503,7 +1666,7 @@ class _PyNativeExecutor:
         """
         self._executor.end_graph(obj, output, *args, *(kwargs.values()))
 
-    def check_run(self, grad, obj, weights, grad_hash_id, *args):
+    def check_run(self, grad, obj, weights, grad_hash_id, *args, **kwargs):
         """
         Whether the forward graph need to construct.
 
@@ -1516,7 +1679,7 @@ class _PyNativeExecutor:
         Return:
             bool, specifies whether the forward graph needs to construct.
         """
-        return self._executor.check_run(grad, obj, weights, grad_hash_id, *args)
+        return self._executor.check_run(grad, obj, weights, grad_hash_id, *args, **kwargs)
 
     def grad(self, obj, grad, weights, grad_position, *args):
         """
@@ -1758,6 +1921,19 @@ class _PyNativeExecutor:
         """
         return self._executor.constant_folding(*args)
 
+    def set_creation_type(self, tensor, creation_type):
+        """
+        Set tensor's view creation type
+
+        Args:
+            tensor (Tensor): input tensor.
+            creation_type (CreationType): The type of view tensor when it is created.
+
+        Return:
+            None.
+        """
+        return self._executor.set_creation_type(tensor, creation_type)
+
 
 class _CellGraphExecutor:
     """
@@ -1834,13 +2010,6 @@ class _CellGraphExecutor:
         else:
             _set_dataset_mode_config('normal')
 
-    @staticmethod
-    def _use_vm_mode():
-        enable_ge = context.get_context("enable_ge")
-        enable_debug_runtime = context.get_context("enable_debug_runtime")
-        exe_mode = context.get_context("mode") == context.PYNATIVE_MODE
-        return not enable_ge or (enable_debug_runtime and exe_mode)
-
     def _build_data_graph(self, obj, phase):
         self._graph_executor.build_data_graph(obj.parameters_dict(), phase)
 
@@ -1872,7 +2041,12 @@ class _CellGraphExecutor:
         obj.__parse_method__ = 'construct'
         if not hasattr(obj, obj.__parse_method__):
             raise AttributeError(
-                'The class {} dose not have method {}'.format(obj.__class__.__name__, obj.__parse_method__))
+                'The class {} does not have method {}'.format(obj.__class__.__name__, obj.__parse_method__))
+        inner_func = inspect.unwrap(obj.construct)
+        if hasattr(get_func(inner_func), ENABLE_DYNAMIC):
+            raise ValueError(
+                "When using set_context(mode=GRAPH_MODE) together with nn.Cell, the 'enable_dynamic' cannot be set!"
+            )
         key_id = str(id(obj)) + str(obj.create_time)
         args = get_auto_dynamic_shape_args(args, key_id)
 
@@ -1883,20 +2057,25 @@ class _CellGraphExecutor:
         self._graph_executor.set_enable_tuple_broaden(self.enable_tuple_broaden)
 
         key = self._graph_executor.generate_arguments_key(obj, args, kwargs, self.enable_tuple_broaden)
-        obj.arguments_key = str(key)
-
-        obj.arguments_key = obj.arguments_key + "." + _get_hook_key(*args, **kwargs)
+        key = str(key)
 
         # When exist parameter in the top graph inputs, need check if the parameter object has changed.
         parameter_ids = _get_parameter_ids(args, kwargs)
         if parameter_ids != "":
-            obj.arguments_key = obj.arguments_key + '.' + parameter_ids
+            key += '.' + parameter_ids
+
+        key += "." + _get_hook_key(*args, **kwargs)
+        key += "." + str(_hook_version())
+
+        obj.arguments_key = key
+
         raw_phase = phase
-        phase = phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+
+        phase = _real_phase(phase, obj)
         obj.phase_cache[raw_phase] = phase
         update_auto_dynamic_shape_phase(args, key_id, phase)
         obj.current_phase = phase
-        if phase in obj.compile_cache and self.has_compiled(phase) and not parameter_hook_updated():
+        if phase in obj.compile_cache and self.has_compiled(phase):
             logger.debug("%r graph has existed.", phase)
             # Release resource should be released when CompileInner won't be executed, such as cur_convert_input_
             # generated in generate_arguments_key.
@@ -1904,7 +2083,7 @@ class _CellGraphExecutor:
             _clear_auto_parallel_context(obj)
             return phase, False
 
-        full_function_name = obj.__class__.__name__ + '.' + str(obj.instance_count) + '.' + str(id(type(obj)))
+        full_function_name = obj.__class__.__name__ + '.' + str(obj.total_instance_count) + '.' + str(id(type(obj)))
         echo_function_name = obj.__class__.__name__
         _check_recompile(obj, args, kwargs, full_function_name, obj.create_time, echo_function_name)
 
@@ -1914,17 +2093,14 @@ class _CellGraphExecutor:
         self._set_compile_cache_dep_files(phase)
 
         self._graph_executor.set_weights_values(obj.parameters_dict())
-        if jit_config_dict:
-            self._graph_executor.set_jit_config(jit_config_dict)
-        else:
+        if not jit_config_dict:
             jit_config_dict = JitConfig().jit_config_dict
-            self._graph_executor.set_jit_config(jit_config_dict)
         gc.collect()
-        result = self._graph_executor.compile(obj, args, kwargs, phase)
+        result = self._graph_executor.compile(
+            obj, args, kwargs, phase, jit_config_dict)
         obj.compile_cache.add(phase)
         if not result:
             raise RuntimeError("Executor compile failed.")
-        set_parameter_hook_updated(False)
         graph = self._graph_executor.get_func_graph(phase)
 
         if graph is None:
@@ -1949,15 +2125,15 @@ class _CellGraphExecutor:
         return self._graph_executor.updata_param_node_default_input(phase, new_param)
 
     def _get_shard_strategy(self, obj):
-        real_phase = obj.phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+        real_phase = _real_phase(obj.phase, obj)
         return self._graph_executor.get_strategy(real_phase)
 
     def _get_num_parallel_ops(self, obj):
-        real_phase = obj.phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+        real_phase = _real_phase(obj.phase, obj)
         return self._graph_executor.get_num_parallel_ops(real_phase)
 
     def _get_allreduce_fusion(self, obj):
-        real_phase = obj.phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+        real_phase = _real_phase(obj.phase, obj)
         return self._graph_executor.get_allreduce_fusion(real_phase)
 
     def __call__(self, obj, *args, phase='predict'):
@@ -2009,10 +2185,10 @@ class _CellGraphExecutor:
             Tensor/Tuple, return execute result.
         """
         if phase == 'save':
-            exe_phase = phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+            exe_phase = _real_phase(phase, obj)
             return self._graph_executor((), exe_phase)
 
-        phase_real = phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+        phase_real = _real_phase(phase, obj)
         if self.has_compiled(phase_real):
             return self._exec_pip(obj, *args, phase=phase_real)
         raise KeyError('{} graph is not exist.'.format(phase_real))
@@ -2039,7 +2215,7 @@ class _CellGraphExecutor:
 
     def get_optimize_graph_proto(self, obj):
         """Return optimize graph binary proto."""
-        exec_id = obj.phase + "." + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+        exec_id = _real_phase(obj.phase, obj)
         if self._graph_executor.has_compiled(exec_id) is False:
             return None
         graph_proto = self._graph_executor.get_optimize_graph_proto(exec_id)
@@ -2121,5 +2297,3 @@ def flops_collection(phase='train'):
 
 _cell_graph_executor = _CellGraphExecutor()
 _pynative_executor = _PyNativeExecutor()
-
-__all__ = ['ms_memory_recycle', 'jit', 'jit_class', 'flops_collection']