Migrate to stablehlo

.dep-versions

@@ -2,7 +2,7 @@
 # To update JAX version alongside compatible dependency tags, run the following script:
 # python3 .github/workflows/set_dep_versions.py {JAX_version}
 jax=0.6.2
-mhlo=1dd2e71331014ae0373f6bf900ce6be393357190
+stablehlo=69d6dae46e1c7de36e6e6973654754f05353cba5
 llvm=f8cb7987c64dcffb72414a40560055cb717dbf74
 enzyme=v0.0.186
 

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

@@ -84,12 +84,12 @@ jobs:
         key: llvm-${{ needs.constants.outputs.llvm_version }}-container-source
         enableCrossOsArchive: True
 
-    - name: Cache MHLO Source
-      id: cache-mhlo-source
+    - name: Cache Stablehlo Source
+      id: cache-stablehlo-source
       uses: actions/cache@v4
       with:
-        path: ${{ github.workspace }}/mlir/mlir-hlo
-        key: mhlo-${{ needs.constants.outputs.mhlo_version }}-container-source
+        path: ${{ github.workspace }}/mlir/stablehlo
+        key: stablehlo-${{ needs.constants.outputs.stablehlo_version }}-container-source
         enableCrossOsArchive: True
 
     - name: Cache Enzyme Source
@@ -109,25 +109,18 @@ jobs:
         path: ${{ github.workspace }}/mlir/llvm-project
 
     - name: Patch LLVM Source
-      if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
+      if: steps.cache-llvm-source.outputs.cache-hit != 'true'
       run: |
         cd $GITHUB_WORKSPACE/mlir/llvm-project
         git apply $GITHUB_WORKSPACE/mlir/patches/llvm-bufferization-segfault.patch
 
-    - name: Clone MHLO Submodule
-      if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
+    - name: Clone Stablehlo Submodule
+      if: steps.cache-stablehlo-source.outputs.cache-hit != 'true'
       uses: actions/checkout@v4
       with:
-        repository: tensorflow/mlir-hlo
-        ref: ${{ needs.constants.outputs.mhlo_version }}
-        path: ${{ github.workspace }}/mlir/mlir-hlo
-
-    - name: Patch MHLO Source
-      if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
-      run: |
-        cd $GITHUB_WORKSPACE/mlir/mlir-hlo
-        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
-        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-rename-sort.patch
+        repository: openxla/stablehlo
+        ref: ${{ needs.constants.outputs.stablehlo_version }}
+        path: ${{ github.workspace }}/mlir/stablehlo
 
     - name: Clone Enzyme Submodule
       if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
@@ -151,12 +144,12 @@ jobs:
         path: ${{ github.workspace }}/llvm-build
         key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{ matrix.python_version }}-wheel-build
 
-    - name: Check MHLO Build Cache
-      id: cache-mhlo-build
+    - name: Check Stablehlo Build Cache
+      id: cache-stablehlo-build
       uses: actions/cache/restore@v4
       with:
-        path: ${{ github.workspace }}/mhlo-build
-        key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
+        path: ${{ github.workspace }}/stablehlo-build
+        key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
         lookup-only: True
 
     - name: Check Enzyme Build Cache
@@ -170,7 +163,7 @@ jobs:
     - name: Install dependencies
       if: |
         steps.cache-llvm-build.outputs.cache-hit != 'true' ||
-        steps.cache-mhlo-build.outputs.cache-hit != 'true' ||
+        steps.cache-stablehlo-build.outputs.cache-hit != 'true' ||
         steps.cache-enzyme-build.outputs.cache-hit != 'true'
       run: |
         cat /etc/dnf.conf | sed "s/\[main\]/\[main\]\ntimeout=5/g" > /etc/dnf.conf
@@ -207,31 +200,29 @@ jobs:
         path: ${{ github.workspace }}/llvm-build
         key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{matrix.python_version}}-wheel-build
 
-    - name: Build MHLO Dialect
-      if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
-      # building with LLD is a strong requirement for mhlo
+    - name: Build Stablehlo Dialect
+      if: steps.cache-stablehlo-build.outputs.cache-hit != 'true'
       run: |
         export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
