fix: codegen for resharded input arrays (#1031)

* fix: codegen for resharded input arrays

* fix: handle unresharding path II

* fix: set

* refactor: add a with_context function

* fix: don't rerun tracing for replicating

* feat: fast-path using hlo_sharding

* fix: missing dispatch

* chore: fix comment

* feat: cache unresharding

* fix: annotate resargs shardings as well

* fix: error on unsupported path

* test: SingleDeviceSharding test

* fix: codegen for cached

* fix: field name

* fix: mlir codegen

Author: avik-pal

src/Compiler.jl

@@ -582,14 +582,7 @@ struct HloSharding{D1,D2,PS} <: AbstractSharding
 end
 
 function Base.convert(::Type{HloSharding}, sharding::NamedSharding)
-    if MLIR.IR._has_context()
-        ctx = MLIR.IR.context()
-    else
-        ctx = MLIR.IR.Context(Reactant.registry[], false)
-        @ccall MLIR.API.mlir_c.RegisterDialects(ctx::MLIR.API.MlirContext)::Cvoid
-    end
-
-    MLIR.IR.context!(ctx) do
+    MLIR.IR.with_context(; allow_use_existing=true) do ctx
         mesh_op = Reactant.Ops.mesh(
             sharding.mesh; mod=MLIR.IR.Module(MLIR.IR.Location(; context=ctx))
         )

src/Sharding.jl

@@ -452,18 +452,33 @@ function make_mlir_fn(
         end
 
         # Ensure the sharding of the mutated arguments is propagated to the results
-        result_not_replicated = falses(length(linear_results))
         for i in mutated_args
             arg = linear_args[i]
-            if has_residx(arg) && haskey(traced_args_to_shardings, arg)
-                residx = findfirst(Base.Fix1(===, arg), linear_results)
-                @assert residx !== nothing
-                result_not_replicated[residx] = true
+
+            if haskey(traced_args_to_shardings, arg) &&
+                (has_residx(arg) || has_resargidx(arg))
+                idx = findfirst(Base.Fix1(===, arg), linear_results)
+                @assert idx !== nothing
                 attr, dialect = linear_arg_shardings[i]
                 if dialect == :sdy
-                    MLIR.API.mlirFuncSetResultAttr(func2, residx - 1, "sdy.sharding", attr)
+                    MLIR.API.mlirFuncSetResultAttr(func2, idx - 1, "sdy.sharding", attr)
+                elseif dialect == :mhlo
+                    MLIR.API.mlirFuncSetResultAttr(func2, idx - 1, "mhlo.sharding", attr)
+                else
+                    error("Unsupported dialect for tensor sharding: $(dialect)")
+                end
+            end
+        end
+
+        for (i, res) in enumerate(linear_results)
+            if has_argidx(res) && haskey(traced_args_to_shardings, res)
+                argidx = findfirst(Base.Fix1(===, res), linear_args)
+                @assert argidx !== nothing
+                attr, dialect = linear_arg_shardings[argidx]
+                if dialect == :sdy
+                    MLIR.API.mlirFuncSetResultAttr(func2, i - 1, "sdy.sharding", attr)
                 elseif dialect == :mhlo
-                    MLIR.API.mlirFuncSetResultAttr(func2, residx - 1, "mhlo.sharding", attr)
+                    MLIR.API.mlirFuncSetResultAttr(func2, i - 1, "mhlo.sharding", attr)
                 else
                     error("Unsupported dialect for tensor sharding: $(dialect)")
                 end
@@ -562,30 +577,6 @@ function push_val!(ad_inputs, x, path)
     return push!(ad_inputs, x)
 end
 
-function get_argidx(x)
-    for path in get_paths(x)
-        if length(path) == 0
-            continue
-        end
-        if path[1] == :args
-            return path[2]::Int, path
-        end
-    end
-    throw(AssertionError("No path found for $x"))
-end
-
-function has_argidx(x)
-    for path in get_paths(x)
-        if length(path) == 0
-            continue
-        end
-        if path[1] == :args
-            return true
-        end
-    end
-    return false
-end
-
 function set!(x, path, tostore; emptypath=false)
     for p in path
         x = Reactant.Compiler.traced_getfield(x, p)
@@ -596,28 +587,33 @@ function set!(x, path, tostore; emptypath=false)
     return emptypath && set_paths!(x, ())
 end
 
-function get_residx(x)
-    for path in get_paths(x)
-        if length(path) == 0
-            continue
-        end
-        if path[1] == :result
-            return path
+for (fn, key) in ((:arg, :args), (:res, :result), (:resarg, :resargs))
+    has_fn = Symbol(:has_, fn, :idx)
+    @eval begin
+        function $(has_fn)(x)
+            for path in get_paths(x)
+                length(path) == 0 && continue
+                path[1] == $(Meta.quot(key)) && return true
+            end
+            return false
         end
     end
-    throw(AssertionError("No path found $x"))
 end
 
-function has_residx(x)
+function get_argidx(x)
+    for path in get_paths(x)
+        length(path) == 0 && continue
+        path[1] == :args && return (path[2]::Int, path)
+    end
+    throw(AssertionError("No path found for $x"))
+end
+
+function get_residx(x)
     for path in get_paths(x)
-        if length(path) == 0
-            continue
-        end
-        if path[1] == :result
-            return true
-        end
+        length(path) == 0 && continue
+        path[1] == :result && return path
     end
-    return false
+    throw(AssertionError("No path found for $x"))
 end
 
 function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}

src/TracedUtils.jl

@@ -79,6 +79,33 @@ function context!(f, ctx::Context)
     end
 end
 
+function with_context(f; allow_use_existing=false)
+    delete_context = false
+    if allow_use_existing && _has_context()
+        ctx = context()
+    else
+        delete_context = true
+        ctx = Context(Reactant.registry[], false)
+        Reactant.Compiler.context_gc_vector[ctx] = Vector{
+            Union{Reactant.TracedRArray,Reactant.TracedRNumber}
+        }(
+            undef, 0
+        )
+        @ccall API.mlir_c.RegisterDialects(ctx::API.MlirContext)::Cvoid
+    end
+
+    activate!(ctx)
+    result = try
+        f(ctx)
+    finally
+        deactivate!(ctx)
+    end
+
+    delete_context && Base.delete!(Reactant.Compiler.context_gc_vector, ctx)
+
+    return result
+end
+
 function enable_multithreading!(enable::Bool=true; context::Context=context())
     API.mlirContextEnableMultithreading(context, enable)
     return context

src/mlir/IR/Context.jl

@@ -223,28 +223,100 @@ function replicate_array_to_all_devices(array::Array, sharding, mesh, size_arr)
         )
     end
 
