feat: initial low-level IFRT API

fix: ifrt HloSharding

refactor: split up into IFRT/PJRT

feat: IFRT Client APIs

feat: IFRT Device API

fix: remove global_ordinals

feat: add devices list abstraction

feat: wrap memory and memory kinds

feat: ifrt::HloSharding now working

fix: use new ABI

chore: run formatter

fix: no finalizer

feat: initial draft of IFRT.Array interface (#774)

* feat: initial draft of IFRT.Array interface

* feat: Base.Array to ifrt::Array

* feat: buffer to host
chore: run formatter

fix: bad rebase

feat: more proxy servers

feat: add ConcreteIFRTArray

feat: add ConcreteIFRTNumber

refactor: rename ConcreteRNumber to ConcretePJRTNumber

revert: concreteifrtarray implementation

chore: run formatter

feat: ifrt loaded executable

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -76,11 +76,11 @@
 
 // IFRT
 #include "xla/python/ifrt/array.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"
 #include "xla/python/ifrt/device_list.h"
-#include "xla/python/ifrt/basic_device_list.h"
 #include "xla/python/ifrt/dtype.h"
 #include "xla/python/ifrt/executable.h"
 #include "xla/python/ifrt/hlo/hlo_program.h"
@@ -1469,6 +1469,10 @@ ifrt_pjrt_loaded_executable_dtor(xla::ifrt::PjRtLoadedExecutable *exec) {
 
 extern "C" void ifrt_array_dtor(HeldIfrtArray *array) { delete array; }
 
+extern "C" void ifrt_loaded_executable_dtor(ifrt::LoadedExecutable *exec) {
+  delete exec;
+}
+
 extern "C" void ifrt_loaded_executable_execute(
     ifrt::LoadedExecutable *exec, int num_args,
     HeldValue<tsl::RCReference<ifrt::Array>> **op_args,
@@ -1538,38 +1542,56 @@ FreeHloModule(HeldValue<std::shared_ptr<xla::HloModule>> *hlo_module) {
 
 #pragma region IfRtClient
 
-// right now only making it available for TPU
-// in the future, we would like this for CPU and GPU PjRt backends too
 extern "C" ifrt::proxy::GrpcServer *
-ifrt_proxy_grpc_server_create_from_ifrt_client_factory_tpu(
-    const char *c_address, const char *tpu_path, const char **error) {
+ifrt_proxy_grpc_server_create_from_ifrt_client_factory_cpu(
+    const char *c_address, uint8_t asynchronous, int node_id, int num_nodes) {
   std::string address = c_address;
 
-  // taken from `MakeTPUClient`
-  std::string tpu_library_path;
-  if (auto path = llvm::sys::Process::GetEnv(kEnvTpuLibraryPath)) {
-    tpu_library_path = *path;
-  } else if (tpu_path) {
-    tpu_library_path = std::string(tpu_path);
-  } else {
-    *error = "Could not find TPU path";
-    return nullptr;
-  }
+  return MyValueOrThrow(
+             ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
+                 address,
+                 [asynchronous, node_id, num_nodes]()
+                     -> absl::StatusOr<std::shared_ptr<ifrt::Client>> {
+                   auto pjrt_client = std::shared_ptr<PjRtClient>(
+                       MakeCPUClient(asynchronous, node_id, num_nodes));
+                   return std::shared_ptr<ifrt::Client>(
+                       xla::ifrt::PjRtClient::Create(pjrt_client).release());
+                 }))
+      .release();
+}
 
-  const PJRT_Api *pluginLoad =
-      LoadPjrtPlugin("tpu", tpu_library_path.c_str(), error);
-  if (pluginLoad == nullptr)
-    return nullptr;
-  auto tpu_status = InitializePjrtPlugin("tpu", error);
-  if (tpu_status)
-    return nullptr;
+extern "C" ifrt::proxy::GrpcServer *
+ifrt_proxy_grpc_server_create_from_ifrt_client_factory_gpu(
+    int node_id, int num_nodes, int *allowed_devices, int num_allowed_devices,
+    double memory_fraction, bool preallocate, const char *platform_name,
+    const char **error) {
+  return MyValueOrThrow(
+             ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
+                 std::string(),
+                 [node_id, num_nodes, allowed_devices, num_allowed_devices,
+                  memory_fraction, preallocate, platform_name,
+                  error]() -> absl::StatusOr<std::shared_ptr<ifrt::Client>> {
+                   auto pjrt_client = std::shared_ptr<PjRtClient>(MakeGPUClient(
+                       node_id, num_nodes, allowed_devices, num_allowed_devices,
+                       memory_fraction, preallocate, platform_name, error));
+                   return std::shared_ptr<ifrt::Client>(
+                       xla::ifrt::PjRtClient::Create(pjrt_client).release());
+                 }))
+      .release();
+}
+
+extern "C" ifrt::proxy::GrpcServer *
+ifrt_proxy_grpc_server_create_from_ifrt_client_factory_tpu(
+    const char *c_address, const char *tpu_path, const char **error) {
+  std::string address = c_address;
 
   return MyValueOrThrow(
              xla::ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
                  address,
-                 []() -> absl::StatusOr<std::shared_ptr<xla::ifrt::Client>> {
-                   auto pjrt_client =
-                       std::shared_ptr<xla::PjRtClient>(GetCApiClient("TPU"));
+                 [tpu_path, error]()
+                     -> absl::StatusOr<std::shared_ptr<xla::ifrt::Client>> {
+                   auto pjrt_client = std::shared_ptr<xla::PjRtClient>(
+                       MakeTPUClient(tpu_path, error));
                    return std::shared_ptr<xla::ifrt::Client>(
                        xla::ifrt::PjRtClient::Create(pjrt_client).release());
                  }))
@@ -1604,28 +1626,28 @@ ifrt_proxy_create_client(const char *c_proxy_server_address,
       nullptr, // callback `on_connection_update`
   };
   return MyValueOrThrow(
-             ifrt::proxy::CreateClient(c_proxy_server_address, options))
+             ifrt::proxy::CreateClient(proxy_server_address, options))
       .release();
 }
 
-extern "C" ifrt::Client *ifrt_make_cpu_client(uint8_t asynchronous, int node_id,
-                                              int num_nodes) {
+extern "C" ifrt::Client *ifrt_make_pjrt_cpu_client(uint8_t asynchronous,
+                                                   int node_id, int num_nodes) {
   return ifrt_pjrt_make_client(
       pjrt_make_cpu_client_shared(asynchronous, node_id, num_nodes));
 }
 
 extern "C" ifrt::Client *
-ifrt_make_gpu_client(int node_id, int num_nodes, int *allowed_devices,
-                     int num_allowed_devices, double memory_fraction,
-                     bool preallocate, const char *platform_name,
-                     const char **error) {
+ifrt_make_pjrt_gpu_client(int node_id, int num_nodes, int *allowed_devices,
+                          int num_allowed_devices, double memory_fraction,
+                          bool preallocate, const char *platform_name,
+                          const char **error) {
   return ifrt_pjrt_make_client(pjrt_make_gpu_client_shared(
       node_id, num_nodes, allowed_devices, num_allowed_devices, memory_fraction,
       preallocate, platform_name, error));
 }
 
-extern "C" ifrt::Client *ifrt_make_tpu_client(const char *tpu_path,
-                                              const char **error) {
+extern "C" ifrt::Client *ifrt_make_pjrt_tpu_client(const char *tpu_path,
+                                                   const char **error) {
   return ifrt_pjrt_make_client(pjrt_make_tpu_client_shared(tpu_path, error));
 }
 
@@ -1815,4 +1837,77 @@ ifrt_hlo_sharding_to_string(ifrt::HloSharding *hlo_sharding) {
   return cstr_from_string(hlo_sharding->DebugString());
 }
 
+extern "C" void
+free_ifrt_sharding(HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding) {
+  delete sharding;
+}
+
+extern "C" HeldValue<std::shared_ptr<ifrt::Sharding>> *
+ifrt_sharding_from_ifrt_hlo_sharding(ifrt::HloSharding *hlo_sharding) {
+  return reactant::capture(std::shared_ptr<ifrt::Sharding>(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) {
+  return ifrt_sharding_from_ifrt_hlo_sharding(
+      ifrt_hlo_sharding_from_xla_hlo_sharding(device_list, memory_kind,
+                                              xla_hlo_sharding));
+}
+
+extern "C" const char *
+ifrt_sharding_to_string(HeldValue<std::shared_ptr<ifrt::Sharding>> *sharding) {
+  return cstr_from_string(sharding->obj()->DebugString());
+}
+
+#pragma endregion
+
+typedef ifrt::Future<> IfRtFutureType;
+
+extern "C" void ifrt_free_future(IfRtFutureType *Future) { delete Future; }
+
+extern "C" uint8_t ifrt_future_is_ready(IfRtFutureType *Future) {
+  return Future->IsReady();
+}
+
+extern "C" void ifrt_future_await(IfRtFutureType *Future) { Future->Await(); }
+
+#pragma region IfRtArray
+
+extern "C" void ifrt_free_array(HeldIfrtArray *array) { delete array; }
+
+extern "C" int64_t *ifrt_array_shape(HeldIfrtArray *array) {
+  absl::Span<const long> dims = array->obj()->shape().dims();
+  int64_t *dims_ptr = new int64_t[dims.size()];
+  std::copy(dims.begin(), dims.end(), dims_ptr);
+  return dims_ptr;
+}
+
+extern "C" int64_t ifrt_array_ndims(HeldIfrtArray *array) {
+  return array->obj()->shape().dims().size();
+}
+
+extern "C" ifrt::DType ifrt_array_eltype(HeldIfrtArray *array) {
+  return array->obj()->dtype();
+}
+
+extern "C" ifrt::Client *ifrt_array_to_client(HeldIfrtArray *array) {
+  return array->obj()->client();
+}
+
+extern "C" HeldValue<std::shared_ptr<const ifrt::Sharding>> *
+ifrt_array_to_sharding(HeldIfrtArray *array) {
+  return reactant::capture(array->obj()->shared_ptr_sharding());
+}
+
+extern "C" void ifrt_array_copy_to_host_buffer(HeldIfrtArray *array,
+                                               void *data) {
+  std::optional<absl::Span<const int64_t>> byte_strides;
+  auto future = array->obj()->CopyToHostBuffer(
+      data, byte_strides, static_cast<ifrt::ArrayCopySemantics>(0));
+  future.Await();
+  return;
+}
+
 #pragma endregion

src/xla/Buffer.jl

@@ -5,6 +5,15 @@ function buffer_on_cpu end
 function to_host end
 function unsafe_buffer_pointer end
 function copy_buffer_to_device end
+function sharding end
+
+Base.convert(::Type{Array}, buffer::AbstractBuffer) = convert(Array{eltype(buffer)}, buffer)
+
+function Base.convert(::Type{<:Array{T}}, buffer::AbstractBuffer) where {T}
+    arr = zeros(T, reverse(size(buffer))...)
+    XLA.to_host(buffer, arr)
+    return arr
+end
 
 @inline function client(
     buffers::Union{Array{<:AbstractBuffer},NTuple{<:Any,AbstractBuffer}}
@@ -19,3 +28,48 @@ end
 )
     return map(synced_buffer, buffers)
 end
+
+function Base.show(io::IO, mime::MIME"text/plain", buffer::B) where {B<:AbstractBuffer}
+    print(io, "$(B) storing ")
+    show(io, mime, convert(Array, buffer))
+    return nothing
+end
+
+# Async Buffers
+abstract type AbstractAsyncBuffer <: AbstractBuffer end
+
+Base.isempty(buffer::AbstractAsyncBuffer) = buffer.buffer.buffer == C_NULL
+
+function Base.convert(T::Type{Array}, buffer::AbstractAsyncBuffer)
+    XLA.await(buffer)
+    return convert(T, buffer.buffer)
+end
+
+function Base.convert(T::Type{<:Array{T1}}, buffer::AbstractAsyncBuffer) where {T1}
+    XLA.await(buffer)
+    return convert(T, buffer.buffer)
+end
+
+for op in (:(Base.ndims), :(Base.size), :(Base.eltype), :device, :client, :sharding)
+    @eval $op(buffer::AbstractAsyncBuffer) = $op(buffer.buffer)
+end
+
+function XLA.synced_buffer(buffer::AbstractAsyncBuffer)
+    XLA.await(buffer)
+    return buffer.buffer
+end
+
+function XLA.await(buffer::AbstractAsyncBuffer)
+    buffer.future === nothing && return nothing
+    future = buffer.future
+    buffer.future = nothing
+    XLA.await(future)
+    return nothing
+end
+
+function XLA.is_ready(buffer::AbstractAsyncBuffer)
+    buffer.future === nothing && return true
+    return XLA.is_ready(buffer.future)
+end
+
+XLA.buffer_on_cpu(buffer::AbstractAsyncBuffer) = XLA.buffer_on_cpu(buffer.buffer)

src/xla/IFRT/Array.jl

@@ -0,0 +1,120 @@
+mutable struct Array <: XLA.AbstractBuffer
+    buffer::Ptr{Cvoid}
+
+    function Array(buffer::Ptr{Cvoid})
+        return finalizer(free_ifrt_array, new(buffer))
+    end
+end
+
+function Array(client::Client, array::Base.Array{T,N}, device::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},
+            XLA.primitive_type(T)::UInt64,
+            N::Csize_t,
+            pointer(sizear)::Ptr{Int64},
+            0::Cint, # kAlwaysCopy
+            device.device::Ptr{Cvoid},
+            string(convert(MemoryKind, XLA.default_memory(device)))::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}
+    sizear = collect(Int64, reverse(size(array)))
+    buffer = GC.@preserve array sizear begin
+        @ccall MLIR.API.mlir_c.ifrt_client_make_array_from_host_buffer(
+            client.client::Ptr{Cvoid},
+            pointer(array)::Ptr{T},
+            XLA.primitive_type(T)::Cint,
+            N::Csize_t,
+            pointer(sizear)::Ptr{Int64},
+            sharding.ptr::Ptr{Cvoid},
+            0::Cint, # kAlwaysCopy
+        )::Ptr{Cvoid}
+    end
+    return Array(buffer)
+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
+    end
+end
+
+function Base.ndims(buffer::Array)
+    GC.@preserve buffer begin
+        return @ccall MLIR.API.mlir_c.ifrt_array_ndims(buffer.buffer::Ptr{Cvoid})::Int64
+    end
+end
+
+function Base.size(buffer::Array)
+    GC.@preserve buffer begin
+        sz = @ccall MLIR.API.mlir_c.ifrt_array_shape(buffer.buffer::Ptr{Cvoid})::Ptr{Int64}
+    end
+    return Tuple(unsafe_wrap(Base.Array, sz, ndims(buffer)))
+end
+
+function Base.eltype(buffer::Array)
+    GC.@preserve buffer begin
+        return XLA.julia_type(
+            @ccall MLIR.API.mlir_c.ifrt_array_eltype(buffer.buffer::Ptr{Cvoid})::Cint
+        )
+    end
+end
+
+function XLA.device(::Array)
+    return error("IFRT.Array can be sharded/replicated across multiple devices. Hence, \
+                  `XLA.device` is not defined.")
+end
+
+function XLA.client(buffer::Array)
+    GC.@preserve buffer begin
+        return Client(
+            @ccall MLIR.API.mlir_c.ifrt_array_to_client(
+                buffer.buffer::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
+XLA.synced_buffer(buffer::Array) = buffer
+
+function XLA.buffer_on_cpu(::Array)
+    return error("IFRT.Array does not support `XLA.buffer_on_cpu`")
+end
+
+function XLA.to_host(buffer::Array, data)
+    GC.@preserve buffer data begin
+        @ccall MLIR.API.mlir_c.ifrt_array_copy_to_host_buffer(
+            buffer.buffer::Ptr{Cvoid}, data::Ptr{Cvoid}
+        )::Cvoid
+    end
+    return nothing
+end
+
+function XLA.unsafe_buffer_pointer(::Array)
+    return error("IFRT.Array does not support `XLA.unsafe_buffer_pointer`")
+end
+
+function XLA.copy_buffer_to_device(::Array, ::Device)
+    return error("IFRT.Array does not support `XLA.copy_buffer_to_device`")
+end
+
+function XLA.sharding(buffer::Array)
+    GC.@preserve buffer begin
+        return Sharding(
+            @ccall MLIR.API.mlir_c.ifrt_array_to_sharding(
+                buffer.buffer::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end

src/xla/IFRT/AsyncArray.jl

@@ -0,0 +1,8 @@
+mutable struct AsyncArray <: XLA.AbstractAsyncBuffer
+    buffer::Array
+    future::Union{Future,Nothing}
+end
+
+const AsyncEmptyArray = AsyncArray(Array(C_NULL), nothing)
+
+AsyncArray(args...; kwargs...) = AsyncArray(Array(args...; kwargs...), nothing)

src/xla/IFRT/Client.jl

@@ -0,0 +1,87 @@
+mutable struct Client <: XLA.AbstractClient
+    client::Ptr{Cvoid}
+
+    function Client(client::Ptr{Cvoid})
+        @assert client != C_NULL
+        return new(client)
+    end
+end
+
+function XLA.free_client(client::Client)
+    GC.@preserve client begin
+        @ccall MLIR.API.mlir_c.ifrt_FreeClient(client.client::Ptr{Cvoid})::Cvoid
+    end
+end
+
+function XLA.num_devices(client::Client)
+    GC.@preserve client begin
+        return @ccall MLIR.API.mlir_c.ifrt_client_device_count(
+            client.client::Ptr{Cvoid}
+        )::Cint
+    end
+end
+
+function XLA.num_addressable_devices(client::Client)
+    GC.@preserve client begin
+        return @ccall MLIR.API.mlir_c.ifrt_client_addressable_device_count(
+            client.client::Ptr{Cvoid}
+        )::Cint
+    end
+end
+
+function XLA.process_index(client::Client)
+    GC.@preserve client begin
+        return @ccall MLIR.API.mlir_c.ifrt_ClientProcessIndex(
+            client.client::Ptr{Cvoid}
+        )::Cint
+    end
+end
+
+function XLA.get_device(client::Client, idx)
+    GC.@preserve client begin
+        return Device(
+            @ccall MLIR.API.mlir_c.ifrt_client_lookup_device(
+                client.client::Ptr{Cvoid}, idx::Cint
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
+function XLA.get_addressable_device(client::Client, idx)
+    GC.@preserve client begin
+        return Device(
+            @ccall MLIR.API.mlir_c.ifrt_client_lookup_addressable_device(
+                client.client::Ptr{Cvoid}, idx::Cint
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
+function XLA.platform_name(client::Client)
+    GC.@preserve client begin
+        str = @ccall MLIR.API.mlir_c.ifrt_ClientGetPlatformName(
+            client.client::Ptr{Cvoid}
+        )::Cstring
+    end
+    return XLA.unsafe_string_and_free(str)
+end
+
+# Different Backends
+const cpu_client_count = Ref(0)
+const gpu_client_count = Ref(0)
+const tpu_client_count = Ref(0)
+
+# XXX: We need other backends to support sharding
+for (backend, fname, counter) in (
+    (:CPUClient, "ifrt_make_pjrt_cpu_client", :cpu_client_count),
+    (:GPUClient, "ifrt_make_pjrt_gpu_client", :gpu_client_count),
+    (:TPUClient, "ifrt_make_pjrt_tpu_client", :tpu_client_count),
+)
+    @eval function $(backend)(args...; checkcount::Bool=true, kwargs...)
+        if checkcount
+            @assert $(counter)[] == 0
+            $(counter)[] += 1
+        end
+        return Client(XLA.$(backend)($(fname), args...; kwargs...))
+    end
+end

src/xla/IFRT/Device.jl

@@ -0,0 +1,111 @@
+struct Device <: XLA.AbstractDevice
+    device::Ptr{Cvoid}
+end
+
+function XLA.client(device::Device)
+    GC.@preserve device begin
+        return Client(
+            @ccall MLIR.API.mlir_c.ifrt_DeviceToClient(
+                device.device::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
+function XLA.device_ordinal(device::Device)
+    GC.@preserve device begin
+        return @ccall MLIR.API.mlir_c.ifrt_DeviceGetGlobalDeviceId(
+            device.device::Ptr{Cvoid}
+        )::Int64
+    end
+end
+
+function XLA.device_kind(device::Device)
+    GC.@preserve device begin
+        str = @ccall MLIR.API.mlir_c.ifrt_DeviceGetKind(device.device::Ptr{Cvoid})::Cstring
+    end
+    return XLA.unsafe_string_and_free(str)
+end
+
+function XLA.get_local_device_id(::Device)
+    return error("Not implemented for ifrt devices")
+end
+
+function XLA.default_memory(device::Device)
+    GC.@preserve device begin
+        return Memory(
+            @ccall MLIR.API.mlir_c.ifrt_DeviceGetDefaultMemory(
+                device.device::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
+function XLA.memories(device::Device)
+    memories_size = Ref{Int32}(0)
+    GC.@preserve device memories_size begin
+        ptr = @ccall MLIR.API.mlir_c.ifrt_DeviceGetMemories(
+            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
+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)
+    if !allequal(default_memory_kinds)
+        error("All devices must have the same default memory")
+    end
+    return first(default_memories)
+end

src/xla/IFRT/Future.jl

@@ -0,0 +1,27 @@
+mutable struct Future <: XLA.AbstractFuture
+    future::Ptr{Cvoid}
+
+    function Future(future::Ptr{Cvoid})
+        @assert future != C_NULL
+        return finalizer(free_future, new(future))
+    end
+end
+
+@inline function free_future(future::Future)
+    @ccall MLIR.API.mlir_c.ifrt_free_future(future.future::Ptr{Cvoid})::Cvoid
+end
+
+function XLA.is_ready(future::Future)
+    GC.@preserve future begin
+        return (@ccall MLIR.API.mlir_c.ifrt_future_is_ready(
+            future.future::Ptr{Cvoid}
+        )::UInt8) != 0
+    end
+end
+
+@inline function XLA.await(future::Future)::Nothing
+    GC.@preserve future begin
+        @ccall MLIR.API.mlir_c.ifrt_future_await(future.future::Ptr{Cvoid})::Cvoid
+    end
+    return nothing
+end

src/xla/IFRT/IFRT.jl

@@ -0,0 +1,15 @@
+module IFRT
+
+using ..Reactant: Reactant, MLIR
+using ..XLA: XLA
+
+include("Client.jl")
+include("Device.jl")
+include("Memory.jl")
+include("Future.jl")
+include("Sharding.jl")
+include("Array.jl")
+include("AsyncArray.jl")
+include("LoadedExecutable.jl")
+
+end

src/xla/IFRT/LoadedExecutable.jl

@@ -0,0 +1,14 @@
+mutable struct LoadedExecutable <: XLA.AbstractLoadedExecutable
+    exec::Ptr{Cvoid}
+
+    function LoadedExecutable(ptr::Ptr{Cvoid})
+        @assert ptr != C_NULL
+        return finalizer(free_exec, new(ptr))
+    end
+end
+
+function free_exec(exec::LoadedExecutable)
+    GC.@preserve exec begin
+        @ccall MLIR.API.mlir_c.ifrt_loaded_executable_dtor(exec.exec::Ptr{Cvoid})::Cvoid
+    end
+end

src/xla/IFRT/Memory.jl

@@ -0,0 +1,47 @@
+mutable struct Memory <: XLA.AbstractMemory
+    ptr::Ptr{Cvoid}
+end
+
+function Base.show(io::IO, memory::Memory)
+    GC.@preserve memory begin
+        str = @ccall MLIR.API.mlir_c.ifrt_MemoryToString(memory.ptr::Ptr{Cvoid})::Cstring
+    end
+    print(io, "XLA.IFRT.Memory(\"", XLA.unsafe_string_and_free(str), "\")")
+    return nothing
+end
+
+mutable struct MemoryKind <: XLA.AbstractMemoryKind
+    ptr::Ptr{Cvoid}
+end
+
+function Base.convert(::Type{MemoryKind}, memory::Memory)
+    GC.@preserve memory begin
+        return MemoryKind(
+            @ccall MLIR.API.mlir_c.ifrt_MemoryGetMemoryKind(
+                memory.ptr::Ptr{Cvoid}
+            )::Ptr{Cvoid}
+        )
+    end
+end
+
+function Base.:(==)(a::MemoryKind, b::MemoryKind)
+    GC.@preserve a b begin
+        return @ccall MLIR.API.mlir_c.ifrt_MemoryKindsAreEqual(
+            a.ptr::Ptr{Cvoid}, b.ptr::Ptr{Cvoid}
+        )::Bool
+    end
+end
+
+function Base.string(memory_kind::MemoryKind)
+    GC.@preserve memory_kind begin
+        str = @ccall MLIR.API.mlir_c.ifrt_MemoryKindToString(
+            memory_kind.ptr::Ptr{Cvoid}
+        )::Cstring
+    end
+    return XLA.unsafe_string_and_free(str)
+end
+
+function Base.show(io::IO, memory_kind::MemoryKind)
+    print(io, "XLA.IFRT.MemoryKind(\"", string(memory_kind), "\")")
+    return nothing
+end

src/xla/IFRT/Sharding.jl

@@ -0,0 +1,83 @@
+# xla::ifrt::HloSharding (distinct from xla::HloSharding)
+mutable struct HloSharding
+    ptr::Ptr{Cvoid}
+
+    function HloSharding(ptr::Ptr{Cvoid})
+        @assert ptr != C_NULL
+        return finalizer(free_hlo_sharding, 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 HloSharding(device_list::BasicDeviceList, 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
+        return HloSharding(
+            @ccall MLIR.API.mlir_c.ifrt_hlo_sharding_from_xla_hlo_sharding(
+                device_list.ptr::Ptr{Cvoid},
+                default_memory_kind.ptr::Ptr{Cvoid},
+                xla_hlo_sharding.ptr::Ptr{Cvoid},
+            )::Ptr{Cvoid}
+        )
+    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))
+    end
+end
+
+function Sharding(device_list::BasicDeviceList, xla_hlo_sharding::XLA.HloSharding)
+    return convert(Sharding, HloSharding(device_list, xla_hlo_sharding))
+end
+
+function free_sharding(sharding::Sharding)
+    @ccall MLIR.API.mlir_c.free_ifrt_sharding(sharding.ptr::Ptr{Cvoid})::Cvoid
+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.string(sharding::Sharding)
+    GC.@preserve sharding begin
+        str = @ccall MLIR.API.mlir_c.ifrt_sharding_to_string(
+            sharding.ptr::Ptr{Cvoid}
+        )::Cstring
+    end
+    return XLA.unsafe_string_and_free(str)
+end
+
+function Base.show(io::IO, ::MIME"text/plain", sharding::Sharding)
+    print(io, "XLA.IFRT.Sharding(\"", string(sharding), "\")")
+    return nothing
+end

src/xla/PJRT/AsyncBuffer.jl

@@ -1,44 +1,8 @@
-mutable struct AsyncBuffer <: XLA.AbstractBuffer
+mutable struct AsyncBuffer <: XLA.AbstractAsyncBuffer
     buffer::Buffer
     future::Union{Future,Nothing}
 end
 
 const AsyncEmptyBuffer = AsyncBuffer(Buffer(C_NULL), nothing)
 
-function AsyncBuffer(client::Client, array::Array{T,N}, device::Device) where {T,N}
-    return AsyncBuffer(Buffer(client, array, device), nothing)
-end
-
-Base.isempty(buffer::AsyncBuffer) = buffer == AsyncEmptyBuffer
-
-function Base.convert(::Type{<:Array{T}}, buffer::AsyncBuffer) where {T}
-    XLA.await(buffer)
-    return convert(Array{T}, buffer.buffer)
-end
-
-for op in (:(Base.ndims), :(Base.size), :device, :client)
-    @eval $op(buffer::AsyncBuffer) = $op(buffer.buffer)
-end
-
-function XLA.synced_buffer(buffer::AsyncBuffer)
-    XLA.await(buffer)
-    return buffer.buffer
-end
-
-function XLA.await(buffer::AsyncBuffer)
-    buffer.future === nothing && return nothing
-    future = buffer.future
-    buffer.future = nothing
-    XLA.await(future)
-    return nothing
-end
-
-function XLA.is_ready(buffer::AsyncBuffer)
-    buffer.future === nothing && return true
-    return XLA.is_ready(buffer.future)
-end
-
-XLA.buffer_on_cpu(buffer::AsyncBuffer) = XLA.buffer_on_cpu(buffer.buffer)
-
-XLA.client(buffer::AsyncBuffer) = XLA.client(buffer.buffer)
-XLA.device(buffer::AsyncBuffer) = XLA.device(buffer.buffer)
+AsyncBuffer(args...; kwargs...) = AsyncBuffer(Buffer(args...; kwargs...), nothing)

src/xla/PJRT/Buffer.jl

@@ -38,7 +38,15 @@ function Base.size(buffer::Buffer)
     GC.@preserve buffer begin
         sz = @ccall MLIR.API.mlir_c.BufferShape(buffer.buffer::Ptr{Cvoid})::Ptr{Int64}
     end
-    return [unsafe_load(sz, i) for i in 1:ndims(buffer)]
+    return Tuple(unsafe_wrap(Array, sz, ndims(buffer)))
+end
+
+function Base.eltype(buffer::Buffer)
+    GC.@preserve buffer begin
+        return XLA.julia_type(
+            @ccall MLIR.API.mlir_c.BufferPrimitiveType(buffer.buffer::Ptr{Cvoid})::Cint
+        )
+    end
 end
 
 function XLA.device(buffer::Buffer)
@@ -65,12 +73,6 @@ function XLA.buffer_on_cpu(buffer::Buffer)
     end
 end
 
-function Base.convert(::Type{<:Array{T}}, buffer::Buffer) where {T}
-    arr = zeros(T, reverse(size(buffer))...)
-    XLA.to_host(buffer, arr)
-    return arr
-end
-
 function XLA.to_host(buffer::Buffer, data)
     GC.@preserve buffer begin
         @ccall MLIR.API.mlir_c.BufferToHost(
@@ -94,3 +96,5 @@ function XLA.copy_buffer_to_device(buffer::Buffer, dev::Device)
         )
     end
 end
+
+XLA.sharding(::Buffer) = error("PJRT Buffers are not sharded.")

src/xla/Utils.jl

@@ -13,37 +13,40 @@ function reactant_err(msg::Cstring)::Cvoid
 end
 
 # https://github.com/openxla/xla/blob/4bfb5c82a427151d6fe5acad8ebe12cee403036a/xla/xla_data.proto#L29
-@inline primitive_type(::Type{Bool}) = 1
-
-@inline primitive_type(::Type{Int8}) = 2
-@inline primitive_type(::Type{UInt8}) = 6
-
-@inline primitive_type(::Type{Int16}) = 3
-@inline primitive_type(::Type{UInt16}) = 7
-
-@inline primitive_type(::Type{Int32}) = 4
-@inline primitive_type(::Type{UInt32}) = 8
-
-@inline primitive_type(::Type{Int64}) = 5
-@inline primitive_type(::Type{UInt64}) = 9
-
-@inline primitive_type(::Type{Float16}) = 10
-@inline primitive_type(::Type{Float32}) = 11
-
-@inline primitive_type(::Type{Reactant.F8E5M2}) = 19
-@inline primitive_type(::Type{Reactant.F8E4M3FN}) = 20
-@inline primitive_type(::Type{Reactant.F8E4M3B11FNUZ}) = 23
-@inline primitive_type(::Type{Reactant.F8E5M2FNUZ}) = 24
-@inline primitive_type(::Type{Reactant.F8E4M3FNUZ}) = 25
+primitive_types_list = [
+    (1, Bool),
+    (2, Int8),
+    (6, UInt8),
+    (3, Int16),
+    (7, UInt16),
+    (4, Int32),
+    (8, UInt32),
+    (5, Int64),
+    (9, UInt64),
+    (10, Float16),
+    (11, Float32),
+    (19, Reactant.F8E5M2),
+    (20, Reactant.F8E4M3FN),
+    (23, Reactant.F8E4M3B11FNUZ),
+    (24, Reactant.F8E5M2FNUZ),
+    (25, Reactant.F8E4M3FNUZ),
+    (12, Float64),
+    (15, Complex{Float32}),
+    (18, Complex{Float64}),
+]
 
 @static if isdefined(Core, :BFloat16)
-    @inline primitive_type(::Type{Core.BFloat16}) = 16
+    push!(primitive_types_list, (16, Core.BFloat16))
 end
 
-@inline primitive_type(::Type{Float64}) = 12
+for (int_val, jl_type) in primitive_types_list
+    @eval begin
+        @inline primitive_type(::Type{$(jl_type)}) = $(int_val)
+        @inline julia_type(::Val{$(int_val)}) = $(jl_type)
+    end
+end
 
-@inline primitive_type(::Type{Complex{Float32}}) = 15
-@inline primitive_type(::Type{Complex{Float64}}) = 18
+@inline julia_type(@nospecialize(x::Integer)) = julia_type(Val(Int64(x)))
 
 function unsafe_string_and_free(str::Cstring, args...)
     str_jl = unsafe_string(str, args...)

src/xla/XLA.jl

@@ -33,6 +33,8 @@ include("Memory.jl")
 
 include("PJRT/PJRT.jl")
 
+include("IFRT/IFRT.jl")
+
 const backends = Dict{String,PJRT.Client}()
 const default_backend = Ref{PJRT.Client}()
 const default_device_idx = Ref{Int}(0)

test/ifrt_manual.jl

@@ -0,0 +1,4 @@
+# Testing manual IFRT buffer creation + compilation + execution
+using Reactant
+using Reactant: XLA
+using Reactant.XLA: IFRT

test/runtests.jl

@@ -62,6 +62,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
             @safetestset "Custom Number Types" include("custom_number_types.jl")
         end
         @safetestset "Sharding" include("sharding.jl")
+        @safetestset "IFRT Low-Level API" include("ifrt_manual.jl")
     end
 
     if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "integration"