Add init interface for customize devices.

paddle/fluid/framework/init.cc

@@ -26,11 +26,29 @@ limitations under the License. */
 namespace paddle {
 namespace framework {
 
+DEFINE_string(devices, "", "The devices to be used.");
+DEFINE_bool(init_p2p, true, "Whether to init p2p.");
+
 std::once_flag gflags_init_flag;
 std::once_flag p2p_init_flag;
 
+using paddle::platform::DeviceContextPool;
+
+void Init(std::vector<std::string> argv) {
+  InitGflags(argv);
+  // init devices
+  std::vector<int> devices;
+  std::string token;
+  std::istringstream tokenStream(FLAGS_devices);
+  while (std::getline(tokenStream, token, ',')) {
+    devices.push_back(std::stoi(token));
+  }
+  InitDevices(FLAGS_init_p2p, devices);
+}
+
 void InitGflags(std::vector<std::string> argv) {
   std::call_once(gflags_init_flag, [&]() {
+    argv.push_back("dummy");
     int argc = argv.size();
     char **arr = new char *[argv.size()];
     std::string line;
@@ -65,13 +83,52 @@ void InitP2P(int count) {
 #endif
 }
 
+void InitP2P(std::vector<int> devices) {
+#ifdef PADDLE_WITH_CUDA
+  std::call_once(p2p_init_flag, [&]() {
+    int count = devices.size();
+    for (int i = 0; i < count; ++i) {
+      for (int j = 0; j < count; ++j) {
+        if (devices[i] == devices[j]) continue;
+        int can_acess = -1;
+        PADDLE_ENFORCE(
+            cudaDeviceCanAccessPeer(&can_acess, devices[i], devices[j]),
+            "Failed to test P2P access.");
+        if (can_acess != 1) {
+          LOG(WARNING) << "Cannot enable P2P access from " << devices[i]
+                       << " to " << devices[j];
+        } else {
+          cudaSetDevice(devices[i]);
+          cudaDeviceEnablePeerAccess(devices[j], 0);
+        }
+      }
+    }
+  });
+#endif
+}
+
 void InitDevices(bool init_p2p) {
-  /*Init all available devices by default */
+/*Init all available devices by default */
+#ifdef PADDLE_WITH_CUDA
+  std::vector<int> devices;
+  try {
+    int count = platform::GetCUDADeviceCount();
+    for (int i = 0; i < count; ++i) {
+      devices.push_back(i);
+    }
+  } catch (const std::exception &exp) {
+    LOG(WARNING) << "Compiled with WITH_GPU, but no GPU found in runtime.";
+  }
+#else
+  LOG(WARNING)
+      << "'CUDA' is not supported, Please re-compile with WITH_GPU option";
+#endif
+  InitDevices(init_p2p, devices);
+}
 
+void InitDevices(bool init_p2p, const std::vector<int> devices) {
   std::vector<platform::Place> places;
-  places.emplace_back(platform::CPUPlace());
   int count = 0;
-
 #ifdef PADDLE_WITH_CUDA
   try {
     count = platform::GetCUDADeviceCount();
@@ -83,12 +140,17 @@ void InitDevices(bool init_p2p) {
       << "'CUDA' is not supported, Please re-compile with WITH_GPU option";
 #endif
 
-  for (int i = 0; i < count; ++i) {
-    places.emplace_back(platform::CUDAPlace(i));
+  for (size_t i = 0; i < devices.size(); ++i) {
+    if (devices[i] >= count || devices[i] < 0) {
+      LOG(WARNING) << "Invalid devices id.";
+      continue;
+    }
+    places.emplace_back(platform::CUDAPlace(devices[i]));
   }
   if (init_p2p) {
-    InitP2P(count);
+    InitP2P(devices);
   }
+  places.emplace_back(platform::CPUPlace());
   platform::DeviceContextPool::Init(places);
 }
 

paddle/fluid/framework/init.h

@@ -22,11 +22,15 @@ limitations under the License. */
 namespace paddle {
 namespace framework {
 
+void Init(std::vector<std::string> argv);
+
 void InitGflags(std::vector<std::string> argv);
 
 void InitGLOG(const std::string &prog_name);
 
 void InitDevices(bool init_p2p);
 
+void InitDevices(bool init_p2p, const std::vector<int> devices);
+
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/inference/io.cc

@@ -24,9 +24,7 @@ limitations under the License. */
 namespace paddle {
 namespace inference {
 
-// Temporarily add this function for exposing framework::InitDevices() when
-// linking the inference shared library.
-void Init(bool init_p2p) { framework::InitDevices(init_p2p); }
+void Init(const std::vector<std::string> argv) { framework::Init(argv); }
 
 void ReadBinaryFile(const std::string& filename, std::string* contents) {
   std::ifstream fin(filename, std::ios::in | std::ios::binary);

paddle/fluid/inference/io.h

@@ -25,7 +25,7 @@ limitations under the License. */
 namespace paddle {
 namespace inference {
 
-void Init(bool init_p2p);
+void Init(const std::vector<std::string> argv);
 
 void LoadPersistables(framework::Executor* executor, framework::Scope* scope,
                       const framework::ProgramDesc& main_program,