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

@@ -1503,6 +1503,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,
@@ -1572,31 +1576,48 @@ 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(
@@ -1636,28 +1657,27 @@ 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, void *distributed_runtime_client) {
+extern "C" ifrt::Client *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, void *distributed_runtime_client) {
   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, distributed_runtime_client));
 }
 
-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));
 }
 
@@ -1847,6 +1867,79 @@ 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
 
 #pragma region PjRtDistributed

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)