fix: correctly set kv_store

Author: Avik Pal

src/xla/Client.jl

@@ -23,17 +23,26 @@ function CPUClient(cfunc, node_id=0, num_nodes=1; asynchronous=true)
     return client
 end
 
-function GPUClient(cfunc, node_id=0, num_nodes=1, platform="gpu"; allowed_devices=nothing)
+function GPUClient(
+    cfunc,
+    node_id=0,
+    num_nodes=1,
+    platform="gpu";
+    allowed_devices=nothing,
+    distributed_runtime_client=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}),
+        (Cint, Cint, Ptr{Cvoid}, Cint, Cdouble, Bool, Cstring, Ptr{Cstring}, Ptr{Cvoid}),
         node_id,
         num_nodes,
         allowed_devices,
@@ -42,6 +51,7 @@ function GPUClient(cfunc, node_id=0, num_nodes=1, platform="gpu"; allowed_device
         false,
         platform,
         refstr,
+        distributed_runtime_client,
     )
     client == C_NULL && throw(AssertionError(unsafe_string(refstr[])))
     LLVMclopts("-nvptx-fma-level=1")

src/xla/XLA.jl

@@ -41,7 +41,7 @@ include("PJRT/PJRT.jl")
 end
 
 function Base.getproperty(bs::BackendState, sym::Symbol)
-    (sym === :initialized  || bs.initialized) && return getfield(bs, sym)
+    (sym === :initialized || bs.initialized) && return getfield(bs, sym)
     initialize_default_clients!(bs)
     return getfield(bs, sym)
 end
@@ -56,6 +56,7 @@ const global_backend_state = BackendState()
 
 client(backend::String) = global_backend_state.clients[backend]
 default_backend() = global_backend_state.default_client
+process_index() = process_index(default_backend())
 
 function set_default_backend(backend::AbstractClient)
     global_backend_state.default_client = backend
@@ -141,6 +142,13 @@ function initialize_default_clients!(state::BackendState)
         else
             if !Reactant.precompiling()
                 try
+                    distributed_runtime_client = if PJRT.global_state.num_processes > 1
+                        @assert PJRT.global_state.client !== nothing
+                        PJRT.global_state.client
+                    else
+                        nothing
+                    end
+
                     if was_initialized && haskey(state.clients, "gpu")
                         XLA.free_client(state.clients["gpu"])
                         XLA.PJRT.gpu_client_count[] -= 1
@@ -149,6 +157,7 @@ function initialize_default_clients!(state::BackendState)
                         PJRT.global_state.process_id,
                         PJRT.global_state.num_processes;
                         allowed_devices=PJRT.global_state.local_device_ids,
+                        distributed_runtime_client,
                     )
                     state.clients["gpu"] = gpu
                     state.default_client = gpu