Skip to content

Commit 449df62

Browse files
paul0403paul0403
authored
Migrate to stablehlo (#1921)
**Context:** Migrate to stablehlo. Catalyst depends on the stablehlo dialect because that is what jax lowers classical primitives (additions, jnp array ops, linalg ops, etc) to. The stablehlo dialect was pulled in from the [tensorflow/mlir-hlo](https://github.com/tensorflow/mlir-hlo) repo, but [that repo is slated for deletion](https://groups.google.com/a/openxla.org/g/openxla-discuss/c/Mppuv1Edv1s) . They suggest users to migrate to [the stablehlo repo](https://github.com/openxla/stablehlo). **Description of the Change:** Use the stablehlo version of the mhlo operations, dialects and passes we are currently using. Remove mhlo build system from regular makefiles and CI scripts. Remove mhlo patches. **Benefits:** No longer depending on a sunsetting external project. [sc-96641]
1 parent 82b076a commit 449df62

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

54 files changed

+596
-848
lines changed

‎.dep-versions‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
# To update JAX version alongside compatible dependency tags, run the following script:
33
# python3 .github/workflows/set_dep_versions.py {JAX_version}
44
jax=0.6.2
5-
mhlo=1dd2e71331014ae0373f6bf900ce6be393357190
5+
stablehlo=69d6dae46e1c7de36e6e6973654754f05353cba5
66
llvm=f8cb7987c64dcffb72414a40560055cb717dbf74
77
enzyme=v0.0.186
88

‎.github/workflows/build-wheel-linux-arm64.yaml‎

Lines changed: 40 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -84,12 +84,12 @@ jobs:
8484
key: llvm-${{ needs.constants.outputs.llvm_version }}-container-source
8585
enableCrossOsArchive: True
8686

87-
- name: Cache MHLO Source
88-
id: cache-mhlo-source
87+
- name: Cache Stablehlo Source
88+
id: cache-stablehlo-source
8989
uses: actions/cache@v4
9090
with:
91-
path: ${{ github.workspace }}/mlir/mlir-hlo
92-
key: mhlo-${{ needs.constants.outputs.mhlo_version }}-container-source
91+
path: ${{ github.workspace }}/mlir/stablehlo
92+
key: stablehlo-${{ needs.constants.outputs.stablehlo_version }}-container-source
9393
enableCrossOsArchive: True
9494

9595
- name: Cache Enzyme Source
@@ -109,25 +109,18 @@ jobs:
109109
path: ${{ github.workspace }}/mlir/llvm-project
110110

111111
- name: Patch LLVM Source
112-
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
112+
if: steps.cache-llvm-source.outputs.cache-hit != 'true'
113113
run: |
114114
cd $GITHUB_WORKSPACE/mlir/llvm-project
115115
git apply $GITHUB_WORKSPACE/mlir/patches/llvm-bufferization-segfault.patch
116116
117-
- name: Clone MHLO Submodule
118-
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
117+
- name: Clone Stablehlo Submodule
118+
if: steps.cache-stablehlo-source.outputs.cache-hit != 'true'
119119
uses: actions/checkout@v4
120120
with:
121-
repository: tensorflow/mlir-hlo
122-
ref: ${{ needs.constants.outputs.mhlo_version }}
123-
path: ${{ github.workspace }}/mlir/mlir-hlo
124-
125-
- name: Patch MHLO Source
126-
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
127-
run: |
128-
cd $GITHUB_WORKSPACE/mlir/mlir-hlo
129-
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
130-
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-rename-sort.patch
121+
repository: openxla/stablehlo
122+
ref: ${{ needs.constants.outputs.stablehlo_version }}
123+
path: ${{ github.workspace }}/mlir/stablehlo
131124

132125
- name: Clone Enzyme Submodule
133126
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
@@ -151,12 +144,12 @@ jobs:
151144
path: ${{ github.workspace }}/llvm-build
152145
key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{ matrix.python_version }}-wheel-build
153146

154-
- name: Check MHLO Build Cache
155-
id: cache-mhlo-build
147+
- name: Check Stablehlo Build Cache
148+
id: cache-stablehlo-build
156149
uses: actions/cache/restore@v4
157150
with:
158-
path: ${{ github.workspace }}/mhlo-build
159-
key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
151+
path: ${{ github.workspace }}/stablehlo-build
152+
key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
160153
lookup-only: True
161154

162155
- name: Check Enzyme Build Cache
@@ -170,7 +163,7 @@ jobs:
170163
- name: Install dependencies
171164
if: |
172165
steps.cache-llvm-build.outputs.cache-hit != 'true' ||
173-
steps.cache-mhlo-build.outputs.cache-hit != 'true' ||
166+
steps.cache-stablehlo-build.outputs.cache-hit != 'true' ||
174167
steps.cache-enzyme-build.outputs.cache-hit != 'true'
175168
run: |
176169
cat /etc/dnf.conf | sed "s/\[main\]/\[main\]\ntimeout=5/g" > /etc/dnf.conf
@@ -207,31 +200,24 @@ jobs:
207200
path: ${{ github.workspace }}/llvm-build
208201
key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{matrix.python_version}}-wheel-build
209202

