fix: Array construction from SingleShards

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -1273,9 +1273,9 @@ extern "C" HeldIfrtArray *ifrt_client_make_single_shard_array_from_host_buffer(
 
 // all arrays are assumed to have same DType
 extern "C" HeldIfrtArray *ifrt_client_assemble_array_from_single_shards(
-    ifrt::Client *client, int ndims, const int64_t *c_shape,
-    HeldValue<std::shared_ptr<const ifrt::Sharding>> *sharding, int narrays,
-    HeldIfrtArray **c_arrays, int c_semantics) {
+    ifrt::Client *client, int32_t ndims, const int64_t *c_shape,
+    HeldValue<std::shared_ptr<const ifrt::Sharding>> *sharding, int32_t narrays,
+    HeldIfrtArray **c_arrays, int32_t c_semantics) {
   auto shape = ifrt::Shape(absl::Span<const int64_t>(c_shape, ndims));
   std::vector<tsl::RCReference<ifrt::Array>> arrays;
   for (int i = 0; i < narrays; i++) {
@@ -1692,6 +1692,10 @@ extern "C" ifrt::Client *ifrt_DeviceToClient(ifrt::Device *device) {
   return device->client();
 }
 
+extern "C" bool ifrt_DeviceIsAddressable(ifrt::Device *device) {
+  return device->IsAddressable();
+}
+
 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);
@@ -1916,6 +1920,14 @@ ifrt_hlo_sharding_to_string(ifrt::HloSharding *hlo_sharding) {
   return cstr_from_string(hlo_sharding->DebugString());
 }
 
+extern "C" ifrt::HloSharding *ifrt_sharding_to_ifrt_hlo_sharding(
+    HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding) {
+  const ifrt::Sharding *val = sharding->obj().get();
+  if (!llvm::isa<ifrt::HloSharding>(val))
+    ReactantThrowError("Expected a HloSharding");
+  return new ifrt::HloSharding(*llvm::dyn_cast<const ifrt::HloSharding>(val));
+}
+
 extern "C" void
 free_ifrt_sharding(HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding) {
   delete sharding;
@@ -1936,11 +1948,35 @@ ifrt_sharding_from_hlo_sharding(ifrt::Client *client,
                                               memory_kind, xla_hlo_sharding));
 }
 
+extern "C" bool ifrt_sharding_is_single_device_sharding(
+    HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding) {
+  return llvm::isa<const ifrt::SingleDeviceSharding>(sharding->obj().get());
+}
+
+extern "C" bool ifrt_sharding_is_fully_replicated(
+    HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding) {
+  return sharding->obj()->IsFullyReplicated();
+}
+
 extern "C" const char *
 ifrt_sharding_to_string(HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding) {
   return cstr_from_string(sharding->obj()->DebugString());
 }
 
+extern "C" int32_t ifrt_sharding_devices_size(
+    HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding) {
+  return sharding->obj()->devices()->size();
+}
+
+extern "C" void ifrt_sharding_to_device_list(
+    HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding,
+    ifrt::Device **devices) {
+  auto device_list = sharding->obj()->devices()->devices();
+  for (int i = 0; i < device_list.size(); i++) {
+    devices[i] = device_list[i];
+  }
+}
+
 #pragma endregion
 
 typedef ifrt::Future<> IfRtFutureType;

src/Compiler.jl

@@ -1107,7 +1107,7 @@ function codegen_flatten!(
                 )
                 push!(flatten_code, :($usbuf = $flatcode))
                 device_to_array_slices = XLA.sharding_to_concrete_array_indices(
-                    condensed_op_sharding, size(carg), mesh
+                    condensed_op_sharding, size(carg), mesh.device_ids
                 )
                 for j in 1:length(mesh)
                     device_id = mesh.device_ids[j]

src/Sharding.jl

@@ -280,7 +280,7 @@ function (sharding::HloSharding)(
     condensed_op_sharding = convert(XLA.CondensedOpSharding, sharding.hlo_sharding)
 
     device_to_array_slices = XLA.sharding_to_concrete_array_indices(
-        condensed_op_sharding, size(x), sharding.mesh
+        condensed_op_sharding, size(x), sharding.mesh.logical_device_ids
     )
 
     data = ntuple(length(sharding.mesh)) do i

src/xla/Device.jl

@@ -10,6 +10,7 @@ function get_local_device_id end
 function device_kind end
 function default_memory end
 function memories end
+function is_addressable end
 
 """
     device_ordinal(device::Device)

src/xla/IFRT/Array.jl

@@ -2,47 +2,91 @@ mutable struct Array <: XLA.AbstractBuffer
     buffer::Ptr{Cvoid}
 
     function Array(buffer::Ptr{Cvoid})
-        return finalizer(free_ifrt_array, new(buffer))
+        # return finalizer(free_ifrt_array, new(buffer))
+        return new(buffer)
     end
 end
 
-function Array(client::Client, array::Base.Array{T,N}, device::Device) where {T,N}
+function Array(
+    client::Client,
+    array::Base.Array{T,N},
+    device::Device,
+    memory_kind::AbstractString=string(convert(MemoryKind, XLA.default_memory(device))),
+) where {T,N}
     sizear = collect(Int64, reverse(size(array)))
     buffer = GC.@preserve array sizear begin
         @ccall MLIR.API.mlir_c.ifrt_client_make_single_shard_array_from_host_buffer(
             client.client::Ptr{Cvoid},
-            pointer(array)::Ptr{T},
+            array::Ptr{T},
             XLA.primitive_type(T)::UInt64,
             N::Csize_t,
-            pointer(sizear)::Ptr{Int64},
+            sizear::Ptr{Int64},
             0::Cint, # kAlwaysCopy
             device.device::Ptr{Cvoid},
-            string(convert(MemoryKind, XLA.default_memory(device)))::Cstring,
+            string(memory_kind)::Cstring,
         )::Ptr{Cvoid}
     end
     return Array(buffer)
 end
 
-function Array(client::Client, array::Base.Array{T,N}, sharding::HloSharding) where {T,N}
-    return Array(client, array, convert(Sharding, sharding))
-end
-
-function Array(client::Client, array::Base.Array{T,N}, sharding::Sharding) where {T,N}
+function Array(
+    client::Client, array::Base.Array{T,N}, sharding::Sharding, logical_device_ids
+) where {T,N}
     sizear = collect(Int64, reverse(size(array)))
-    buffer = GC.@preserve array sizear begin
-        @ccall MLIR.API.mlir_c.ifrt_client_make_array_from_host_buffer(
+
+    if is_single_device_sharding(sharding) || is_fully_replicated(sharding)
+        buffer = GC.@preserve array sizear begin
+            @ccall MLIR.API.mlir_c.ifrt_client_make_array_from_host_buffer(
+                client.client::Ptr{Cvoid},
+                array::Ptr{T},
+                XLA.primitive_type(T)::Cint,
+                N::Csize_t,
+                sizear::Ptr{Int64},
+                sharding.ptr::Ptr{Cvoid},
+                0::Cint, # kAlwaysCopy
+            )::Ptr{Cvoid}
+        end
+        return Array(buffer)
+    end
+
+    all_devices = XLA.devices(sharding)
+    array_slices = XLA.sharding_to_concrete_array_indices(
+        convert(XLA.HloSharding, sharding), size(array), logical_device_ids
+    )
+    array_shape = collect(Int64, reverse(size(array)))
+    arrays_list = [
+        Array(client, array[slice...], device).buffer for
+        (device, slice) in zip(all_devices, array_slices) if XLA.is_addressable(device)
+    ]
+
+    buffer = GC.@preserve client arrays_list array_shape sharding begin
+        @ccall MLIR.API.mlir_c.ifrt_client_assemble_array_from_single_shards(
             client.client::Ptr{Cvoid},
-            pointer(array)::Ptr{T},
-            XLA.primitive_type(T)::Cint,
-            N::Csize_t,
-            pointer(sizear)::Ptr{Int64},
+            Int32(length(array_shape))::Int32,
+            array_shape::Ptr{Int64},
             sharding.ptr::Ptr{Cvoid},
-            0::Cint, # kAlwaysCopy
+            Int32(length(arrays_list))::Int32,
+            arrays_list::Ptr{Ptr{Cvoid}},
+            2::Cint, # kDonateInput
         )::Ptr{Cvoid}
     end
+
     return Array(buffer)
 end
 
+function Array(client::Client, array::Base.Array{T,N}, sharding) where {T,N}
+    @assert sharding isa Reactant.Sharding.AbstractSharding
+    if !(sharding isa Reactant.Sharding.HloSharding)
+        sharding = convert(Reactant.Sharding.HloSharding, sharding)
+    end
+
+    (; hlo_sharding, mesh) = sharding
+    devices = XLA.get_device.((client,), mesh.device_ids)
+    ifrt_sharding = Sharding([devices...], hlo_sharding)
+
+    return Array(client, array, ifrt_sharding, mesh.logical_device_ids)
+end
+
 @inline function free_ifrt_array(buffer::Array)
     sbuffer = buffer.buffer
     if sbuffer != C_NULL

src/xla/IFRT/Device.jl

@@ -63,3 +63,11 @@ function XLA.client(device_list::AbstractVector{Device})
     @assert allequal(clients) "All devices must have the same client"
     return first(clients)
 end
+
+function XLA.is_addressable(device::Device)
+    GC.@preserve device begin
+        return @ccall MLIR.API.mlir_c.ifrt_DeviceIsAddressable(
+            device.device::Ptr{Cvoid}
+        )::Bool
+    end
+end

src/xla/IFRT/Sharding.jl

@@ -13,10 +13,21 @@ function free_hlo_sharding(hlo_sharding::HloSharding)
     @ccall MLIR.API.mlir_c.free_ifrt_hlo_sharding(hlo_sharding.ptr::Ptr{Cvoid})::Cvoid
 end
 
+function Base.convert(::Type{XLA.HloSharding}, sharding::HloSharding)
+    GC.@preserve sharding begin
+        return XLA.HloSharding(
+            @ccall MLIR.API.mlir_c.ifrt_hlo_sharding_to_xla_hlo_sharding(
+                sharding.ptr::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
 function HloSharding(
     device_list::AbstractVector{<:Device}, xla_hlo_sharding::XLA.HloSharding
 )
-    default_memory_kind = convert(MemoryKind, XLA.default_memory(device_list))
+    addressable_devices = filter(XLA.is_addressable, device_list)
+    default_memory_kind = convert(MemoryKind, XLA.default_memory(addressable_devices))
     return HloSharding(device_list, xla_hlo_sharding, default_memory_kind)
 end
 
@@ -89,6 +100,21 @@ function free_sharding(sharding::Sharding)
     @ccall MLIR.API.mlir_c.free_ifrt_sharding(sharding.ptr::Ptr{Cvoid})::Cvoid
 end
 
+function XLA.devices(sharding::Sharding)
+    GC.@preserve sharding begin
+        ndevices = @ccall MLIR.API.mlir_c.ifrt_sharding_devices_size(
+            sharding.ptr::Ptr{Cvoid}
+        )::Int32
+    end
+    devices = Ref{NTuple{Int64(ndevices),Ptr{Cvoid}}}()
+    GC.@preserve sharding devices begin
+        @ccall MLIR.API.mlir_c.ifrt_sharding_to_device_list(
+            sharding.ptr::Ptr{Cvoid}, devices::Ptr{Ptr{Cvoid}}
+        )::Cvoid
+    end
+    return [Device(device) for device in devices[]]
+end
+
 function Base.convert(::Type{Sharding}, hlo_sharding::HloSharding)
     GC.@preserve hlo_sharding begin
         return Sharding(
@@ -99,6 +125,20 @@ function Base.convert(::Type{Sharding}, hlo_sharding::HloSharding)
     end
 end
 
+function Base.convert(::Type{HloSharding}, sharding::Sharding)
+    GC.@preserve sharding begin
+        return HloSharding(
+            @ccall MLIR.API.mlir_c.ifrt_sharding_to_ifrt_hlo_sharding(
+                sharding.ptr::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
+function Base.convert(::Type{XLA.HloSharding}, sharding::Sharding)
+    return convert(XLA.HloSharding, convert(HloSharding, sharding))
+end
+
 function Base.string(sharding::Sharding)
     GC.@preserve sharding begin
         str = @ccall MLIR.API.mlir_c.ifrt_sharding_to_string(
@@ -108,6 +148,22 @@ function Base.string(sharding::Sharding)
     return XLA.unsafe_string_and_free(str)
 end
 
+function is_fully_replicated(sharding::Sharding)
+    GC.@preserve sharding begin
+        return @ccall MLIR.API.mlir_c.ifrt_sharding_is_fully_replicated(
+            sharding.ptr::Ptr{Cvoid}
+        )::Bool
+    end
+end
+
+function is_single_device_sharding(sharding::Sharding)
+    GC.@preserve sharding begin
+        return @ccall MLIR.API.mlir_c.ifrt_sharding_is_single_device_sharding(
+            sharding.ptr::Ptr{Cvoid}
+        )::Bool
+    end
+end
+
 function Base.show(io::IO, ::MIME"text/plain", sharding::Sharding)
     print(io, "XLA.IFRT.Sharding(\"", string(sharding), "\")")
     return nothing