Refactor GetNumDevices and create GetNumGlobalDevices (#9184)

Author: pgmoka

torch_xla/csrc/init_python_bindings.cpp

@@ -1487,9 +1487,12 @@ void InitXlaModuleBindings(py::module m) {
     if (UseVirtualDevice()) {
       return 1;
     } else {
-      return runtime::GetComputationClient()->GetNumDevices();
+      return runtime::GetComputationClient()->GetNumLocalDevices();
     }
   });
+  m.def("_xla_num_global_devices", []() -> int64_t {
+    return runtime::GetComputationClient()->GetNumDevices();
+  });
   m.def("_xla_get_all_devices", []() {
     std::vector<std::string> all_devices =
         runtime::GetComputationClient()->GetAllDevices();
@@ -1505,7 +1508,7 @@ void InitXlaModuleBindings(py::module m) {
   m.def("_xla_get_runtime_devices",
         []() { return runtime::GetComputationClient()->GetLocalDevices(); });
   m.def("_xla_num_runtime_devices", []() -> int64_t {
-    return runtime::GetComputationClient()->GetNumDevices();
+    return runtime::GetComputationClient()->GetNumLocalDevices();
   });
   m.def("_xla_get_all_runtime_devices", []() {
     std::vector<std::string> all_devices =

torch_xla/csrc/runtime/computation_client.h

@@ -374,6 +374,8 @@ class ComputationClient {
 
   virtual std::intptr_t GetCudaStreamForDevice(int local_device_id) const = 0;
 
+  virtual size_t GetNumLocalDevices() const = 0;
+
   virtual size_t GetNumDevices() const = 0;
 
   virtual std::vector<std::string> GetLocalDevices() const = 0;

torch_xla/csrc/runtime/ifrt_computation_client.cpp

@@ -622,10 +622,14 @@ IfrtComputationClient::ExecuteReplicated(
   return data_handles;
 }
 
-size_t IfrtComputationClient::GetNumDevices() const {
+size_t IfrtComputationClient::GetNumLocalDevices() const {
   return client_->addressable_device_count();
 }
 
+size_t IfrtComputationClient::GetNumDevices() const {
+  return client_->device_count();
+}
+
 std::string IfrtComputationClient::GetDefaultDevice() const {
   return IfrtDeviceToString(client_->addressable_devices()[0]);
 }

torch_xla/csrc/runtime/ifrt_computation_client.h

@@ -79,6 +79,8 @@ class IfrtComputationClient : public ComputationClient {
       absl::Span<const std::string> devices,
       const ExecuteReplicatedOptions& options) override;
 
+  size_t GetNumLocalDevices() const override;
+
   size_t GetNumDevices() const override;
 
   std::string GetDefaultDevice() const override;

torch_xla/csrc/runtime/pjrt_computation_client.cpp

@@ -934,10 +934,14 @@ PjRtComputationClient::ExecuteReplicated(
   return data_handles;
 }
 
-size_t PjRtComputationClient::GetNumDevices() const {
+size_t PjRtComputationClient::GetNumLocalDevices() const {
   return client_->addressable_device_count();
 }
 
+size_t PjRtComputationClient::GetNumDevices() const {
+  return client_->device_count();
+}
+
 std::string PjRtComputationClient::GetDefaultDevice() const {
   return PjRtDeviceToString(client_->addressable_devices()[0]);
 }

torch_xla/csrc/runtime/pjrt_computation_client.h

@@ -86,6 +86,8 @@ class PjRtComputationClient : public ComputationClient {
       absl::Span<const std::string> devices,
       const ExecuteReplicatedOptions& options) override;
 
+  size_t GetNumLocalDevices() const override;
+
   size_t GetNumDevices() const override;
 
   std::string GetDefaultDevice() const override;

torch_xla/csrc/tensor_impl.cpp

@@ -57,7 +57,7 @@ struct XLAGuardImpl : public c10::impl::DeviceGuardImplInterface {
       return 0;
     }
 
-    return client->GetNumDevices();
+    return client->GetNumLocalDevices();
   }
 };