mindspore 2.7.0rc1__cp310-cp310-win_amd64.whl → 2.7.1__cp310-cp310-win_amd64.whl

mindspore/common/api.py

@@ -44,18 +44,20 @@ 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, MSContext, TensorPy as Tensor
+    _run_jit_pipeline, _ms_memory_recycle, _bind_device_ctx, TensorPy as Tensor, dump_func_graph, _GraphFragment_
 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
@@ -74,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"""
 
@@ -201,6 +208,11 @@ def _handle_func_args(func, *args, **kwargs):
         args = bound_arguments.args
         kwargs = bound_arguments.kwargs
 
+    return args, kwargs
+
+
+def _check_func_args(func, *args):
+    """Check the *args inputs of the function"""
     positional_args = 0
     default_args = 0
     has_var = False
@@ -214,14 +226,13 @@ def _handle_func_args(func, *args, **kwargs):
                 default_args += 1
 
     if has_var:
-        return args, kwargs
+        return
 
     if len(args) < positional_args:
         raise TypeError(f"Function {func.__name__} needs {positional_args} positional argument, but got {len(args)}.")
     if len(args) > positional_args + default_args:
         raise TypeError(f"Function {func.__name__} needs {positional_args} positional argument and {default_args} "
                         f"default argument, total {positional_args + default_args}, but got {len(args)}.")
-    return args, kwargs
 
 
 sys_path = list(sys.path)
@@ -342,7 +353,7 @@ def _get_parameter_layout():
     return layout
 
 
-def _handle_arg(obj, arg, has_mutable_arg):
+def _handle_arg(obj, arg, has_mutable_arg, is_predict):
     """Handle arg for runtime .If need handle the arg, return True"""
     from mindspore._extends.parse import compile_config
     if isinstance(arg, PythonTensor):
@@ -357,7 +368,7 @@ def _handle_arg(obj, arg, has_mutable_arg):
         if isinstance(arg, list) and not arg:
             return None
         return arg
-    elif (context.get_context("grad_for_scalar") or str(compile_config.GRAD_FOR_SCALAR) == '1') and \
+    elif not is_predict and (context.get_context("grad_for_scalar") or str(compile_config.GRAD_FOR_SCALAR) == '1') and \
             isinstance(arg, (int, float)):
         return arg
     elif hasattr(obj, "enable_tuple_broaden") and obj.enable_tuple_broaden and isinstance(arg, tuple) and \
@@ -387,17 +398,16 @@ def _handle_arg_predict(obj, arg, has_mutable_arg):
     return arg
 
 
-def _get_args_for_run(obj, args, kwargs, has_mutable_args_list, is_predict):
+def _get_args_for_run(obj, args, kwargs, has_mutable_args_list, is_predict=False):
     """Get the actual input args and kwargs for runtime."""
     new_args = []
-    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)
+        new_arg = _handle_arg(obj, arg, has_mutable_arg, is_predict)
         if new_arg is not None:
             new_args.append(new_arg)
 
     for _, value in kwargs.items():
-        new_value = fn(obj, value, None)
+        new_value = _handle_arg(obj, value, None, is_predict)
         if new_value is not None:
             new_args.append(new_value)
 
@@ -538,10 +548,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 = ""
@@ -588,6 +600,8 @@ class _JitExecutor:
 
         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
@@ -598,7 +612,7 @@ class _JitExecutor:
         else:
             self._graph_executor = GraphExecutor_.get_instance()
         self._create_time = ms_create_time
-        self._compile_args = None
+        self._mutable_flags = 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
@@ -623,18 +637,8 @@ class _JitExecutor:
             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()
-            compile_args = get_auto_dynamic_shape_args_with_check_input_signature(
-                compile_args,
-                key_id,
-                self.input_signature,
-                self._enable_auto_dynamic
-            )
-            self._compile_args = compile_args
 
-        new_inputs = self._generate_run_args(args_list, kwargs)
+        new_inputs = self._generate_run_args(args_list, kwargs, is_predict=True)
         if self.jit_config_dict:
             jit_config_dict = self.jit_config_dict
         else:
@@ -647,11 +651,25 @@ class _JitExecutor:
         res = _convert_python_data(output)
         return True, res
 
+    def compile_frontend(self, *args, **kwargs):
+        """Only compile to the frontend graph."""
+        args_list = args
+        if self.obj is not None:
+            args_list = args_list[1:]
+        os.environ['MS_DEV_PRECOMPILE_ONLY'] = '1'
+        phase = ""
+        _pynative_executor.set_jit_compile_phase(phase)
+        phase = self.compile(self.fn.__name__, *args_list, **kwargs)
+        _pynative_executor.set_jit_compile_phase(phase)
+        os.unsetenv('MS_DEV_PRECOMPILE_ONLY')
+        return self._graph_executor.get_func_graph(phase), self._mutable_flags, phase, self.enable_tuple_broaden
+
     @_wrap_func
     def __call__(self, *args, **kwargs):
         predict, res = self._predict(*args, **kwargs)
         if predict:
             return res
+        _check_func_args(self.fn, *args)
         if jit_context() and jit_context().is_nested():
             return jit_context().run_graph("", None, *())
         args_list = args
@@ -659,9 +677,9 @@ class _JitExecutor:
             args_list = args_list[1:]
         phase = ""
         try:
-            _pynative_executor.set_jit_compile_status(True, phase)
+            _pynative_executor.set_jit_compile_phase(phase)
             phase = self.compile(self.fn.__name__, *args_list, **kwargs)
-            _pynative_executor.set_jit_compile_status(False, phase)
+            _pynative_executor.set_jit_compile_phase(phase)
         except Exception as err:
             _pynative_executor.clear_res()
             raise err
@@ -684,24 +702,24 @@ class _JitExecutor:
 
     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.
+        self._graph_executor.set_real_args(args, kwargs)
         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.
         # 2) Args contains sequence with gradient tensor.
         compile_args = _add_mutable_attr(args, compile_args, _pynative_executor.requires_grad())
-        self._compile_args = compile_args
+        mutable_flags = _get_mutable_flags(compile_args)
+        self._mutable_flags = mutable_flags
+        # Store the _mutable_flags in the cell obj for incremental inference.
+        if self.obj is not None:
+            self.obj._mutable_flags = mutable_flags
         generate_name, echo_function_name = self._get_generate_name()
         # The full Function name
         full_function_name = generate_name
@@ -735,20 +753,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 = str(key) + "." + _get_hook_key(*args, **kwargs)
+        key += "." + _get_hook_key(*args, **kwargs)
+        key += "." + str(_hook_version())
 
-        phase = generate_name + '.' + str(key)
+        phase = generate_name + '.' + key
 
-        update_auto_dynamic_shape_phase_with_check_input_signature(compile_args, key_id, phase, self.input_signature)
+        if self.input_signature is None:
+            update_auto_dynamic_shape_phase(compile_args, key_id, phase)
 
         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()
@@ -765,16 +786,9 @@ class _JitExecutor:
 
         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, jit_config_dict)
-            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())
@@ -783,7 +797,6 @@ class _JitExecutor:
 
         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
@@ -831,43 +844,73 @@ 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}")
+        Validator.check_symbolic_shape(compile_args, args_list)
+        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):
+    def _generate_run_args(self, args_list, kwargs, is_predict=False):
         """
         Generate input args, which are required for running.
 
