mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.3.0__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-2023 Huawei Technologies Co., Ltd
+# Copyright 2020-2024 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.
@@ -40,7 +40,7 @@ from mindspore.common.sparse_tensor import COOTensor as PythonCOOTensor
 from mindspore.common.sparse_tensor import RowTensor as PythonRowTensor
 from mindspore._c_expression import GraphExecutor_, Tensor, CSRTensor, RowTensor, COOTensor, \
     PyNativeExecutor_, verify_inputs_signature, init_exec_dataset, _set_dataset_mode_config, init_pipeline, \
-    _ms_memory_recycle, _bind_device_ctx
+    _ms_memory_recycle, _bind_device_ctx, jit_mode_pi_enable, jit_mode_pi_compile
 from mindspore.parallel._ps_context import _is_role_sched
 from mindspore.parallel._utils import _check_full_batch, _get_parameter_broadcast, _is_pynative_parallel, \
     _is_in_auto_parallel_mode
@@ -62,6 +62,68 @@ function_phases = dict()
 BROADCAST_PHASE = "_broadcast_"
 _PYNATIVE_PARALLEL_FUNC_NAME = "after_shard"
 
+ARG_SPECIFIED = "arg_specified_infos"
+TOTAL_ARG_LEN = "total_arg_length"
+
+
+def _check_recompile_args(compile_args, kwargs):
+    """Check recompile of graph"""
+
+    def _check_constant_tensor_arg(arg):
+        if hasattr(arg, "__ms_mutable__"):
+            return False
+        if isinstance(arg, (list, tuple)):
+            return any(_check_constant_tensor_arg(x) for x in arg)
+        return isinstance(arg, Tensor)
+
+    for v in kwargs.values():
+        compile_args += (v,)
+    for arg in compile_args:
+        if not isinstance(arg, tuple) and not isinstance(arg, list):
+            continue
+        if _check_constant_tensor_arg(arg):
+            logger.warning(f"Constant value tensor are detected in tuple or list, which might cause recompiling "
+                           f"when tensor value changes. You can use mutable(Tensor) or mutable(tuple(Tensor)) "
+                           f"to set tensor's value as variable to to avoid recompiling. The tuple or list arg "
+                           f"is: {arg} .")
+            return
+
+
+def _check_recompile(obj, compile_args, kwargs, full_function_name, create_time, echo_function_name):
+    """Warning when the function has been compiled."""
+    ignore_dirs = ["mindspore/ops", "mindspore/nn"]
+    if any((lambda x: x in full_function_name)(x) for x in ignore_dirs):
+        return
+
+    if full_function_name in function_phases:
+        warning_times = 1
+        if len(function_phases[full_function_name]) >= warning_times \
+                and create_time not in function_phases[full_function_name]:
+            if isinstance(obj, ms.nn.Cell):
+                tips = f"Please try to create {echo_function_name} instance only once to avoid recompiling. "
+                logger.info(f"The {echo_function_name} has been compiled again. "
+                            f"{tips} ")
+            else:
+                tips = "Try to decorate the function with @jit(hash_args=...) " \
+                       "or @jit(compile_once=True) to reduce the compile time. " \
+                       "For more details, get instructions about `jit` at " \
+                       "https://www.mindspore.cn/search?inputValue=jit."
+                logger.warning(f"The {echo_function_name} has been compiled again. "
+                               f"{tips} ")
+        else:
+            _check_recompile_args(compile_args, kwargs)
+    else:
+        function_phases[full_function_name] = set()
+    function_phases[full_function_name].add(create_time)
+
+
+def _ms_adapter_tensor_as_parameter_output(data):
+    """Check whether the data is an output from a parameter which is a ms_adapter tensor.
+       Pylint: disable=unidiomatic-typecheck.
+    """
+    return ms_adapter_registry.is_registered and isinstance(data, ms_adapter_registry.tensor) \
+        and hasattr(data, "__ms_parameter_output__") and getattr(data, "__ms_parameter_output__")
+
 
 def _convert_python_data(data):
     """
@@ -73,8 +135,10 @@ def _convert_python_data(data):
     Returns:
         data, a data convert C++ to python
     """
-    if isinstance(data, Tensor) and data.adapter_flag:
+    if isinstance(data, (Tensor, PythonTensor)) and data.adapter_flag:
         return ms_adapter_registry.tensor(data)
