mindspore 1.10.0__cp37-cp37m-win_amd64.whl → 2.0.0rc1__cp37-cp37m-win_amd64.whl

mindspore/rewrite/symbol_tree_builder.py

@@ -13,19 +13,54 @@
 # limitations under the License.
 # ============================================================================
 """SymbolTree builder."""
-from __future__ import absolute_import
 from copy import copy
 from typing import Optional
 import ast
 import inspect
 
 from mindspore.nn import Cell
-from mindspore.rewrite.symbol_tree import SymbolTree
-from mindspore.rewrite.node import TreeNode
-from mindspore.rewrite.parser_register import ParserRegister
-from mindspore.rewrite.parser import Parser
-from mindspore.rewrite.ast_transformers import FlattenRecursiveStmt
-from mindspore.rewrite.ast_helpers import AstFinder, AstModifier
+from .symbol_tree import SymbolTree
+from .node import TreeNode
+from .parser_register import ParserRegister
+from .parser import Parser
+from .ast_transformers import FlattenRecursiveStmt
+from .ast_helpers import AstModifier
+from .ast_helpers import AstFinder
+
+
+class FunctionSymbolTreeBuilder:
+    """Create function SymbolTree"""
+    def __init__(self, network: Cell, ast_root):
+        self._origin_net = network
+        self._ast_root: ast.Module = ast_root
+        self._root_tree: Optional[SymbolTree] = None
+
+    @staticmethod
+    def _ast_transform(ast_root: ast.AST) -> ast.AST:
+        """
+        Optimize ast before parse.
+
+        Args:
+             ast_root (ast.AST): An instance of ast to be optimized.
+
+        Returns:
+             An instance of ast been optimized.
+        """
+        transform_list = [FlattenRecursiveStmt()]
+        for transformer in transform_list:
+            ast_root = transformer.transform(ast_root)
+        return ast_root
+
+    def build(self) -> SymbolTree:
+        """
+        Build SymbolTree.
+
+        Returns:
+             An instance of SymbolTree.
+        """
+        self._root_tree: SymbolTree = SymbolTree(self._origin_net, self._ast_root)
+        self._root_tree.finish_build()
+        return self._root_tree
 
 
 class SymbolTreeBuilder:
@@ -43,6 +78,8 @@ class SymbolTreeBuilder:
         network_str = inspect.getsource(type(network))
         self._ast_root: ast.Module = ast.parse(network_str)
         self._root_tree: Optional[SymbolTree] = None
