@@ -1387,6 +1387,7 @@ function compile_mlir!(
1387
1387
donated_args:: Symbol = :auto , # :auto | :none
1388
1388
optimize_then_pad:: Bool = true ,
1389
1389
runtime:: Union{Val{:PJRT},Val{:IFRT}} ,
1390
+ legalize_chlo_to_stablehlo:: Bool = false ,
1390
1391
kwargs... ,
1391
1392
)
1392
1393
@assert donated_args ∈ (:auto , :none )
@@ -1551,6 +1552,9 @@ function compile_mlir!(
1551
1552
" canonicalize" ,
1552
1553
" remove-unnecessary-enzyme-ops" ,
1553
1554
" enzyme-simplify-math" ,
1555
+ (
1556
+ legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []
1557
+ ). .. ,
1554
1558
opt_passes2,
1555
1559
lower_enzymexla_linalg_pass,
1556
1560
jit,
@@ -1566,6 +1570,9 @@ function compile_mlir!(
1566
1570
" canonicalize" ,
1567
1571
" remove-unnecessary-enzyme-ops" ,
1568
1572
" enzyme-simplify-math" ,
1573
+ (
1574
+ legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []
1575
+ ). .. ,
1569
1576
opt_passes2,
1570
1577
kern,
1571
1578
raise_passes,
@@ -1594,6 +1601,9 @@ function compile_mlir!(
1594
1601
" canonicalize" ,
1595
1602
" remove-unnecessary-enzyme-ops" ,
1596
1603
" enzyme-simplify-math" ,
1604
+ (
1605
+ legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []
1606
+ ). .. ,
1597
1607
opt_passes2,
1598
1608
]
1599
1609
end ,
@@ -1618,6 +1628,9 @@ function compile_mlir!(
1618
1628
" canonicalize" ,
1619
1629
" remove-unnecessary-enzyme-ops" ,
1620
1630
" enzyme-simplify-math" ,
1631
+ (
1632
+ legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []
1633
+ ). .. ,
1621
1634
opt_passes2,
1622
1635
]
1623
1636
else
@@ -1631,6 +1644,9 @@ function compile_mlir!(
1631
1644
" canonicalize" ,
1632
1645
" remove-unnecessary-enzyme-ops" ,
1633
1646
" enzyme-simplify-math" ,
1647
+ (
1648
+ legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []
1649
+ ). .. ,
1634
1650
opt_passes2,
1635
1651
kern,
1636
1652
raise_passes,
@@ -1657,6 +1673,9 @@ function compile_mlir!(
1657
1673
" canonicalize" ,
1658
1674
" remove-unnecessary-enzyme-ops" ,
1659
1675
" enzyme-simplify-math" ,
1676
+ (
1677
+ legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []
1678
+ ). .. ,
1660
1679
opt_passes2,
1661
1680
kern,
1662
1681
]
@@ -1679,6 +1698,7 @@ function compile_mlir!(
1679
1698
" canonicalize" ,
1680
1699
" remove-unnecessary-enzyme-ops" ,
1681
1700
" enzyme-simplify-math" ,
1701
+ (legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []). .. ,
1682
1702
opt_passes2,
1683
1703
],
1684
1704
' ,' ,
@@ -1715,6 +1735,9 @@ function compile_mlir!(
1715
1735
" canonicalize" ,
1716
1736
" remove-unnecessary-enzyme-ops" ,
1717
1737
" enzyme-simplify-math" ,
1738
+ (
1739
+ legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []
1740
+ ). .. ,
1718
1741
opt_passes2,
1719
1742
lower_enzymexla_linalg_pass,
1720
1743
jit,
@@ -1727,6 +1750,9 @@ function compile_mlir!(
1727
1750
" canonicalize" ,
1728
1751
" remove-unnecessary-enzyme-ops" ,
1729
1752
" enzyme-simplify-math" ,
1753
+ (
1754
+ legalize_chlo_to_stablehlo ? [" chlo-legalize-to-stablehlo" ] : []
1755
+ ). .. ,
1730
1756
opt_passes2,
1731
1757
kern,
1732
1758
raise_passes,
@@ -2190,6 +2216,7 @@ function get_common_compile_options()
2190
2216
:optimize_then_pad => true ,
2191
2217
:optimize_communications => true ,
2192
2218
:cudnn_hlo_optimize => false ,
2219
+ :legalize_chlo_to_stablehlo => false ,
2193
2220
)
2194
2221
end
2195
2222
@@ -2241,6 +2268,8 @@ const COMMON_COMPILE_OPTIONS_DOCS = """
2241
2268
- `cudnn_hlo_optimize`: Run cuDNN specific HLO optimizations. This is only relevant for
2242
2269
GPU backends and is `false` by default. **Experimental and not heavily tested.**
2243
2270
_(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.
2244
2273
"""
2245
2274
2246
2275
const SYNC_DOCS = """
0 commit comments