From d60d028a401ce9fef5b33c0bbe14393e3dedc370 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 27 Jun 2022 14:42:38 +0800
Subject: [PATCH] feat(mge/device): enable to get cuda/cudnn/tensorrt version

GitOrigin-RevId: 5864c61d10972aa30923cf87e457e9819e8ea52d
---
 imperative/python/megengine/device.py | 33 +++++++++++++++++++++++++++++++++
 imperative/python/src/common.cpp      |  5 +++++
 src/core/impl/version.cpp             | 25 +++++++++++++++++++++++++
 src/core/include/megbrain/version.h   |  3 +++
 4 files changed, 66 insertions(+)

diff --git a/imperative/python/megengine/device.py b/imperative/python/megengine/device.py
index 3566dae8..fa4b5db4 100644
--- a/imperative/python/megengine/device.py
+++ b/imperative/python/megengine/device.py
@@ -4,7 +4,10 @@ import re
 from typing import Optional
 
 from .core._imperative_rt.common import CompNode, DeviceType
+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
+from .core._imperative_rt.common import get_tensorrt_version as _get_tensorrt_version
 from .core._imperative_rt.common import set_prealloc_config as _set_prealloc_config
 from .core._imperative_rt.common import what_is_xpu as _what_is_xpu
 from .core._imperative_rt.utils import _try_coalesce_all_free_memory
@@ -17,6 +20,9 @@ __all__ = [
     "get_mem_status_bytes",
     "get_cuda_compute_capability",
     "get_cuda_device_property",
+    "get_cuda_version",
+    "get_cudnn_version",
+    "get_tensorrt_version",
     "get_allocated_memory",
     "get_reserved_memory",
     "get_max_reserved_memory",
@@ -263,3 +269,30 @@ def coalesce_free_memory():
        * This function may do nothing if there are no chunks that can be freed.
     """
     return _try_coalesce_all_free_memory()
+
+
+def get_cuda_version():
+    r"""Gets the CUDA version used when compiling MegEngine.
+
+    Returns:
+        a version number, indicating `CUDA_VERSION_MAJOR * 1000 + CUDA_VERSION_MINOR * 10`.
+    """
+    return _get_cuda_version()
+
+
+def get_cudnn_version():
+    r"""Get the Cudnn version used when compiling MegEngine.
+
+    Returns:
+        a version number, indicating `CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL`.
+    """
+    return _get_cudnn_version()
+
+
+def get_tensorrt_version():
+    r"""Get the TensorRT version used when compiling MegEngine.
+
+    Returns:
+        a version number, indicating `NV_TENSORRT_MAJOR * 1000 + NV_TENSORRT_MINOR * 100 + NV_TENSORRT_PATCH`.
+    """
+    return _get_tensorrt_version()
diff --git a/imperative/python/src/common.cpp b/imperative/python/src/common.cpp
index 61fa8e0a..4edee731 100644
--- a/imperative/python/src/common.cpp
+++ b/imperative/python/src/common.cpp
@@ -8,6 +8,7 @@
 #include "megbrain/comp_node.h"
 #include "megbrain/graph.h"
 #include "megbrain/imperative/physical_tensor.h"
+#include "megbrain/version.h"
 #if MGB_ENABLE_OPR_MM
 #include "megbrain/opr/mm_handler.h"
 #endif
@@ -264,6 +265,10 @@ void init_common(py::module m) {
         return mge_gen_code;
     });
 
+    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("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 b49c95db..586854cf 100644
--- a/src/core/impl/version.cpp
+++ b/src/core/impl/version.cpp
@@ -26,4 +26,29 @@ Version mgb::get_version() {
 #endif
 }
 
+#if MGB_CUDA
+#include "NvInfer.h"
+#include "cuda.h"
+#include "cudnn.h"
+int mgb::get_cuda_version() {
+    return CUDA_VERSION;
+}
+int mgb::get_cudnn_version() {
+    return CUDNN_VERSION;
+}
+int mgb::get_tensorrt_version() {
+    return NV_TENSORRT_VERSION;
+}
+#else
+int mgb::get_cuda_version() {
+    return -1;
+}
+int mgb::get_cudnn_version() {
+    return -1;
+}
+int mgb::get_tensorrt_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 c3b21e7a..258708d2 100644
--- a/src/core/include/megbrain/version.h
+++ b/src/core/include/megbrain/version.h
@@ -20,6 +20,9 @@ struct Version {
 };
 
 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();
 }  // namespace mgb
 
 // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}