Skip to content

feat: xla/mhlo export passes #1463

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jul 20, 2025
Merged

feat: xla/mhlo export passes #1463

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f5d456d
feat: run mhlo passes before legalize
avik-pal Jul 18, 2025
9bc14b8
feat: allow direct conversion of shlo to mhlo
avik-pal Jul 18, 2025
d4b52b6
feat: before optimizations HLO
avik-pal Jul 18, 2025
3856fa2
chore: run fmt
avik-pal Jul 18, 2025
838e474
chore: bump jll
avik-pal Jul 20, 2025
0c8f270
fix: tests
avik-pal Jul 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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"
Expand Down
161 changes: 78 additions & 83 deletions src/Compiler.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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,
]
Expand All @@ -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,
],
',',
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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...
)
(;
Expand Down
8 changes: 8 additions & 0 deletions src/xla/HloModule.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions test/buffer_donation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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)

Expand Down
16 changes: 16 additions & 0 deletions test/compile.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading