feat(imperative): add user environment information collector

GitOrigin-RevId: 068c79d30dc7fab69ff220ddd5de4c0427994e6f

imperative/python/megengine/__init__.py

@@ -42,8 +42,6 @@ def _exit(code):
 
 def _atexit(handler):
     _exit_handlers.append(handler)
-
-
 _atexit(_close)
 
 _persistent_cache = _PersistentCacheOnServer()

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()

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; });

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}}}

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}}}