feat: initial IFRT integration (#764)

* feat: Low-Level XLA.IFRT integration

refactor: rework how OpSharding works

feat: generate_device_list

feat: add placeholder code to simplify future sharding logic

fixup

fix: store results as HloSharding

docs: fix duplicate docs

feat: compile with logical device ids

fix: use correct global device ids

feat: use a global state to setup pjrt distributed runtime

fix: devices are not necessarily from 0 to N-1

fix: initialize clients on first use rather than on init

fix: make device selection consistent with clients

feat: add OMPI cluster detection

fix: correctly set kv_store

refactor: Distributed setup is not PJRT specific

refactor: OMPI detection doesn't need to be in an extension

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

feat: construct IFRT clients with distributed options

refactor: remove BasicDevicesList

fix: use global device ids

feat: sharding annotations across nodes now working

fix: Array construction from SingleShards

feat: support to_host for distributed cases

feat: add Gloo/MPI collectives for distributed CPU client

feat: low level compile API

feat: low-level IFRT compile + execute working

* test: low level IFRT tests

* docs: add warning on distributed cases

Author: avik-pal

src/Compiler.jl

@@ -508,13 +508,6 @@ function compile_mlir(f, args; client=nothing, kwargs...)
 
         mlir_fn_res = compile_mlir!(mod, f, args; backend, kwargs...)
 
-        client, _ = __resolve_device_and_client(
-            client,
-            mlir_fn_res.seen_args,
-            mlir_fn_res.linear_args,
-            mlir_fn_res.is_sharded,
-        )
-
         # Attach a name, and partitioning attributes to the module
         __add_mhlo_attributes_and_name!(
             mod, f; mlir_fn_res.num_partitions, mlir_fn_res.num_replicas
@@ -1509,6 +1502,8 @@ function compile_xla(f, args; client=nothing, kwargs...)
             num_parameters=length(mlir_fn_res.linear_args),
             mlir_fn_res.is_sharded,
             global_device_ids,
+            mlir_fn_res.num_replicas,
+            mlir_fn_res.num_partitions,
         )
 
         return mod, exec, mlir_fn_res, device, client

src/Distributed.jl

@@ -8,24 +8,24 @@ function initialize(;
     coordinator_address::Union{Nothing,String}=nothing,
     num_processes::Union{Nothing,Integer}=nothing,
     process_id::Union{Nothing,Integer}=nothing,
-    local_device_ids::Union{Nothing,Vector{Int}}=nothing,
+    local_gpu_device_ids::Union{Nothing,Vector{Int}}=nothing,
     initialization_timeout_in_seconds::Integer=300,
     kwargs...,
 )
     @assert !initialized[] "`Distributed.initialize` has already been called"
 
-    (coordinator_address, num_processes, process_id, local_device_ids) = auto_detect_unset_distributed_params(;
+    (coordinator_address, num_processes, process_id, local_gpu_device_ids) = auto_detect_unset_distributed_params(;
         coordinator_address,
         num_processes,
         process_id,
-        local_device_ids,
+        local_gpu_device_ids,
         initialization_timeout_in_seconds,
     )
 
-    @debug "Detected Reactant distributed params" coordinator_address num_processes process_id local_device_ids
+    @debug "Detected Reactant distributed params" coordinator_address num_processes process_id local_gpu_device_ids
 
     Reactant.XLA.update_global_state!(;
-        coordinator_address, num_processes, process_id, local_device_ids, kwargs...
+        coordinator_address, num_processes, process_id, local_gpu_device_ids, kwargs...
     )
 
     @debug "New Global State" Reactant.XLA.global_state
@@ -57,14 +57,14 @@ function auto_detect_unset_distributed_params(;
     coordinator_address::Union{Nothing,String}=nothing,
     num_processes::Union{Nothing,Integer}=nothing,
     process_id::Union{Nothing,Integer}=nothing,
-    local_device_ids::Union{Nothing,Vector{Int}}=nothing,
+    local_gpu_device_ids::Union{Nothing,Vector{Int}}=nothing,
     initialization_timeout_in_seconds::Integer=300,
 )
     if all(
         Base.Fix2(!==, nothing),
-        (coordinator_address, num_processes, process_id, local_device_ids),
+        (coordinator_address, num_processes, process_id, local_gpu_device_ids),
     )
-        return coordinator_address, num_processes, process_id, local_device_ids
+        return coordinator_address, num_processes, process_id, local_gpu_device_ids
     end
 
     idx = findfirst(is_env_present, detector_list)
@@ -91,11 +91,11 @@ function auto_detect_unset_distributed_params(;
         process_id = get_process_id(detector)
     end
 
-    if local_device_ids === nothing
-        local_device_ids = [get_local_process_id(detector)]
+    if local_gpu_device_ids === nothing
+        local_gpu_device_ids = [get_local_process_id(detector)]
     end
 
-    return coordinator_address, num_processes, process_id, local_device_ids
+    return coordinator_address, num_processes, process_id, local_gpu_device_ids
 end
 
 # OpenMPIORTEEnvDetector & OpenMPIPMIXEnvDetector

src/Types.jl

@@ -135,11 +135,11 @@ function ConcretePJRTArray(
             if idx === nothing
                 device = XLA.default_device(client)
             else
-                device = XLA.get_addressable_device(client, idx)
+                device = XLA.get_device(client, idx)
             end
         else
             if idx !== nothing
-                device_from_idx = XLA.get_addressable_device(client, idx)
+                device_from_idx = XLA.get_device(client, idx)
                 @assert device_from_idx == device "If both `idx` and `device` are \
                                                    specified, `idx` must match `device`"
             end

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/Client.jl

@@ -14,55 +14,3 @@ function get_addressable_device end
 function platform_name end
 
 default_device(client::AbstractClient) = first(addressable_devices(client))
-
-# Clients for Different Backends
-function CPUClient(cfunc, node_id=0, num_nodes=1; asynchronous=true)
-    f = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, string(cfunc))
-    client = ccall(f, Ptr{Cvoid}, (UInt, Cint, Cint), asynchronous, node_id, num_nodes)
-    LLVMclopts("-nvptx-fma-level=1")
-    return client
-end
-
-function GPUClient(
-    cfunc,
-    node_id=0,
-    num_nodes=1,
-    platform="gpu";
-    allowed_devices::Union{Nothing,Vector{Int}}=nothing,
-    distributed_runtime_client::Union{Nothing,DistributedRuntimeClient}=nothing,
-)
-    f = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, string(cfunc))
-    refstr = Ref{Cstring}()
-
-    num_allowed_devices = allowed_devices === nothing ? 0 : length(allowed_devices)
-    allowed_devices = allowed_devices === nothing ? C_NULL : allowed_devices
-    distributed_runtime_client =
-        distributed_runtime_client === nothing ? C_NULL : distributed_runtime_client.client
-
-    client = ccall(
-        f,
-        Ptr{Cvoid},
-        (Cint, Cint, Ptr{Cvoid}, Cint, Cdouble, Bool, Cstring, Ptr{Cstring}, Ptr{Cvoid}),
-        node_id,
-        num_nodes,
-        allowed_devices,
-        num_allowed_devices,
-        XLA_REACTANT_GPU_MEM_FRACTION[],
-        false,
-        platform,
-        refstr,
-        distributed_runtime_client,
-    )
-    client == C_NULL && throw(AssertionError(unsafe_string(refstr[])))
-    LLVMclopts("-nvptx-fma-level=1")
-    return client
-end
-
-function TPUClient(cfunc, tpu_path::String)
-    f = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, string(cfunc))
-    refstr = Ref{Cstring}()
-    client = ccall(f, Ptr{Cvoid}, (Cstring, Ptr{Cstring}), tpu_path, refstr)
-    client == C_NULL && throw(AssertionError(unsafe_string(refstr[])))
-    LLVMclopts("-nvptx-fma-level=1")
-    return client
-end

src/xla/Device.jl

@@ -10,19 +10,15 @@ function get_local_device_id end
 function device_kind end
 function default_memory end
 function memories end
+function is_addressable end
 
 """
     device_ordinal(device::Device)
-    device_ordinal(client::XLA.AbstractClient, local_device_id::Int)
 
-Given the device or local device id, return the corresponding global device ordinal in the client.
+Given the device, return the corresponding global device ordinal in the client.
 """
 function device_ordinal end
 
-function device_ordinal(client::AbstractClient, local_device_id::Integer)
-    return device_ordinal(get_addressable_device(client, local_device_id))
-end
-
 function Base.string(device::AbstractDevice)
     client = XLA.client(device)
     pname = XLA.platform_name(client)

src/xla/Distributed.jl

@@ -106,7 +106,7 @@ end
 @kwdef mutable struct State
     process_id::Int = 0
     num_processes::Int = 1
-    local_device_ids::Union{Nothing,Vector{Int}} = nothing
+    local_gpu_device_ids::Union{Nothing,Vector{Int}} = nothing
     service::Union{Nothing,DistributedRuntimeService} = nothing
     client::Union{Nothing,DistributedRuntimeClient} = nothing
     coordinator_address::Union{Nothing,String} = nothing
@@ -129,7 +129,7 @@ function update!(
     coordinator_address::String,
     num_processes::Int,
     process_id::Int,
-    local_device_ids::Vector{Int},
+    local_gpu_device_ids::Vector{Int},
     coordinator_bind_address::Union{Nothing,String}=nothing,
     cluster_register_timeout_in_minutes::Integer=60,
     rpc_timeout_in_seconds::Integer=120,
@@ -141,7 +141,7 @@ function update!(
     @assert 0 ≤ process_id < num_processes
 
     state.coordinator_address = coordinator_address
-    state.local_device_ids = local_device_ids
+    state.local_gpu_device_ids = local_gpu_device_ids
     state.process_id = process_id
     state.num_processes = num_processes