210-
- name: Build MHLO Dialect
211-
if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
212-
# building with LLD is a strong requirement for mhlo
203+
- name: Build Stablehlo Dialect
204+
if: steps.cache-stablehlo-build.outputs.cache-hit != 'true'
213205
run: |
214-
export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
215-
216-
cmake -S mlir/mlir-hlo -B $GITHUB_WORKSPACE/mhlo-build -G Ninja \
217-
-DCMAKE_BUILD_TYPE=Release \
218-
-DLLVM_ENABLE_ASSERTIONS=ON \
219-
-DMLIR_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir" \
220-
-DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
221-
-DLLVM_ENABLE_ZLIB=FORCE_ON \
222-
-DLLVM_ENABLE_ZSTD=OFF \
223-
-DCMAKE_CXX_VISIBILITY_PRESET=default \
224-
-DLLVM_ENABLE_LLD=ON
225-
226-
LIT_FILTER_OUT="chlo_legalize_to_mhlo" cmake --build $GITHUB_WORKSPACE/mhlo-build --target check-mlir-hlo
206+
C_COMPILER=$(which gcc) \
207+
CXX_COMPILER=$(which g++) \
208+
LLVM_BUILD_DIR="$(pwd)/llvm-build" \
209+
STABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
210+
COMPILER_LAUNCHER="" \
211+
ENABLE_LLD=OFF \
212+
make stablehlo
227213
228-
- name: Save MHLO Build
229-
id: save-mhlo-build
230-
if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
214+
- name: Save Stablehlo Build
215+
id: save-stablehlo-build
216+
if: steps.cache-stablehlo-build.outputs.cache-hit != 'true'
231217
uses: actions/cache/save@v4
232218
with:
233-
path: ${{ github.workspace }}/mhlo-build
234-
key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
219+
path: ${{ github.workspace }}/stablehlo-build
220+
key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
235221

236222
- name: Build Enzyme
237223
if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
@@ -305,21 +291,21 @@ jobs:
305291
key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-3.11-wheel-build
306292
fail-on-cache-miss: True
307293

308-
- name: Get Cached MHLO Source
309-
id: cache-mhlo-source
294+
- name: Get Cached Stablehlo Source
295+
id: cache-stablehlo-source
310296
uses: actions/cache/restore@v4
311297
with:
312-
path: ${{ github.workspace }}/mlir/mlir-hlo
313-
key: mhlo-${{ needs.constants.outputs.mhlo_version }}-container-source
298+
path: ${{ github.workspace }}/mlir/stablehlo
299+
key: stablehlo-${{ needs.constants.outputs.stablehlo_version }}-container-source
314300
enableCrossOsArchive: True
315301
fail-on-cache-miss: True
316302

317-
- name: Get Cached MHLO Build
318-
id: cache-mhlo-build
303+
- name: Get Cached Stablehlo Build
304+
id: cache-stablehlo-build
319305
uses: actions/cache/restore@v4
320306
with:
321-
path: ${{ github.workspace }}/mhlo-build
322-
key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
307+
path: ${{ github.workspace }}/stablehlo-build
308+
key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
323309
fail-on-cache-miss: True
324310

325311
- name: Get Cached Enzyme Source
@@ -372,8 +358,8 @@ jobs:
372358
-DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
373359
-DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())") \
374360
-DMLIR_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir" \
375-
-DMHLO_DIR="$GITHUB_WORKSPACE/mhlo-build/lib/cmake/mlir-hlo" \
376-
-DMHLO_BINARY_DIR="$GITHUB_WORKSPACE/mhlo-build/bin" \
361+
-DSTABLEHLO_DIR="$GITHUB_WORKSPACE/mlir/stablehlo" \
362+
-DSTABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
377363
-DEnzyme_DIR="$GITHUB_WORKSPACE/enzyme-build" \
378364
-DENZYME_SRC_DIR="$GITHUB_WORKSPACE/mlir/Enzyme" \
379365
-DLLVM_ENABLE_ZLIB=FORCE_ON \
@@ -396,7 +382,7 @@ jobs:
396382
run: |
397383
PYTHON=python${{ matrix.python_version }} \
398384
LLVM_BUILD_DIR="$GITHUB_WORKSPACE/llvm-build" \
399-
MHLO_BUILD_DIR="$GITHUB_WORKSPACE/mhlo-build" \
385+
STABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
400386
DIALECTS_BUILD_DIR="$GITHUB_WORKSPACE/quantum-build" \
401387
RT_BUILD_DIR="$GITHUB_WORKSPACE/runtime-build" \
402388
OQC_BUILD_DIR="$GITHUB_WORKSPACE/oqc-build" \

