fix: sharding improvements for GB (#919)

* fix: warn

* fix: auto-raise if we are using sharding

* feat: fix codegen for mixed sharding + nosharidng

* fix(IFRT): correct device to host transfer

Author: avik-pal

src/Compiler.jl

@@ -42,14 +42,34 @@ end
     @nospecialize(obj::AbstractArray{T}), field, val
 ) where {T}
     ancestor_obj = ancestor(obj)
-    (isbitstype(T) || ancestor_obj isa RArray) && return Base.setfield!(obj, field, val)
+    (isbitstype(T) || ancestor_obj isa RArray) && return setfield_carray!(obj, field, val)
     return Base.setindex!(obj, val, field)
 end
 
 @inline function traced_setfield!(@nospecialize(obj::Dict), field, val)
     return Base.setindex!(obj, field, val)
 end
 
+# fallback
+@inline function setfield_carray!(obj, field, val)
+    return Base.setfield!(obj, field, val)
+end
+
+@inline function setfield_carray!(obj::ConcretePJRTArray, field, val)
+    if field !== :data || typeof(val) == typeof(getfield(obj, field))
+        return Base.setfield!(obj, field, val)
+    end
+
+    # This case is triggered if the user had provided an unsharded input (NoSharding), but
+    # we had to replicate it before feeding it to XLA
+    @assert !Reactant.Sharding.is_sharded(obj) "Expected unsharded input. Open an issue on \
+                                                Reactant.jl with a MWE."
+    devices = Reactant.XLA.device.(val)
+    device = Reactant.XLA.device(only(obj.data))
+    idx = findfirst(isequal(device), devices)
+    return Base.setfield!(obj, field, (val[idx],))
+end
+
 function create_result(
     tocopy::T, path, result_stores, path_to_shard_info, sharding_mesh
 ) where {T}
@@ -737,10 +757,25 @@ function compile_mlir!(
         MLIR.IR.deactivate!(MLIR.IR.body(mod))
         MLIR.IR.deactivate!(mod)
     end
-    (; fnwrapped, traced_result, seen_args, ret, linear_args, in_tys, linear_results) =
-        mlir_fn_res
+    (;
+        fnwrapped,
+        traced_result,
+        seen_args,
+        ret,
+        linear_args,
+        in_tys,
+        linear_results,
+        is_sharded,
+    ) = mlir_fn_res
     compiled_f = mlir_fn_res.f
 
+    # Custom Kernels without Raising will lead to suboptimal codegen for is_sharded, force
+    # raising
+    if is_sharded
+        is_raising = true
+        raise isa Bool && (raise = true)
+    end
+
     concrete_seen = OrderedIdDict()
 
     concrete_result = make_tracer(
@@ -1274,20 +1309,40 @@ function codegen_flatten!(
                     device_to_array_slices, _ = XLA.sharding_to_concrete_array_indices(
                         condensed_op_sharding, size(carg), mesh.logical_device_ids
                     )
+
+                    # Extract the buffer_slice
+                    buf_slice = Dict{eltype(device_to_array_slices),Symbol}()
+                    counter = 0
                     for j in 1:length(mesh)
-                        device_id = mesh.logical_device_ids[j]
-                        buf = Symbol(:buf_, i, :_, device_id)
+                        sliced_buf = Symbol(:sliced_buf_, i, :_, counter)
                         slice = device_to_array_slices[j]
+                        haskey(buf_slice, slice) && continue
+                        counter += 1
                         push!(
                             flatten_code,
-                            :($buf = XLA.synced_buffer(only($usbuf[$(slice)...].data))),
+                            :(
+                                $sliced_buf = only(
+                                    Reactant._fast_slice($usbuf, $(slice...)).data
+                                )
+                            ),
                         )
+                        buf_slice[slice] = sliced_buf
+                    end
+
+                    for j in 1:length(mesh)
+                        device_id = mesh.logical_device_ids[j]
+                        buf = Symbol(:buf_, i, :_, device_id)
+                        slice = device_to_array_slices[j]
                         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)),
+                            :(
+                                $sbuf = XLA.copy_buffer_to_device(
+                                    XLA.synced_buffer($(buf_slice[slice])),
+                                    $(XLA.get_device(client, device_id)),
+                                )
+                            ),
                         )
                     end
                 end

src/ConcreteRArray.jl

@@ -251,6 +251,18 @@ function Base.getindex(a::ConcreteIFRTArray, args::Vararg{Int,N}) where {N}
     return convert(Array, a)[args...]
 end
 
+# This doesn't follow the semantics of getindex with ranges. It is mostly meant to be used
+# inside Compiler.jl
+@inline function _fast_slice(
+    a::AbstractConcreteArray{T,N}, args::Vararg{UnitRange,N}
+) where {T,N}
+    # Avoid slicing all-together
+    args == ntuple(Base.Fix1(UnitRange, 1) ∘ Base.Fix1(size, a), N) && return a
+    # For all other cases do a compile
+    fn = compile(getindex, (a, args...))
+    return fn(a, args...)
+end
+
 function mysetindex!(a, v, args::Vararg{Any,N}) where {N}
     setindex!(a, v, args...)
     return nothing

src/TracedUtils.jl

