[IFRT] Change tsl::RCReference<xla::ifrt::DeviceList> to xla::ifrt::DeviceListRef

IFRT will introduce a wrapper around `tsl::RCReference<xla::ifrt::DeviceList>`
that is more concise and is also safer to use (e.g., equality and hash compares
the dereferenced value, not the `tsl::RCReference`.

This migration will happen in 3 steps:
1. Define `xla::ifrt::DeviceListRef` alias that is interchangeable with `tsl::RCReference<xla::ifrt::DeviceList>`.
2. Change `tsl::RCReference<xla::ifrt::DeviceList>` to `xla::ifrt::DeviceListRef` in IFRT API, implementations, and user code. (current step)
3. Introduce a real wrapper `xla::ifrt::DeviceListRef`, replacing the alias.

PiperOrigin-RevId: 730231976

Author: hyeontaek

xla/python/ifrt/array_impl_test_lib.cc

@@ -536,8 +536,7 @@ TEST(ArrayImplTest, CopyToSameDevices) {
 
 TEST(ArrayImplTest, CopyToDifferentDevice) {
   TF_ASSERT_OK_AND_ASSIGN(auto client, test_util::GetClient());
-  tsl::RCReference<DeviceList> devices =
-      client->MakeDeviceList(client->addressable_devices());
+  DeviceListRef devices = client->MakeDeviceList(client->addressable_devices());
 
   DType dtype(DType::kF32);
   Shape shape({2, 3});

xla/python/ifrt/basic_device_list.cc

@@ -37,17 +37,15 @@ namespace ifrt {
 
 char BasicDeviceList::ID = 0;
 
-tsl::RCReference<DeviceList> BasicDeviceList::Create(Devices devices) {
-  return tsl::MakeRef<BasicDeviceList>(std::move(devices));
+DeviceListRef BasicDeviceList::Create(Devices devices) {
+  return DeviceListRef(tsl::MakeRef<BasicDeviceList>(std::move(devices)));
 }
 
-tsl::RCReference<DeviceList> BasicDeviceList::Create(
-    absl::Span<Device* const> devices) {
+DeviceListRef BasicDeviceList::Create(absl::Span<Device* const> devices) {
   return Create(Devices(devices.begin(), devices.end()));
 }
 
-tsl::RCReference<DeviceList> BasicDeviceList::Create(
-    std::initializer_list<Device*> devices) {
+DeviceListRef BasicDeviceList::Create(std::initializer_list<Device*> devices) {
   return Create(Devices(devices.begin(), devices.end()));
 }
 

xla/python/ifrt/basic_device_list.h

@@ -54,10 +54,9 @@ class BasicDeviceList : public llvm::RTTIExtends<BasicDeviceList, DeviceList> {
   using Devices = absl::InlinedVector<Device*, kInlineDeviceSize>;
 
   // Constructor with a pre-populated `devices`.
-  static tsl::RCReference<DeviceList> Create(Devices devices);
-  static tsl::RCReference<DeviceList> Create(absl::Span<Device* const> devices);
-  static tsl::RCReference<DeviceList> Create(
-      std::initializer_list<Device*> devices);
+  static DeviceListRef Create(Devices devices);
+  static DeviceListRef Create(absl::Span<Device* const> devices);
+  static DeviceListRef Create(std::initializer_list<Device*> devices);
 
   ~BasicDeviceList() override = default;
 
@@ -95,7 +94,7 @@ class BasicDeviceList : public llvm::RTTIExtends<BasicDeviceList, DeviceList> {
   struct AddressableDeviceListCache {
     absl::once_flag once_flag;
     DeviceList* device_list = nullptr;
-    tsl::RCReference<DeviceList> device_list_holder;
+    DeviceListRef device_list_holder;
   };
   mutable AddressableDeviceListCache addressable_device_list_cache_;
 

xla/python/ifrt/client.h

@@ -153,7 +153,7 @@ class Client : public llvm::RTTIExtends<Client, llvm::RTTIRoot> {
   // device.
   virtual absl::StatusOr<std::vector<tsl::RCReference<Array>>> CopyArrays(
       absl::Span<tsl::RCReference<Array>> arrays,
-      std::optional<tsl::RCReference<DeviceList>> devices,
+      std::optional<DeviceListRef> devices,
       std::optional<MemoryKind> memory_kind, ArrayCopySemantics semantics) = 0;
 
   // Remaps shards across input `Array`s to create new `Array`s based on `plan`.
@@ -237,7 +237,7 @@ class Client : public llvm::RTTIExtends<Client, llvm::RTTIRoot> {
       int local_hardware_id) const = 0;
 
   // Creates a device list from the given list of devices.
-  virtual tsl::RCReference<DeviceList> MakeDeviceList(
+  virtual DeviceListRef MakeDeviceList(
       absl::Span<Device* const> devices) const = 0;
 
   // TODO(hyeontaek): Potentially remove this method to encourage supporting
@@ -246,7 +246,7 @@ class Client : public llvm::RTTIExtends<Client, llvm::RTTIRoot> {
 
   // Returns a topology that covers the provided devices.
   virtual absl::StatusOr<std::shared_ptr<Topology>> GetTopologyForDevices(
-      const tsl::RCReference<DeviceList>& devices) const = 0;
+      const DeviceListRef& devices) const = 0;
 
   // Returns the default layout on `device` with `memory_kind` for a buffer with
   // `dtype` and single-shard dimensions `dims`.

xla/python/ifrt/custom_call_program.h

@@ -37,8 +37,7 @@ struct CustomCallProgram
   // Specification for a single array. The sharding of all input and output
   // specs must use only the devices in `devices`.
   CustomCallProgram(std::string type, std::string name,
-                    absl::Cord serialized_program_text,
-                    tsl::RCReference<DeviceList> devices,
+                    absl::Cord serialized_program_text, DeviceListRef devices,
                     std::vector<ArraySpec> input_specs,
                     std::vector<ArraySpec> output_specs)
       : type(std::move(type)),
@@ -62,7 +61,7 @@ struct CustomCallProgram
   absl::Cord serialized_program_text;
 
   // List of devices to compile and run the custom call program on.
-  tsl::RCReference<DeviceList> devices;
+  DeviceListRef devices;
 
   // Specification for input and output arrays. The custom call program must
   // expect to receive input arrays and return output arrays both following the

xla/python/ifrt/custom_call_program_serdes.cc

@@ -81,7 +81,7 @@ class CustomCallProgramSerDes
           "Failed to parse serialized CustomCallProgramProto");
     }
     TF_ASSIGN_OR_RETURN(
-        tsl::RCReference<DeviceList> devices,
+        DeviceListRef devices,
         DeviceList::FromProto(deserialize_program_options->client,
                               proto.devices()));
     std::vector<ArraySpec> input_specs;

xla/python/ifrt/custom_call_program_serdes_test.cc

@@ -51,7 +51,7 @@ class CustomCallProgramSerDesTest : public test_util::DeviceTest {};
 TEST_P(CustomCallProgramSerDesTest, RoundTrip) {
   Shape shape0({10, 20});
   Shape shard_shape0({5, 20});
-  tsl::RCReference<DeviceList> devices = GetDevices({0, 1});
+  DeviceListRef devices = GetDevices({0, 1});
   std::shared_ptr<const Sharding> sharding0 =
       ConcreteEvenSharding::Create(devices, MemoryKind(),
                                    /*shape=*/shape0,

xla/python/ifrt/device_list.cc

@@ -31,7 +31,7 @@ namespace ifrt {
 
 char DeviceList::ID = 0;
 
-absl::StatusOr<tsl::RCReference<DeviceList>> DeviceList::FromProto(
+absl::StatusOr<DeviceListRef> DeviceList::FromProto(
     xla::ifrt::Client* client, const DeviceListProto& proto) {
   absl::InlinedVector<Device*, 1> devices;
   devices.reserve(proto.device_ids_size());
@@ -52,8 +52,7 @@ DeviceListProto DeviceList::ToProto() const {
   return proto;
 }
 
-std::vector<DeviceId> GetDeviceIds(
-    const tsl::RCReference<DeviceList>& device_list) {
+std::vector<DeviceId> GetDeviceIds(const DeviceListRef& device_list) {
   std::vector<DeviceId> ids;
   ids.reserve(device_list->devices().size());
   for (const Device* device : device_list->devices()) {

xla/python/ifrt/device_list.h

@@ -104,8 +104,7 @@ class DeviceList : public tsl::ReferenceCounted<DeviceList>,
 using DeviceListRef = tsl::RCReference<DeviceList>;
 
 // Returns the id of each device in `device_list`.
-std::vector<DeviceId> GetDeviceIds(
-    const tsl::RCReference<DeviceList>& device_list);
+std::vector<DeviceId> GetDeviceIds(const DeviceListRef& device_list);
 
 // Hash function for `DeviceList`. Assumes that every unique device has a unique
 // `Device` object, not duplicate `Device` objects ("d1 == d2 if d1->id() ==

xla/python/ifrt/device_test_util.cc

@@ -132,13 +132,11 @@ std::shared_ptr<MockClient> MakeDeviceTestClient(int num_devices,
         return it->second.get();
       });
   ON_CALL(*client, MakeDeviceList)
-      .WillByDefault([](absl::Span<Device* const> devices)
-                         -> tsl::RCReference<DeviceList> {
+      .WillByDefault([](absl::Span<Device* const> devices) -> DeviceListRef {
         return BasicDeviceList::Create(devices);
       });
   ON_CALL(*client, GetTopologyForDevices)
-      .WillByDefault(
-          [](const tsl::RCReference<DeviceList>&) { return nullptr; });
+      .WillByDefault([](const DeviceListRef&) { return nullptr; });
   return client;
 }
 
@@ -149,12 +147,11 @@ void DeviceTest::SetUp() {
   client_ = MakeDeviceTestClient(num_devices, num_addressable_devices);
 }
 
-tsl::RCReference<DeviceList> DeviceTest::GetDevices(
-    absl::Span<const int> device_indices) {
+DeviceListRef DeviceTest::GetDevices(absl::Span<const int> device_indices) {
   return test_util::GetDevices(client_.get(), device_indices).value();
 }
 
-tsl::RCReference<DeviceList> DeviceTest::GetAddressableDevices(
+DeviceListRef DeviceTest::GetAddressableDevices(
     absl::Span<const int> device_indices) {
   return test_util::GetAddressableDevices(client_.get(), device_indices)
       .value();

xla/python/ifrt/device_test_util.h

@@ -46,13 +46,12 @@ class DeviceTest : public testing::TestWithParam<DeviceTestParam> {
   // Returns `DeviceList` containing devices at given indexes (not ids) within
   // `client.devices()`.
   // REQUIRES: 0 <= device_indices[i] < num_devices
-  tsl::RCReference<DeviceList> GetDevices(absl::Span<const int> device_indices);
+  DeviceListRef GetDevices(absl::Span<const int> device_indices);
 
   // Returns `DeviceList` containing devices at given indexes (not ids) within
   // `client.addressable_devices()`.
   // REQUIRES: 0 <= device_indices[i] < num_addressable_devices
-  tsl::RCReference<DeviceList> GetAddressableDevices(
-      absl::Span<const int> device_indices);
+  DeviceListRef GetAddressableDevices(absl::Span<const int> device_indices);
 
  private:
   std::shared_ptr<Client> client_;

xla/python/ifrt/executable.h

@@ -234,7 +234,7 @@ class LoadedExecutable
   // API).
   virtual absl::StatusOr<ExecuteResult> Execute(
       absl::Span<tsl::RCReference<Array>> args, const ExecuteOptions& options,
-      std::optional<tsl::RCReference<DeviceList>> devices) = 0;
+      std::optional<DeviceListRef> devices) = 0;
 
   // Deletes the executable from the devices. The operation may be asynchronous.
   // The returned future will have the result of the deletion on the devices.

xla/python/ifrt/mock.cc

@@ -155,7 +155,7 @@ MockClient::MockClient(std::unique_ptr<xla::ifrt::Client> delegated)
           });
   ON_CALL(*this, CopyArrays)
       .WillByDefault([this](absl::Span<tsl::RCReference<Array>> arrays,
-                            std::optional<tsl::RCReference<DeviceList>> devices,
+                            std::optional<DeviceListRef> devices,
                             std::optional<MemoryKind> memory_kind,
                             ArrayCopySemantics semantics) {
         return delegated_->CopyArrays(arrays, std::move(devices), memory_kind,
@@ -229,10 +229,9 @@ MockClient::MockClient(std::unique_ptr<xla::ifrt::Client> delegated)
     return delegated_->GetDefaultCompiler();
   });
   ON_CALL(*this, GetTopologyForDevices)
-      .WillByDefault(
-          [this](const tsl::RCReference<xla::ifrt::DeviceList>& devices) {
-            return delegated_->GetTopologyForDevices(devices);
-          });
+      .WillByDefault([this](const DeviceListRef& devices) {
+        return delegated_->GetTopologyForDevices(devices);
+      });
   ON_CALL(*this, GetDefaultLayout)
       .WillByDefault([this](xla::ifrt::DType dtype,
                             absl::Span<const int64_t> dims,

xla/python/ifrt/mock.h

@@ -148,7 +148,7 @@ class MockClient : public llvm::RTTIExtends<MockClient, Client> {
               (final));
   MOCK_METHOD(absl::StatusOr<std::vector<tsl::RCReference<Array>>>, CopyArrays,
               (absl::Span<tsl::RCReference<Array>> arrays,
-               std::optional<tsl::RCReference<DeviceList>> devices,
+               std::optional<DeviceListRef> devices,
                std::optional<MemoryKind> memory_kind,
                ArrayCopySemantics semantics),
               (final));
@@ -180,12 +180,11 @@ class MockClient : public llvm::RTTIExtends<MockClient, Client> {
               (const, final));
   MOCK_METHOD(absl::StatusOr<Device*>, LookupAddressableDevice,
               (int local_hardware_id), (const, final));
-  MOCK_METHOD(tsl::RCReference<DeviceList>, MakeDeviceList,
+  MOCK_METHOD(DeviceListRef, MakeDeviceList,
               (absl::Span<Device* const> devices), (const));
   MOCK_METHOD(Compiler*, GetDefaultCompiler, (), (final));
   MOCK_METHOD(absl::StatusOr<std::shared_ptr<Topology>>, GetTopologyForDevices,
-              (const tsl::RCReference<xla::ifrt::DeviceList>& devices),
-              (const, final));
+              (const xla::ifrt::DeviceListRef& devices), (const, final));
   MOCK_METHOD(absl::StatusOr<std::shared_ptr<const PjRtLayout>>,
               GetDefaultLayout,
               (xla::ifrt::DType dtype, absl::Span<const int64_t> dims,
@@ -319,7 +318,7 @@ class MockLoadedExecutable
   MOCK_METHOD(absl::StatusOr<ExecuteResult>, Execute,
               (absl::Span<tsl::RCReference<Array>> args,
                const ExecuteOptions& options,
-               std::optional<tsl::RCReference<DeviceList>> devices),
+               std::optional<DeviceListRef> devices),
               (final));
   MOCK_METHOD(Future<>, Delete, (), (final));
   MOCK_METHOD(bool, IsDeleted, (), (const, final));
@@ -357,7 +356,7 @@ class MockSharding : public llvm::RTTIExtends<MockSharding, Sharding> {
             BasicDeviceList::Create({}), MemoryKind(),
             /*is_fully_replicated=*/false) {}
 
-  MockSharding(tsl::RCReference<DeviceList> devices, MemoryKind memory_kind,
+  MockSharding(DeviceListRef devices, MemoryKind memory_kind,
                bool is_fully_replicated)
       : llvm::RTTIExtends<MockSharding, Sharding>(devices, memory_kind,
                                                   is_fully_replicated) {}
@@ -395,7 +394,7 @@ class MockSharding : public llvm::RTTIExtends<MockSharding, Sharding> {
   MOCK_METHOD(bool, HasSamePartitioning, (const Sharding& other),
               (const final));
   MOCK_METHOD(absl::StatusOr<std::unique_ptr<Sharding>>, WithDeviceAssignment,
-              (std::optional<tsl::RCReference<DeviceList>> devices,
+              (std::optional<DeviceListRef> devices,
                std::optional<MemoryKind> memory_kind),
               (const final));
   MOCK_METHOD(void, Hash, (absl::HashState), (const final));

xla/python/ifrt/remap_impl_test_lib.cc

@@ -69,7 +69,7 @@ absl::StatusOr<ArraySpec> CreateArraySpec(Client* client,
                                           absl::Span<const int> device_indices,
                                           Shape shard_shape = Shape({2, 3}),
                                           DType dtype = DType(DType::kS32)) {
-  TF_ASSIGN_OR_RETURN(tsl::RCReference<DeviceList> device_list,
+  TF_ASSIGN_OR_RETURN(DeviceListRef device_list,
                       test_util::GetAddressableDevices(client, device_indices));
   TF_ASSIGN_OR_RETURN(Shape shape,
                       GetShape(device_indices.size(), shard_shape));

xla/python/ifrt/remap_plan_test.cc

@@ -70,7 +70,7 @@ TEST_P(RemapPlanTest, ToFromProto) {
 
   Shape shape({20, 20});
   Shape shard_shape({5, 20});
-  tsl::RCReference<DeviceList> devices = GetDevices({0, 1, 2, 3});
+  DeviceListRef devices = GetDevices({0, 1, 2, 3});
   std::shared_ptr<const Sharding> sharding =
       ConcreteEvenSharding::Create(devices, MemoryKind(), /*shape=*/shape,
                                    /*shard_shape=*/shard_shape);