fix: run finalizer on array (#1076)

* fix: run finalizer on array

* fix: more missing finalizer

* fix: avoid double free of HloSharding

* fix: remove the ifrt::HloSharding middleman

* Update Project.toml

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.59"
+version = "0.2.60"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -87,7 +87,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.9"
-Reactant_jll = "0.0.108"
+Reactant_jll = "0.0.109"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/xla/IFRT/Array.jl

@@ -2,8 +2,7 @@ mutable struct Array <: XLA.AbstractBuffer
     buffer::Ptr{Cvoid}
 
     function Array(buffer::Ptr{Cvoid})
-        # return finalizer(free_ifrt_array, new(buffer))
-        return new(buffer)
+        return finalizer(free_ifrt_array, new(buffer))
     end
 end
 
@@ -128,9 +127,8 @@ function Array(
 end
 
 @inline function free_ifrt_array(buffer::Array)
-    sbuffer = buffer.buffer
-    if sbuffer != C_NULL
-        @ccall MLIR.API.mlir_c.ifrt_free_array(sbuffer::Ptr{Cvoid})::Cvoid
+    if buffer.buffer != C_NULL
+        @ccall MLIR.API.mlir_c.ifrt_free_array(buffer.buffer::Ptr{Cvoid})::Cvoid
     end
 end
 
@@ -243,7 +241,7 @@ function disassemble_into_single_device_arrays(array::Array, only_addressable_de
     arrays = GC.@preserve array begin
         @ccall MLIR.API.mlir_c.ifrt_array_disassemble_into_single_device_arrays(
             array.buffer::Ptr{Cvoid},
-            Int32(0)::Int32,
+            0::Int32,
             c_single_device_shard_semantics::Int32,
             narrays::Ptr{Int32},
         )::Ptr{Ptr{Cvoid}}

src/xla/IFRT/Sharding.jl

@@ -1,53 +1,28 @@
-# xla::ifrt::HloSharding (distinct from xla::HloSharding)
-mutable struct HloSharding
+mutable struct Sharding
     ptr::Ptr{Cvoid}
 
-    function HloSharding(ptr::Ptr{Cvoid})
+    function Sharding(ptr::Ptr{Cvoid})
         @assert ptr != C_NULL
-        # return finalizer(free_hlo_sharding, new(ptr))
-        return new(ptr)
-    end
-end
-
-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}
-        )
+        return finalizer(free_sharding, new(ptr))
     end
 end
 
-function HloSharding(
-    device_list::AbstractVector{<:Device}, xla_hlo_sharding::XLA.HloSharding
-)
+function Sharding(device_list::AbstractVector{<:Device}, xla_hlo_sharding::XLA.HloSharding)
     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
-
-function HloSharding(
-    device_list::AbstractVector{<:Device},
-    xla_hlo_sharding::XLA.HloSharding,
-    memoy_kind::AbstractString,
-)
-    return HloSharding(device_list, xla_hlo_sharding, MemoryKind(memoy_kind))
+    return Sharding(device_list, xla_hlo_sharding, default_memory_kind)
 end
 
-function HloSharding(
+function Sharding(
     device_list::AbstractVector{<:Device},
     xla_hlo_sharding::XLA.HloSharding,
-    memory_kind::MemoryKind,
+    memory_kind::Union{AbstractString,MemoryKind},
 )
+    memory_kind isa AbstractString && (memory_kind = MemoryKind(memory_kind))
     client = XLA.client(device_list)
     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(
+        return Sharding(
+            @ccall MLIR.API.mlir_c.ifrt_sharding_from_xla_hlo_sharding(
                 client.client::Ptr{Cvoid},
                 [d.device for d in device_list]::Ptr{Ptr{Cvoid}},
                 length(device_list)::Int32,
@@ -58,87 +33,39 @@ function HloSharding(
     end
 end
 
-function Base.string(hlo_sharding::HloSharding)
-    GC.@preserve hlo_sharding begin
-        str = @ccall MLIR.API.mlir_c.ifrt_hlo_sharding_to_string(
-            hlo_sharding.ptr::Ptr{Cvoid}
-        )::Cstring
-    end
-    return XLA.unsafe_string_and_free(str)
-end
-
-function Base.show(io::IO, ::MIME"text/plain", hlo_sharding::HloSharding)
-    print(io, "XLA.IFRT.HloSharding(\"", string(hlo_sharding), "\")")
-    return nothing
-end
-
-# HloSharding is more specific than Sharding. But Sharding is a neater way to deal with
-# most of the IFRT APIs.
-mutable struct Sharding
-    ptr::Ptr{Cvoid}
-
-    function Sharding(ptr::Ptr{Cvoid})
-        @assert ptr != C_NULL
-        # 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,
-    memory_kind::Union{AbstractString,MemoryKind},
-)
-    return convert(Sharding, HloSharding(device_list, xla_hlo_sharding, memory_kind))
-end
-
 function free_sharding(sharding::Sharding)
     @ccall MLIR.API.mlir_c.free_ifrt_sharding(sharding.ptr::Ptr{Cvoid})::Cvoid
 end
 
-function XLA.devices(sharding::Sharding)
+function XLA.num_devices(sharding::Sharding)
     GC.@preserve sharding begin
-        ndevices = @ccall MLIR.API.mlir_c.ifrt_sharding_devices_size(
+        return @ccall MLIR.API.mlir_c.ifrt_sharding_devices_size(
             sharding.ptr::Ptr{Cvoid}
         )::Int32
     end
+end
+
+function XLA.devices(sharding::Sharding)
+    ndevices = XLA.num_devices(sharding)
     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[]]
+    return map(Device, devices[])
 end
 
-function Base.convert(::Type{Sharding}, hlo_sharding::HloSharding)
-    GC.@preserve hlo_sharding begin
-        return Sharding(
-            @ccall MLIR.API.mlir_c.ifrt_sharding_from_ifrt_hlo_sharding(
-                hlo_sharding.ptr::Ptr{Cvoid}
-            )::Ptr{Cvoid}
-        )
-    end
-end
-
-function Base.convert(::Type{HloSharding}, sharding::Sharding)
+function Base.convert(::Type{XLA.HloSharding}, sharding::Sharding)
     GC.@preserve sharding begin
-        return HloSharding(
-            @ccall MLIR.API.mlir_c.ifrt_sharding_to_ifrt_hlo_sharding(
+        return XLA.HloSharding(
+            @ccall MLIR.API.mlir_c.ifrt_sharding_to_xla_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(