@@ -226,9 +226,7 @@ function make_mlir_fn(
     # Insert meshes for the sharded arguments
     traced_args_to_shardings = OrderedIdDict()
     for (k, v) in seen_args
-        if (
-            k isa Reactant.AbstractConcreteNumber || k isa Reactant.AbstractConcreteArray
-        ) && hasfield(typeof(k), :sharding)
+        if k isa Reactant.AbstractConcreteNumber || k isa Reactant.AbstractConcreteArray
             if Reactant.Sharding.is_sharded(k)
                 Reactant.Ops.mesh(k.sharding.mesh)
                 traced_args_to_shardings[v] = k.sharding
@@ -373,6 +371,18 @@ function make_mlir_fn(
 
         linear_arg_shardings = Vector{MLIR.IR.Attribute}(undef, length(linear_args))
 
+        # If an argument is mutated but is not sharded (aka sharding is NoSharding), we
+        # need to force a replicated sharding.
+        for i in mutated_args
+            arg = linear_args[i]
+            if !haskey(traced_args_to_shardings, arg)
+                # Force a replicated sharding
+                traced_args_to_shardings[arg] = Reactant.Sharding.NamedSharding(
+                    sharding_mesh, ntuple(Returns(nothing), ndims(arg))
+                )
+            end
+        end
+
         # Attach `sdy.sharding` attribute to the argument
         for (i, arg) in enumerate(linear_args)
             if haskey(traced_args_to_shardings, arg)

src/Types.jl

@@ -163,8 +163,10 @@ function ConcretePJRTArray(
         sdata, sharding = sharding(client, device, data)
         return ConcretePJRTArray{T,N,1,typeof(sharding)}(sdata, size(data), sharding)
     end
-    @assert device === nothing && idx === nothing "If `sharding` is not `NoSharding`, \
-                                                   `device` and `idx` cannot be specified!"
+    if device !== nothing || idx !== nothing
+        @warn "`device` and `idx` specified for non-`NoSharding` sharding. These arguments \
+               will be ignored."
+    end
     sharded_data, sharding = sharding(client, nothing, data)
     return ConcretePJRTArray{T,N,length(sharded_data),typeof(sharding)}(
         sharded_data, size(data), sharding
@@ -282,9 +284,10 @@ function ConcreteIFRTArray(
             end
         end
     else
-        @assert device === nothing && idx === nothing "If `sharding` is not `NoSharding`, \
-                                                       `device` and `idx` cannot be \
-                                                       specified!"
+        if device !== nothing || idx !== nothing
+            @warn "`device` and `idx` specified for non-`NoSharding` sharding. These \
+                   arguments will be ignored."
+        end
     end
     sharded_data, sharding = sharding(client, device, data)
     return ConcreteIFRTArray{T,N}(sharded_data, size(data), sharding)

src/xla/Buffer.jl

@@ -23,12 +23,6 @@ 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
 

src/xla/IFRT/Array.jl

@@ -140,10 +140,15 @@ end
 function XLA.to_host(buffer::Array, data, reactant_sharding)
     reactant_sharding = Reactant.Sharding.unwrap_shardinfo(reactant_sharding)
 
+    # While some client implementations might support directly copying to host, but we
+    # avoid the complexity of supporting that for now.
+    single_device_arrays = disassemble_into_single_device_arrays(buffer, true)
+
     if reactant_sharding isa Reactant.Sharding.NoSharding
-        GC.@preserve buffer data begin
+        data_buffer = first(single_device_arrays)
+        GC.@preserve data_buffer data begin
             @ccall MLIR.API.mlir_c.ifrt_array_copy_to_host_buffer(
-                buffer.buffer::Ptr{Cvoid}, data::Ptr{Cvoid}
+                data_buffer.buffer::Ptr{Cvoid}, data::Ptr{Cvoid}
             )::Cvoid
         end
         return data
@@ -159,10 +164,6 @@ function XLA.to_host(buffer::Array, data, reactant_sharding)
                untouched."
     end
 
-    # While some client implementations might support directly copying to host, but we 
-    # avoid the complexity of supporting that for now.
-    single_device_arrays = disassemble_into_single_device_arrays(buffer, true)
-
     array_slices, _ = XLA.sharding_to_concrete_array_indices(
         convert(XLA.CondensedOpSharding, reactant_sharding.hlo_sharding),
         size(data),

test/sharding.jl

@@ -262,3 +262,34 @@ end
         [2, 1],        #=iota_transpose_perm=#
     ) == [0, 2, 4, 6, 1, 3, 5, 7]
 end
+
+@testset "Sharding with Mutation" begin
+    if length(addressable_devices) ≥ 8
+        mesh = Sharding.Mesh(reshape(Reactant.addressable_devices(), 2, 2, 2), (:x, :y, :z))
+
+        x_ra = Reactant.to_rarray(
+            randn(Float32, 4, 5); sharding=Sharding.NamedSharding(mesh, ((:x, :y), :z))
+        )
+
+        y_ra = Reactant.to_rarray(randn(Float32, 5, 4); sharding=Sharding.NoSharding())
+
+        function fn(x, y)
+            z = x * y
+            y[1:2, 1:2] .= 1
+            return z
+        end
+
+        y_ra_arr = Array(y_ra)
+        x_ra_arr = Array(x_ra)
+        z_ra_arr = fn(x_ra_arr, y_ra_arr)
+
+        z_ra = @jit fn(x_ra, y_ra)
+        y_ra_final = Array(y_ra)
+
+        @test z_ra_arr ≈ Array(z_ra)
+        @test y_ra_final[1:2, 1:2] ≈ y_ra_arr[1:2, 1:2]
+        @test all(y_ra_final[1:2, 1:2] .== 1)
+    else
+        @warn "Not enough addressable devices to run sharding tests"
+    end
+end