-    shard_info = Reactant.Sharding.ShardInfo(
-        reactant_sharding,
-        Reactant.Sharding.sharding_to_array_slices(reactant_sharding, size_arr),
-    )
-    sharding_constraint = Reactant.Sharding.NamedSharding(
+    XLA.is_replicated(hlo_sharding) && return array
+
+    output_sharding = Reactant.Sharding.NamedSharding(
         mesh, ntuple(Returns(nothing), length(size_arr))
     )
 
-    data = Reactant.ConcreteIFRTArray{eltype(array),length(size_arr),typeof(shard_info)}(
-        AsyncArray(array, nothing), size_arr, shard_info
+    # Manually write the MLIR for resharding resharding
+    ctx = MLIR.IR.Context(Reactant.registry[], false)
+    Reactant.Compiler.context_gc_vector[ctx] = Vector{
+        Union{Reactant.TracedRArray,Reactant.TracedRNumber}
+    }(
+        undef, 0
     )
+    @ccall MLIR.API.mlir_c.RegisterDialects(ctx::MLIR.API.MlirContext)::Cvoid
+    MLIR.IR.activate!(ctx)
+
+    sdycache = IdDict{
+        Reactant.Sharding.Mesh,
+        @NamedTuple{
+            sym_name::MLIR.IR.Attribute,
+            mesh_attr::MLIR.IR.Attribute,
+            mesh_op::MLIR.IR.Operation,
+        }
+    }()
+    Reactant.Compiler.activate_sdycache!(sdycache)
+
+    output_buffer = try
+        data_mlir_type = [MLIR.IR.TensorType(reverse(size_arr), MLIR.IR.Type(eltype(array)))]
+        mod = MLIR.IR.Module(MLIR.IR.Location(; context=ctx))
+
+        (; sym_name, mesh_attr) = Reactant.Ops.mesh(mesh; mod=mod)
+        common_args = (ctx, sym_name, mesh_attr, size_arr)
+        common_kwargs = (; dialect=:sdy, do_transpose=true)
+        input_tensor_sharding_attr, _ = Reactant.Sharding.get_tensor_sharding_attribute(
+            reactant_sharding, common_args...; common_kwargs...
+        )
+        output_tensor_sharding_attr, _ = Reactant.Sharding.get_tensor_sharding_attribute(
+            output_sharding, common_args...; common_kwargs...
+        )
 
-    # TODO: Directly write the MLIR for this part??
-    fn_compiled = Reactant.compile(
-        identity,
-        (data,);
-        shardy_passes=:to_mhlo_shardings,
-        optimize=false,
-        output_shardings=Dict(1 => sharding_constraint),
-    )
+        func = MLIR.Dialects.func.func_(;
+            sym_name="main",
+            function_type=MLIR.IR.FunctionType(data_mlir_type, data_mlir_type),
+            no_inline=true,
+            body=MLIR.IR.Region(),
+        )
+        fnbody = MLIR.IR.Block(data_mlir_type, [MLIR.IR.Location()])
+        push!(MLIR.IR.region(func, 1), fnbody)
+        MLIR.IR.activate!(fnbody)
+        try
+            MLIR.Dialects.func.return_([MLIR.IR.argument(fnbody, 1)])
+        finally
+            MLIR.IR.deactivate!(fnbody)
+        end
+        push!(MLIR.IR.body(mod), func)
+
+        MLIR.API.mlirFuncSetArgAttr(func, 0, "sdy.sharding", input_tensor_sharding_attr)
+        MLIR.API.mlirFuncSetResultAttr(func, 0, "sdy.sharding", output_tensor_sharding_attr)
+
+        Reactant.Compiler.run_pass_pipeline!(
+            mod,
+            join(
+                [
+                    "sdy-propagation-pipeline",
+                    "sdy-close-shardings",
+                    "xla-sdy-stablehlo-export-pipeline",
+                    "canonicalize",
+                    "cse",
+                ],
+                ",",
+            ),
+        )
+
+        exec = XLA.compile(
+            XLA.client(array),
+            nothing,
+            mod;
+            is_sharded=true,
+            global_device_ids=vec(mesh.device_ids),
+            num_outputs=1,                # unused
+            num_parameters=1,             # unused
+            num_replicas=-1,              # unused
+            num_partitions=-1,            # unused
+            use_shardy_partitioner=false, # unused
+        )
+
+        only(XLA.execute(exec, (array.buffer,), (UInt8(0),), Val(1)))
+    finally
+        Reactant.Compiler.deactivate_sdycache!(sdycache)
+        MLIR.IR.deactivate!(ctx)
+    end
+    delete!(Reactant.Compiler.context_gc_vector, ctx)
 
-    return fn_compiled(data).data.buffer
+    return output_buffer
 end
 
 function XLA.unsafe_buffer_pointer(::Array)

src/xla/IFRT/Array.jl

@@ -18,3 +18,10 @@ function replicate_array_to_all_devices(array::AsyncArray, args...)
     wait(array)
     return replicate_array_to_all_devices(array.buffer, args...)
 end
+
+function XLA.to_host(array::AsyncArray, data, reactant_sharding)
+    wait(array)
+    return XLA.to_host(array.buffer, data, reactant_sharding)
+end
+
+XLA.sharding(x::AsyncArray) = XLA.sharding(x.buffer)

src/xla/IFRT/AsyncArray.jl

@@ -19,10 +19,7 @@ end
 end
 
 function test()
-    ctx = MLIR.IR.Context(Reactant.registry[], false)
-    @ccall MLIR.API.mlir_c.RegisterDialects(ctx::MLIR.API.MlirContext)::Cvoid
-
-    MLIR.IR.context!(ctx) do
+    MLIR.IR.with_context() do ctx
         mod = MLIR.IR.Module(MLIR.IR.Location())
         modbody = MLIR.IR.body(mod)
 

test/bcast.jl

@@ -349,3 +349,31 @@ end
         @warn "Not enough addressable devices to run sharding tests"
     end
 end
+
+@testset "Bad Codegen for Resharded Inputs: #1027" begin
+    if length(addressable_devices) ≥ 12 && Reactant.XLA.runtime() isa Val{:IFRT}
+        x_ra = Reactant.to_rarray(
+            randn(Float32, 32, 32);
+            sharding=Sharding.NamedSharding(
+                Sharding.Mesh(reshape(0:11, 3, 4), (:x, :y)), (:x, :y)
+            ),
+        )
+
+        z_ra = Reactant.to_rarray(ones(Float32, 32, 32))
+
+        function test1!(x, z)
+            y = x .+ x'
+            x .+= y
+            z .= x
+            return z
+        end
+
+        @jit test1!(x_ra, z_ra)
+
+        @test contains(
+            string(Reactant.XLA.sharding(z_ra.data.buffer)), "SingleDeviceSharding"
+        )
+    else
+        @warn "Not enough addressable devices to run sharding tests"
+    end
+end