@@ -878,7 +921,11 @@ class _JitExecutor:
         Returns:
             new_inputs, new input args, which are required for running.
         """
-        return _get_args_for_run(self, args_list, kwargs, _get_mutable_flags(self._compile_args), False)
+        if self.obj is not None and hasattr(self.obj, '_mutable_flags'):
+            mutable_flags = self.obj._mutable_flags
+        else:
+            mutable_flags = self._mutable_flags
+        return _get_args_for_run(self, args_list, kwargs, mutable_flags, is_predict)
 
     def _get_func_graph_proto(self, obj, exec_id, ir_type="onnx_ir", use_prefix=False, incremental=False):
         """Get graph proto from pipeline."""
@@ -950,7 +997,7 @@ def _check_option_backend(option, backend):
         'ge_options': ['GE'],
         'infer_boost': ['ms_backend'],
     }
-    if option in option_backend_cfgs and backend not in option_backend_cfgs[option]:
+    if option in option_backend_cfgs and backend != '' and backend not in option_backend_cfgs[option]:
         logger.warning(f"For 'jit(options)', the option '{option}' is only support backend in "
                        f"'{option_backend_cfgs[option]}', but got '{backend}', ignore it.")
 
@@ -1077,10 +1124,7 @@ def _jit_ast(hash_obj, dynamic, jit_config, jit_graph_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())
+                setattr(get_func(func), "amp_strategy", get_curr_amp_strategy())
 
             jit_graph_name = ''
             if hasattr(staging_specialize, "__jit_graph_name__"):
@@ -1088,6 +1132,8 @@ def _jit_ast(hash_obj, dynamic, jit_config, 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
@@ -1127,28 +1173,26 @@ def jit(
 
     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/docs/en/master/features/compile/graph_construction.html#ast>`_ :
-              Parse Python ast to build graph.
-            - `bytecode <https://www.mindspore.cn/docs/en/master/features/compile/graph_construction.html#bytecode>`_ :
-              Parse Python bytecode to build graph at runtime. This is an experimental prototype that is subject to
-              change and/or deletion.
-            - `trace <https://www.mindspore.cn/docs/en/master/features/compile/graph_construction.html#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 ms_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
@@ -1156,12 +1200,16 @@ def jit(
             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 only supports GRAPH Mode in Ascend,
+              only supports whole graph sinking or sub graph sinking in pipeline parallel, and does not support
+              dynamic shape scenes. In addition, this backend incurs additional compilation costs and is difficult to
+              debug and tune.
 
         **options (dict): A dictionary of options to pass to the compilation backend.
 
@@ -1184,11 +1232,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).
 
@@ -1199,11 +1247,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
@@ -1306,9 +1354,8 @@ def jit(
     jit_level = Validator.check_string(jit_level, ["O0", "O1"], "jit_level", "jit")
     dynamic = Validator.check_int_range(dynamic, 0, 1, Validator.INC_BOTH, "dynamic", "jit")
     fullgraph = Validator.check_bool(fullgraph, "fullgraph", "jit")
-    if backend == "":
-        backend = "GE" if MSContext.get_instance().get_ascend_soc_version() == "ascend910" else "ms_backend"
-    backend = Validator.check_string(backend, ["ms_backend", "GE"], "backend", "jit")
+    if backend != "":
+        backend = Validator.check_string(backend, ["ms_backend", "GE"], "backend", "jit")
     jit_syntax_level = "LAX" if fullgraph is False else "STRICT"
     hash_obj = _get_hash_obj(options)
     _check_options(options, backend)
@@ -1323,7 +1370,7 @@ def jit(
     elif capture_mode == "bytecode":
         wrap_func = PIJitCaptureContext(fullgraph=fullgraph, jit_config=jit_config)
     else:
-        wrap_func = _jit_trace()
+        wrap_func = _jit_trace(jit_config)
 
     if function is not None:
         return wrap_func(function)
@@ -1530,6 +1577,20 @@ def _parameter_broadcast(obj):
     _build_broadcast_graph(broadcast_params_dict, broadcast_phase)
 
 
+def _run_in_jit():
+    """In jit, this function always returns true. Otherwise, returns false."""
+    def _temp_func():
+        return 0
+
+    from mindspore.ops.primitive import constexpr
+
+    @constexpr(check=False)
+    def _check_func(func):
+        return func is None
+
+    return _check_func(_temp_func)
+
+
 class _no_grad(contextlib.ContextDecorator):
     """
     Context Manager to disable gradient calculation. When enter this context, we will disable calculate