+        if isinstance(network, Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     @staticmethod
     def merge_module_of_subtree(main_tree: SymbolTree, sub_stree: SymbolTree):
@@ -140,7 +177,7 @@ class SymbolTreeBuilder:
         """
 
         for node in self._root_tree.nodes():
-            if isinstance(node, TreeNode):
+            if isinstance(node, TreeNode) and node.get_instance():
                 SymbolTreeBuilder.merge_module_of_subtree(self._root_tree, node.symbol_tree)
 
     def _reduce_redundant_import(self):

mindspore/rewrite/symbol_tree_dumper.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ============================================================================
 """SymbolTree dumper."""
-from __future__ import absolute_import
 import inspect
 
 from .node import Node

mindspore/rewrite/topological_manager.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ============================================================================
 """SymbolTree topological-relationship manager."""
-from __future__ import absolute_import
 from typing import Tuple
 
 from .api.scoped_value import ScopedValue
@@ -79,7 +78,7 @@ class TopoManager(Observable):
         if consumers is None:
             self._target_consumer[product] = [(consumer, index)]
         else:
-            self._target_consumer[product].append((consumer, index))
+            self._target_consumer.get(product).append((consumer, index))
 
     def _erase_provider(self, product: ScopedValue):
         """

mindspore/run_check/_check_version.py

@@ -26,16 +26,16 @@ from abc import abstractmethod, ABCMeta
 from packaging import version
 import numpy as np
 from mindspore import log as logger
+from mindspore._c_expression import MSContext, ms_ctx_param
 from ..version import __version__
-from ..default_config import __package_name__
 
 
 class EnvChecker(metaclass=ABCMeta):
     """basic class for environment check"""
 
     @abstractmethod
-    def check_env(self, e):
-        pass
+    def check_env(self):
+        """check dependency"""
 
     @abstractmethod
     def set_env(self):
@@ -46,12 +46,35 @@ class EnvChecker(metaclass=ABCMeta):
         pass
 
 
+class CPUEnvChecker(EnvChecker):
+    """CPU environment check."""
+
+    def __init__(self, library_path):
+        self.library_path = library_path
+
+    def check_env(self):
+        pass
+
+    def check_version(self):
+        pass
+
+    def set_env(self):
+        """set env for cpu"""
+        plugin_dir = os.path.dirname(self.library_path)
+        akg_dir = os.path.join(plugin_dir, "plugin/cpu")
+        if os.getenv('LD_LIBRARY_PATH'):
+            os.environ['LD_LIBRARY_PATH'] = akg_dir + ":" + os.environ['LD_LIBRARY_PATH']
+        else:
+            os.environ['LD_LIBRARY_PATH'] = akg_dir
+
+
 class GPUEnvChecker(EnvChecker):
     """GPU environment check."""
 
-    def __init__(self):
+    def __init__(self, library_path):
         self.version = ["10.1", "11.1", "11.6"]
-        self.lib_key_to_lib_name = {'libcu': 'libcuda.so'}
+        self.lib_key_to_lib_name = {'libcu': 'libcuda.so', 'libcudnn': 'libcudnn.so'}
+        self.library_path = library_path
         # env
         self.path = os.getenv("PATH")
         self.ld_lib_path = os.getenv("LD_LIBRARY_PATH")
@@ -62,8 +85,8 @@ class GPUEnvChecker(EnvChecker):
         self.cuda_bin_path = self._get_bin_path("cuda")
         self.cudnn_lib_path = self._get_lib_path("libcudnn")
 
-    def check_env(self, e):
-        raise e
+    def check_env(self):
+        pass
 
     def check_version(self):
         """Check cuda version."""
@@ -72,8 +95,8 @@ class GPUEnvChecker(EnvChecker):
             version_match = True
         if not version_match:
             if self.v == "0":
-                logger.warning("Can not found cuda libs, please confirm that the correct "
-                               "cuda version has been installed, you can refer to the "
+                logger.warning("Can not found cuda libs. Please confirm that the correct "
+                               "cuda version has been installed. Refer to the "
                                "installation guidelines: https://www.mindspore.cn/install")
             else:
                 logger.warning(f"MindSpore version {__version__} and cuda version {self.v} does not match, "
@@ -83,22 +106,44 @@ class GPUEnvChecker(EnvChecker):
         nvcc_version = self._get_nvcc_version(False)
         if nvcc_version and (nvcc_version not in self.version):
             logger.warning(f"MindSpore version {__version__} and nvcc(cuda bin) version {nvcc_version} "
-                           "does not match, please refer to the installation guide for version matching "
+                           "does not match. Please refer to the installation guide for version matching "
                            "information: https://www.mindspore.cn/install")
         cudnn_version = self._get_cudnn_version()
         if cudnn_version and int(cudnn_version) < 760:
             logger.warning(f"MindSpore version {__version__} and cudDNN version {cudnn_version} "
-                           "does not match, please refer to the installation guide for version matching "
+                           "does not match. Please refer to the installation guide for version matching "
                            "information: https://www.mindspore.cn/install. The recommended version is "
-                           "CUDA10.1 with cuDNN7.6.x and CUDA11.1 with cuDNN8.0.x")
+                           "CUDA10.1 with cuDNN7.6.x, CUDA11.1 with cuDNN8.0.x and CUDA11.6 with cuDNN8.5.x.")
         if cudnn_version and int(cudnn_version) < 800 and int(str(self.v).split('.')[0]) > 10:
             logger.warning(f"CUDA version {self.v} and cuDNN version {cudnn_version} "
-                           "does not match, please refer to the installation guide for version matching "
+                           "does not match. Please refer to the installation guide for version matching "
                            "information: https://www.mindspore.cn/install. The recommended version is "
-                           "CUDA11.1 with cuDNN8.0.x")
+                           "CUDA11.1 with cuDNN8.0.x or CUDA11.6 with cuDNN8.5.x.")
+
+    def get_cudart_version(self):
+        """Get cuda runtime version by libcudart.so."""
+        for path in self.cuda_lib_path:
+            real_path = glob.glob(path + "/lib*/libcudart.so.*.*.*")
+            if real_path == []:
+                continue
+            ls_cudart = subprocess.run(["ls", real_path[0]], timeout=10, text=True,
+                                       capture_output=True, check=False)
+            if ls_cudart.returncode == 0:
+                self.v = ls_cudart.stdout.split('/')[-1].strip('libcudart.so.').strip()
+                break
+        return self.v
 
     def set_env(self):
-        return
+        """set env for gpu"""
+        v = self.get_cudart_version()
+        v = version.parse(v)
+        v_str = str(v.major) + "." + str(v.minor)
+        plugin_dir = os.path.dirname(self.library_path)
+        akg_dir = os.path.join(plugin_dir, "gpu" + v_str)
+        if os.getenv('LD_LIBRARY_PATH'):
+            os.environ['LD_LIBRARY_PATH'] = akg_dir + ":" + os.environ['LD_LIBRARY_PATH']
+        else:
+            os.environ['LD_LIBRARY_PATH'] = akg_dir
 
     def _get_bin_path(self, bin_name):
         """Get bin path by bin name."""
@@ -110,7 +155,7 @@ class GPUEnvChecker(EnvChecker):
         """Get cuda bin path by lib path."""
         path_list = []
         for path in self.cuda_lib_path:
-            path = os.path.abspath(path.strip()+"/bin/")
+            path = os.path.abspath(path.strip() + "/bin/")
             if Path(path).is_dir():
                 path_list.append(path)
         return np.unique(path_list)
@@ -150,22 +195,9 @@ class GPUEnvChecker(EnvChecker):
         version_str = ''.join(cudnn_version)
         return version_str[0:3]
 
-    def _get_cudart_version(self):
-        """Get cuda runtime version by libcudart.so."""
-        for path in self.cuda_lib_path:
-            real_path = glob.glob(path + "/lib*/libcudart.so.*.*.*")
-            if real_path == []:
-                continue
-            ls_cudart = subprocess.run(["ls", real_path[0]], timeout=10, text=True,
-                                       capture_output=True, check=False)
-            if ls_cudart.returncode == 0:
-                self.v = ls_cudart.stdout.split('/')[-1].strip('libcudart.so.').strip()
-                break
-        return self.v
-
     def _check_version(self):
         """Check cuda version"""
-        v = self._get_cudart_version()
+        v = self.get_cudart_version()
         v = version.parse(v)
         v_str = str(v.major) + "." + str(v.minor)
         if v_str not in self.version:
@@ -178,23 +210,23 @@ class GPUEnvChecker(EnvChecker):
         current_path = os.path.split(os.path.realpath(__file__))[0]
         mindspore_path = os.path.join(current_path, "../lib/plugin")
         try:
-            real_path = glob.glob(mindspore_path + "/libmindspore_gpu.so")
-            if real_path == []:
-                logger.error(f"{self.lib_key_to_lib_name[lib_name]} (need by mindspore-gpu) is not found, please "
+            real_path = self.library_path
+            if real_path is None or real_path == []:
+                logger.error(f"{self.lib_key_to_lib_name[lib_name]} (need by mindspore-gpu) is not found. Please "
                              f"confirm that libmindspore_gpu.so is in directory:{mindspore_path} and the correct cuda "
                              "version has been installed, you can refer to the installation "
                              "guidelines: https://www.mindspore.cn/install")
                 return path_list
-            ldd_r = subprocess.Popen(['ldd', real_path[0]], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            ldd_r = subprocess.Popen(['ldd', self.library_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
             ldd_result = subprocess.Popen(['/bin/grep', lib_name], stdin=ldd_r.stdout, stdout=subprocess.PIPE)
             result = ldd_result.communicate(timeout=5)[0].decode()
             for i in result.split('\n'):
                 path = i.partition("=>")[2]
                 if path.lower().find("not found") > 0:
-                    logger.error(f"Cuda {self.version} version({lib_name}*.so need by mindspore-gpu) is not found, "
-                                 "please confirm that the path of cuda is set to the env LD_LIBRARY_PATH, or check "
+                    logger.error(f"Cuda {self.version} version({lib_name}*.so need by mindspore-gpu) is not found. "
+                                 "Please confirm that the path of cuda is set to the env LD_LIBRARY_PATH, or check "
                                  "whether the CUDA version in wheel package and the CUDA runtime in current device "
-                                 "matches, please refer to the installation guidelines: "
+                                 "matches. Please refer to the installation guidelines: "
                                  "https://www.mindspore.cn/install")
                     continue
                 path = path.partition(lib_name)[0]
@@ -202,8 +234,8 @@ class GPUEnvChecker(EnvChecker):
                     path_list.append(os.path.abspath(path.strip() + "../"))
             return np.unique(path_list)
         except subprocess.TimeoutExpired:
-            logger.warning("Failed to check cuda version due to the ldd command timeout, please confirm that "
-                           "the correct cuda version has been installed, you can refer to the "
+            logger.warning("Failed to check cuda version due to the ldd command timeout. Please confirm that "
+                           "the correct cuda version has been installed. For details, refer to the "
                            "installation guidelines: https://www.mindspore.cn/install")
             return path_list
 
@@ -221,8 +253,9 @@ class GPUEnvChecker(EnvChecker):
 class AscendEnvChecker(EnvChecker):
     """ascend environment check"""
 
-    def __init__(self):
-        self.version = ["1.84"]
+    def __init__(self, library_path):
+        self.library_path = library_path
+        self.version = ["6.11", "6.3"]
         atlas_nnae_version = "/usr/local/Ascend/nnae/latest/compiler/version.info"
         atlas_toolkit_version = "/usr/local/Ascend/ascend-toolkit/latest/compiler/version.info"
         hisi_fwk_version = "/usr/local/Ascend/latest/compiler/version.info"
@@ -278,15 +311,14 @@ class AscendEnvChecker(EnvChecker):
         self.ascend_opp_path_check = "/op"
         self.v = ""
 
-    def check_env(self, e):
+    def check_env(self):
         self._check_env()
-        raise e
 
     def check_version(self):
         if not Path(self.fwk_version).is_file():
             logger.warning("Using custom Ascend AI software package (Ascend Data Center Solution) path, package "
-                           "version checking is skipped, please make sure Ascend AI software package (Ascend Data "
-                           "Center Solution) version is supported, you can reference to the installation guidelines "
+                           "version checking is skipped. Please make sure Ascend AI software package (Ascend Data "
+                           "Center Solution) version is supported. For details, refer to the installation guidelines "
                            "https://www.mindspore.cn/install")
             return
 
@@ -295,31 +327,54 @@ class AscendEnvChecker(EnvChecker):
             v_list = str([x for x in self.version])
             logger.warning(f"MindSpore version {__version__} and Ascend AI software package (Ascend Data Center "
                            f"Solution)version {v} does not match, the version of software package expect one of "
-                           f"{v_list}, please reference to the match info on: https://www.mindspore.cn/install")
+                           f"{v_list}. Please refer to the match info on: https://www.mindspore.cn/install")
 
     def check_deps_version(self):
         """
-            te, topi, hccl wheel package version check
+            te and hccl wheel package version check
             in order to update the change of 'LD_LIBRARY_PATH' env, run a sub process
         """
-        input_args = ["--mindspore_version=" + __version__]
-        for v in self.version:
-            input_args.append("--supported_version=" + v)
-        deps_version_checker = os.path.join(os.path.split(os.path.realpath(__file__))[0],
-                                            "_check_deps_version.py")
-        call_cmd = [sys.executable, deps_version_checker] + input_args
+
+        mindspore_version = __version__
+        supported_version = self.version
+        attention_warning = False
         try:
-            process = subprocess.run(call_cmd, timeout=3, text=True, capture_output=True, check=False)
-            if process.stdout.strip() != "":
-                logger.warning(process.stdout.strip())
-                warning_countdown = 3
-                for i in range(warning_countdown, 0, -1):
-                    logger.warning(f"Please pay attention to the above warning, countdown: {i}")
-                    time.sleep(1)
-        except subprocess.TimeoutExpired:
-            logger.info("Package te, topi, hccl version check timed out, skip.")
+            from te import version as tever
+            v = '.'.join(tever.version.split('.')[0:2])
+            if v not in supported_version:
+                attention_warning = True
+                logger.warning(f"MindSpore version {mindspore_version} and \"te\" wheel package version {v} does not "
+                               "match. For details, refer to the installation guidelines: "
+                               "https://www.mindspore.cn/install")
+            from hccl import sys_version as hccl_version
+            v = '.'.join(hccl_version.__sys_version__.split('.')[0:2])
+            if v not in supported_version:
+                attention_warning = True
+                logger.warning(f"MindSpore version {mindspore_version} and \"hccl\" wheel package version {v} does not "
+                               "match. For details, refer to the installation guidelines: "
+                               "https://www.mindspore.cn/install")
+        # DO NOT modify exception type to any other, you DO NOT know what kind of exceptions the te will throw.
+        # pylint: disable=broad-except
+        except Exception as e:
+            logger.error("CheckFailed:", e.args)
+            logger.error("MindSpore relies on whl packages of \"te\" and \"hccl\" in the \"latest\" "
+                         "folder of the Ascend AI software package (Ascend Data Center Solution). Please check whether"
+                         " they are installed correctly or not, refer to the match info on: "
+                         "https://www.mindspore.cn/install")
+        if attention_warning:
+            warning_countdown = 3
+            for i in range(warning_countdown, 0, -1):
+                logger.warning(f"Please pay attention to the above warning, countdown: {i}")
+                time.sleep(1)
 
     def set_env(self):
+        plugin_dir = os.path.dirname(self.library_path)
+        akg_dir = os.path.join(plugin_dir, "ascend")
+        if os.getenv('LD_LIBRARY_PATH'):
+            os.environ['LD_LIBRARY_PATH'] = akg_dir + ":" + os.environ['LD_LIBRARY_PATH']
+        else:
+            os.environ['LD_LIBRARY_PATH'] = akg_dir
+
         if not self.tbe_path:
             self._check_env()
             return
@@ -331,14 +386,12 @@ class AscendEnvChecker(EnvChecker):
         except Exception:
             sys.path = deepcopy(origin_path)
             if Path(self.tbe_path).is_dir():
-                if os.getenv('LD_LIBRARY_PATH'):
-                    os.environ['LD_LIBRARY_PATH'] = self.tbe_path + ":" + os.environ['LD_LIBRARY_PATH']
-                else:
-                    os.environ['LD_LIBRARY_PATH'] = self.tbe_path
+                os.environ['LD_LIBRARY_PATH'] = self.tbe_path + ":" + os.environ['LD_LIBRARY_PATH']
             else:
-                raise EnvironmentError(
-                    f"No such directory: {self.tbe_path}, Please check if Ascend AI software package (Ascend Data "
+                logger.error(
+                    f"No such directory: {self.tbe_path}. Please check if Ascend AI software package (Ascend Data "
                     "Center Solution) is installed correctly.")
+                return
 
         # check te version after set te env
         self.check_deps_version()
@@ -354,57 +407,63 @@ class AscendEnvChecker(EnvChecker):
 
             os.environ['TBE_IMPL_PATH'] = self.op_impl_path
         else:
-            raise EnvironmentError(
-                f"No such directory: {self.op_impl_path}, Please check if Ascend AI software package (Ascend Data "
+            logger.error(
+                f"No such directory: {self.op_impl_path}. Please check if Ascend AI software package (Ascend Data "
                 "Center Solution) is installed correctly.")
+            return
 
         if Path(self.cce_path).is_dir():
             os.environ['PATH'] = self.cce_path + ":" + os.environ['PATH']
         else:
-            raise EnvironmentError(
-                f"No such directory: {self.cce_path}, Please check if Ascend AI software package (Ascend Data Center "
+            logger.error(
+                f"No such directory: {self.cce_path}. Please check if Ascend AI software package (Ascend Data Center "
                 "Solution) is installed correctly.")
+            return
 
         if self.op_path is None:
             pass
         elif Path(self.op_path).is_dir():
             os.environ['ASCEND_OPP_PATH'] = self.op_path
         else:
-            raise EnvironmentError(
-                f"No such directory: {self.op_path}, Please check if Ascend AI software package (Ascend Data Center "
+            logger.error(
+                f"No such directory: {self.op_path}. Please check if Ascend AI software package (Ascend Data Center "
                 "Solution) is installed correctly.")
+            return
 
         if self.aicpu_path is None:
             pass
         elif Path(self.aicpu_path).is_dir():
             os.environ['ASCEND_AICPU_PATH'] = self.aicpu_path
         else:
-            raise EnvironmentError(
-                f"No such directory: {self.aicpu_path}, Please check if Ascend AI software package (Ascend Data Center"
+            logger.error(
+                f"No such directory: {self.aicpu_path}. Please check if Ascend AI software package (Ascend Data Center"
                 " Solution) is installed correctly.")
+            return
 
     def _check_env(self):
         """ascend dependence path check"""
         if self.path is None or self.path_check not in self.path:
-            logger.warning("Can not find ccec_compiler(need by mindspore-ascend), please check if you have set env "
-                           "PATH, you can reference to the installation guidelines https://www.mindspore.cn/install")
+            logger.warning("Can not find ccec_compiler(need by mindspore-ascend). Please check whether the Environment "
+                           "Variable PATH is set. For details, refer to the installation guidelines: "
+                           "https://www.mindspore.cn/install")
 
         if self.python_path is None or self.python_path_check not in self.python_path:
             logger.warning(
-                "Can not find tbe op implement(need by mindspore-ascend), please check if you have set env "
-                "PYTHONPATH, you can reference to the installation guidelines "
+                "Can not find the tbe operator implementation(need by mindspore-ascend). Please check whether the "
+                "Environment Variable PYTHONPATH is set. For details, refer to the installation guidelines: "
                 "https://www.mindspore.cn/install")
 
         if self.ld_lib_path is None or not (self.ld_lib_path_check_fwk in self.ld_lib_path and
                                             self.ld_lib_path_check_addons in self.ld_lib_path):
-            logger.warning("Can not find driver so(need by mindspore-ascend), please check if you have set env "
-                           "LD_LIBRARY_PATH, you can reference to the installation guidelines "
-                           "https://www.mindspore.cn/install")
+            logger.warning("Can not find driver so(need by mindspore-ascend). Please check whether the "
+                           "Environment Variable LD_LIBRARY_PATH is set. For details, refer to the installation "
+                           "guidelines: https://www.mindspore.cn/install")
 
         if self.ascend_opp_path is None or self.ascend_opp_path_check not in self.ascend_opp_path:
             logger.warning(
-                "Can not find opp path (need by mindspore-ascend), please check if you have set env ASCEND_OPP_PATH, "
-                "you can reference to the installation guidelines https://www.mindspore.cn/install")
+                "Can not find opp path (need by mindspore-ascend). Please check whether the Environment Variable "
+                "ASCEND_OPP_PATH is set. For details, refer to the installation guidelines: "
+                "https://www.mindspore.cn/install")
 
     def _read_version(self, file_path):
         """get ascend version info"""
@@ -418,34 +477,59 @@ class AscendEnvChecker(EnvChecker):
         return self.v
 
 
+def check_env(device, _):
+    """callback function for checking environment variables"""
+    if device.lower() == "ascend":
+        env_checker = AscendEnvChecker(None)
+        env_checker.check_version()
+    elif device.lower() == "gpu":
+        env_checker = GPUEnvChecker(None)
+    else:
+        logger.info(f"Device {device} does not need to check any environment variable, skipping.")
+        return
+    env_checker.check_env()
+
+
+def set_env(device, library_path):
+    """callback function for setting environment variables"""
+    if not os.getenv("MS_DEV_CLOSE_VERSION_CHECK") is None:
+        if device in os.environ["MS_DEV_CLOSE_VERSION_CHECK"]:
+            return
+        os.environ["MS_DEV_CLOSE_VERSION_CHECK"] = os.environ["MS_DEV_CLOSE_VERSION_CHECK"] + ":" + device
+    else:
+        os.environ["MS_DEV_CLOSE_VERSION_CHECK"] = device
+
+    if device.lower() == "ascend":
+        env_checker = AscendEnvChecker(library_path)
+    elif device.lower() == "gpu":
+        env_checker = GPUEnvChecker(library_path)
+    elif device.lower() == "cpu":
+        env_checker = CPUEnvChecker(library_path)
+    else:
+        logger.info(f"Device {device} does not need to check any environment variable, skipping.")
+        return
+
+    env_checker.check_version()
+    env_checker.set_env()
+
+
 def check_version_and_env_config():
     """check version and env config"""
-    if __package_name__.lower() == "mindspore-ascend":
-        env_checker = AscendEnvChecker()
+    if platform.system().lower() == 'linux':
         # Note: pre-load libgomp.so to solve error like "cannot allocate memory in statis TLS block"
         try:
             import ctypes
             ctypes.cdll.LoadLibrary("libgomp.so.1")
         except OSError:
-            logger.warning(
-                "Pre-Load Lirary libgomp.so.1 failed, this might cause cannot allocate TLS memory problem, "
-                "if so find solution in FAQ in https://www.mindspore.cn/docs/en/r1.10/faq/installation.html.")
-    elif __package_name__.lower() == "mindspore-gpu":
-        env_checker = GPUEnvChecker()
-    else:
-        logger.info(f"Package version {__package_name__} does not need to check any environment variable, skipping.")
-        return
-    if os.getenv("MS_DEV_CLOSE_VERSION_CHECK") == "ON":
-        return
-    os.environ["MS_DEV_CLOSE_VERSION_CHECK"] = "ON"
-
-    try:
-        # check version of ascend site or cuda
-        env_checker.check_version()
-        import mindspore._c_expression as _c_expression  # pylint: disable=unused-import
-        env_checker.set_env()
-    except ImportError as e:
-        env_checker.check_env(e)
+            logger.warning("Pre-Load Library libgomp.so.1 failed, which might cause TLS memory allocation failure. If "
+                           "the failure occurs, please refer to the FAQ for a solution: "
+                           "https://www.mindspore.cn/docs/en/r2.0/faq/installation.html.")
+        if not os.getenv("MS_DEV_CLOSE_VERSION_CHECK") is None:
+            return
+        MSContext.get_instance().register_check_env_callback(check_env)
+        MSContext.get_instance().register_set_env_callback(set_env)
+        MSContext.get_instance().set_param(ms_ctx_param.device_target,
+                                           MSContext.get_instance().get_param(ms_ctx_param.device_target))
 
 
 def _set_pb_env():
@@ -454,7 +538,7 @@ def _set_pb_env():
         logger.info("Current env variable `PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp`. "
                     "When the checkpoint file is too large, "
                     "it may cause memory limit error during load checkpoint file. "
-                    "This can be solved by set env `PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python`.")
+                    "This can be solved by setting env `PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python`.")
     elif os.getenv("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION") is None:
         logger.info("Setting the env `PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python` to prevent memory overflow "
                     "during save or load checkpoint file.")
@@ -462,13 +546,26 @@ def _set_pb_env():
 
 
 def _add_cuda_path():
+    """add cuda path on windows."""
     if platform.system().lower() == 'windows':
-        if __package_name__.lower() == "mindspore_gpu":
-            cuda_home = os.environ.get('CUDA_PATH')
-            if cuda_home is None:
-                logger.error("mindspore-gpu on windows need CUDA_PATH, but not set it now")
+        cuda_home = os.environ.get('CUDA_PATH')
+        if cuda_home is None:
+            pass
+        else:
+            cuda_bin_path = os.path.join(os.environ['CUDA_PATH'], 'bin')
+            if sys.version_info >= (3, 8):
+                os.add_dll_directory(cuda_bin_path)
+            else:
+                os.environ['PATH'] += os.pathsep + cuda_bin_path
+        cudann_home = os.environ.get('CUDNN_HOME')
+        if cudann_home is None:
+            pass
+        else:
+            cuda_home_bin_path = os.path.join(os.environ['CUDNN_HOME'], 'bin')
+            if sys.version_info >= (3, 8):
+                os.add_dll_directory(cuda_home_bin_path)
             else:
-                os.add_dll_directory(os.path.join(os.environ['CUDA_PATH'], 'bin'))
+                os.environ['PATH'] += os.pathsep + cuda_home_bin_path
 
 
 check_version_and_env_config()

mindspore/run_check/run_check.py

@@ -40,7 +40,8 @@ def _check_mul():
     input_y = ms.Tensor(np.array([4.0, 5.0, 6.0]), ms.float32)
     mul = ms.ops.Mul()
     mul(input_x, input_y)
-    print(f"The result of multiplication calculation is correct, MindSpore has been installed successfully!")
+    print("The result of multiplication calculation is correct, MindSpore has been installed on platform "
+          f"[{ms.get_context('device_target')}] successfully!")
 
 
 def run_check():