diff --git a/src/Tracing.jl b/src/Tracing.jl index 288020cfa9..7a945a744c 100644 --- a/src/Tracing.jl +++ b/src/Tracing.jl @@ -1520,7 +1520,10 @@ Base.@nospecializeinfer function make_tracer( (XLA.PJRT.AsyncEmptyBuffer,), size(prev), Sharding.NoShardInfo() ) else - error("TODO(#2230): implement sharding") + ndev = Sharding.ndevices(sharding) + res = ConcretePJRTArray{T,N,ndev}( + ntuple(Returns(XLA.PJRT.AsyncEmptyBuffer), ndev), size(prev), Sharding.NoShardInfo() + ) end seen[prev] = res return res @@ -1531,7 +1534,9 @@ Base.@nospecializeinfer function make_tracer( XLA.IFRT.AsyncEmptyArray, size(prev), Sharding.NoShardInfo() ) else - error("TODO(#2230): implement sharding") + res = ConcreteIFRTArray{T,N}( + XLA.IFRT.AsyncEmptyArray, size(prev), Sharding.NoShardInfo() + ) end seen[prev] = res return res diff --git a/src/compiler/Compiler.jl b/src/compiler/Compiler.jl index 7bcc804a87..f4724de4fd 100644 --- a/src/compiler/Compiler.jl +++ b/src/compiler/Compiler.jl @@ -1198,8 +1198,14 @@ function compile_mlir!( end end + sharding = if mlir_fn_res.is_sharded && !isempty(result_shardings_after_masking) + first(result_shardings_after_masking) + else + Sharding.NoSharding() + end + concrete_result = make_tracer( - OrderedIdDict(), traced_result, ("result",), TracedToConcrete; runtime + OrderedIdDict(), traced_result, ("result",), TracedToConcrete; runtime, sharding ) return Reactant.TracedUtils.CompiledMlirFnResult( diff --git a/test/core/sharding/sharding.jl b/test/core/sharding/sharding.jl index 0f563857b3..00a1f9e888 100644 --- a/test/core/sharding/sharding.jl +++ b/test/core/sharding/sharding.jl @@ -663,3 +663,24 @@ end end end =# + +struct WrapTest{A} + data::A +end + +wrap_f(a) = WrapTest(a .* a) + +@testset "Wrapped Sharded Array" begin + if length(addressable_devices) ≥ 2 + ndev = length(addressable_devices) + mesh = Sharding.Mesh(reshape(addressable_devices, ndev), (:dev,)) + sh = Sharding.DimsSharding(mesh, (1,), (:dev,)) + + x = Reactant.to_rarray(collect(Float64, 1:64); sharding = sh) + r = @jit(wrap_f(x)) + @test r isa WrapTest + @test Array(r.data) ≈ collect(Float64, 1:64) .^ 2 + else + @warn "Not enough addressable devices to run wrapped sharded array tests" + end +end