+    if _ms_adapter_tensor_as_parameter_output(data) and hasattr(data, "tensor"):
+        return data.tensor
     if isinstance(data, Tensor) and not isinstance(data, PythonTensor):
         return PythonTensor(data, internal=True)
     if isinstance(data, CSRTensor) and not isinstance(data, PythonCSRTensor):
@@ -83,7 +147,7 @@ def _convert_python_data(data):
         return PythonCOOTensor(coo_tensor=data)
     if isinstance(data, RowTensor) and not isinstance(data, PythonRowTensor):
         return PythonRowTensor(row_tensor=data)
-    if isinstance(data, tuple):
+    if data.__class__ is tuple:
         # Handle namedtuple since its type is tuple.
         if hasattr(data, "_fields"):
             type_name = data.__class__.__name__
@@ -91,12 +155,12 @@ def _convert_python_data(data):
             fields = data_dict.keys()
             return namedtuple(type_name, fields)(**_convert_python_data(data_dict))
         return tuple(_convert_python_data(x) for x in data)
-    if isinstance(data, list):
+    if data.__class__ is list:
         # Keep list object not change for inplace operation.
         for i in range(len(data)):
             data[i] = _convert_python_data(data[i])
         return data
-    if isinstance(data, dict):
+    if data.__class__ is dict:
         # Keep the dict object not change.
         keys = tuple(data.keys())
         for key in keys:
@@ -167,6 +231,7 @@ def _handle_func_args(func, *args, **kwargs):
 
 sys_path = list(sys.path)
 # Get the entry script path.
+entry_script_path = None
 if sys.argv and sys.argv[0] != '':
     entry_script_path = os.path.realpath(sys.argv[0])
     entry_script_path_dir = os.path.split(entry_script_path)[0]
@@ -260,7 +325,7 @@ def _get_parameter_layout():
     return layout
 
 
-def _handle_arg(obj, arg):
+def _handle_arg(obj, arg, compile_arg):
     """Handle arg for runtime .If need handle the arg, return True"""
     if isinstance(arg, PythonTensor):
         if arg.has_init:
@@ -269,7 +334,7 @@ def _handle_arg(obj, arg):
             return arg
     elif isinstance(arg, (Tensor, CSRTensor, COOTensor)):
         return arg
-    elif hasattr(arg, "__ms_mutable__") and getattr(arg, "__ms_mutable__"):
+    elif compile_arg is not None and hasattr(compile_arg, "__ms_mutable__") and getattr(compile_arg, "__ms_mutable__"):
         # mutable([]) will be eliminated by FuncGraphSpecializer, and empty list is not supported by backend.
         if isinstance(arg, list) and not arg:
             return None
@@ -282,22 +347,172 @@ def _handle_arg(obj, arg):
     return None
 
 
-def _get_args_for_run(obj, args, kwargs):
+def _handle_arg_predict(obj, arg, compile_arg):
+    """Handle arg for runtime .If need handle the arg, return True"""
+    if arg is None:
+        return None
+
+    if isinstance(arg, (int, float)):
+        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__"):
+            # mutable([]) will be eliminated by FuncGraphSpecializer, and empty list is not supported by backend.
+            if isinstance(arg, list) and not arg:
+                return None
+            return arg
+        if hasattr(obj, "enable_tuple_broaden") and obj.enable_tuple_broaden and isinstance(arg, tuple) and \
+                _check_all_tensor(arg):
+            return arg
+        return None
+    return arg
+
+
+def _get_args_for_run(obj, args, kwargs, compile_args):
+    """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)
+        if new_arg is not None:
+            new_args.append(new_arg)
+
+    for _, value in kwargs.items():
+        new_value = _handle_arg(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."""
     new_args = []
-    for arg in args:
-        new_arg = _handle_arg(obj, arg)
+    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(obj, value)
+        new_value = _handle_arg_predict(obj, value, None)
         if new_value is not None:
             new_args.append(new_value)
 
     return new_args
 
 