@@ -1799,17 +1860,16 @@ class _PyNativeExecutor:
         """
         return self._executor.requires_grad()
 
-    def set_jit_compile_status(self, status, phase):
+    def set_jit_compile_phase(self, phase):
         """
-        Set jit is compiling
+        Set jit phase
 
         Args:
-            status(bool): jit compile status
             phase (str): The phase of cell/function instance.
         Return:
             None.
         """
-        self._executor.set_jit_compile_status(status, phase)
+        self._executor.set_jit_compile_phase(phase)
 
     def set_is_run_recompute(self, status):
         """
@@ -1894,6 +1954,32 @@ 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)
+
+    def queue_backward_final_callback(self, callback):
+        """
+        add backward final callback
+
+        Args:
+            callback(Function): callback function.
+
+        Return:
+            None.
+        """
+        return self._executor.queue_backward_final_callback(callback)
+
+
 
 class _CellGraphExecutor:
     """
@@ -2002,6 +2088,11 @@ class _CellGraphExecutor:
         if not hasattr(obj, obj.__parse_method__):
             raise AttributeError(
                 '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)
 
@@ -2012,20 +2103,27 @@ 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():
+        obj._add_attr("compile_phase", phase)
+        obj.compile_phase = phase
+        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.
@@ -2051,7 +2149,6 @@ class _CellGraphExecutor:
         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:
@@ -2075,16 +2172,20 @@ class _CellGraphExecutor:
         new_param = {x.name: replace[x] for x in replace if id(x) != id(replace[x])}
         return self._graph_executor.updata_param_node_default_input(phase, new_param)
 
+    def set_real_args(self, args, kwargs):
+        """Set real arguments to graph executor."""
+        self._graph_executor.set_real_args(args, kwargs)
+
     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'):
@@ -2136,10 +2237,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))
@@ -2164,9 +2265,22 @@ class _CellGraphExecutor:
             return None
         return self._graph_executor.get_func_graph_proto(exec_id, ir_type, incremental)
 
+    def _get_onnx_func_graph_proto(self, obj, exec_id, use_prefix=False, input_names=None, output_names=None,
+                                   opset_version=11, export_params=True, keep_initializers_as_inputs=False,
+                                   dynamic_axes=None, extra_save_params=False, save_file_dir=None):
+        """Get graph proto from pipeline."""
+        if use_prefix:
+            exec_id = exec_id + '.' + obj.arguments_key
+        if self._graph_executor.has_compiled(exec_id) is False:
+            return None
+
+        return self._graph_executor.get_onnx_func_graph_proto(exec_id, input_names, output_names, opset_version,
+                                                              export_params, keep_initializers_as_inputs, dynamic_axes,
+                                                              extra_save_params, save_file_dir)
+
     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)
@@ -2246,5 +2360,190 @@ def flops_collection(phase='train'):
     return _cell_graph_executor.flops_collection(phase)
 
 
+class _ScriptGraph:
+    """Store the graph compiled by the frontend compiler."""
+    def __init__(self, func_graph, func, origin_cell, mutable_flags, phase, enable_tuple_broaden):
+        self.func_graph = func_graph
+        self.func = func
+        self.origin_cell = origin_cell
+        self.mutable_flags = mutable_flags
+        self.phase = phase
+        self.enable_tuple_broaden = enable_tuple_broaden
+
+    def print(self):
+        """Print the MindIR of the frontend graph."""
+        graph_str = dump_func_graph(self.func_graph)
+        print(graph_str, flush=True)
+
+
+def _frontend_compile_ast(dynamic, jit_config, jit_graph_name=''):
+    """Return the wrapped function for ast mode jit."""
+    def wrap_func(func):
+        if hasattr(func, "construct") and isinstance(func, ms.nn.Cell):
+            # Bound the cell object to get the self arg.
+            return types.MethodType(_frontend_compile_ast(dynamic, jit_config,
+                                                          func._jit_graph_name)(func.construct.__func__), func)
+
+        if isinstance(func, types.MethodType):
+            return types.MethodType(_frontend_compile_ast(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
+
+        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)
+            func_graph, mutable_flags, phase, enable_tuple_broaden = jit_executor.compile_frontend(*args, **kwargs)
+            return _ScriptGraph(func_graph, func, process_obj, mutable_flags, phase, enable_tuple_broaden)
+
+        # `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, "__jit_graph_name__", jit_graph_name)
+        return staging_specialize
+
+    return wrap_func
+
+
+def _frontend_compile(function: Callable,
+                      *,
+                      dynamic: int = 0,
+                      fullgraph: bool = False):
+    """
+    Create a frontend MindSpore graph from a Python function by the ast capture mode.
+
+    Args:
+        function (Callable, optional): The Python function or Cell instance that will be compiled as a frontend graph.
+            Default: ``None``.
+
+    Keyword Args:
+        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.
+
+        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``
+            or ``bytecode``. Default: ``False``.
+
+    Returns:
+        a :class:`_ScriptGraph` object.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import Tensor
+        >>> from mindspore import ops
+        >>> from mindspore.common.api import _frontend_compile
+        ...
+        >>> x = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
+        >>> y = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
+        ...
+        >>> def tensor_add(x, y):
+        ...     z = x + y
+        ...     return z
+        ...
+        >>> tensor_add_graph = _frontend_compile(tensor_add)(x, y)
+        >>> tensor_add_graph.print()
+        ...
+    """
+
+    dynamic = Validator.check_int_range(dynamic, 0, 1, Validator.INC_BOTH, "dynamic", "jit")
+    fullgraph = Validator.check_bool(fullgraph, "fullgraph", "jit")
+    jit_syntax_level = "LAX" if fullgraph is False else "STRICT"
+    jit_config = JitConfig(jit_syntax_level=jit_syntax_level)
+    return _frontend_compile_ast(dynamic, jit_config)(function)
+
+
+class _GraphFragment(_GraphFragment_):
+    """
+    Represents the output by backend graph split.
+    """
+    def __init__(self, frag):
+        if frag is None or not isinstance(frag, _GraphFragment_):
+            raise TypeError(f"Expect input `frag` to be a _GraphFragment_, but got {type(frag)}")
+        _GraphFragment_.__init__(self, frag)
+
+    def __call__(self, *args):
+        return super().__call__(args)
+
+    def __repr__(self):
+        return self.__str__()
+
+    def id(self):
+        return self.id_()
+
+    def is_graph(self):
+        return self.is_graph_()
+
+    def py_key(self):
+        return self.py_key_()
+
+    def args_list(self):
+        return self.args_list_()
+
+
+def _graph_split(script_graph):
+    """
+    Split the script_graph into several fragments according to the nodes with the split op attribute.
+
+    Args:
+        a :class:`_ScriptGraph` object.
+
+    Returns:
+        several :class:`_GraphFragment` object.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import Tensor
+        >>> from mindspore import ops
+        >>> from mindspore.common.api import _frontend_compile, _graph_split
+        ...
+        >>> x = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
+        >>> y = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))
+        >>> add = ops.Add().add_prim_attr("split_op", True).add_prim_attr("func_id", "add_func")
+        ...
+        >>> def tensor_add(x, y):
+        ...     z1 = x + y
+        ...     z2 = add(z1, x)
+        ...     return z2
+        ...
+        >>> tensor_add_graph = _frontend_compile(tensor_add)(x, y)
+        >>> frags = _graph_split(tensor_add_graph)
+        >>> print(frags)
+        ...
+    """
+    outputs = JitExecutor_.get_instance().split_graph(script_graph.func_graph)
+    fragments = []
+    for arg in outputs:
+        fragments.append(_GraphFragment(arg))
+    return fragments
+
 _cell_graph_executor = _CellGraphExecutor()
 _pynative_executor = _PyNativeExecutor()