Cherry-pick PR #1408 into the RC (#1423)

PennyLaneAI · Jan 8, 2025 · f7d0603 · f7d0603
1 parent 07ffe5b
commit f7d0603
Show file tree

Hide file tree

Showing 11 changed files with 234 additions and 491 deletions.
diff --git a/.github/workflows/build-enzyme-v0.0.130.yaml b/.github/workflows/build-enzyme-v0.0.130.yaml
diff --git a/.github/workflows/build-wheel-linux-arm64.yaml b/.github/workflows/build-wheel-linux-arm64.yaml
@@ -62,23 +62,23 @@ jobs:
       id: cache-llvm-source
       uses: actions/cache@v4
       with:
-        path:  mlir/llvm-project
+        path: ${{ github.workspace }}/mlir/llvm-project
         key: llvm-${{ needs.constants.outputs.llvm_version }}-default-source
         enableCrossOsArchive: True
 
     - name: Cache MHLO Source
       id: cache-mhlo-source
       uses: actions/cache@v4
       with:
-        path: mlir/mlir-hlo
+        path: ${{ github.workspace }}/mlir/mlir-hlo
         key: mhlo-${{ needs.constants.outputs.mhlo_version }}-default-source
         enableCrossOsArchive: True
 
     - name: Cache Enzyme Source
       id: cache-enzyme-source
       uses: actions/cache@v4
       with:
-        path:  mlir/Enzyme
+        path: ${{ github.workspace }}/mlir/Enzyme
         key: enzyme-${{ needs.constants.outputs.enzyme_version }}-default-source
         enableCrossOsArchive: True
 
@@ -88,60 +88,48 @@ jobs:
       with:
         repository: llvm/llvm-project
         ref: ${{ needs.constants.outputs.llvm_version }}
-        path: mlir/llvm-project
+        path: ${{ github.workspace }}/mlir/llvm-project
 
     - name: Clone MHLO Submodule
       if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
       uses: actions/checkout@v4
       with:
         repository: tensorflow/mlir-hlo
         ref: ${{ needs.constants.outputs.mhlo_version }}
-        path: mlir/mlir-hlo
+        path: ${{ github.workspace }}/mlir/mlir-hlo
 
     - name: Clone Enzyme Submodule
       if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
       uses: actions/checkout@v4
       with:
         repository: EnzymeAD/Enzyme
         ref: ${{ needs.constants.outputs.enzyme_version }}
-        path: mlir/Enzyme
+        path: ${{ github.workspace }}/mlir/Enzyme
 
     # Cache external project builds
     - name: Restore LLVM Build
       id: cache-llvm-build
       uses: actions/cache/restore@v4
       with:
-        path:  llvm-build
+        path: ${{ github.workspace }}/llvm-build
         key: ${{ matrix.container_name }}-llvm-${{ needs.constants.outputs.llvm_version }}-wheel-build
 
     - name: Restore MHLO Build
       id: cache-mhlo-build
       uses: actions/cache/restore@v4
       with:
-        path:  mhlo-build
+        path: ${{ github.workspace }}/mhlo-build
         key: ${{ matrix.container_name }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
         lookup-only: True
 
     - name: Restore Enzyme Build
       id: cache-enzyme-build
       uses: actions/cache/restore@v4
       with:
-        path: enzyme-build
+        path: ${{ github.workspace }}/enzyme-build
         key: ${{ matrix.container_name }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-build
         lookup-only: True
 
-    - name: Build LLD
-      if: steps.cache-llvm-build.outputs.cache-hit != 'true'
-      run: |
-        set -x
-        # With GCC 13, LLVM fails some tests, then we use GCC 12 instead (Copied from below)
-        export GCC_VERSION=12
-        docker run --rm --platform linux/aarch64 \
-            -v /var/run/docker.sock:/var/run/docker.sock \
-            -v `pwd`:/catalyst \
-            -i ${{ matrix.container_img }} \
-            bash /catalyst/.github/workflows/scripts/linux_arm64/rh8/build_lld.sh $GCC_VERSION ${{ matrix.python_version.major_minor }} ${{ matrix.python_version.patch }} ${{ matrix.python_version.package }}
-
     - name: Build LLVM / MLIR
       if: steps.cache-llvm-build.outputs.cache-hit != 'true'
       run: |
@@ -159,7 +147,7 @@ jobs:
       if: steps.cache-llvm-build.outputs.cache-hit != 'true'
       uses: actions/cache/save@v4
       with:
-        path:  llvm-build
+        path: ${{ github.workspace }}/llvm-build
         key: ${{ matrix.container_name }}-llvm-${{ needs.constants.outputs.llvm_version }}-wheel-build
 
     - name: Build MHLO Dialect
@@ -178,7 +166,7 @@ jobs:
       if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
       uses: actions/cache/save@v4
       with:
-        path:  mhlo-build
+        path: ${{ github.workspace }}/mhlo-build
         key: ${{ matrix.container_name }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
 
     - name: Build Enzyme
@@ -197,7 +185,7 @@ jobs:
       if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
       uses: actions/cache/save@v4
       with:
-        path: enzyme-build
+        path: ${{ github.workspace }}/enzyme-build
         key: ${{ matrix.container_name }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-build
 
   catalyst-linux-wheels-arm64:
@@ -230,7 +218,7 @@ jobs:
       id: cache-llvm-source
       uses: actions/cache/restore@v4
       with:
-        path: mlir/llvm-project
+        path: ${{ github.workspace }}/mlir/llvm-project
         key: llvm-${{ needs.constants.outputs.llvm_version }}-default-source
         enableCrossOsArchive: True
         fail-on-cache-miss: True
@@ -239,15 +227,15 @@ jobs:
       id: cache-llvm-build
       uses: actions/cache/restore@v4
       with:
-        path:  llvm-build
+        path: ${{ github.workspace }}/llvm-build
         key: ${{ matrix.container_name }}-llvm-${{ needs.constants.outputs.llvm_version }}-wheel-build
         fail-on-cache-miss: True
 
     - name: Get Cached MHLO Source
       id: cache-mhlo-source
       uses: actions/cache/restore@v4
       with:
-        path: mlir/mlir-hlo
+        path: ${{ github.workspace }}/mlir/mlir-hlo
         key: mhlo-${{ needs.constants.outputs.mhlo_version }}-default-source
         enableCrossOsArchive: True
         fail-on-cache-miss: True
@@ -256,15 +244,15 @@ jobs:
       id: cache-mhlo-build
       uses: actions/cache/restore@v4
       with:
-        path: mhlo-build
+        path: ${{ github.workspace }}/mhlo-build
         key: ${{ matrix.container_name }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
         fail-on-cache-miss: True
 
     - name: Get Cached Enzyme Source
       id: cache-enzyme-source
       uses: actions/cache/restore@v4
       with:
-        path: mlir/Enzyme
+        path: ${{ github.workspace }}/mlir/Enzyme
         key: enzyme-${{ needs.constants.outputs.enzyme_version }}-default-source
         enableCrossOsArchive: True
         fail-on-cache-miss: True
@@ -273,7 +261,7 @@ jobs:
       id: cache-enzyme-build
       uses: actions/cache/restore@v4
       with:
-        path: enzyme-build
+        path: ${{ github.workspace }}/enzyme-build
         key: ${{ matrix.container_name }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-build
         fail-on-cache-miss: True
 
@@ -292,7 +280,7 @@ jobs:
       uses: actions/upload-artifact@v4
       with:
         name: catalyst-linux_arm64-wheel-py-${{ matrix.python_version.major_minor}}.zip
-        path: wheel/
+        path: ${{ github.workspace }}/wheel/
         retention-days: 14
 
   test-wheels:
@@ -326,16 +314,7 @@ jobs:
       uses: actions/download-artifact@v4
       with:
         name: catalyst-linux_arm64-wheel-py-${{ matrix.python_version.major_minor }}.zip
-        path: dist
-
-    # Needed for accessing llvm-symbolizer
-    - name: Get Cached LLVM Build
-      id: cache-llvm-build
-      uses: actions/cache@v4
-      with:
-        path:  llvm-build
-        key: ${{ matrix.container_name }}-llvm-${{ needs.constants.outputs.llvm_version }}-wheel-build
-        fail-on-cache-miss: True
+        path: ${{ github.workspace }}/dist
 
     - name: Run Python Pytest Tests
       run: |