+def _is_args_fullmode(args, is_init=True):
+    """Check whether the arguments is for incremental-mode.
+
+    Args:
+        args (Union[list, tuple, dict, Tensor]): Given arguments.
+        is_init (bool): Is check in argument initialization phase.
+
+    Raises:
+        RuntimeError: loss necessary keys and values for incremental-mode.
+
+    Returns:
+        bool: Fullmode or not.
+    """
+    if not isinstance(args, dict):
+        return True
+    if not is_init and (args.get(ARG_SPECIFIED, None) is None or args.get(TOTAL_ARG_LEN, None) is None):
+        raise RuntimeError(
+            "The incremental inputs should be processed(with \"%s\" and \"%s\"), but got %s." %
+            (ARG_SPECIFIED, TOTAL_ARG_LEN, str(args)))
+    return False
+
+
+def _process_dyn_args(fn, dyn_args):
+    """Process the dynamic arguments, return the necessary data for latter processing.
+
+    Args:
+        fn (Function): The root function to compile.
+        dyn_args (Union[dict, list, tuple, None]): Given arguments for dynamic compilation.
+            None for nothing, list or tuple for fullmode setting, dict for incremental configuration.
+
+    Returns:
+        A dict which contains args for dynamic compilation. None for nothing dynamic.
+    """
+    if dyn_args is None:
+        # nothing should be done for None.
+        return dyn_args
+
+    if isinstance(dyn_args, dict) and ARG_SPECIFIED in dyn_args:
+        return dyn_args
+
+    args_sig = inspect.signature(fn)
+    if _is_args_fullmode(dyn_args):
+        if not isinstance(dyn_args, (list, tuple)):
+            temp_dyn_args = (dyn_args,)
+        else:
+            temp_dyn_args = dyn_args
+
+        # If dyn_args is fullmode, it should be apply directly.
+        args_sig_parameters = list(args_sig.parameters.values())
+        if not args_sig_parameters:
+            return ()
+
+        # fn may be Cell's construct while the first input is 'self'.
+        if args_sig_parameters[0].name == "self" and (len(temp_dyn_args) + 1) == len(args_sig_parameters):
+            bound_args = args_sig.bind(None, *temp_dyn_args)
+            bound_args.apply_defaults()
+            return bound_args.args[1:]
+
+        bound_args = args_sig.bind(*temp_dyn_args)
+        bound_args.apply_defaults()
+        return bound_args.args
+
+    # The dyn_args is not fullmode, a real compilation arguments should be assembled by latter procession...
+    arg_names = []
+    args_sig_parameters = list(args_sig.parameters.values())
+    for arg_p in args_sig_parameters:
+        if arg_p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD):
+            arg_names.append(arg_p.name)
+        else:
+            raise TypeError("Dynamic arguments is not accepted for VAR_POSITIONAL or VAR_KEYWORD parameters!")
+
+    offset = -1 if fn.__name__ == 'construct' and args_sig_parameters[0].name == "self" else 0
+    meet_index = set()
+
+    def _check_index_valid(index):
+        if index >= len(arg_names):
+            raise ValueError("For dict mode, valid index is \"0\"-\"%d\", but got %s!" % (len(arg_names) - 1, index))
+        if index in meet_index:
+            raise ValueError("For dict mode, there are more than one same specified key for real index: %d!" % index)
+        meet_index.add(index)
+
+    arg_handler_infos = []
+    for k, v in dyn_args.items():
+        if not isinstance(k, str):
+            raise TypeError("For dict mode, only string key is accepted, but got %s!" % k)
+        if k in arg_names:
+            cur_id = arg_names.index(k)
+            _check_index_valid(cur_id)
+            arg_handler_infos.append([cur_id + offset, v])
+        else:
+            raise ValueError("For dict mode, valid key is %s, but got %s!" % (arg_names, k))
+    return {ARG_SPECIFIED: arg_handler_infos, TOTAL_ARG_LEN: len(args_sig_parameters)}
+
+
+def _generate_dyn_compile_args(compile_args, dyn_args):
+    """Generate the dynamic compile arguments."""
+    if not dyn_args:
+        return compile_args
+    if _is_args_fullmode(dyn_args, False):
+        if not isinstance(dyn_args, (list, tuple)):
+            return (dyn_args,)
+        return dyn_args
+    arg_specified_infos = dyn_args.get(ARG_SPECIFIED, None)
+    if arg_specified_infos is None:
+        raise RuntimeError("For dict mode, a key with \"%s\" should exist, but got %s!" %
+                           (ARG_SPECIFIED, str(dyn_args)))
+    new_compile_args = list(compile_args)
+    for index, arg in arg_specified_infos:
+        new_compile_args[index] = arg
+    return tuple(new_compile_args)
+
+
 class _MindsporeFunctionExecutor:
     """
     Represents a function compiled by graph compiler.
@@ -315,6 +530,7 @@ class _MindsporeFunctionExecutor:
     Returns:
         The result of pipeline running in graph mode.
     """
