Build wheels on macOS arm64

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

@@ -0,0 +1,341 @@
+name: Build Catalyst Wheel on macOS (arm64)
+
+# TODO: remove on pull_request before merging
+on:
+  pull_request:
+  push:
+    branches: [ main ]
+  workflow_dispatch:
+
+env:
+  MACOSX_DEPLOYMENT_TARGET: 11.0
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  constants:
+    name: "Set build matrix"
+    uses: ./.github/workflows/constants.yaml
+
+  build-dependencies:
+    needs: [constants]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [macos-latest-xlarge]
+        arch: [arm64]
+        python_version: ${{ fromJson(needs.constants.outputs.python_versions) }}
+
+    name: Build Dependencies (Python ${{ matrix.python_version }})
+    runs-on: macos-latest-xlarge
+
+    steps:
+    - name: Checkout Catalyst repo
+      uses: actions/checkout@v3
+
+    # Cache external project sources
+    - name: Cache LLVM Source
+      id: cache-llvm-source
+      uses: actions/cache@v3
+      with:
+        path:  mlir/llvm-project
+        key: ${{ runner.os }}-${{ runner.arch }}-llvm-${{ needs.constants.outputs.llvm_version }}-generic-source
+        enableCrossOsArchive: True
+
+    - name: Cache MHLO Source
+      id: cache-mhlo-source
+      uses: actions/cache@v3
+      with:
+        path: mlir/mlir-hlo
+        key: ${{ runner.os }}-${{ runner.arch }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-generic-source
+        enableCrossOsArchive: True
+
+    - name: Cache Enzyme Source
+      id: cache-enzyme-source
+      uses: actions/cache@v3
+      with:
+        path:  mlir/Enzyme
+        key: ${{ runner.os }}-${{ runner.arch }}-enzyme-${{ needs.constants.outputs.enzyme_version }}-generic-source
+        enableCrossOsArchive: True
+
+    - name: Clone LLVM Submodule
+      if: steps.cache-llvm-source.outputs.cache-hit != 'true'
+      uses: actions/checkout@v3
+      with:
+        repository: llvm/llvm-project
+        ref: ${{ needs.constants.outputs.llvm_version }}
+        path: mlir/llvm-project
+
+    - name: Clone MHLO Submodule
+      if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
+      uses: actions/checkout@v3
+      with:
+        repository: tensorflow/mlir-hlo
+        ref: ${{ needs.constants.outputs.mhlo_version }}
+        path: mlir/mlir-hlo
+
+    - name: Clone Enzyme Submodule
+      if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
+      uses: actions/checkout@v3
+      with:
+        repository: EnzymeAD/Enzyme
+        ref: ${{ needs.constants.outputs.enzyme_version }}
+        path: mlir/Enzyme
+
+    # Cache external project builds
+    - name: Cache LLVM Build
+      id: cache-llvm-build
+      uses: actions/cache@v3
+      with:
+        path:  llvm-build
+        key: ${{ runner.os }}-${{ runner.arch }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{matrix.python_version}}-generic-build
+
+    - name: Cache MHLO Build
+      id: cache-mhlo-build
+      uses: actions/cache@v3
+      with:
+        path:  mhlo-build
+        key: ${{ runner.os }}-${{ runner.arch }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-generic-build
+
+    - name: Cache Enzyme Build
+      id: cache-enzyme-build
+      uses: actions/cache@v3
+      with:
+        path:  enzyme-build
+        key: ${{ runner.os }}-${{ runner.arch }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-generic-build
+
+    - name: Set up Python ${{ matrix.python_version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python_version }}
+
+    - name: Install Dependencies (Python)
+      run: |
+        python${{ matrix.python_version }} -m pip install wheel numpy pybind11 cmake ninja
+
+    - name: Build LLVM / MLIR
+      if: steps.cache-llvm-build.outputs.cache-hit != 'true'
+      run: |
+        cmake -S mlir/llvm-project/llvm -B llvm-build -G Ninja \
+              -DCMAKE_BUILD_TYPE=Release \
+              -DLLVM_BUILD_EXAMPLES=OFF \
+              -DLLVM_TARGETS_TO_BUILD="host" \
+              -DLLVM_ENABLE_PROJECTS="mlir" \
+              -DLLVM_ENABLE_ASSERTIONS=ON \
+              -DLLVM_INSTALL_UTILS=ON \
+              -DLLVM_ENABLE_ZLIB=OFF \
+              -DLLVM_ENABLE_ZSTD=OFF \
+              -DLLVM_ENABLE_LLD=OFF \
+              -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
+              -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
+              -DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())")
+
+        # TODO: when updating LLVM, test to see if mlir/unittests/Bytecode/BytecodeTest.cpp:55 is passing
+        # and remove filter
+        # This tests fails on CI/CD not locally.
+        LIT_FILTER_OUT="Bytecode" cmake --build llvm-build --target check-mlir
+
+    - name: Build MHLO Dialect
+      if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
+      run: |
+        export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH
+        cmake -S mlir/mlir-hlo -B mhlo-build -G Ninja \
+              -DCMAKE_BUILD_TYPE=Release \
+              -DLLVM_ENABLE_ASSERTIONS=ON \
+              -DMLIR_DIR=$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir \
+              -DLLVM_ENABLE_LLD=OFF \
+              -DLLVM_ENABLE_ZLIB=OFF \
+              -DLLVM_ENABLE_ZSTD=OFF
+
+        cmake --build mhlo-build --target check-mlir-hlo
+
+    - name: Build Enzyme
+      if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
+      run: |
+        cmake -S mlir/Enzyme/enzyme -B enzyme-build -G Ninja \
+              -DCMAKE_BUILD_TYPE=Release \
+              -DLLVM_DIR=$GITHUB_WORKSPACE/llvm-build/lib/cmake/llvm \
+              -DENZYME_STATIC_LIB=ON
+
+        cmake --build enzyme-build --target EnzymeStatic-18
+
+  catalyst-macos-wheels-x86-64:
+    needs: [constants, build-dependencies]
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [macos-latest-xlarge]
+        arch: [arm64]
+        python_version: ${{ fromJson(needs.constants.outputs.python_versions) }}
+
+    name: Build Wheels (Python ${{ matrix.python_version }})
+    runs-on: ${{ matrix.os }}
+
+    steps:
+    - name: Checkout Catalyst repo
+      uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python_version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python_version }}
+
+    - name: Install Dependencies (Python)
+      run: |
+        python${{ matrix.python_version }} -m pip install wheel numpy pybind11 cmake ninja
+
+    - name: Get Cached LLVM Source
+      id: cache-llvm-source
+      uses: actions/cache@v3
+      with:
+        path: mlir/llvm-project
+        key: ${{ runner.os }}-${{ runner.arch }}-llvm-${{ needs.constants.outputs.llvm_version }}-generic-source
+        enableCrossOsArchive: True
+        fail-on-cache-miss: True
+
+    - name: Get Cached LLVM Build
+      id: cache-llvm-build
+      uses: actions/cache@v3
+      with:
+        path:  llvm-build
+        key: ${{ runner.os }}-${{ runner.arch }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{matrix.python_version}}-generic-build
+        fail-on-cache-miss: True
+
+    - name: Get Cached MHLO Source
+      id: cache-mhlo-source
+      uses: actions/cache@v3
+      with:
+        path: mlir/mlir-hlo
+        key: ${{ runner.os }}-${{ runner.arch }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-generic-source
+        enableCrossOsArchive: True
+        fail-on-cache-miss: True
+
+    - name: Get Cached MHLO Build
+      id: cache-mhlo-build
+      uses: actions/cache@v3
+      with:
+        path:  mhlo-build
+        key: ${{ runner.os }}-${{ runner.arch }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-generic-build
+        fail-on-cache-miss: True
+
+    - name: Get Cached Enzyme Source
+      id: cache-enzyme-source
+      uses: actions/cache@v3
+      with:
+        path:  mlir/Enzyme
+        key: ${{ runner.os }}-${{ runner.arch }}-enzyme-${{ needs.constants.outputs.enzyme_version }}-generic-source
+        fail-on-cache-miss: True
+
+    - name: Get Cached Enzyme Build
+      id: cache-enzyme-build
+      uses: actions/cache@v3
+      with:
+        path:  enzyme-build
+        key: ${{ runner.os }}-${{ runner.arch }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-generic-build
+        fail-on-cache-miss: True
+
+    # Build Catalyst-Runtime
+    - name: Build Catalyst-Runtime
+      run: |
+        cmake -S runtime -B runtime-build -G Ninja \
+              -DCMAKE_BUILD_TYPE=Release \
+              -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$GITHUB_WORKSPACE/runtime-build/lib \
+              -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
+              -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
+              -DENABLE_LIGHTNING_KOKKOS=ON \
+              -DLIGHTNING_GIT_TAG="latest_release" \
+              -DKokkos_ENABLE_SERIAL=ON \
+              -DKokkos_ENABLE_COMPLEX_ALIGN=OFF \
+              -DBUILD_QIR_STDLIB_FROM_SRC=ON \
+              -DENABLE_OPENQASM=ON
+
+        cmake --build runtime-build --target rt_capi
+
+    # Build Quantum and Gradient Dialects
+    - name: Build MLIR Dialects
+      run: |
+        cmake -S mlir -B quantum-build -G Ninja \
+              -DCMAKE_BUILD_TYPE=Release \
+              -DLLVM_ENABLE_ASSERTIONS=ON \
+              -DQUANTUM_ENABLE_BINDINGS_PYTHON=ON \
+              -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 \
+              -DEnzyme_DIR=$GITHUB_WORKSPACE/enzyme-build \
+              -DENZYME_SRC_DIR=$GITHUB_WORKSPACE/mlir/Enzyme \
+              -DLLVM_ENABLE_ZLIB=OFF \
+              -DLLVM_ENABLE_ZSTD=OFF \
+              -DLLVM_ENABLE_LLD=OFF
+
+        cmake --build quantum-build --target check-dialects compiler_driver
+
+    - name: Build wheel
+      run: |
+        PYTHON=python${{ matrix.python_version }} \
+        LLVM_BUILD_DIR=$GITHUB_WORKSPACE/llvm-build \
+        MHLO_BUILD_DIR=$GITHUB_WORKSPACE/mhlo-build \
+        DIALECTS_BUILD_DIR=$GITHUB_WORKSPACE/quantum-build \
+        RT_BUILD_DIR=$GITHUB_WORKSPACE/runtime-build \
+        ENZYME_BUILD_DIR=$GITHUB_WORKSPACE/enzyme-build \
+        make wheel
+
+    - name: Upload Wheel Artifact
+      uses: actions/upload-artifact@v3
+      with:
+        name: catalyst-macos_arm64-wheel-py-${{ matrix.python_version }}.zip
+        path: dist/
+        retention-days: 14
+
+  test-wheels:
+    needs: [constants, catalyst-macos-wheels-x86-64]
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [macos-latest-xlarge]
+        arch: [arm64]
+        python_version: ${{ fromJson(needs.constants.outputs.python_versions) }}
+
+    # To check all wheels for supported python3 versions
+    name: Test Wheels (Python ${{ matrix.python_version }}) on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+
+    steps:
+    - name: Checkout Catalyst repo
+      uses: actions/checkout@v3
+
+    - name: Download Wheel Artifact
+      uses: actions/download-artifact@v3
+      with:
+        name: catalyst-macos_arm64-wheel-py-${{ matrix.python_version }}.zip
+        path: dist
+
+    - name: Set up Python ${{ matrix.python_version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python_version }}
+
+    - name: Install Python dependencies
+      run: |
+        # tensorflow-cpu is not distributed for macOS ARM
+        python${{ matrix.python_version }} -m pip install pytest pytest-xdist
+        python${{ matrix.python_version }} -m pip install tensorflow  # for autograph tests
+
+    - name: Install Catalyst
+      run: |
+        python${{ matrix.python_version }} -m pip install $GITHUB_WORKSPACE/dist/*.whl
+
+    - name: Install PennyLane Plugins
+      run: |
+        python${{ matrix.python_version }} -m pip install PennyLane-Lightning-Kokkos
+        python${{ matrix.python_version }} -m pip install amazon-braket-pennylane-plugin
+
+    - name: Run Python Pytest Tests
+      run: |
+        python${{ matrix.python_version }} -m pytest $GITHUB_WORKSPACE/frontend/test/pytest -n auto
+        python${{ matrix.python_version }} -m pytest $GITHUB_WORKSPACE/frontend/test/pytest --backend="lightning.kokkos" -n auto
+        python${{ matrix.python_version }} -m pytest $GITHUB_WORKSPACE/frontend/test/pytest --runbraket=LOCAL -n auto