build-wheel-macos-arm64.yaml

name: Build Catalyst Wheel on macOS (arm64)

on:
  pull_request:
    types:
      - opened
      - reopened
      - synchronize
      - labeled
      - ready_for_review
  push:
    branches: [ main ]
  workflow_dispatch:
  workflow_call:

env:
  MACOSX_DEPLOYMENT_TARGET: 14.0

concurrency:
  group: Build Catalyst Wheel on macOS (arm64)-${{ github.ref }}
  cancel-in-progress: true

jobs:
  check_if_wheel_build_required:
    if: github.event.pull_request.draft == false
    uses: ./.github/workflows/check-for-wheel-build.yml

  constants:
    needs: [check_if_wheel_build_required]
    if: needs.check_if_wheel_build_required.outputs.build-wheels == 'true'
    name: "Set build matrix"
    uses: ./.github/workflows/constants.yaml

  build-dependencies:
    needs: [constants, check_if_wheel_build_required]

    strategy:
      fail-fast: false
      matrix:
        python_version: [3.9]

    name: Build Dependencies (Python ${{ matrix.python_version }})
    runs-on: macos-latest

    if: needs.check_if_wheel_build_required.outputs.build-wheels == 'true'

    steps:
    - name: Checkout Catalyst repo
      uses: actions/checkout@v4

    # Cache external project sources
    # Hopefully these can be shared with the main check-catalyst action since we don't run this
    # build in a container.
    - name: Cache LLVM Source
      id: cache-llvm-source
      uses: actions/cache@v4
      with:
        path: 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
        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
        key: enzyme-${{ needs.constants.outputs.enzyme_version }}-default-source
        enableCrossOsArchive: True

    - name: Clone LLVM Submodule
      if: steps.cache-llvm-source.outputs.cache-hit != 'true'
      uses: actions/checkout@v4
      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@v4
      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@v4
      with:
        repository: EnzymeAD/Enzyme
        ref: ${{ needs.constants.outputs.enzyme_version }}
        path: mlir/Enzyme

    # Cache external project builds
    - name: Restore LLVM Build
      id: cache-llvm-build
      uses: actions/cache/restore@v4
      with:
        path:  llvm-build
        key: ${{ runner.os }}-${{ runner.arch }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{matrix.python_version}}-wheel-build

    - name: Restore MHLO Build
      id: cache-mhlo-build
      uses: actions/cache/restore@v4
      with:
        path:  mhlo-build
        key: ${{ runner.os }}-${{ runner.arch }}-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
        key: ${{ runner.os }}-${{ runner.arch }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-build
        lookup-only: True

    - name: Setup Python version
      # There are multiple Python versions installed on the GitHub image, 3.9 - 3.12 is already
      # available under /Library/Frameworks/Python.framework/Versions/, but homebrew also provides
      # 3.11 and 3.12. Make sure to consistently use the system versions.
      run: |
        echo /Library/Frameworks/Python.framework/Versions/${{ matrix.python_version }}/bin >> $GITHUB_PATH

    - name: Install Dependencies (Python)
      run: |
        python${{ matrix.python_version }} -m pip install numpy pybind11 PyYAML 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=FORCE_ON \
              -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())") \
              -DCMAKE_CXX_VISIBILITY_PRESET=hidden

        # 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: Save LLVM Build
      id: save-llvm-build
      if: steps.cache-llvm-build.outputs.cache-hit != 'true'
      uses: actions/cache/save@v4
      with:
        path:  llvm-build
        key: ${{ runner.os }}-${{ runner.arch }}-llvm-${{ needs.constants.outputs.llvm_version }}-${{matrix.python_version}}-wheel-build

    - name: Build MHLO Dialect
      if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
      run: |
        export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH

        export TARGET_FILE=mlir/mlir-hlo/mhlo/transforms/CMakeLists.txt
        export PATCH_FILE=mlir/patches/mhlo-Add-PassesIncGen-in-transforms-CMakeList.patch
        if patch --dry-run -p1 -N $TARGET_FILE $PATCH_FILE > /dev/null 2>&1; then patch -p1 $TARGET_FILE $PATCH_FILE; fi

        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 \
              -DPython3_EXECUTABLE=$(which python${{ matrix.python_version }}) \
              -DLLVM_ENABLE_LLD=OFF \
              -DLLVM_ENABLE_ZLIB=OFF \
              -DLLVM_ENABLE_ZSTD=FORCE_ON \
              -DCMAKE_CXX_VISIBILITY_PRESET=hidden

        cmake --build mhlo-build --target check-mlir-hlo

    - name: Save MHLO Build
      id: save-mhlo-build
      if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
      uses: actions/cache/save@v4
      with:
        path:  mhlo-build
        key: ${{ runner.os }}-${{ runner.arch }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build

    - 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 \
              -DCMAKE_CXX_VISIBILITY_PRESET=hidden

        cmake --build enzyme-build --target EnzymeStatic-19

    - name: Save Enzyme Build
      id: save-enzyme-build
      if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
      uses: actions/cache/save@v4
      with:
        path: enzyme-build
        key: ${{ runner.os }}-${{ runner.arch }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-build

  catalyst-macos-wheels-arm64:
    needs: [constants, build-dependencies]
    strategy:
      fail-fast: false
      matrix:
        python_version: ${{ fromJson(needs.constants.outputs.python_versions) }}

    name: Build Wheels (Python ${{ matrix.python_version }})
    runs-on: macos-latest

    steps:
    - name: Checkout Catalyst repo
      uses: actions/checkout@v4

    - name: Install Dependencies (System)
      run: |
        brew install libomp

    - name: Setup Python version
      # There are multiple Python versions installed on the GitHub image, 3.9 - 3.12 is already
      # available under /Library/Frameworks/Python.framework/Versions/, but homebrew also provides
      # 3.11 and 3.12. Make sure to consistently use the system versions.
      run: |
        echo /Library/Frameworks/Python.framework/Versions/${{ matrix.python_version }}/bin >> $GITHUB_PATH

    - name: Install Dependencies (Python)
      run: |
        python${{ matrix.python_version }} -m pip install numpy pybind11 PyYAML cmake ninja delocate

    - name: Get Cached LLVM Source
      id: cache-llvm-source
      uses: actions/cache/restore@v4
      with:
        path: mlir/llvm-project
        key: llvm-${{ needs.constants.outputs.llvm_version }}-default-source
        enableCrossOsArchive: True
        fail-on-cache-miss: True

    - name: Get Cached LLVM Build
      id: cache-llvm-build
      uses: actions/cache/restore@v4
      with:
        path:  llvm-build
        key: ${{ runner.os }}-${{ runner.arch }}-llvm-${{ needs.constants.outputs.llvm_version }}-3.9-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
        key: mhlo-${{ needs.constants.outputs.mhlo_version }}-default-source
        enableCrossOsArchive: True
        fail-on-cache-miss: True

    - name: Get Cached MHLO Build
      id: cache-mhlo-build
      uses: actions/cache/restore@v4
      with:
        path: mhlo-build
        key: ${{ runner.os }}-${{ runner.arch }}-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
        key: enzyme-${{ needs.constants.outputs.enzyme_version }}-default-source
        enableCrossOsArchive: True
        fail-on-cache-miss: True

    - name: Get Cached Enzyme Build
      id: cache-enzyme-build
      uses: actions/cache/restore@v4
      with:
        path: enzyme-build
        key: ${{ runner.os }}-${{ runner.arch }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-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())") \
              -DLIGHTNING_GIT_TAG=352e7d1 \
              -DENABLE_WARNINGS=OFF \
              -DENABLE_OPENQASM=ON \
              -DENABLE_OPENMP=OFF \
              -DLQ_ENABLE_KERNEL_OMP=OFF

        cmake --build runtime-build --target rt_capi rtd_lightning rtd_openqasm rtd_dummy

    - name: Test Catalyst-Runtime
      run: |
        python${{ matrix.python_version }} -m pip install 'amazon-braket-pennylane-plugin>1.27.1'
        cmake --build runtime-build --target runner_tests_lightning runner_tests_openqasm
        ./runtime-build/tests/runner_tests_lightning
        ./runtime-build/tests/runner_tests_openqasm

    # Build OQC-Runtime
    - name: Build OQC-Runtime
      run: |
          OQC_BUILD_DIR=$GITHUB_WORKSPACE/oqc-build \
          RT_BUILD_DIR=$GITHUB_WORKSPACE/runtime-build \
          PYTHON=$(which python${{ matrix.python_version }}) \
          make oqc

    # 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=FORCE_ON \
              -DLLVM_ENABLE_LLD=OFF \
              -DLLVM_DIR=$GITHUB_WORKSPACE/llvm-build/lib/cmake/llvm

        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 \
        OQC_BUILD_DIR=$GITHUB_WORKSPACE/oqc-build \
        ENZYME_BUILD_DIR=$GITHUB_WORKSPACE/enzyme-build \
        make wheel

    - name: Repair wheel using delocate-wheel
      run: |
        # ignore-missing-dependencies only ignores libopenblas.dylib
        delocate-wheel --require-archs=arm64 -w ./wheel -v dist/*.whl --ignore-missing-dependencies -vv

    - name: Upload Wheel Artifact
      uses: actions/upload-artifact@v4
      with:
        name: catalyst-macos_arm64-wheel-py-${{ matrix.python_version }}.zip
        path: wheel/
        retention-days: 14

  test-wheels:
    needs: [constants, catalyst-macos-wheels-arm64]
    strategy:
      fail-fast: false
      matrix:
        python_version: ${{ fromJson(needs.constants.outputs.python_versions) }}

    # To check all wheels for supported python3 versions
    name: Test Wheels (Python ${{ matrix.python_version }}) on Mac arm64
    runs-on: macos-latest

    steps:
    - name: Checkout Catalyst repo
      uses: actions/checkout@v4

    - name: Download Wheel Artifact
      uses: actions/download-artifact@v4
      with:
        name: catalyst-macos_arm64-wheel-py-${{ matrix.python_version }}.zip
        path: dist

    - name: Setup Python version
      # There are multiple Python versions installed on the GitHub image, 3.9 - 3.12 is already
      # available under /Library/Frameworks/Python.framework/Versions/, but homebrew also provides
      # 3.11 and 3.12. Make sure to consistently use the system versions.
      run: |
        echo /Library/Frameworks/Python.framework/Versions/${{ matrix.python_version }}/bin >> $GITHUB_PATH

    - name: Install Python dependencies
      run: |
        python${{ matrix.python_version }} -m pip install pytest pytest-xdist pytest-mock

    - 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>1.27.1'

    - name: Install OQC client
      run: |
        python${{ matrix.python_version }} -m pip install oqc-qcaas-client

    - name: Install Catalyst
      run: |
        python${{ matrix.python_version }} -m pip install dist/*.whl --extra-index-url https://test.pypi.org/simple

    - name: Run Python Pytest Tests
      run: |
        python${{ matrix.python_version }} -m pytest frontend/test/pytest -n auto
        python${{ matrix.python_version }} -m pytest frontend/test/pytest --backend="lightning.kokkos" -n auto
        python${{ matrix.python_version }} -m pytest frontend/test/async_tests
        python${{ matrix.python_version }} -m pytest frontend/test/pytest --runbraket=LOCAL -n auto
        python${{ matrix.python_version }} -m pytest frontend/test/test_oqc/oqc -n auto