Obtain compile-time cuda version

Summary: Pull Request resolved: #86

Differential Revision: D17947428

Pulled By: ppwwyyxx

fbshipit-source-id: f7d2c9b96a1478d2d7b9438b15ba942db0e4d06b

detectron2/layers/csrc/cuda_version.cu

@@ -0,0 +1,7 @@
+#include <cuda_runtime_api.h>
+
+namespace detectron2 {
+int get_cudart_version() {
+  return CUDART_VERSION;
+}
+} // namespace detectron2

detectron2/layers/csrc/vision.cpp

@@ -1,4 +1,5 @@
 // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
 #include "ROIAlign/ROIAlign.h"
 #include "ROIAlignRotated/ROIAlignRotated.h"
 #include "box_iou_rotated/box_iou_rotated.h"
@@ -7,6 +8,29 @@
 
 namespace detectron2 {
 
+#ifdef WITH_CUDA
+extern int get_cudart_version();
+#endif
+
+std::string get_cuda_version() {
+#ifdef WITH_CUDA
+  std::ostringstream oss;
+
+  // copied from
+  // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/cuda/detail/CUDAHooks.cpp#L231
+  auto printCudaStyleVersion = [&](int v) {
+    oss << (v / 1000) << "." << (v / 10 % 100);
+    if (v % 10 != 0) {
+      oss << "." << (v % 10);
+    }
+  };
+  printCudaStyleVersion(get_cudart_version());
+  return oss.str();
+#else
+  return std::string("not available");
+#endif
+}
+
 // similar to
 // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Version.cpp
 std::string get_compiler_version() {
@@ -32,6 +56,7 @@ std::string get_compiler_version() {
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("get_compiler_version", &get_compiler_version, "get_compiler_version");
+  m.def("get_cuda_version", &get_cuda_version, "get_cuda_version");
 
   m.def("box_iou_rotated", &box_iou_rotated, "IoU for rotated boxes");
 

detectron2/utils/collect_env.py

@@ -40,6 +40,7 @@ def collect_env_info():
         data.append(("detectron2._C", "failed to import"))
     else:
         data.append(("Detectron2 Compiler", _C.get_compiler_version()))
+        data.append(("Detectron2 CUDA Compiler", _C.get_cuda_version()))
 
     data.append(get_env_module())
     data.append(("PyTorch", torch.__version__))

setup.py

@@ -17,7 +17,9 @@ def get_extensions():
 
     main_source = os.path.join(extensions_dir, "vision.cpp")
     sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp"))
-    source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu"))
+    source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu")) + glob.glob(
+        os.path.join(extensions_dir, "*.cu")
+    )
 
     sources = [main_source] + sources