mindspore 2.6.0rc1__cp39-cp39-win_amd64.whl → 2.7.0rc1__cp39-cp39-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,7 +44,7 @@ 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
@@ -58,7 +60,7 @@ 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 = {}
@@ -134,8 +136,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 +278,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 +342,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 +352,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 +366,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 +375,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 +387,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
 
 
@@ -586,7 +580,8 @@ 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))
@@ -606,9 +601,57 @@ 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()
+            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)
+        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 +659,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,15 +670,16 @@ 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):
@@ -674,7 +715,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)
@@ -705,6 +746,8 @@ class _JitExecutor:
 
         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():
             # Release resource should be released when CompileInner won't be executed, such as cur_convert_input_
             # generated in generate_arguments_key.
@@ -716,10 +759,9 @@ 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.
@@ -727,7 +769,8 @@ class _JitExecutor:
                 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)
+            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:
@@ -735,12 +778,15 @@ class _JitExecutor:
         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
 
@@ -832,7 +878,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 +1039,68 @@ 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"):
+                if isinstance(func, types.MethodType):
+                    setattr(func.__func__, "amp_strategy", get_curr_amp_strategy())
+                else:
+                    setattr(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)
+            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,22 +1123,22 @@ 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`` .
 
-            - `ast <https://www.mindspore.cn/tutorials/en/master/compile/static_graph.html>`_ :
+            - `ast <https://www.mindspore.cn/docs/en/master/features/compile/graph_construction.html#ast>`_ :
               Parse Python ast to build graph.
-            - `bytecode` :
+            - `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` : Trace the execution of Python code to build graph. This is an experimental prototype that is
+            - `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.
 
         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
@@ -1045,8 +1153,8 @@ def jit(
         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.
@@ -1114,29 +1222,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 +1318,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 +1639,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 +1652,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):
         """
@@ -1834,13 +1970,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 +2001,7 @@ 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__))
         key_id = str(id(obj)) + str(obj.create_time)
         args = get_auto_dynamic_shape_args(args, key_id)
 
@@ -1904,7 +2033,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,13 +2043,11 @@ 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.")
@@ -2121,5 +2248,3 @@ def flops_collection(phase='train'):
 
 _cell_graph_executor = _CellGraphExecutor()
 _pynative_executor = _PyNativeExecutor()
-
-__all__ = ['ms_memory_recycle', 'jit', 'jit_class', 'flops_collection']