Add initial IFRT Julia bindings

src/xla/Buffer.jl

@@ -1,16 +1,39 @@
 # Buffer
-@inline function free_buffer(buffer)
-    sbuffer = buffer.buffer
-    if sbuffer != C_NULL
-        @ccall MLIR.API.mlir_c.PjRtBufferFree(sbuffer::Ptr{Cvoid})::Cvoid
+mutable struct HeldBuffer
+    ptr::Ptr{Cvoid}
+
+    function HeldBuffer(ptr::Ptr{Cvoid})
+        return finalizer(release_buffer, new(ptr))
     end
 end
 
+@inline function release_buffer(held_buffer::HeldBuffer)
+    @ccall MLIR.API.mlir_c.reactant_release_pjrtbuffer(held_buffer.ptr::Ptr{Cvoid})::Cvoid
+end
+
 mutable struct Buffer
     buffer::Ptr{Cvoid}
+    held::Union{Nothing,HeldBuffer}
+
     function Buffer(buffer::Ptr{Cvoid})
-        return finalizer(free_buffer, new(buffer))
+        return finalizer(free_buffer, new(buffer, nothing))
+    end
+end
+
+@inline function free_buffer(buffer)
+    if isnothing(buffer.holded) && buffer.buffer != C_NULL
+        @ccall MLIR.API.mlir_c.PjRtBufferFree(buffer.buffer::Ptr{Cvoid})::Cvoid
+    end
+end
+
+function hold!(buffer::Buffer)
+    if buffer.holded == C_NULL
+        sbuffer = buffer.buffer
+        buffer.holded = HeldBuffer(
+            @ccall MLIR.API.mlir_c.reactant_hold_pjrtbuffer(sbuffer::Ptr{Cvoid})::Ptr{Cvoid}
+        )
     end
+    return buffer
 end
 
 function Base.ndims(buffer::Buffer)

src/xla/Client.jl

@@ -1,12 +1,13 @@
 mutable struct Client
     client::Ptr{Cvoid}
     global_ordinals::Vector{Cint}
+    holded::Ptr{Cvoid}
 
     function Client(client::Ptr{Cvoid})
         @assert client != C_NULL
         global_ordinals = Cint[]
 
