feat: xla/mhlo export passes (#1463)

* feat: run mhlo passes before legalize

* feat: allow direct conversion of shlo to mhlo

* feat: before optimizations HLO

* chore: run fmt

* chore: bump jll

* fix: tests

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.145"
+version = "0.2.146"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -90,7 +90,7 @@ PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.15"
-Reactant_jll = "0.0.218"
+Reactant_jll = "0.0.219"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/Compiler.jl

@@ -1475,6 +1475,22 @@ function get_optimize_comms_passes(options::OptimizeCommunicationOptions)
     return res
 end
 
+function get_stablehlo_to_hlo_passes(; stablehlo_to_mhlo::Bool=true)
+    passes = (
+        "func.func(stablehlo-ext-chlo-recompose-ops)",
+        "symbol-dce",
+        "func.func(chlo-legalize-to-high-level-mhlo)",
+        "func.func(chlo-legalize-to-stablehlo)",
+    )
+    if stablehlo_to_mhlo
+        passes = (passes..., "stablehlo-legalize-to-hlo")
+    end
+    passes = (
+        passes..., "canonicalize", "func.func(stablehlo-ext-sink-constants-to-control-flow)"
+    )
+    return passes
+end
+
 function compile_mlir!(
     mod,
     f,
@@ -1485,6 +1501,7 @@ function compile_mlir!(
     fn_kwargs=(),
     backend="gpu",
     runtime::Union{Val{:PJRT},Val{:IFRT}},
+    legalize_stablehlo_to_mhlo::Bool=false,
     kwargs...,
 )
     # Explicitly don't use block! to avoid creating a closure, which creates
@@ -1624,6 +1641,13 @@ function compile_mlir!(
     lower_enzymexla_linalg_pass = "lower-enzymexla-linalg{backend=$backend \
                                    blas_int_width=$blas_int_width}"
 
+    legalize_chlo_to_stablehlo =
+        if legalize_stablehlo_to_mhlo || compile_options.legalize_chlo_to_stablehlo
+            get_stablehlo_to_hlo_passes(; stablehlo_to_mhlo=legalize_stablehlo_to_mhlo)
+        else
+            ()
+        end
+
     if compile_options.optimization_passes === :all
         run_pass_pipeline!(
             mod,
@@ -1641,13 +1665,7 @@ function compile_mlir!(
                         "canonicalize",
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
-                        (
-                            if compile_options.legalize_chlo_to_stablehlo
-                                ["func.func(chlo-legalize-to-stablehlo)"]
-                            else
-                                []
-                            end
-                        )...,
+                        legalize_chlo_to_stablehlo...,
                         opt_passes2,
                         lower_enzymexla_linalg_pass,
                         jit,
@@ -1663,13 +1681,7 @@ function compile_mlir!(
                         "canonicalize",
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
-                        (
-                            if compile_options.legalize_chlo_to_stablehlo
-                                ["func.func(chlo-legalize-to-stablehlo)"]
-                            else
-                                []
-                            end
-                        )...,
+                        legalize_chlo_to_stablehlo...,
                         opt_passes2,
                         kern,
                         raise_passes,
@@ -1698,13 +1710,7 @@ function compile_mlir!(
                         "canonicalize",
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
-                        (
-                            if compile_options.legalize_chlo_to_stablehlo
-                                ["func.func(chlo-legalize-to-stablehlo)"]
-                            else
-                                []
-                            end
-                        )...,
+                        legalize_chlo_to_stablehlo...,
                         opt_passes2,
                     ]
                 end,
@@ -1729,13 +1735,7 @@ function compile_mlir!(
                         "canonicalize",
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
-                        (
-                            if compile_options.legalize_chlo_to_stablehlo
-                                ["func.func(chlo-legalize-to-stablehlo)"]
-                            else
-                                []
-                            end
-                        )...,
+                        legalize_chlo_to_stablehlo...,
                         opt_passes2,
                     ]
                 else
@@ -1749,13 +1749,7 @@ function compile_mlir!(
                         "canonicalize",
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
-                        (
-                            if compile_options.legalize_chlo_to_stablehlo
-                                ["func.func(chlo-legalize-to-stablehlo)"]
-                            else
-                                []
-                            end
-                        )...,
+                        legalize_chlo_to_stablehlo...,
                         opt_passes2,
                         kern,
                         raise_passes,
@@ -1782,13 +1776,7 @@ function compile_mlir!(
                         "canonicalize",
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
-                        (
-                            if compile_options.legalize_chlo_to_stablehlo
-                                ["func.func(chlo-legalize-to-stablehlo)"]
-                            else
-                                []
-                            end
-                        )...,
+                        legalize_chlo_to_stablehlo...,
                         opt_passes2,
                         kern,
                     ]
@@ -1811,13 +1799,7 @@ function compile_mlir!(
                     "canonicalize",
                     "remove-unnecessary-enzyme-ops",
                     "enzyme-simplify-math",
-                    (
-                        if compile_options.legalize_chlo_to_stablehlo
-                            ["func.func(chlo-legalize-to-stablehlo)"]
-                        else
-                            []
-                        end
-                    )...,
+                    legalize_chlo_to_stablehlo...,
                     opt_passes2,
                 ],
                 ',',
@@ -1854,13 +1836,7 @@ function compile_mlir!(
                         "canonicalize",
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
-                        (
-                            if compile_options.legalize_chlo_to_stablehlo
-                                ["func.func(chlo-legalize-to-stablehlo)"]
-                            else
-                                []
-                            end
-                        )...,
+                        legalize_chlo_to_stablehlo...,
                         opt_passes2,
                         lower_enzymexla_linalg_pass,
                         jit,
@@ -1873,13 +1849,7 @@ function compile_mlir!(
                         "canonicalize",
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
-                        (
-                            if compile_options.legalize_chlo_to_stablehlo
-                                ["func.func(chlo-legalize-to-stablehlo)"]
-                            else
-                                []
-                            end
-                        )...,
+                        legalize_chlo_to_stablehlo...,
                         opt_passes2,
                         kern,
                         raise_passes,
@@ -2406,7 +2376,13 @@ See also [`@code_xla`](@ref), [`@code_hlo`](@ref).
 """
 macro code_mhlo(args...)
     compile_expr, (; compiled) = compile_call_expr(
-        __module__, compile_xla, get_common_compile_options(), args...
+        __module__,
+        compile_mlir,
+        merge(
+            get_common_compile_options(),
+            Dict{Symbol,Any}(:legalize_stablehlo_to_mhlo => true),
+        ),
+        args...,
     )
     #! format: off
     return esc(
@@ -2427,20 +2403,25 @@ This is the post optimizations XLA HLO module.
 ## Options
 
 $(COMMON_COMPILE_OPTIONS_DOCS)
+  - `before_xla_optimizations`: If `true`, return the `before_optimizations` HLO module.
 
 See also [`@code_mhlo`](@ref), [`@code_hlo`](@ref).
 """
 macro code_xla(args...)
     compile_expr, (; compiled) = compile_call_expr(
-        __module__, compile_xla, get_common_compile_options(), args...
+        __module__,
+        compile_xla,
+        merge(
+            get_common_compile_options(),
+            Dict{Symbol,Any}(:before_xla_optimizations => false),
+        ),
+        args...,
     )
     #! format: off
     return esc(
         :(
             $(compile_expr);
-            exec = $(compiled)[2];
-            hlo_modules = $(XLA.get_hlo_modules)(exec);
-            length(hlo_modules) == 1 ? only(hlo_modules) : hlo_modules
+            $(compiled)[3]
         )
     )
     #! format: on
@@ -3374,7 +3355,14 @@ function __resolve_device_and_client(client, seen_args, linear_args, is_sharded)
     return (client, device)
 end
 
-function compile_xla(f, args; client=nothing, serializable::Bool=false, kwargs...)
+function compile_xla(
+    f,
+    args;
+    before_xla_optimizations::Bool=false,
+    client=nothing,
+    serializable::Bool=false,
+    kwargs...,
+)
     # register MLIR dialects
     ctx = MLIR.IR.Context(Reactant.registry[], false)
     context_gc_vector[ctx] = Vector{Union{TracedRArray,TracedRNumber}}(undef, 0)
@@ -3430,20 +3418,27 @@ function compile_xla(f, args; client=nothing, serializable::Bool=false, kwargs..
             module_string = ""
         end
 
-        exec = XLA.compile(
-            client,
-            device,
-            mod;
-            num_outputs=length(mlir_fn_res.linear_results),
-            num_parameters=length(mlir_fn_res.linear_args),
-            mlir_fn_res.is_sharded,
-            global_device_ids,
-            mlir_fn_res.num_replicas,
-            mlir_fn_res.num_partitions,
-            mlir_fn_res.use_shardy_partitioner,
-        )
+        if before_xla_optimizations
+            exec = nothing
+            hlo_modules = XLA.HloModule(mod)
+        else
+            exec = XLA.compile(
+                client,
+                device,
+                mod;
+                num_outputs=length(mlir_fn_res.linear_results),
+                num_parameters=length(mlir_fn_res.linear_args),
+                mlir_fn_res.is_sharded,
+                global_device_ids,
+                mlir_fn_res.num_replicas,
+                mlir_fn_res.num_partitions,
+                mlir_fn_res.use_shardy_partitioner,
+            )
+            hlo_modules = XLA.get_hlo_modules(exec)
+            hlo_modules = length(hlo_modules) == 1 ? only(hlo_modules) : hlo_modules
+        end
 
-        return mod, exec, mlir_fn_res, device, client, module_string
+        return mod, exec, hlo_modules, mlir_fn_res, device, client, module_string
     finally
         MLIR.IR.deactivate!(ctx)
     end
@@ -3459,7 +3454,7 @@ const __thunk_rev_body_cache = Dict{Expr,Symbol}()
 function compile(f, args; kwargs...)
     compile_options, kwargs = __get_compile_options_and_kwargs(; kwargs...)
 
-    _, exec, mlir_fn_res, device, client, str = compile_xla(
+    _, exec, _, mlir_fn_res, device, client, str = compile_xla(
         f, args; compile_options, kwargs...
     )
     (;

src/xla/HloModule.jl

@@ -11,6 +11,14 @@ function free_hlo_module(hlo_module)
     @ccall MLIR.API.mlir_c.FreeHloModule(hlo_module.ptr::Ptr{Cvoid})::Cvoid
 end
 
+function HloModule(mod::MLIR.IR.Module)
+    return HloModule(
+        @ccall MLIR.API.mlir_c.convertMlirModuleToHloModule(
+            mod::MLIR.API.MlirModule
+        )::Ptr{Cvoid}
+    )
+end
+
 function Base.show(io::IO, hlo_module::HloModule)
     GC.@preserve hlo_module begin
         str = @ccall MLIR.API.mlir_c.HloModuleToString(hlo_module.ptr::Ptr{Cvoid})::Cstring

test/buffer_donation.jl

@@ -25,7 +25,7 @@ end
     hlo = @code_hlo(donate_fill_x_with_2(a, b))
     @test length(findall("tf.aliasing_output = 0", repr(hlo))) == 1
 
-    (; preserved_args) = Reactant.Compiler.compile_xla(donate_fill_x_with_2, (a, b))[3]
+    (; preserved_args) = Reactant.Compiler.compile_xla(donate_fill_x_with_2, (a, b))[4]
     preserved_args_idx = last.(preserved_args)
     @test preserved_args_idx == [1] # only `y`(i.e. `b`) is preserved
 
@@ -36,7 +36,7 @@ end
     hlo = @code_hlo(donate_inplace_mul(a, b))
     @test length(findall("tf.aliasing_output = 0", repr(hlo))) == 1
 
-    (; preserved_args) = Reactant.Compiler.compile_xla(donate_inplace_mul, (a, b))[3]
+    (; preserved_args) = Reactant.Compiler.compile_xla(donate_inplace_mul, (a, b))[4]
     preserved_args_idx = last.(preserved_args)
     @test preserved_args_idx == [1] # only `y`(i.e. `b`) is preserved
 
@@ -71,7 +71,7 @@ end
     z = Reactant.to_rarray(ones(3))
 
     @code_hlo assert_nonallocating = true update_inplace!(x, y, z)
-    (; preserved_args) = Reactant.Compiler.compile_xla(update_inplace!, (x, y, z))[3]
+    (; preserved_args) = Reactant.Compiler.compile_xla(update_inplace!, (x, y, z))[4]
     preserved_args_idx = last.(preserved_args)
     @test preserved_args_idx == [1, 2] # y and z are both preserved (preserved_args is 0-indexed)
 

test/compile.jl

@@ -227,3 +227,19 @@ end
     @test y.x isa Reactant.RArray
     @test y.x == fcustom_path(x).x
 end
+
+# CHLO legalize options
+# test that we are running some mhlo passes first before legalizing, else we will end up
+# decomposing some necessary ops
+function fn_test(x)
+    y = Reactant.Ops.top_k(x, 16).values
+    y_complex = Complex.(y, -y .+ 1)
+    conj!(y_complex)
+    return y_complex
+end
+
+@testset "chlo legalize" begin
+    x_ra = Reactant.to_rarray(rand(Float32, 128))
+    hlo = @code_hlo legalize_chlo_to_stablehlo = true fn_test(x_ra)
+    @test occursin("mhlo.topk", repr(hlo))
+end