-
-        cmake -S mlir/mlir-hlo -B $GITHUB_WORKSPACE/mhlo-build -G Ninja \
+        cmake -S $GITHUB_WORKSPACE/mlir/stablehlo -B $GITHUB_WORKSPACE/stablehlo-build -G Ninja \
               -DCMAKE_BUILD_TYPE=Release \
               -DLLVM_ENABLE_ASSERTIONS=ON \
               -DMLIR_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir" \
-              -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
-              -DLLVM_ENABLE_ZLIB=FORCE_ON \
-              -DLLVM_ENABLE_ZSTD=OFF \
-              -DCMAKE_CXX_VISIBILITY_PRESET=default \
-              -DLLVM_ENABLE_LLD=ON
+              -DSTABLEHLO_ENABLE_LLD=ON \
+              -DSTABLEHLO_ENABLE_BINDINGS_PYTHON=OFF \
+              -DSTABLEHLO_ENABLE_SPLIT_DWARF=ON \
+              -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
+              -DCMAKE_CXX_VISIBILITY_PRESET=default
 
-        LIT_FILTER_OUT="chlo_legalize_to_mhlo" cmake --build $GITHUB_WORKSPACE/mhlo-build --target check-mlir-hlo
+        cmake --build $GITHUB_WORKSPACE/stablehlo-build
 
-    - name: Save MHLO Build
-      id: save-mhlo-build
-      if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
+    - name: Save Stablehlo Build
+      id: save-stablehlo-build
+      if: steps.cache-stablehlo-build.outputs.cache-hit != 'true'
       uses: actions/cache/save@v4
       with:
-        path: ${{ github.workspace }}/mhlo-build
-        key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
+        path: ${{ github.workspace }}/stablehlo-build
+        key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
 
     - name: Build Enzyme
       if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
@@ -305,21 +296,21 @@ jobs:
         key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-3.11-wheel-build
         fail-on-cache-miss: True
 
-    - name: Get Cached MHLO Source
-      id: cache-mhlo-source
+    - name: Get Cached Stablehlo Source
+      id: cache-stablehlo-source
       uses: actions/cache/restore@v4
       with:
-        path: ${{ github.workspace }}/mlir/mlir-hlo
-        key: mhlo-${{ needs.constants.outputs.mhlo_version }}-container-source
+        path: ${{ github.workspace }}/mlir/stablehlo
+        key: stablehlo-${{ needs.constants.outputs.stablehlo_version }}-container-source
         enableCrossOsArchive: True
         fail-on-cache-miss: True
 
-    - name: Get Cached MHLO Build
-      id: cache-mhlo-build
+    - name: Get Cached Stablehlo Build
+      id: cache-stablehlo-build
       uses: actions/cache/restore@v4
       with:
-        path: ${{ github.workspace }}/mhlo-build
-        key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
+        path: ${{ github.workspace }}/stablehlo-build
+        key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
         fail-on-cache-miss: True
 
     - name: Get Cached Enzyme Source
@@ -372,8 +363,8 @@ jobs:
               -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())") \
               -DMLIR_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir" \
-              -DMHLO_DIR="$GITHUB_WORKSPACE/mhlo-build/lib/cmake/mlir-hlo" \
-              -DMHLO_BINARY_DIR="$GITHUB_WORKSPACE/mhlo-build/bin" \
+              -DSTABLEHLO_DIR="$GITHUB_WORKSPACE/mlir/stablehlo" \
+              -DSTABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
               -DEnzyme_DIR="$GITHUB_WORKSPACE/enzyme-build" \
               -DENZYME_SRC_DIR="$GITHUB_WORKSPACE/mlir/Enzyme" \
               -DLLVM_ENABLE_ZLIB=FORCE_ON \
@@ -396,7 +387,7 @@ jobs:
       run: |
         PYTHON=python${{ matrix.python_version }} \
         LLVM_BUILD_DIR="$GITHUB_WORKSPACE/llvm-build" \
-        MHLO_BUILD_DIR="$GITHUB_WORKSPACE/mhlo-build" \
+        STABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
         DIALECTS_BUILD_DIR="$GITHUB_WORKSPACE/quantum-build" \
         RT_BUILD_DIR="$GITHUB_WORKSPACE/runtime-build" \
         OQC_BUILD_DIR="$GITHUB_WORKSPACE/oqc-build" \

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

@@ -103,12 +103,12 @@ jobs:
         key: llvm-${{ needs.constants.outputs.llvm_version }}-container-source
         enableCrossOsArchive: True
 
-    - name: Cache MHLO Source
-      id: cache-mhlo-source
+    - name: Cache Stablehlo Source
+      id: cache-stablehlo-source
       uses: actions/cache@v4
       with:
-        path: ${{ github.workspace }}/mlir/mlir-hlo
-        key: mhlo-${{ needs.constants.outputs.mhlo_version }}-container-source
+        path: ${{ github.workspace }}/mlir/stablehlo
+        key: stablehlo-${{ needs.constants.outputs.stablehlo_version }}-container-source
         enableCrossOsArchive: True
 
     - name: Cache Enzyme Source
@@ -128,25 +128,18 @@ jobs:
         path: ${{ github.workspace }}/mlir/llvm-project
 
     - name: Patch LLVM Source
-      if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
+      if: steps.cache-llvm-source.outputs.cache-hit != 'true'
       run: |
         cd $GITHUB_WORKSPACE/mlir/llvm-project
         git apply $GITHUB_WORKSPACE/mlir/patches/llvm-bufferization-segfault.patch
 
-    - name: Clone MHLO Submodule
-      if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
+    - name: Clone Stablehlo Submodule
+      if: steps.cache-stablehlo-source.outputs.cache-hit != 'true'
       uses: actions/checkout@v4
       with:
-        repository: tensorflow/mlir-hlo
-        ref: ${{ needs.constants.outputs.mhlo_version }}
-        path: ${{ github.workspace }}/mlir/mlir-hlo
-
-    - name: Patch MHLO Source
-      if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
-      run: |
-        cd $GITHUB_WORKSPACE/mlir/mlir-hlo
-        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
-        git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-rename-sort.patch
+        repository: openxla/stablehlo
+        ref: ${{ needs.constants.outputs.stablehlo_version }}
+        path: ${{ github.workspace }}/mlir/stablehlo
 
     - name: Clone Enzyme Submodule
       if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
@@ -170,12 +163,12 @@ jobs:
         path: ${{ github.workspace }}/llvm-build
         key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{ matrix.python_version }}-wheel-build
 
-    - name: Check MHLO Build Cache
-      id: cache-mhlo-build
+    - name: Check Stablehlo Build Cache
+      id: cache-stablehlo-build
       uses: actions/cache/restore@v4
       with:
-        path: ${{ github.workspace }}/mhlo-build
-        key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
+        path: ${{ github.workspace }}/stablehlo-build
+        key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
         lookup-only: True
 
     - name: Check Enzyme Build Cache
@@ -189,7 +182,7 @@ jobs:
     - name: Install dependencies (AlmaLinux)
       if: |
         steps.cache-llvm-build.outputs.cache-hit != 'true' ||
-        steps.cache-mhlo-build.outputs.cache-hit != 'true' ||
+        steps.cache-stablehlo-build.outputs.cache-hit != 'true' ||
         steps.cache-enzyme-build.outputs.cache-hit != 'true'
       run: |
         # Reduce wait time for repos not responding
@@ -205,7 +198,6 @@ jobs:
         PYTHON_BINS=$(find /opt/_internal/cpython-${{ matrix.python_version }}.*/bin -maxdepth 1 -type d | tr '\n' ':' | sed 's/:$//')
         echo $PYTHON_BINS >> $GITHUB_PATH
 
-    # LLD is required for MHLO builds.
     # (Don't forget to add the build directory to PATH in subsequent steps, so
     # other tools can find it, in particular collect2 invoked by gcc.)
     - name: Build LLVM / MLIR
@@ -230,31 +222,29 @@ jobs:
         path: ${{ github.workspace }}/llvm-build
         key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{matrix.python_version}}-wheel-build
 
-    - name: Build MHLO Dialect
-      if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
-      # building with LLD is a strong requirement for mhlo
+    - name: Build Stablehlo Dialect
+      if: steps.cache-stablehlo-build.outputs.cache-hit != 'true'
       run: |
         export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
-
-        cmake -S mlir/mlir-hlo -B $GITHUB_WORKSPACE/mhlo-build -G Ninja \
+        cmake -S $GITHUB_WORKSPACE/mlir/stablehlo -B $GITHUB_WORKSPACE/stablehlo-build -G Ninja \
               -DCMAKE_BUILD_TYPE=Release \
               -DLLVM_ENABLE_ASSERTIONS=ON \
               -DMLIR_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir" \
-              -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
-              -DLLVM_ENABLE_ZLIB=FORCE_ON \
-              -DLLVM_ENABLE_ZSTD=OFF \
-              -DCMAKE_CXX_VISIBILITY_PRESET=default \
-              -DLLVM_ENABLE_LLD=ON
+              -DSTABLEHLO_ENABLE_LLD=ON \
+              -DSTABLEHLO_ENABLE_BINDINGS_PYTHON=OFF \
+              -DSTABLEHLO_ENABLE_SPLIT_DWARF=ON \
+              -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
+              -DCMAKE_CXX_VISIBILITY_PRESET=default
 
-        LIT_FILTER_OUT="chlo_legalize_to_mhlo" cmake --build $GITHUB_WORKSPACE/mhlo-build --target check-mlir-hlo
+        cmake --build $GITHUB_WORKSPACE/stablehlo-build
 
-    - name: Save MHLO Build
-      id: save-mhlo-build
-      if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
+    - name: Save Stablehlo Build
+      id: save-stablehlo-build
+      if: steps.cache-stablehlo-build.outputs.cache-hit != 'true'
       uses: actions/cache/save@v4
       with:
-        path: ${{ github.workspace }}/mhlo-build
-        key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
+        path: ${{ github.workspace }}/stablehlo-build
+        key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
 
     - name: Build Enzyme
       if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
@@ -328,21 +318,21 @@ jobs:
         key: ${{ matrix.container_img }}-llvm-${{ needs.constants.outputs.llvm_version }}-3.11-wheel-build
         fail-on-cache-miss: True
 
-    - name: Get Cached MHLO Source
-      id: cache-mhlo-source
+    - name: Get Cached Stablehlo Source
+      id: cache-stablehlo-source
       uses: actions/cache/restore@v4
       with:
-        path: ${{ github.workspace }}/mlir/mlir-hlo
-        key: mhlo-${{ needs.constants.outputs.mhlo_version }}-container-source
+        path: ${{ github.workspace }}/mlir/stablehlo
+        key: stablehlo-${{ needs.constants.outputs.stablehlo_version }}-container-source
         enableCrossOsArchive: True
         fail-on-cache-miss: True
 
-    - name: Get Cached MHLO Build
-      id: cache-mhlo-build
+    - name: Get Cached Stablehlo Build
+      id: cache-stablehlo-build
       uses: actions/cache/restore@v4
       with:
-        path: ${{ github.workspace }}/mhlo-build
-        key: ${{ matrix.container_img }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
+        path: ${{ github.workspace }}/stablehlo-build
+        key: ${{ matrix.container_img }}-stablehlo-${{ needs.constants.outputs.stablehlo_version }}-wheel-build
         fail-on-cache-miss: True
 
     - name: Get Cached Enzyme Source
@@ -397,8 +387,8 @@ jobs:
               -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())") \
               -DMLIR_DIR="$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir" \
-              -DMHLO_DIR="$GITHUB_WORKSPACE/mhlo-build/lib/cmake/mlir-hlo" \
-              -DMHLO_BINARY_DIR="$GITHUB_WORKSPACE/mhlo-build/bin" \
+              -DSTABLEHLO_DIR="$GITHUB_WORKSPACE/mlir/stablehlo" \
+              -DSTABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
               -DEnzyme_DIR="$GITHUB_WORKSPACE/enzyme-build" \
               -DENZYME_SRC_DIR="$GITHUB_WORKSPACE/mlir/Enzyme" \
               -DLLVM_ENABLE_ZLIB=FORCE_ON \
@@ -421,7 +411,7 @@ jobs:
       run: |
         PYTHON=python${{ matrix.python_version }} \
         LLVM_BUILD_DIR="$GITHUB_WORKSPACE/llvm-build" \
-        MHLO_BUILD_DIR="$GITHUB_WORKSPACE/mhlo-build" \
+        STABLEHLO_BUILD_DIR="$GITHUB_WORKSPACE/stablehlo-build" \
         DIALECTS_BUILD_DIR="$GITHUB_WORKSPACE/quantum-build" \
         RT_BUILD_DIR="$GITHUB_WORKSPACE/runtime-build" \
         OQC_BUILD_DIR="$GITHUB_WORKSPACE/oqc-build" \