-        client = new(client, global_ordinals)
+        client = new(client, global_ordinals, C_NULL)
 
         # https://github.com/pytorch/xla/blob/8b2414094578e829b99a8383877c86d357eeb682/torch_xla/csrc/runtime/pjrt_computation_client.cc#L127
         devices = [
@@ -29,7 +30,20 @@ end
 Base.:(==)(a::Client, b::Client) = a.client == b.client
 
 @inline function free_client(client::Client)
-    @ccall MLIR.API.mlir_c.FreeClient(client.client::Ptr{Cvoid})::Cvoid
+    if client.holded == C_NULL
+        @ccall MLIR.API.mlir_c.FreeClient(client.client::Ptr{Cvoid})::Cvoid
+    else
+        @ccall MLIR.API.mlir_c.reactant_release_pjrtclient(client.holded::Ptr{Cvoid})::Cvoid
+    end
+end
+
+function hold!(client::Client)
+    if client.holded == C_NULL
+        client.holded = @ccall MLIR.API.mlir_c.reactant_hold_pjrtclient(
+            client.client::Ptr{Cvoid}
+        )::Ptr{Cvoid}
+    end
+    return client
 end
 
 function ClientNumDevices(client::Client)

src/xla/IFRT/Array.jl

@@ -0,0 +1,38 @@
+@cenum ArrayCopySemantics::UInt32 begin
+    AlwaysCopy = 0
+    ReuseInput = 1
+    DonateInput = 2
+end
+
+# currently, only supports IFRT-PjRt
+mutable struct Array
+    ptr::Ptr{Cvoid}
+
+    function Array(ptr::Ptr{Cvoid})
+        @assert ptr != C_NULL
+        return finalizer(free_array, new(ptr))
+    end
+end
+
+function free_array(array)
+    @ccall MLIR.API.mlir_c.reactant_release_ifrt_array(array.ptr::Ptr{Cvoid})::Cvoid
+end
+
+function Array(client::Client, buffer::XLA.Buffer)
+    hold!(buffer)
+    GC.@preserve client buffer begin
+        return Array(
+            @ccall MLIR.API.mlir_c.ifrt_pjrt_ArrayFromHostBuffer(
+                client.ptr::Ptr{Cvoid}, buffer.holded::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
+function CopyArrayToHostBuffer(array::Array, data)
+    GC.@preserve array data begin
+        @ccall MLIR.API.mlir_c.ifrt_CopyArrayToHostBuffer(
+            array.ptr::Ptr{Cvoid}, data::Ptr{Cvoid}, AlwaysCopy::Cuint
+        )::Cvoid
+    end
+end

src/xla/IFRT/Client.jl

@@ -0,0 +1,31 @@
+# currently, only supports IFRT-PjRt
+mutable struct Client
+    ptr::Ptr{Cvoid}
+
+    function Client(ptr::Ptr{Cvoid})
+        @assert ptr != C_NULL
+        return finalizer(free_client, new(ptr))
+    end
+end
+
+function Client(pjrt_client::XLA.Client)
+    # it needs a `std::shared_ptr<xla::PjRtClient>`
+    hold!(pjrt_client)
+    return Client(
+        @ccall MLIR.API.mlir_c.ifrt_pjrt_MakeClient(
+            pjrt_client.holded::Ptr{Cvoid}
+        )::Ptr{Cvoid}
+    )
+end
+
+function free_client(client)
+    @ccall MLIR.API.mlir_c.ifrt_FreeClient(client.ptr::Ptr{Cvoid})::Cvoid
+end
+
+function compile(client::Client, code::MLIR.IR.Module)
+    return LoadedExecutable(
+        @ccall MLIR.API.mlir_c.ifrt_ClientCompile(
+            client.ptr::Ptr{Cvoid}, code.module_::MLIR.API.MlirModule
+        )::Ptr{Cvoid}
+    )
+end

src/xla/IFRT/IFRT.jl

@@ -0,0 +1,13 @@
+module IFRT
+
+using CEnum
+
+import ..XLA
+import .XLA: hold!
+import ..MLIR
+
+include("LoadedExecutable.jl")
+include("Client.jl")
+include("Array.jl")
+
+end

src/xla/IFRT/LoadedExecutable.jl

@@ -0,0 +1,48 @@
+# currently, only supports IFRT-PjRt
+mutable struct LoadedExecutable
+    ptr::Ptr{Cvoid}
+
+    function LoadedExecutable(ptr::Ptr{Cvoid})
+        @assert ptr != C_NULL
+        return finalizer(free_exec, new(ptr))
+    end
+end
+
+@inline function free_exec(exec)
+    @ccall MLIR.API.mlir_c.ifrt_pjrt_FreeLoadedExecutable(exec.ptr::Ptr{Cvoid})::Cvoid
+end
+
+function execute(
+    exec::LoadedExecutable,
+    args::NTuple{N,Ptr{Cvoid}},
+    donated_mask::NTuple{N,UInt8},
+    ::Val{n_results},
+) where {N,n_results}
+    results = Ref{NTuple{n_results,Ptr{Cvoid}}}()
+    has_future = Ref{UInt8}()
+    status = Ref{NTuple{1,Ptr{Cvoid}}}() # unused right now
+
+    args = Base.RefValue(args)
+    donated_mask = Base.RefValue(donated_mask)
+
+    GC.@preserve exec args donated_mask results has_future status begin
+        @ccall MLIR.API.mlir_c.ifrt_Execute(
+            exec.ptr::Ptr{Cvoid},
+            N::Cint,
+            args::Ptr{Cvoid},
+            donated_mask::Ptr{Cvoid},
+            n_results::Cint,
+            Base.unsafe_convert(Ptr{Cvoid}, results)::Ptr{Cvoid},
+            has_future::Ptr{Cvoid},
+            status::Ptr{Cvoid},
+        )::Cvoid
+    end
+
+    @assert has_future[] == true
+
+    results = results[]
+
+    return ntuple(Val(n_results)) do i
+        return Array(results[i])
+    end
+end

src/xla/XLA.jl

@@ -96,4 +96,6 @@ function __init__()
     return nothing
 end
 
+include("IFRT/IFRT.jl")
+
 end