Skip to content

Commit c4f0ab6

Browse files
committed
feat: add an option for chlo to stablehlo
1 parent a645d94 commit c4f0ab6

File tree

2 files changed

+31
-2
lines changed

2 files changed

+31
-2
lines changed

Project.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "Reactant"
22
uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
33
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]>"]
4-
version = "0.2.132"
4+
version = "0.2.133"
55

66
[deps]
77
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -90,7 +90,7 @@ PythonCall = "0.9"
9090
Random = "1.10"
9191
Random123 = "1.7"
9292
ReactantCore = "0.1.12"
93-
Reactant_jll = "0.0.198"
93+
Reactant_jll = "0.0.200"
9494
ScopedValues = "1.3.0"
9595
Scratch = "1.2"
9696
Sockets = "1.10"

src/Compiler.jl

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1387,6 +1387,7 @@ function compile_mlir!(
13871387
donated_args::Symbol=:auto, # :auto | :none
13881388
optimize_then_pad::Bool=true,
13891389
runtime::Union{Val{:PJRT},Val{:IFRT}},
1390+
legalize_chlo_to_stablehlo::Bool=false,
13901391
kwargs...,
13911392
)
13921393
@assert donated_args (:auto, :none)
@@ -1551,6 +1552,9 @@ function compile_mlir!(
15511552
"canonicalize",
15521553
"remove-unnecessary-enzyme-ops",
15531554
"enzyme-simplify-math",
1555+
(
1556+
legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : []
1557+
)...,
15541558
opt_passes2,
15551559
lower_enzymexla_linalg_pass,
15561560
jit,
@@ -1566,6 +1570,9 @@ function compile_mlir!(
15661570
"canonicalize",
15671571
"remove-unnecessary-enzyme-ops",
15681572
"enzyme-simplify-math",
1573+
(
1574+
legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : []
1575+
)...,
15691576
opt_passes2,
15701577
kern,
15711578
raise_passes,
@@ -1594,6 +1601,9 @@ function compile_mlir!(
15941601
"canonicalize",
15951602
"remove-unnecessary-enzyme-ops",
15961603
"enzyme-simplify-math",
1604+
(
1605+
legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : []
1606+
)...,
15971607
opt_passes2,
15981608
]
15991609
end,
@@ -1618,6 +1628,9 @@ function compile_mlir!(
16181628
"canonicalize",
16191629
"remove-unnecessary-enzyme-ops",
16201630
"enzyme-simplify-math",
1631+
(
1632+
legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : []
1633+
)...,
16211634
opt_passes2,
16221635
]
16231636
else
@@ -1631,6 +1644,9 @@ function compile_mlir!(
16311644
"canonicalize",
16321645
"remove-unnecessary-enzyme-ops",
16331646
"enzyme-simplify-math",
1647+
(
1648+
legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : []
1649+
)...,
16341650
opt_passes2,
16351651
kern,
16361652
raise_passes,
@@ -1657,6 +1673,9 @@ function compile_mlir!(
16571673
"canonicalize",
16581674
"remove-unnecessary-enzyme-ops",
16591675
"enzyme-simplify-math",
1676+
(
1677+
legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : []
1678+
)...,
16601679
opt_passes2,
16611680
kern,
16621681
]
@@ -1679,6 +1698,7 @@ function compile_mlir!(
16791698
"canonicalize",
16801699
"remove-unnecessary-enzyme-ops",
16811700
"enzyme-simplify-math",
1701+
(legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : [])...,
16821702
opt_passes2,
16831703
],
16841704
',',
@@ -1715,6 +1735,9 @@ function compile_mlir!(
17151735
"canonicalize",
17161736
"remove-unnecessary-enzyme-ops",
17171737
"enzyme-simplify-math",
1738+
(
1739+
legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : []
1740+
)...,
17181741
opt_passes2,
17191742
lower_enzymexla_linalg_pass,
17201743
jit,
@@ -1727,6 +1750,9 @@ function compile_mlir!(
17271750
"canonicalize",
17281751
"remove-unnecessary-enzyme-ops",
17291752
"enzyme-simplify-math",
1753+
(
1754+
legalize_chlo_to_stablehlo ? ["chlo-legalize-to-stablehlo"] : []
1755+
)...,
17301756
opt_passes2,
17311757
kern,
17321758
raise_passes,
@@ -2190,6 +2216,7 @@ function get_common_compile_options()
21902216
:optimize_then_pad => true,
21912217
:optimize_communications => true,
21922218
:cudnn_hlo_optimize => false,
2219+
:legalize_chlo_to_stablehlo => false,
21932220
)
21942221
end
21952222

@@ -2241,6 +2268,8 @@ const COMMON_COMPILE_OPTIONS_DOCS = """
22412268
- `cudnn_hlo_optimize`: Run cuDNN specific HLO optimizations. This is only relevant for
22422269
GPU backends and is `false` by default. **Experimental and not heavily tested.**
22432270
_(Only for CUDA backend)_
2271+
- `legalize_chlo_to_stablehlo`: If `true`, `chlo` dialect ops will be converted to
2272+
`stablehlo` ops. This is `false` by default.
22442273
"""
22452274

22462275
const SYNC_DOCS = """

0 commit comments

Comments
 (0)