Merge pull request #3590 from junshi15/GPUUtilities

Add functions to check and grab GPU

include/caffe/common.hpp

@@ -153,6 +153,11 @@ 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);
+  // Search from start_id to the highest possible device ordinal,
+  // return the ordinal of the first available device.
+  static int FindDevice(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; }

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::FindDevice(const int start_id) {
+  NO_GPU;
+  return -1;
+}
 
 class Caffe::RNG::Generator {
  public:
@@ -192,6 +201,39 @@ void Caffe::DeviceQuery() {
   return;
 }
 
+bool Caffe::CheckDevice(const int device_id) {
+  // This function checks the availability of GPU #device_id.
+  // It attempts to create a context on the device by calling cudaFree(0).
+  // cudaSetDevice() alone is not sufficient to check the availability.
+  // It lazily records device_id, however, does not initialize a
+  // context. So it does not know if the host thread has the permission to use
+  // the device or not.
+  //
+  // In a shared environment where the devices are set to EXCLUSIVE_PROCESS
+  // or EXCLUSIVE_THREAD mode, cudaSetDevice() returns cudaSuccess
+  // even if the device is exclusively occupied by another process or thread.
+  // Cuda operations that initialize the context are needed to check
+  // the permission. cudaFree(0) is one of those with no side effect,
+  // except the context initialization.
+  bool r = ((cudaSuccess == cudaSetDevice(device_id)) &&
+            (cudaSuccess == cudaFree(0)));
+  // reset any error that may have occurred.
+  cudaGetLastError();
+  return r;
+}
+
+int Caffe::FindDevice(const int start_id) {
+  // This function finds the first available device by checking devices with
+  // ordinal from start_id to the highest available value. In the
+  // EXCLUSIVE_PROCESS or EXCLUSIVE_THREAD mode, if it succeeds, it also
+  // claims the device due to the initialization of the context.
+  int count = 0;
+  CUDA_CHECK(cudaGetDeviceCount(&count));
+  for (int i = start_id; i < count; i++) {
+    if (CheckDevice(i)) return i;
+  }
+  return -1;
+}
 
 class Caffe::RNG::Generator {
  public: