From f0a3ab97c056bdb0ba36516ff131ce7bb02f51f1 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Thu, 4 May 2023 16:33:25 +0800
Subject: [PATCH] feat(imperative): add user environment information collector

GitOrigin-RevId: 068c79d30dc7fab69ff220ddd5de4c0427994e6f
---
 imperative/python/megengine/__init__.py |  2 --
 imperative/python/megengine/device.py   | 13 +++++++++++++
 imperative/python/src/common.cpp        |  1 +
 src/core/impl/version.cpp               | 17 +++++++++++++++++
 src/core/include/megbrain/version.h     |  1 +
 5 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py
index 218bdd098..56ec95f91 100644
--- a/imperative/python/megengine/__init__.py
+++ b/imperative/python/megengine/__init__.py
@@ -42,8 +42,6 @@ def _exit(code):
 
 def _atexit(handler):
     _exit_handlers.append(handler)
-
-
 _atexit(_close)
 
 _persistent_cache = _PersistentCacheOnServer()
diff --git a/imperative/python/megengine/device.py b/imperative/python/megengine/device.py
index 90007f887..147cf990a 100644
--- a/imperative/python/megengine/device.py
+++ b/imperative/python/megengine/device.py
@@ -4,6 +4,9 @@ import re
 from typing import Optional
 
 from .core._imperative_rt.common import CompNode, DeviceType
+from .core._imperative_rt.common import (
+    get_cuda_driver_version as _get_cuda_driver_version,
+)
 from .core._imperative_rt.common import get_cuda_version as _get_cuda_version
 from .core._imperative_rt.common import get_cudnn_version as _get_cudnn_version
 from .core._imperative_rt.common import get_device_prop as _get_device_prop
@@ -26,6 +29,7 @@ __all__ = [
     "get_cuda_version",
     "get_cudnn_version",
     "get_tensorrt_version",
+    "get_cuda_driver_version",
     "get_allocated_memory",
     "get_reserved_memory",
     "get_max_reserved_memory",
@@ -299,3 +303,12 @@ def get_tensorrt_version():
         a version number, indicating `NV_TENSORRT_MAJOR * 1000 + NV_TENSORRT_MINOR * 100 + NV_TENSORRT_PATCH`.
     """
     return _get_tensorrt_version()
+
+
+def get_cuda_driver_version():
+    r"""Gets the latest version of CUDA supported by the driver.
+
+    Returns:
+        a supported latest cuda version number, indicating `CUDA_VERSION_MAJOR * 1000 + CUDA_VERSION_MINOR * 10`.
+    """
+    return _get_cuda_driver_version()
diff --git a/imperative/python/src/common.cpp b/imperative/python/src/common.cpp
index 4edee7316..4d5dc90a4 100644
--- a/imperative/python/src/common.cpp
+++ b/imperative/python/src/common.cpp
@@ -268,6 +268,7 @@ void init_common(py::module m) {
     m.def("get_cuda_version", []() { return mgb::get_cuda_version(); });
     m.def("get_cudnn_version", []() { return mgb::get_cudnn_version(); });
     m.def("get_tensorrt_version", []() { return mgb::get_tensorrt_version(); });
+    m.def("get_cuda_driver_version", []() { return mgb::get_cuda_driver_version(); });
 
     m.def("what_is_xpu",
           [] { return CompNode::Locator::parse("xpux").to_physical().type; });
diff --git a/src/core/impl/version.cpp b/src/core/impl/version.cpp
index 469acaa7b..4dff54e5f 100644
--- a/src/core/impl/version.cpp
+++ b/src/core/impl/version.cpp
@@ -63,4 +63,21 @@ int mgb::get_cudnn_version() {
 }
 #endif
 
+#if __has_include("cuda.h") && MGB_CUDA
+#include "cuda.h"
+int mgb::get_cuda_driver_version() {
+    int driver_version = -1;
+    auto error_code = cudaDriverGetVersion(&driver_version);
+    if (error_code != cudaSuccess) {
+        mgb_log_warn("cudaDriverGetVersion failed, error code: %d", error_code);
+        return -1;
+    }
+    return driver_version;
+}
+#else
+int mgb::get_cuda_driver_version() {
+    return -1;
+}
+#endif
+
 // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}
diff --git a/src/core/include/megbrain/version.h b/src/core/include/megbrain/version.h
index 258708d27..325e542d6 100644
--- a/src/core/include/megbrain/version.h
+++ b/src/core/include/megbrain/version.h
@@ -23,6 +23,7 @@ MGE_WIN_DECLSPEC_FUC Version get_version();
 MGE_WIN_DECLSPEC_FUC int get_cuda_version();
 MGE_WIN_DECLSPEC_FUC int get_cudnn_version();
 MGE_WIN_DECLSPEC_FUC int get_tensorrt_version();
+MGE_WIN_DECLSPEC_FUC int get_cuda_driver_version();
 }  // namespace mgb
 
 // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}
-- 
GitLab