Make local backend default for GPU. (#1109)

Author: pschuh

Sources/x10/xla_client/local_device.cc

@@ -539,16 +539,20 @@ std::unique_ptr<ComputationClient::Device> MakeLocalDeviceFromClient(
 std::vector<std::unique_ptr<ComputationClient::Device>>
 GetAllLocalDevicesForPlatform(const char* platform_name,
                               const char* device_prefix) {
+  auto platform = xla::PlatformUtil::GetPlatform(platform_name);
+  if (!platform.ok()) return {};
   xla::LocalClientOptions options;
-  options.set_platform(
-      xla::PlatformUtil::GetPlatform(platform_name).ValueOrDie());
-  xla::LocalClient* client =
-      xla::ClientLibrary::GetOrCreateLocalClient(options).ValueOrDie();
+  options.set_platform(platform.ValueOrDie());
+  auto local_client_statusor =
+      xla::ClientLibrary::GetOrCreateLocalClient(options);
+  if (!local_client_statusor.ok()) return {};
+  xla::LocalClient* client = local_client_statusor.ValueOrDie();
   std::vector<std::unique_ptr<ComputationClient::Device>> devices;
   devices.reserve(client->device_count());
   for (int i = 0; i < client->device_count(); ++i) {
     devices.push_back(MakeLocalDeviceFromClient(
-        absl::StrCat(device_prefix, ":", i), client, i, i, true));
+        absl::StrCat(device_prefix, ":", i), client, i, i,
+            std::string(platform_name) == "cpu"));
   }
   return devices;
 }

Sources/x10/xla_client/xrt_computation_client.cc

@@ -35,6 +35,7 @@
 #include "tensorflow/compiler/xla/xla_client/util.h"
 #include "tensorflow/compiler/xla/xla_client/xla_util.h"
 #include "tensorflow/compiler/xla/xla_client/xrt_local_service.h"
+#include "tensorflow/compiler/xla/xla_client/local_device.h"
 #include "tensorflow/compiler/xrt/xrt_util.h"
 #include "tensorflow/cc/ops/const_op.h"
 #include "tensorflow/compiler/xla/shape_util.h"
@@ -534,7 +535,7 @@ bool ParseEnvDeviceCounts(XrtComputationClient::Options* options) {
 }
 
 bool ParseEnvDevices(XrtComputationClient::Options* options) {
-  std::string device = GpuIsAvailable() ? "GPU" : "CPU";
+  std::string device = "CPU";
   std::string default_device_spec = absl::StrFormat(
       "%s:0;/job:localservice/replica:0/task:0/device:XLA_%s:0", device,
       device);
@@ -629,6 +630,11 @@ XrtComputationClient::XrtComputationClient(
   for (const auto& dev_target : options_.global_device_map) {
     AddDevice(std::make_unique<XrtDevice>(dev_target.first, this));
   }
+
+  for (auto& device : GetAllLocalDevicesForPlatform("gpu", "GPU")) {
+    options_.default_device = "GPU:0";
+    AddDevice(std::move(device));
+  }
 }
 
 std::vector<size_t> XrtComputationClient::PartitionTransferToServer(

Sources/x10/xla_client/xrt_local_service.cc

@@ -67,6 +67,8 @@ XrtLocalService::XrtLocalService(const std::string& cluster_spec,
                                  const std::string& job_name, int task_index) {
   tensorflow::ServerDef server_def;
   FillServerDef(cluster_spec, job_name, task_index, &server_def);
+  (*server_def.mutable_default_session_config()
+      ->mutable_device_count())["GPU"] = 0;
   TF_CHECK_OK(tensorflow::NewServer(server_def, &server_));
 }