+
     def __init__(self, fn, ms_create_time, input_signature=None, obj=None, jit_config=None):
         init_pipeline()
         if not isinstance(fn, (types.FunctionType, types.MethodType)):
@@ -329,9 +545,9 @@ class _MindsporeFunctionExecutor:
         self.enable_tuple_broaden = False
         self._graph_executor = GraphExecutor_.get_instance()
         self._create_time = ms_create_time
+        self._compile_args = None
         self.jit_config_dict = jit_config.jit_config_dict if jit_config else None
 
-
     @_wrap_func
     def __call__(self, *args, **kwargs):
         args_list = args
@@ -359,7 +575,6 @@ class _MindsporeFunctionExecutor:
 
         return output
 
-
     def compile(self, method_name, *args, **kwargs):
         """Returns pipeline for the given args."""
         # Check whether hook function registered on Cell object.
@@ -376,6 +591,7 @@ class _MindsporeFunctionExecutor:
 
         # Restore the mutable attr for every arg.
         compile_args = _restore_mutable_attr(args, compile_args)
+        self._compile_args = compile_args
         generate_name, echo_function_name = self._get_generate_name()
         # The full Function name
         full_function_name = generate_name
@@ -414,9 +630,12 @@ class _MindsporeFunctionExecutor:
         update_auto_dynamic_shape_phase_with_check_input_signature(compile_args, key_id, phase, self.input_signature)
 
         if phase in ms_compile_cache:
+            # 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()
             return phase
 
-        self._check_recompile(full_function_name, create_time, echo_function_name)
+        _check_recompile(self.obj, compile_args, kwargs, full_function_name, create_time, echo_function_name)
 
         # If enable compile cache, get the dependency files list and set to graph executor.
         self._set_compile_cache_dep_files()
@@ -448,29 +667,6 @@ class _MindsporeFunctionExecutor:
 
         return phase
 
-    def _check_recompile(self, full_function_name, create_time, echo_function_name):
-        """Warning when the function has been compiled."""
-        ignore_dirs = ["mindspore/ops", "mindspore/nn"]
-        if any((lambda x: x in full_function_name)(x) for x in ignore_dirs):
-            return
-
-        if full_function_name in function_phases:
-            warning_times = 1
-            if len(function_phases[full_function_name]) >= warning_times \
-                    and create_time not in function_phases[full_function_name]:
-                tips = "Try to decorate the function with @jit(hash_args=...) " \
-                       "or @jit(compile_once=True) to reduce the compile time. " \
-                       "For more details, get instructions about `jit` at " \
-                       "https://www.mindspore.cn/search?inputValue=jit."
-
-                logger.warning(f"The {echo_function_name} has been compiled again. "
-                               f"{tips} ")
-        else:
-            function_phases[full_function_name] = set()
-
-        function_phases[full_function_name].add(create_time)
-
-
     @staticmethod
     def _optimizer_state_init(opt_states):
         """set data for all optimizer states in case it is executed in graph mode"""
@@ -481,7 +677,6 @@ class _MindsporeFunctionExecutor:
             if opt_param.has_init and (prefix in prefix_list or opt_param.name == "global_step"):
                 opt_param.init_data()
 
-
     def _get_key_id(self):
         """get key id."""
         if isinstance(self.obj, ms.nn.Cell):
@@ -493,7 +688,6 @@ class _MindsporeFunctionExecutor:
             key_id = key_id + ".grad"
         return key_id
 
