feat: sharding annotations across nodes now working

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -1891,6 +1891,10 @@ hlo_sharding_to_string(const xla::HloSharding *hlo_sharding) {
   return cstr_from_string(hlo_sharding->ToString(true));
 }
 
+extern "C" ifrt::MemoryKind *ifrt_memory_kind_from_string(const char *c_str) {
+  return new ifrt::MemoryKind(std::string(c_str));
+}
+
 extern "C" ifrt::HloSharding *ifrt_hlo_sharding_from_xla_hlo_sharding(
     ifrt::Client *client, ifrt::Device **device_list, int32_t num_devices,
     ifrt::MemoryKind *memory_kind, xla::HloSharding *xla_hlo_sharding) {

src/xla/IFRT/Device.jl

@@ -48,19 +48,13 @@ function XLA.memories(device::Device)
             device.device::Ptr{Cvoid}, memories_size::Ptr{Int32}
         )::Ptr{Ptr{Cvoid}}
     end
-    memories = Vector{Memory}(undef, memories_size[])
-    for i in 1:memories_size[]
-        memories[i] = Memory(unsafe_load(ptr, i))
-    end
-    return memories
+    return [Memory(unsafe_load(ptr, i)) for i in 1:memories_size[]]
 end
 
 function XLA.default_memory(device_list::AbstractVector{Device})
     default_memories = XLA.default_memory.(device_list)
     default_memory_kinds = convert.(MemoryKind, default_memories)
-    if !allequal(default_memory_kinds)
-        error("All devices must have the same default memory")
-    end
+    @assert allequal(default_memory_kinds) "All devices must have the same default memory"
     return first(default_memories)
 end
 

src/xla/IFRT/Memory.jl

@@ -14,6 +14,15 @@ mutable struct MemoryKind <: XLA.AbstractMemoryKind
     ptr::Ptr{Cvoid}
 end
 
+function MemoryKind(str::AbstractString)
+    str = string(str)
+    GC.@preserve str begin
+        return MemoryKind(
+            @ccall MLIR.API.mlir_c.ifrt_memory_kind_from_string(str::Cstring)::Ptr{Cvoid}
+        )
+    end
+end
+
 function Base.convert(::Type{MemoryKind}, memory::Memory)
     GC.@preserve memory begin
         return MemoryKind(

src/xla/IFRT/Sharding.jl

@@ -4,7 +4,8 @@ mutable struct HloSharding
 
     function HloSharding(ptr::Ptr{Cvoid})
         @assert ptr != C_NULL
-        return finalizer(free_hlo_sharding, new(ptr))
+        # return finalizer(free_hlo_sharding, new(ptr))
+        return new(ptr)
     end
 end
 
@@ -16,14 +17,30 @@ function HloSharding(
     device_list::AbstractVector{<:Device}, xla_hlo_sharding::XLA.HloSharding
 )
     default_memory_kind = convert(MemoryKind, XLA.default_memory(device_list))
+    return HloSharding(device_list, xla_hlo_sharding, default_memory_kind)
+end
+
+function HloSharding(
+    device_list::AbstractVector{<:Device},
+    xla_hlo_sharding::XLA.HloSharding,
+    memoy_kind::AbstractString,
+)
+    return HloSharding(device_list, xla_hlo_sharding, MemoryKind(memoy_kind))
+end
+
+function HloSharding(
+    device_list::AbstractVector{<:Device},
+    xla_hlo_sharding::XLA.HloSharding,
+    memory_kind::MemoryKind,
+)
     client = XLA.client(device_list)
-    GC.@preserve device_list default_memory_kind xla_hlo_sharding client begin
+    GC.@preserve device_list memory_kind xla_hlo_sharding client begin
         return HloSharding(
             @ccall MLIR.API.mlir_c.ifrt_hlo_sharding_from_xla_hlo_sharding(
                 client.client::Ptr{Cvoid},
                 [d.device for d in device_list]::Ptr{Ptr{Cvoid}},
                 length(device_list)::Int32,
-                default_memory_kind.ptr::Ptr{Cvoid},
+                memory_kind.ptr::Ptr{Cvoid},
                 xla_hlo_sharding.ptr::Ptr{Cvoid},
             )::Ptr{Cvoid}
         )
@@ -51,14 +68,23 @@ mutable struct Sharding
 
     function Sharding(ptr::Ptr{Cvoid})
         @assert ptr != C_NULL
-        return finalizer(free_sharding, new(ptr))
+        # return finalizer(free_sharding, new(ptr))
+        return new(ptr)
     end
 end
 
 function Sharding(device_list::AbstractVector{<:Device}, xla_hlo_sharding::XLA.HloSharding)
     return convert(Sharding, HloSharding(device_list, xla_hlo_sharding))
 end
 
+function Sharding(
+    device_list::AbstractVector{<:Device},
+    xla_hlo_sharding::XLA.HloSharding,
+    memoy_kind::Union{AbstractString,MemoryKind},
+)
+    return convert(Sharding, HloSharding(device_list, xla_hlo_sharding, memoy_kind))
+end
+
 function free_sharding(sharding::Sharding)
     @ccall MLIR.API.mlir_c.free_ifrt_sharding(sharding.ptr::Ptr{Cvoid})::Cvoid
 end