refactor: remove BasicDevicesList

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -82,7 +82,6 @@
 // IFRT
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/attribute_map.h"
-#include "xla/python/ifrt/basic_device_list.h"
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt/compiler.h"
 #include "xla/python/ifrt/device.h"
@@ -1800,26 +1799,10 @@ extern "C" ifrt::Client *ifrt_DeviceToClient(ifrt::Device *device) {
   return device->client();
 }
 
-extern "C" HeldValue<tsl::RCReference<ifrt::DeviceList>> *
-ifrt_CreateBasicDeviceListFromDevices(ifrt::Device **device_list,
-                                      int32_t num_devices) {
+tsl::RCReference<ifrt::DeviceList> ifrt_CreateDeviceListFromDevices(
+    ifrt::Client *client, ifrt::Device **device_list, int32_t num_devices) {
   absl::Span<ifrt::Device *const> devices(device_list, num_devices);
-  return reactant::capture(ifrt::BasicDeviceList::Create(devices));
-}
-
-extern "C" const char *ifrt_BasicDeviceListToString(
-    HeldValue<tsl::RCReference<ifrt::DeviceList>> *device_list) {
-  return cstr_from_string(device_list->obj()->DebugString());
-}
-
-extern "C" int ifrt_BasicDeviceListSize(
-    HeldValue<tsl::RCReference<ifrt::DeviceList>> *device_list) {
-  return device_list->obj()->size();
-}
-
-extern "C" ifrt::Device *const ifrt_BasicDeviceListGetDevice(
-    HeldValue<tsl::RCReference<ifrt::DeviceList>> *device_list, int32_t index) {
-  return device_list->obj()->devices()[index];
+  return client->MakeDeviceList(devices);
 }
 
 extern "C" ifrt::Memory *ifrt_DeviceGetDefaultMemory(ifrt::Device *device) {
@@ -1889,10 +1872,11 @@ hlo_sharding_to_string(const xla::HloSharding *hlo_sharding) {
 }
 
 extern "C" ifrt::HloSharding *ifrt_hlo_sharding_from_xla_hlo_sharding(
-    HeldValue<tsl::RCReference<ifrt::DeviceList>> *device_list,
+    ifrt::Client *client, ifrt::Device **device_list, int32_t num_devices,
     ifrt::MemoryKind *memory_kind, xla::HloSharding *xla_hlo_sharding) {
-  return ifrt::HloSharding::Create(device_list->obj(), *memory_kind,
-                                   *xla_hlo_sharding)
+  return ifrt::HloSharding::Create(
+             ifrt_CreateDeviceListFromDevices(client, device_list, num_devices),
+             *memory_kind, *xla_hlo_sharding)
       .release();
 }
 
@@ -1919,12 +1903,13 @@ ifrt_sharding_from_ifrt_hlo_sharding(ifrt::HloSharding *hlo_sharding) {
 }
 
 extern "C" HeldValue<std::shared_ptr<ifrt::Sharding>> *
-ifrt_sharding_from_hlo_sharding(
-    HeldValue<tsl::RCReference<ifrt::DeviceList>> *device_list,
-    ifrt::MemoryKind *memory_kind, xla::HloSharding *xla_hlo_sharding) {
+ifrt_sharding_from_hlo_sharding(ifrt::Client *client,
+                                ifrt::Device **device_list, int32_t num_devices,
+                                ifrt::MemoryKind *memory_kind,
+                                xla::HloSharding *xla_hlo_sharding) {
   return ifrt_sharding_from_ifrt_hlo_sharding(
-      ifrt_hlo_sharding_from_xla_hlo_sharding(device_list, memory_kind,
-                                              xla_hlo_sharding));
+      ifrt_hlo_sharding_from_xla_hlo_sharding(client, device_list, num_devices,
+                                              memory_kind, xla_hlo_sharding));
 }
 
 extern "C" const char *

src/xla/IFRT/Device.jl

@@ -55,52 +55,6 @@ function XLA.memories(device::Device)
     return memories
 end
 
-# Device List
-## TODO: This is semi-deprecated in openxla. At some point we want to just replace this with
-##       a simple vector of devices
-struct BasicDeviceList <: AbstractVector{Device}
-    ptr::Ptr{Cvoid}
-
-    function BasicDeviceList(devices::AbstractVector{Device})
-        GC.@preserve devices begin
-            ptr = @ccall MLIR.API.mlir_c.ifrt_CreateBasicDeviceListFromDevices(
-                [d.device for d in devices]::Ptr{Ptr{Cvoid}}, length(devices)::Int32
-            )::Ptr{Cvoid}
-        end
-        return new(ptr)
-    end
-end
-
-function Base.getindex(device_list::BasicDeviceList, index::Integer)
-    if !(1 ≤ index ≤ length(device_list))
-        throw(BoundsError(device_list, index))
-    end
-    GC.@preserve device_list begin
-        device_ptr = @ccall MLIR.API.mlir_c.ifrt_BasicDeviceListGetDevice(
-            device_list.ptr::Ptr{Cvoid}, (index - 1)::Int32
-        )::Ptr{Cvoid}
-    end
-    return Device(device_ptr)
-end
-
-function Base.size(device_list::BasicDeviceList)
-    GC.@preserve device_list begin
-        len = @ccall MLIR.API.mlir_c.ifrt_BasicDeviceListSize(
-            device_list.ptr::Ptr{Cvoid}
-        )::Int32
-    end
-    return (len,)
-end
-
-function Base.string(device_list::BasicDeviceList)
-    GC.@preserve device_list begin
-        str = @ccall MLIR.API.mlir_c.ifrt_BasicDeviceListToString(
-            device_list.ptr::Ptr{Cvoid}
-        )::Cstring
-    end
-    return XLA.unsafe_string_and_free(str)
-end
-
 function XLA.default_memory(device_list::AbstractVector{Device})
     default_memories = XLA.default_memory.(device_list)
     default_memory_kinds = convert.(MemoryKind, default_memories)
@@ -109,3 +63,9 @@ function XLA.default_memory(device_list::AbstractVector{Device})
     end
     return first(default_memories)
 end
+
+function XLA.client(device_list::AbstractVector{Device})
+    clients = XLA.client.(device_list)
+    @assert allequal(clients) "All devices must have the same client"
+    return first(clients)
+end

src/xla/IFRT/Sharding.jl

@@ -12,12 +12,17 @@ function free_hlo_sharding(hlo_sharding::HloSharding)
     @ccall MLIR.API.mlir_c.free_ifrt_hlo_sharding(hlo_sharding.ptr::Ptr{Cvoid})::Cvoid
 end
 
-function HloSharding(device_list::BasicDeviceList, xla_hlo_sharding::XLA.HloSharding)
+function HloSharding(
+    device_list::AbstractVector{<:Device}, xla_hlo_sharding::XLA.HloSharding
+)
     default_memory_kind = convert(MemoryKind, XLA.default_memory(device_list))
-    GC.@preserve device_list default_memory_kind xla_hlo_sharding begin
+    client = XLA.client(device_list)
+    GC.@preserve device_list default_memory_kind xla_hlo_sharding client begin
         return HloSharding(
             @ccall MLIR.API.mlir_c.ifrt_hlo_sharding_from_xla_hlo_sharding(
-                device_list.ptr::Ptr{Cvoid},
+                client.client::Ptr{Cvoid},
+                [d.device for d in device_list]::Ptr{Ptr{Cvoid}},
+                length(device_list)::Int32,
                 default_memory_kind.ptr::Ptr{Cvoid},
                 xla_hlo_sharding.ptr::Ptr{Cvoid},
             )::Ptr{Cvoid}
@@ -50,7 +55,7 @@ mutable struct Sharding
     end
 end
 
-function Sharding(device_list::BasicDeviceList, xla_hlo_sharding::XLA.HloSharding)
+function Sharding(device_list::AbstractVector{<:Device}, xla_hlo_sharding::XLA.HloSharding)
     return convert(Sharding, HloSharding(device_list, xla_hlo_sharding))
 end