-
     def _get_generate_name(self):
         """get generate name."""
         generate_name = self.fn.__module__ + "." + self.fn.__name__ + "." + self.fn.__code__.co_filename + "." + str(
@@ -506,54 +700,47 @@ class _MindsporeFunctionExecutor:
             generate_name = generate_name[:generate_name.rfind(str(id(self.fn)))] + str(id(self.shard_parent_obj))
         return generate_name, echo_function_name
 
-
     def _set_compile_cache_dep_files(self):
         # If enable compile cache, get the dependency files list
         enable_compile_cache = context.get_context("enable_compile_cache")
-        if enable_compile_cache is not True and enable_compile_cache != "1":
+        if enable_compile_cache is None:
             enable_compile_cache = os.getenv('MS_COMPILER_CACHE_ENABLE')
         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(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 `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 = 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)}.")
-            for i, elem in enumerate(compile_args):
-                if isinstance(elem, PythonTensor):
-                    Validator.check_dynamic_shape(compile_args[i], args_list[i], i)
+            self._graph_executor.check_argument_consistency(compile_args, args_list, "input_signature")
+            Validator.check_symbolic_shape(compile_args, 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:
-            if not isinstance(self.input_signature, (tuple, list)):
-                self.input_signature = (self.input_signature,)
-            self.input_signature = list(self.input_signature)
-            dyn_shape = False
-            for i, elem in enumerate(self.input_signature):
-                if isinstance(elem, PythonTensor) and is_shape_unknown(elem.shape):
-                    Validator.check_dynamic_shape(self.input_signature[i], args_list[i], i)
-                    dyn_shape = True
+            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(self.input_signature) == len(args_list) - 1:
+                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(self.input_signature)}'. The last actual args may "
+                                   f"of input_signature args '{len(compile_args)}'. The last actual args may "
                                    f"be 'sens' and added it to compile args.")
-                    self.input_signature.append(args_list[-1])
-                compile_args = tuple(self.input_signature)
+                    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(self.input_signature, args_list):
+                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
 
@@ -568,7 +755,7 @@ class _MindsporeFunctionExecutor:
         Returns:
             new_inputs, new input args, which are required for running.
         """
-        return _get_args_for_run(self, args_list, kwargs)
+        return _get_args_for_run(self, args_list, kwargs, self._compile_args)
 
 
 # The attributes used to identify a given object.
@@ -596,31 +783,58 @@ def _get_jit_hash(hash_input):
     return _get_obj_id(hash_input)
 
 
-def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_once=False):
+def _update_graph_executor_config(jit_config):
+    """Update GraphExecutor jit_config"""
+    if isinstance(jit_config, JitConfig):
+        jit_config = jit_config.jit_config_dict
+    if not isinstance(jit_config, dict):
+        return
+    valid_config = dict()
+    for k, v in jit_config.items():
+        valid_config[str(k)] = str(v)
+    GraphExecutor_.get_instance().set_jit_config(JitConfig(**valid_config).jit_config_dict)
+
+
+def jit(fn=None, mode="PSJit", input_signature=None, hash_args=None, jit_config=None, compile_once=False):
     """
     Create a callable MindSpore graph from a Python function.
 
     This allows the MindSpore runtime to apply optimizations based on graph.
 
-    Note:
-        If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
-        will not accept `**kwargs`.
-
     Args:
         fn (Function): The Python function that will be run as a graph. Default: ``None`` .
-        input_signature (Tensor): The Tensor which describes the input arguments. The shape and dtype of the Tensor
-            will be supplied to this function. If input_signature is specified, each input to `fn` must be a `Tensor`.
-            And the input parameters of `fn` cannot accept `**kwargs`. The shape and dtype of actual inputs should
-            keep the same as input_signature. Otherwise, TypeError will be raised. Default: ``None`` .
+        mode (str): The type of jit used, the value of mode should be ``PIJit`` or ``PSJit``. Default: ``PSJit`` .
+
+            - `PSJit <https://www.mindspore.cn/docs/en/master/note/static_graph_syntax_support.html>`_ :
+              Parse python ast to build graph.
+            - `PIJit <https://www.mindspore.cn/docs/en/master/design/dynamic_graph_and_static_graph.html>`_ :
+              Parse python bytecode to build graph at runtime.
+
+        input_signature (Union[Tuple, List, Dict, Tensor]): The Tensor which describes the input arguments. The
+            shape and dtype of the Tensor will be supplied to this function. If `input_signature` is specified, the
+            input parameters of `fn` cannot accept `**kwargs`, and the shape and dtype of actual inputs should keep the
+            same as `input_signature`. Otherwise, TypeError will be raised. There are two mode for `input_signature`:
+
+            - Full mode: Arguments is a Tuple, List or a Tensor, and they will be used as all compile inputs
+              for graph-compiling.
+            - Incremental mode: Argument is a Dict, and they will set to some of the graph inputs, which will be
+              substituted into the input at the corresponding position for graph-compiling.
+
+            Default: ``None`` .
+
         hash_args (Union[Object, List or Tuple of Objects]): The local free variables used inside `fn`,
             like functions or objects of class defined outside `fn`. Calling `fn` again with change of `hash_args`
             will trigger recompilation. Default: ``None`` .
         jit_config (JitConfig): Jit config for compile. Default: ``None`` .
         compile_once(bool): ``True``: The function would be compiled once when it was created many times.
             But it may be wrong if the free variables were changed. ``False`` : It would be recompiled when
-            it was created again
+            it was created again.
             Default: ``False`` .
 
+    Note:
+        If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
+        will not accept `**kwargs`.
+
     Returns:
         Function, if `fn` is not None, returns a callable function that will execute the compiled function; If `fn` is
         None, returns a decorator and when this decorator invokes with a single `fn` argument, the callable function is
@@ -663,6 +877,13 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_
         ...
         >>> out = tensor_add_with_sig(x, y)
         ...
+        >>> @jit(input_signature={"y": Tensor(np.ones([1, 1, 3, 3]).astype(np.float32))})
+        ... def tensor_add_with_sig_1(x, y):
+        ...     z = x + y
+        ...     return z
+        ...
+        >>> out1 = tensor_add_with_sig_1(x, y)
+        ...
         ... # Set hash_args as fn, otherwise cache of compiled closure_fn will not be reused.
         ... # While fn differs during calling again, recompilation will be triggered.
         >>> def func(x):
@@ -702,6 +923,8 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_
         else:
             hash_obj = int(time.time() * 1e9)
 
+        dyn_args = _process_dyn_args(func, input_signature)
+
         @wraps(func)
         def staging_specialize(*args, **kwargs):
             if os.getenv("MS_JIT") == '0':
@@ -715,14 +938,49 @@ def jit(fn=None, input_signature=None, hash_args=None, jit_config=None, compile_
             # only the function or cell instance wrapped by shard will fall into this branch
             if _is_pynative_parallel() and func.__name__ == _PYNATIVE_PARALLEL_FUNC_NAME:
                 process_obj = hash_args
-            out = _MindsporeFunctionExecutor(func, hash_obj, input_signature, process_obj, jit_config)(*args, **kwargs)
+            out = _MindsporeFunctionExecutor(func, hash_obj, dyn_args, process_obj, jit_config)(*args, **kwargs)
             return out
 
         return staging_specialize
 
+    def pi_wrap_mindspore(decorated):
+        func = decorated
+        if isinstance(func, ms.nn.Cell):
+            func = func.construct
+        if isinstance(func, type) and issubclass(func, ms.nn.Cell):
+            func = func.construct
+        if isinstance(func, types.MethodType):
+            func = func.__func__
+        if not isinstance(func, types.FunctionType):
+            logger.warning("only support function and mindspore.nn.Cell instance")
+            return decorated
+
+        # generator, coroutine, awaitable and a function that return them is unsupported
+        UNSUPPORTED_CODE_TYPE = (inspect.CO_GENERATOR | inspect.CO_COROUTINE |
+                                 inspect.CO_ASYNC_GENERATOR | inspect.CO_ITERABLE_COROUTINE)
+        if func.__code__.co_flags & UNSUPPORTED_CODE_TYPE:
+            return decorated
+
+        _update_graph_executor_config(jit_config)
+        config = dict()
+        if isinstance(jit_config, JitConfig):
+            config.update(jit_config.jit_config_dict)
+        elif jit_config is not None:
+            config.update(jit_config)
+        jit_mode_pi_enable()
+
+        if jit_mode_pi_compile(func, config, input_signature) is False:
+            logger.warning('add fn {} to compile failed '.format(func))
+
+        return decorated
+
+    wrap_func = wrap_mindspore
+    if mode == "PIJit":
+        wrap_func = pi_wrap_mindspore
+
     if fn is not None:
-        return wrap_mindspore(fn)
-    return wrap_mindspore
+        return wrap_func(fn)
+    return wrap_func
 
 
 def ms_function(fn=None, input_signature=None, hash_args=None, jit_config=None):
@@ -732,15 +990,14 @@ def ms_function(fn=None, input_signature=None, hash_args=None, jit_config=None):
     This allows the MindSpore runtime to apply optimizations based on graph.
 
     Note:
-        `ms_function` will be deprecated and removed in a future version. Please use `jit` instead.
-        If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
-        will not accept `**kwargs`.
+        - `ms_function` will be deprecated and removed in a future version. Please use :func:`mindspore.jit` instead.
+        - If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
+          will not accept `**kwargs`.
 
     Args:
         fn (Function): The Python function that will be run as a graph. Default: ``None`` .
         input_signature (Tensor): The Tensor which describes the input arguments. The shape and dtype of the Tensor
-            will be supplied to this function. If input_signature is specified, each input to `fn` must be a `Tensor`.
-            And the input parameters of `fn` cannot accept `**kwargs`. The shape and dtype of actual inputs should
+            will be supplied to this function. The shape and dtype of actual inputs of `fn` should
             keep the same as input_signature. Otherwise, TypeError will be raised. Default: ``None`` .
         hash_args (Union[Object, List or Tuple of Objects]): The local free variables used inside `fn`,
             like functions or objects of class defined outside `fn`. Calling `fn` again with change of `hash_args`
@@ -909,7 +1166,7 @@ def ms_class(cls):
     This allows MindSpore to identify user-defined classes and thus obtain their attributes and methods.
 
     Note:
-        `ms_class` will be deprecated and removed in a future version. Please use `jit_class` instead.
+        `ms_class` will be deprecated and removed in a future version. Please use :func:`mindspore.jit_class` instead.
 
     Args:
         cls (Class): User-defined class.
@@ -1037,6 +1294,8 @@ def set_adapter_config(config):
             ms_adapter_registry.register_parameter(value)
         elif key == "convert_object_map":
             ms_adapter_registry.register_convert_map(value)
+        elif key == "convert_adapter_tensor_map":
+            ms_adapter_registry.register_convert_adapter_tensor_map(value)
         else:
             raise ValueError(f"Unsupported key in adapter config: {key}")
 
@@ -1135,16 +1394,6 @@ class _PyNativeExecutor:
         self._executor = PyNativeExecutor_.get_instance()
         self._executor.set_py_exe_path(sys.executable)
         self._executor.set_kernel_build_server_dir(os.path.split(kernel_build_server.__file__)[0] + os.sep)
-        self._top_cell = None
-
-    def __call__(self):
-        """
-        PyNative executor run grad graph.
-
-        Return:
-            The return object after running grad graph.
-        """
-        return self._executor()
 
     @staticmethod
     def parameter_broadcast(obj, phase):
@@ -1246,7 +1495,7 @@ class _PyNativeExecutor:
         Return:
             None.
         """
-        self._executor.grad_net(grad, obj, weights, grad_position, *args, *(kwargs.values()))
+        return self._executor.grad(grad, obj, weights, grad_position, *args, *(kwargs.values()))
 
     def clear_res(self):
         """
@@ -1300,6 +1549,18 @@ class _PyNativeExecutor:
         """
         self._executor.set_grad_flag(flag)
 
+    def set_async_for_graph(self, flag):
+        """
+        Set the flag for graph async run.
+
+        Args:
+            flag (bool): Specifying whether enable graph async run.
+
+        Return:
+            None.
+        """
+        self._executor.set_async_for_graph(flag)
+
     def enable_grad(self):
         """
         The global flag whether needing to calculate gradient.
@@ -1333,6 +1594,17 @@ class _PyNativeExecutor:
         """
         self._executor.set_jit_compile_status(status, phase)
 
+    def set_is_run_recompute(self, status):
+        """
+        Set recompute grad is compiling
+
+        Args:
+            status(bool): grad is in recompute status
+        Return:
+            None.
+        """
+        self._executor.set_is_run_recompute(status)
+
     def set_dynamic_input(self, obj, *args):
         """
         Set dynamic shape tensor of input arguments.
@@ -1380,15 +1652,6 @@ class _PyNativeExecutor:
         """
         self._executor.set_hook_changed(cell)
 
-    def get_top_cell(self):
-        """
-        Get the top cell object.
-
-        Return:
-            The top cell object.
-        """
-        return self._top_cell
-
     def constant_folding(self, *args):
         """
         Get value by infer value.
@@ -1491,9 +1754,9 @@ class _CellGraphExecutor:
     def _set_compile_cache_dep_files(self, phase):
         # If enable compile cache, get the dependency files list
         enable_compile_cache = context.get_context("enable_compile_cache")
-        if enable_compile_cache is not True and enable_compile_cache != "1":
+        if enable_compile_cache is None:
             enable_compile_cache = os.getenv('MS_COMPILER_CACHE_ENABLE')
-        if "train" in phase and (enable_compile_cache is True or enable_compile_cache == "1"):
+        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 compile(self, obj, *args, phase='predict', do_convert=True, jit_config_dict=None, **kwargs):
@@ -1522,17 +1785,26 @@ class _CellGraphExecutor:
         self.enable_tuple_broaden = False
         if hasattr(obj, "enable_tuple_broaden"):
             self.enable_tuple_broaden = obj.enable_tuple_broaden
-        logger.debug("Convert the network.", do_convert)
+        logger.debug(f"Convert the network: {do_convert}.")
         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)
+        raw_phase = phase
         phase = phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
+        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):
             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.
+            self._graph_executor.clear_compile_arguments_resource()
             return phase, False
 
+        full_function_name = obj.__class__.__name__ + '.' + str(obj.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)
+
         obj.check_names()
         _check_full_batch()
         self._set_dataset_mode(obj)
@@ -1560,7 +1832,6 @@ class _CellGraphExecutor:
         elif 'skip_auto_parallel_compile' not in obj.get_flags().keys():
             obj.parameter_layout_dict = self._graph_executor.get_parameter_layout(phase)
             obj.parallel_parameter_name_list = self._graph_executor.get_parallel_parameter_name_list(phase)
-
         if "export.air" in phase:
             self._build_data_graph(obj, phase)
         elif BROADCAST_PHASE not in phase and _get_parameter_broadcast():
@@ -1600,6 +1871,18 @@ class _CellGraphExecutor:
         """
         return self._graph_executor.has_compiled(phase)
 
+    def flops_collection(self, phase='train'):
+        """
+        Specify whether have been compiled.
+
+        Args:
+            phase (str): The phase name. Default: 'predict'.
+
+        Returns:
+            bool, specifies whether the specific graph has been compiled.
+        """
+        return self._graph_executor.flops_collection(phase)
+
     @_wrap_func
     def _exec_pip(self, obj, *args, phase=''):
         """Execute the generated pipeline."""
@@ -1632,6 +1915,14 @@ class _CellGraphExecutor:
         """Clear the memory resource of a network."""
         self._graph_executor.del_net_res(obj, net_id)
 
+    def inc_graph_cell_count(self):
+        """Increase the count of GraphCell instance."""
+        self._graph_executor.inc_graph_cell_count()
+
+    def dec_graph_cell_count(self):
+        """Decrease the count of GraphCell instance."""
+        self._graph_executor.dec_graph_cell_count()
+
     def _get_branch_control_input(self):
         if ('obf_ratio' not in self.obfuscate_config.keys()) or (
                 'obf_random_seed' not in self.obfuscate_config.keys()):
@@ -1738,7 +2029,21 @@ def _bind_device_context():
     _bind_device_ctx()
 
 
+def flops_collection(phase='train'):
+    """
+    Recycle memory used by MindSpore.
+    When train multi Neural network models in one process, memory used by MindSpore is very large,
+    this is because MindSpore cached runtime memory for every model.
+    To recycle these cached memory, users can call this function after training of one model.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.ms_memory_recycle()
+    """
+    return _cell_graph_executor.flops_collection(phase)
+
+
 _cell_graph_executor = _CellGraphExecutor()
 _pynative_executor = _PyNativeExecutor()
 
-__all__ = ['ms_function', 'ms_memory_recycle', 'ms_class', 'jit', 'jit_class']
+__all__ = ['ms_function', 'ms_memory_recycle', 'ms_class', 'jit', 'jit_class', 'flops_collection']