Merge 45bde08 into 93e7c19

BVLC · Feb 18, 2016 · 1048576 · 1048576
2 parents 93e7c19 + 45bde08
commit 1048576
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/include/caffe/common.hpp b/include/caffe/common.hpp
@@ -153,6 +153,10 @@ class Caffe {
   static void SetDevice(const int device_id);
   // Prints the current GPU status.
   static void DeviceQuery();
+  // Check if specified device is available
+  static bool CheckDevice(const int device_id);
+  // Get the first available device id since start_id
+  static int GrabDevice(const int start_id = 0);
   // Parallel training info
   inline static int solver_count() { return Get().solver_count_; }
   inline static void set_solver_count(int val) { Get().solver_count_ = val; }

diff --git a/src/caffe/common.cpp b/src/caffe/common.cpp
@@ -70,6 +70,15 @@ void Caffe::DeviceQuery() {
   NO_GPU;
 }
 
+bool Caffe::CheckDevice(const int device_id) {
+  NO_GPU;
+  return false;
+}
+
+int Caffe::GrabDevice(const int start_id) {
+  NO_GPU;
+  return -1;
+}
 
 class Caffe::RNG::Generator {
  public:
@@ -192,6 +201,26 @@ void Caffe::DeviceQuery() {
   return;
 }
 
+bool Caffe::CheckDevice(const int device_id) {
+  bool r =  ((cudaSuccess == cudaSetDevice(device_id)) &&
+            (cudaSuccess == cudaFree(0)));
+  cudaGetLastError();
+  return r;
+}
+
+int Caffe::GrabDevice(const int start_id) {
+  int count = 0;
+  int r = -1;
+  if (cudaSuccess == cudaGetDeviceCount(&count)) {
+    for (int i = start_id; i < count; i++) {
+      if (CheckDevice(i)) {
+        r = i;
+        break;
+      }
+    }  // for
+  }
+  return r;
+}
 
 class Caffe::RNG::Generator {
  public: