feat: run shardy passes on julia end (#926)

* feat: run shardy passes on julia end

fix: conditionally apply the passes

fix: expose the passes as an option

* feat: defaults based on entry-point

* chore: bump jll version

Author: avik-pal

Project.toml

@@ -86,7 +86,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.5"
-Reactant_jll = "0.0.89"
+Reactant_jll = "0.0.90"
 Scratch = "1.2"
 Sockets = "1.10"
 SpecialFunctions = "2.4"

src/Compiler.jl

@@ -722,6 +722,8 @@ function compile_mlir!(
         }
     }();
     optimize::Union{Bool,Symbol}=true,
+    # default refers to letting XLA handle the shardy inport/propagation/export
+    shardy_passes::Symbol=:to_mhlo_shardings, # [:default, :to_mhlo_shardings]
     no_nan::Bool=false,
     backend="gpu",
     fn_kwargs=(),
@@ -977,6 +979,34 @@ function compile_mlir!(
         error("Invalid optimize option: $(Meta.quot(optimize))")
     end
 
+    # shardy passes
+    use_shardy_partitioner = false
+    if is_sharded
+        if shardy_passes == :default
+            # If `:default` is passed in, we will run a pass to export the sharding
+            # inside the corresponding compile function for IFRT/PJRT. This keeps the
+            # sharding readable.
+            use_shardy_partitioner = true
+        elseif shardy_passes == :to_mhlo_shardings
+            # Convert all shardy ops to corresponding mhlo attrs/ops that can be consumed by
+            # XLA (note we need to set `use_shardy_partitioner` to `false` in the options)
+            # TODO: Use https://github.com/openxla/shardy/blob/01d3205086132d1bdf0867e911c05f489918431d/shardy/dialect/sdy/transforms/propagation/propagation_pipeline.cc#L28 to pass in the options
+            run_pass_pipeline!(
+                mod,
+                join(
+                    ["sdy-propagation-pipeline", "xla-sdy-stablehlo-export-pipeline"], ','
+                ),
+            )
+
+            # Run our optimization passes here -- we need to be careful to not apply folding
+            # here since that violates the semantics of `sdy.constant` which was converted to
+            # `stablehlo.constant` by the previous pass.
+            run_pass_pipeline!(mod, join(["canonicalize", "cse"], ','))
+        else
+            error("Invalid shardy_passes option: $(Meta.quot(shardy_passes))")
+        end
+    end
+
     preserved_args = Tuple{TracedType,Int}[]
     results = [MLIR.IR.operand(ret, i) for i in 1:MLIR.IR.noperands(ret)]
     nresults = MLIR.IR.Value[]
@@ -1040,6 +1070,7 @@ function compile_mlir!(
         concrete_result,
         mlir_fn_res.sharding_mesh,
         mlir_fn_res.mutated_args,
+        use_shardy_partitioner,
     )
 end
 
@@ -1050,7 +1081,11 @@ See also [`@code_xla`](@ref), [`@code_mhlo`](@ref).
 """
 macro code_hlo(args...)
     default_options = Dict{Symbol,Any}(
-        :optimize => true, :no_nan => false, :client => nothing, :raise => false
+        :optimize => true,
+        :no_nan => false,
+        :client => nothing,
+        :raise => false,
+        :shardy_passes => :(:default),
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_mlir, default_options, args...
@@ -1074,7 +1109,11 @@ See also [`@code_xla`](@ref), [`@code_hlo`](@ref).
 """
 macro code_mhlo(args...)
     default_options = Dict{Symbol,Any}(
-        :optimize => true, :no_nan => false, :client => nothing, :raise => false
+        :optimize => true,
+        :no_nan => false,
+        :client => nothing,
+        :raise => false,
+        :shardy_passes => :(:default),
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_xla, default_options, args...
@@ -1098,7 +1137,11 @@ See also [`@code_mhlo`](@ref), [`@code_hlo`](@ref).
 """
 macro code_xla(args...)
     default_options = Dict{Symbol,Any}(
-        :optimize => true, :no_nan => false, :client => nothing, :raise => false
+        :optimize => true,
+        :no_nan => false,
+        :client => nothing,
+        :raise => false,
+        :shardy_passes => :(:to_mhlo_shardings),
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_xla, default_options, args...
@@ -1125,6 +1168,7 @@ macro compile(args...)
         :no_nan => false,
         :client => nothing,
         :raise => false,
+        :shardy_passes => :(:to_mhlo_shardings),
     )
     return esc(first(compile_call_expr(__module__, compile, default_options, args...)))
 end
@@ -1141,6 +1185,7 @@ macro jit(args...)
         :no_nan => false,
         :client => nothing,
         :raise => false,
+        :shardy_passes => :(:to_mhlo_shardings),
     )
     compile_expr, (; compiled, args) = compile_call_expr(
         __module__, compile, default_options, args...
@@ -1166,6 +1211,7 @@ function compile_call_expr(mod, compiler, options::Dict, args...)
             options[option_name] = option.args[2]
         end
     end
+
     call = only(args)
     f_symbol = gensym(:f)
     args_symbol = gensym(:args)
@@ -1207,18 +1253,20 @@ function compile_call_expr(mod, compiler, options::Dict, args...)
         error("Invalid function call: $(call)")
     end
 
-    return quote
-        $(f_symbol) = $(fname)
-        $(args_symbol) = $(args_rhs)
-        $(kwargs_symbol) = (; $(kwargs_rhs...))
-        $(compiled_symbol) = $(compiler)(
-            $(f_symbol),
-            $(args_symbol);
-            fn_kwargs=$(kwargs_symbol),
-            $(Expr.(:kw, keys(options), values(options))...),
-        )
-    end,
-    (; compiled=compiled_symbol, args=args_symbol)
+    return (
+        quote
+            $(f_symbol) = $(fname)
+            $(args_symbol) = $(args_rhs)
+            $(kwargs_symbol) = (; $(kwargs_rhs...))
+            $(compiled_symbol) = $(compiler)(
+                $(f_symbol),
+                $(args_symbol);
+                fn_kwargs=$(kwargs_symbol),
+                $(Expr.(:kw, keys(options), values(options))...),
+            )
+        end,
+        (; compiled=compiled_symbol, args=args_symbol),
+    )
 end
 
 """
@@ -1765,6 +1813,7 @@ function compile_xla(f, args; client=nothing, kwargs...)
             global_device_ids,
             mlir_fn_res.num_replicas,
             mlir_fn_res.num_partitions,
+            mlir_fn_res.use_shardy_partitioner,
         )
 
         return mod, exec, mlir_fn_res, device, client

src/TracedUtils.jl

@@ -148,6 +148,7 @@ mutable struct CompiledMlirFnResult{
     concrete_result::CR
     sharding_mesh::M
     mutated_args::MA
+    use_shardy_partitioner::Bool
 end
 
 function make_mlir_fn(
@@ -434,6 +435,7 @@ function make_mlir_fn(
         nothing,
         sharding_mesh,
         mutated_args,
+        true,
     )
 end
 

src/xla/IFRT/LoadedExecutable.jl

@@ -79,6 +79,7 @@ function XLA.compile(
     num_parameters::Int64,
     num_replicas::Int64,
     num_partitions::Int64,
+    use_shardy_partitioner::Bool,
 )
     device_id = is_sharded ? Int64(-1) : Int64(XLA.device_ordinal(device))
     GC.@preserve client mod begin
@@ -90,6 +91,7 @@ function XLA.compile(
             global_device_ids::Ptr{Clong},
             length(global_device_ids)::Clong,
             XLA.CUDA_DATA_DIR[]::Cstring,
+            use_shardy_partitioner::Bool,
         )::Ptr{Cvoid}
     end
     return LoadedExecutable(

src/xla/PJRT/LoadedExecutable.jl

@@ -73,6 +73,7 @@ function XLA.compile(
     num_parameters::Int64,
     num_replicas::Int64,
     num_partitions::Int64,
+    use_shardy_partitioner::Bool,
 )
     device_id = is_sharded ? Int64(-1) : Int64(XLA.device_ordinal(device))
     GC.@preserve client mod begin
@@ -84,6 +85,7 @@ function XLA.compile(
             global_device_ids::Ptr{Clong},
             length(global_device_ids)::Clong,
             XLA.CUDA_DATA_DIR[]::Cstring,
+            use_shardy_partitioner::Bool,
         )::Ptr{Cvoid}
     end
     return LoadedExecutable(

test/sharding.jl

@@ -1,4 +1,3 @@
-# Currently an extremely simple test
 using Reactant, Test
 
 const addressable_devices = Reactant.addressable_devices()
@@ -113,6 +112,9 @@ end
         )
         @test Array(@jit fn_test2(x_ra)) ≈ fn_test2(x)
         @test Reactant.to_number(@jit sum(x_ra)) ≈ sum(x)
+
+        @test Array(@jit shardy_passes = :to_mhlo_shardings fn_test3(x_ra)) ≈ fn_test3(x)
+        @test Reactant.to_number(@jit shardy_passes = :to_mhlo_shardings sum(x_ra)) ≈ sum(x)
     else
         @warn "Not enough addressable devices to run sharding tests"
     end
@@ -125,8 +127,11 @@ end
         x_ra = Reactant.to_rarray(
             x; sharding=Sharding.NamedSharding(mesh, (("data", "model"), nothing))
         )
-        @test Array(@jit fn_test2(x_ra)) ≈ fn_test2(x)
-        @test Reactant.to_number(@jit sum(x_ra)) ≈ sum(x)
+        @test Array(@jit shardy_passes = :default fn_test2(x_ra)) ≈ fn_test2(x)
+        @test Reactant.to_number(@jit shardy_passes = :default sum(x_ra)) ≈ sum(x)
+
+        @test Array(@jit shardy_passes = :to_mhlo_shardings fn_test3(x_ra)) ≈ fn_test3(x)
+        @test Reactant.to_number(@jit shardy_passes = :to_mhlo_shardings sum(x_ra)) ≈ sum(x)
     else
         @warn "Not enough addressable devices to run sharding tests"
     end
@@ -142,8 +147,12 @@ end
                 mesh, ("model", nothing); is_closed=(false, false)
             ),
         )
-        @test Array(@jit fn_test2(x_ra)) ≈ fn_test2(x)
-        @test Reactant.to_number(@jit sum(x_ra)) ≈ sum(x)
+
+        @test Array(@jit shardy_passes = :default fn_test2(x_ra)) ≈ fn_test2(x)
+        @test Reactant.to_number(@jit shardy_passes = :default sum(x_ra)) ≈ sum(x)
+
+        @test Array(@jit shardy_passes = :to_mhlo_shardings fn_test3(x_ra)) ≈ fn_test3(x)
+        @test Reactant.to_number(@jit shardy_passes = :to_mhlo_shardings sum(x_ra)) ≈ sum(x)
     else
         @warn "Not enough addressable devices to run sharding tests"
     end
@@ -193,6 +202,9 @@ end
 
         hlo = @code_hlo fn_with_constraint(x_ra)
         @test contains(repr(hlo), "sharding_constraint")
+        hlo = @code_hlo shardy_passes = :to_mhlo_shardings fn_with_constraint(x_ra)
+        @test !contains(repr(hlo), "sharding_constraint")
+        @test length(collect(eachmatch(r"mhlo.sharding", repr(hlo)))) == 3
 
         z = Reactant.to_rarray(x; sharding=constraint)
         res = @jit fn_with_constraint(x_ra)
@@ -208,6 +220,11 @@ end
 
         hlo = @code_hlo fn_with_constraint(x_ra_no_sharding)
         @test contains(repr(hlo), "sharding_constraint")
+        hlo = @code_hlo shardy_passes = :to_mhlo_shardings fn_with_constraint(
+            x_ra_no_sharding
+        )
+        @test !contains(repr(hlo), "sharding_constraint")
+        @test length(collect(eachmatch(r"mhlo.sharding", repr(hlo)))) == 3
 
         res = @jit fn_with_constraint(x_ra_no_sharding)
         @test x .+ x ≈ Array(res)
@@ -228,14 +245,17 @@ end
             x; sharding=Sharding.NamedSharding(mesh, ("data", "model"))
         )
 
-        @test Array(@jit sum(x_ra; dims=2)) ≈ sum(x; dims=2)
+        @test Array(@jit shardy_passes = :default sum(x_ra; dims=2)) ≈ sum(x; dims=2)
+        @test Array(@jit shardy_passes = :to_mhlo_shardings sum(x_ra; dims=2)) ≈
+            sum(x; dims=2)
 
         x = reshape(collect(Float32, 1:25), 5, 5)
         x_ra = Reactant.to_rarray(
             x; sharding=Sharding.NamedSharding(mesh, ("data", "model"))
         )
 
-        @test Array(@jit fn_test2(x_ra)) ≈ fn_test2(x)
+        @test Array(@jit shardy_passes = :default fn_test2(x_ra)) ≈ fn_test2(x)
+        @test Array(@jit shardy_passes = :to_mhlo_shardings fn_test2(x_ra)) ≈ fn_test2(x)
     else
         @warn "Not enough addressable devices to run sharding tests"
     end
@@ -271,24 +291,32 @@ end
             randn(Float32, 4, 5); sharding=Sharding.NamedSharding(mesh, ((:x, :y), :z))
         )
 
-        y_ra = Reactant.to_rarray(randn(Float32, 5, 4); sharding=Sharding.NoSharding())
+        y_ra_arr = randn(Float32, 5, 4)
+        y_ra = Reactant.to_rarray(y_ra_arr; sharding=Sharding.NoSharding())
+        y_ra_2 = Reactant.to_rarray(y_ra_arr; 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)
+        z_ra = @jit shardy_passes = :default 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)
+
+        z_ra2 = @jit shardy_passes = :to_mhlo_shardings fn(x_ra, y_ra_2)
+        y_ra_final2 = Array(y_ra_2)
+
+        @test z_ra_arr ≈ Array(z_ra2)
+        @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