‎.github/workflows/build-wheel-linux-x86_64.yaml‎

Lines changed: 40 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -103,12 +103,12 @@ jobs:
103103
key: llvm-${{ needs.constants.outputs.llvm_version }}-container-source
104104
enableCrossOsArchive: True
105105

106-
- name: Cache MHLO Source
107-
id: cache-mhlo-source
106+
- name: Cache Stablehlo Source
107+
id: cache-stablehlo-source
108108
uses: actions/cache@v4
109109
with:
110-
path: ${{ github.workspace }}/mlir/mlir-hlo
111-
key: mhlo-${{ needs.constants.outputs.mhlo_version }}-container-source
110+
path: ${{ github.workspace }}/mlir/stablehlo
111+
key: stablehlo-${{ needs.constants.outputs.stablehlo_version }}-container-source
112112
enableCrossOsArchive: True
113113

114114
- name: Cache Enzyme Source
@@ -128,25 +128,18 @@ jobs:
128128
path: ${{ github.workspace }}/mlir/llvm-project
129129

130130
- name: Patch LLVM Source
131-
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
131+
if: steps.cache-llvm-source.outputs.cache-hit != 'true'
132132
run: |
133133
cd $GITHUB_WORKSPACE/mlir/llvm-project
134134
git apply $GITHUB_WORKSPACE/mlir/patches/llvm-bufferization-segfault.patch
135135
136-
- name: Clone MHLO Submodule
137-
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
136+
- name: Clone Stablehlo Submodule
137+
if: steps.cache-stablehlo-source.outputs.cache-hit != 'true'
138138
uses: actions/checkout@v4
139139
with:
140-
repository: tensorflow/mlir-hlo
141-
ref: ${{ needs.constants.outputs.mhlo_version }}
142-
path: ${{ github.workspace }}/mlir/mlir-hlo
143-
144-
- name: Patch MHLO Source
145-
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
146-
run: |
147-
cd $GITHUB_WORKSPACE/mlir/mlir-hlo
148-
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
149-
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-rename-sort.patch
140+
repository: openxla/stablehlo
141+
ref: ${{ needs.constants.outputs.stablehlo_version }}
142+
path: ${{ github.workspace }}/mlir/stablehlo
150143

151144
- name: Clone Enzyme Submodule
152145
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
@@ -170,12 +163,12 @@ jobs:
170163
path: ${{ github.workspace }}/llvm-build
171164
key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{ matrix.python_version }}-wheel-build
172165

173-
- name: Check MHLO Build Cache
174-
id: cache-mhlo-build
166+
- name: Check Stablehlo Build Cache
167+
id: cache-stablehlo-build
175168
uses: actions/cache/restore@v4
176169
with:
177-
path: ${{ github.workspace }}/mhlo-build
178-
key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
170+
path: ${{ github.workspace }}/stablehlo-build
171+
key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
179172
lookup-only: True
180173

181174
- name: Check Enzyme Build Cache
@@ -189,7 +182,7 @@ jobs:
189182
- name: Install dependencies (AlmaLinux)
190183
if: |
191184
steps.cache-llvm-build.outputs.cache-hit != 'true' ||
192-
steps.cache-mhlo-build.outputs.cache-hit != 'true' ||
185+
steps.cache-stablehlo-build.outputs.cache-hit != 'true' ||
193186
steps.cache-enzyme-build.outputs.cache-hit != 'true'
194187
run: |
195188
# Reduce wait time for repos not responding
@@ -205,7 +198,6 @@ jobs:
205198
PYTHON_BINS=$(find /opt/_internal/cpython-${{ matrix.python_version }}.*/bin -maxdepth 1 -type d | tr '\n' ':' | sed 's/:$//')
206199
echo $PYTHON_BINS >> $GITHUB_PATH
207200
208-
# LLD is required for MHLO builds.
209201
# (Don't forget to add the build directory to PATH in subsequent steps, so
210202
# other tools can find it, in particular collect2 invoked by gcc.)
211203
- name: Build LLVM / MLIR
@@ -230,31 +222,24 @@ jobs:
230222
path: ${{ github.workspace }}/llvm-build
231223
key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{matrix.python_version}}-wheel-build
232224

