fix: use global device ids

Author: avik-pal

src/Compiler.jl

@@ -1101,15 +1101,15 @@ function codegen_flatten!(
                     condensed_op_sharding, size(carg), mesh
                 )
                 for j in 1:length(mesh)
-                    local_device_id = device_ids[j]
-                    buf = Symbol(:buf_, i, :_, local_device_id)
+                    device_id = device_ids[j]
+                    buf = Symbol(:buf_, i, :_, device_id)
                     slice = device_to_array_slices[j]
                     push!(
                         flatten_code,
                         :($buf = XLA.synced_buffer(only($usbuf[$(slice)...].data))),
                     )
-                    sbuf = Symbol(:sbuf_, i, :_, local_device_id)
-                    device = XLA.get_addressable_device(client, local_device_id)
+                    sbuf = Symbol(:s, buf)
+                    device = XLA.get_device(client, device_id)
                     push!(flatten_names, sbuf)
                     push!(flatten_code, :($sbuf = XLA.copy_buffer_to_device($buf, $device)))
                 end
@@ -1457,7 +1457,7 @@ function compile_xla(f, args; client=nothing, kwargs...)
         )
 
         # compile MLIR module to XLA executable
-        local_device_ids = if mlir_fn_res.is_sharded
+        global_device_ids = if mlir_fn_res.is_sharded
             collect(Int64, mlir_fn_res.sharding_mesh.sorted_device_ids)
         else
             Int64[]
@@ -1471,7 +1471,7 @@ function compile_xla(f, args; client=nothing, kwargs...)
             num_outputs=length(mlir_fn_res.linear_results),
             num_parameters=length(mlir_fn_res.linear_args),
             mlir_fn_res.is_sharded,
-            local_device_ids,
+            global_device_ids,
         )
 
         return mod, exec, mlir_fn_res, device, client

src/Sharding.jl

@@ -8,14 +8,14 @@ struct Mesh{D,ND}
     shape::Dims{D}
     axis_names::NTuple{D,Symbol}
 
-    function Mesh(devices::AbstractArray{XLA.AbstractDevice}, axis_names)
-        return Mesh(XLA.get_local_device_id.(devices), axis_names)
+    function Mesh(devices::AbstractArray{<:XLA.AbstractDevice}, axis_names)
+        return Mesh(XLA.device_ordinal.(devices), axis_names)
     end
 
     function Mesh(
-        devices::NTuple{D,XLA.AbstractDevice}, shape::Dims{D}, axis_names
+        devices::NTuple{D,<:XLA.AbstractDevice}, shape::Dims{D}, axis_names
     ) where {D}
-        return Mesh(XLA.get_local_device_id.(devices), shape, axis_names)
+        return Mesh(XLA.device_ordinal.(devices), shape, axis_names)
     end
 
     function Mesh(
@@ -114,7 +114,7 @@ function (sharding::NamedSharding)(
         XLA.PJRT.AsyncBuffer(
             client,
             x[device_to_array_slices[i]...],
-            XLA.get_addressable_device(client, mesh.sorted_device_ids[i]),
+            XLA.get_device(client, mesh.sorted_device_ids[i]),
         )
     end
 
@@ -199,7 +199,7 @@ function (sharding::LazySharding)(
     client::XLA.PJRT.Client, ::Nothing, x::Union{AbstractArray,Number}
 )
     data = XLA.PJRT.AsyncBuffer(
-        client, x, XLA.get_addressable_device(client, vec(sharding.sharding.mesh)[1])
+        client, x, XLA.get_device(client, vec(sharding.sharding.mesh)[1])
     )
 
     return (data,), ShardInfo(sharding, (ntuple(i -> 1:size(x, i), ndims(x)),))

src/Types.jl

@@ -123,11 +123,11 @@ function ConcreteRArray(
             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/Device.jl

@@ -13,16 +13,11 @@ function memories 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/PJRT/LoadedExecutable.jl

@@ -81,12 +81,11 @@ function XLA.compile(
     device::Union{Device,Nothing},
     mod::MLIR.IR.Module;
     is_sharded::Bool=false,
-    local_device_ids::Vector{Int64}=Int64[],
+    global_device_ids::Vector{Int64}=Int64[],
     num_outputs::Int64,
     num_parameters::Int64,
 )
     device_id = is_sharded ? Int64(-1) : Int64(XLA.device_ordinal(device))
-    global_device_ids = Int64.(XLA.device_ordinal.((client,), local_device_ids))
     GC.@preserve client mod begin
         exec = @ccall MLIR.API.mlir_c.ClientCompile(
             client.client::Ptr{Cvoid},