233-
- name: Build MHLO Dialect
234-
if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
235-
# building with LLD is a strong requirement for mhlo
225+
- name: Build Stablehlo Dialect
226+
if: steps.cache-stablehlo-build.outputs.cache-hit != 'true'
236227
run: |
237-
export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
238-
239-
cmake -S mlir/mlir-hlo -B $GITHUB_WORKSPACE/mhlo-build -G Ninja \
240-
-DCMAKE_BUILD_TYPE=Release \
241-
-DLLVM_ENABLE_ASSERTIONS=ON \
242-
-DMLIR_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir" \
243-
-DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
244-
-DLLVM_ENABLE_ZLIB=FORCE_ON \
245-
-DLLVM_ENABLE_ZSTD=OFF \
246-
-DCMAKE_CXX_VISIBILITY_PRESET=default \
247-
-DLLVM_ENABLE_LLD=ON
248-
249-
LIT_FILTER_OUT="chlo_legalize_to_mhlo" cmake --build $GITHUB_WORKSPACE/mhlo-build --target check-mlir-hlo
228+
C_COMPILER=$(which gcc) \
229+
CXX_COMPILER=$(which g++) \
230+
LLVM_BUILD_DIR="$(pwd)/llvm-build" \
231+
STABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
232+
COMPILER_LAUNCHER="" \
233+
ENABLE_LLD=OFF \
234+
make stablehlo
250235
251-
- name: Save MHLO Build
252-
id: save-mhlo-build
253-
if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
236+
- name: Save Stablehlo Build
237+
id: save-stablehlo-build
238+
if: steps.cache-stablehlo-build.outputs.cache-hit != 'true'
254239
uses: actions/cache/save@v4
255240
with:
256-
path: ${{ github.workspace }}/mhlo-build
257-
key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
241+
path: ${{ github.workspace }}/stablehlo-build
242+
key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
258243

259244
- name: Build Enzyme
260245
if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
@@ -328,21 +313,21 @@ jobs:
328313
key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-3.11-wheel-build
329314
fail-on-cache-miss: True
330315

331-
- name: Get Cached MHLO Source
332-
id: cache-mhlo-source
316+
- name: Get Cached Stablehlo Source
317+
id: cache-stablehlo-source
333318
uses: actions/cache/restore@v4
334319
with:
335-
path: ${{ github.workspace }}/mlir/mlir-hlo
336-
key: mhlo-${{ needs.constants.outputs.mhlo_version }}-container-source
320+
path: ${{ github.workspace }}/mlir/stablehlo
321+
key: stablehlo-${{ needs.constants.outputs.stablehlo_version }}-container-source
337322
enableCrossOsArchive: True
338323
fail-on-cache-miss: True
339324

340-
- name: Get Cached MHLO Build
341-
id: cache-mhlo-build
325+
- name: Get Cached Stablehlo Build
326+
id: cache-stablehlo-build
342327
uses: actions/cache/restore@v4
343328
with:
344-
path: ${{ github.workspace }}/mhlo-build
345-
key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
329+
path: ${{ github.workspace }}/stablehlo-build
330+
key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
346331
fail-on-cache-miss: True
347332

348333
- name: Get Cached Enzyme Source
@@ -397,8 +382,8 @@ jobs:
397382
-DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
398383
-DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())") \
399384
-DMLIR_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir" \
400-
-DMHLO_DIR="$GITHUB_WORKSPACE/mhlo-build/lib/cmake/mlir-hlo" \
401-
-DMHLO_BINARY_DIR="$GITHUB_WORKSPACE/mhlo-build/bin" \
385+
-DSTABLEHLO_DIR="$GITHUB_WORKSPACE/mlir/stablehlo" \
386+
-DSTABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
402387
-DEnzyme_DIR="$GITHUB_WORKSPACE/enzyme-build" \
403388
-DENZYME_SRC_DIR="$GITHUB_WORKSPACE/mlir/Enzyme" \
404389
-DLLVM_ENABLE_ZLIB=FORCE_ON \
@@ -421,7 +406,7 @@ jobs:
421406
run: |
422407
PYTHON=python${{ matrix.python_version }} \
423408
LLVM_BUILD_DIR="$GITHUB_WORKSPACE/llvm-build" \
424-
MHLO_BUILD_DIR="$GITHUB_WORKSPACE/mhlo-build" \
409+
STABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
425410
DIALECTS_BUILD_DIR="$GITHUB_WORKSPACE/quantum-build" \
426411
RT_BUILD_DIR="$GITHUB_WORKSPACE/runtime-build" \
427412
OQC_BUILD_DIR="$GITHUB_WORKSPACE/oqc-build" \

0 commit comments

Comments
 (0)