build-wheel-macos-arm64.yaml

name: Build Catalyst Wheel on macOS (arm64)

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

env:
  MACOSX_DEPLOYMENT_TARGET: 14.0

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

jobs:
  check_if_wheel_build_required:
    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: ${{ fromJson(needs.constants.outputs.python_versions) }}

    name: Build Dependencies (Python ${{ matrix.python_version }})
    runs-on:
      group: 'Office 24th floor M2 Mac'

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

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

    - name: Setup Runner Environment
      id: setup_env
      uses: ./.github/workflows/utils/setup_self_hosted_macos_env
      with:
        python_version: ${{ matrix.python_version }}

    # Cache external project sources
    - name: Cache LLVM Source
      id: cache-llvm-source
      uses: actions/cache@v3
      with:
        path:  ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/llvm-project
        key: Linux-llvm-${{ needs.constants.outputs.llvm_version }}-generic-source
        enableCrossOsArchive: True

    - name: Cache MHLO Source
      id: cache-mhlo-source
      uses: actions/cache@v3
      with:
        path: ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/mlir-hlo
        key: Linux-mhlo-${{ needs.constants.outputs.mhlo_version }}-generic-source
        enableCrossOsArchive: True

    - name: Cache Enzyme Source
      id: cache-enzyme-source
      uses: actions/cache@v3
      with:
        path:  ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/Enzyme
        key: Linux-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

      # This step is needed because actions/checkout cannot clone to outside GITHUB_WORKSPACE
      # https://github.com/actions/checkout/issues/197
    - name: Copy Submodule to tmp cache directory
      run: |
        if [ ! -d "${{ steps.setup_env.outputs.dependency_build_dir }}/mlir" ]; then
          echo 'Creating directory at ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir'
          mkdir -p ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir
        fi

        if [ "${{ steps.cache-llvm-source.outputs.cache-hit }}" != 'true' ]; then
          echo 'Copying mlir/llvm-project to ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/'
          mkdir -p ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/llvm-project
          mv mlir/llvm-project/* ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/llvm-project
        fi

        if [ "${{ steps.cache-mhlo-source.outputs.cache-hit }}" != 'true' ]; then
          echo 'Copying mlir/mlir-hlo to ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/'
          mkdir -p ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/mlir-hlo
          mv mlir/mlir-hlo/* ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/mlir-hlo
        fi

        if [ "${{ steps.cache-enzyme-source.outputs.cache-hit }}" != 'true' ]; then
          echo 'Copying mlir/Enzyme to ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/'
          mkdir -p ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/Enzyme
          mv mlir/Enzyme/* ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/Enzyme
        fi

        ls ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/

    # Cache external project builds
    - name: Cache LLVM Build
      id: cache-llvm-build
      uses: actions/cache@v3
      with:
        path:  ${{ steps.setup_env.outputs.dependency_build_dir }}/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:  ${{ steps.setup_env.outputs.dependency_build_dir }}/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: ${{ steps.setup_env.outputs.dependency_build_dir }}/enzyme-build
        key: ${{ runner.os }}-${{ runner.arch }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-generic-build

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

    - name: Build LLVM / MLIR
      if: steps.cache-llvm-build.outputs.cache-hit != 'true'
      run: |
        cmake -S ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/llvm-project/llvm -B ${{ steps.setup_env.outputs.dependency_build_dir }}/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 ${{ steps.setup_env.outputs.dependency_build_dir }}/llvm-build --target check-mlir

    - name: Build MHLO Dialect
      if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
      run: |
        export PATH=${{ steps.setup_env.outputs.dependency_build_dir }}/llvm-build/bin:$PATH
        cmake -S ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/mlir-hlo -B ${{ steps.setup_env.outputs.dependency_build_dir }}/mhlo-build -G Ninja \
              -DCMAKE_BUILD_TYPE=Release \
              -DLLVM_ENABLE_ASSERTIONS=ON \
              -DMLIR_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/llvm-build/lib/cmake/mlir \
              -DLLVM_ENABLE_LLD=OFF \
              -DLLVM_ENABLE_ZLIB=OFF \
              -DLLVM_ENABLE_ZSTD=FORCE_ON \
              -DCMAKE_CXX_VISIBILITY_PRESET=hidden

        cmake --build ${{ steps.setup_env.outputs.dependency_build_dir }}/mhlo-build --target check-mlir-hlo

    - name: Build Enzyme
      if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
      run: |
        cmake -S ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/Enzyme/enzyme -B ${{ steps.setup_env.outputs.dependency_build_dir }}/enzyme-build -G Ninja \
              -DCMAKE_BUILD_TYPE=Release \
              -DLLVM_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/llvm-build/lib/cmake/llvm \
              -DENZYME_STATIC_LIB=ON \
              -DCMAKE_CXX_VISIBILITY_PRESET=hidden

        cmake --build ${{ steps.setup_env.outputs.dependency_build_dir }}/enzyme-build --target EnzymeStatic-18

  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:
      group: 'Office 24th floor M2 Mac'

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

    - name: Setup Runner Environment
      id: setup_env
      uses: ./.github/workflows/utils/setup_self_hosted_macos_env
      with:
        python_version: ${{ matrix.python_version }}

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

    - name: Get Cached LLVM Source
      id: cache-llvm-source
      uses: actions/cache@v3
      with:
        path: ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/llvm-project
        key: Linux-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:  ${{ steps.setup_env.outputs.dependency_build_dir }}/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: ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/mlir-hlo
        key: Linux-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: ${{ steps.setup_env.outputs.dependency_build_dir }}/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: ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/Enzyme
        key: Linux-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: ${{ steps.setup_env.outputs.dependency_build_dir }}/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
    # TODO: Update git tag to v0.35.0 after the next lightning release.
    - 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="2716864521c539429bd382b92151a918d1a076ce" \
              -DKokkos_ENABLE_SERIAL=ON \
              -DKokkos_ENABLE_OPENMP=ON \
              -DKokkos_ENABLE_COMPLEX_ALIGN=OFF \
              -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

    - name: Test Catalyst-Runtime
      env:
        VENV_SITE_PACKAGES: ${{ steps.setup_env.outputs.venv_site_packages }}
      run: |
        export PYTHONPATH="$VENV_SITE_PACKAGES:$PYTHONPATH"
        python${{ matrix.python_version }} -m pip install amazon-braket-pennylane-plugin
        cmake --build runtime-build --target runner_tests_lightning runner_tests_openqasm
        ./runtime-build/tests/runner_tests_lightning
        ./runtime-build/tests/runner_tests_openqasm

    # 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=${{ steps.setup_env.outputs.dependency_build_dir }}/llvm-build/lib/cmake/mlir \
              -DMHLO_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/mhlo-build/lib/cmake/mlir-hlo \
              -DMHLO_BINARY_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/mhlo-build/bin \
              -DEnzyme_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/enzyme-build \
              -DENZYME_SRC_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/Enzyme \
              -DLLVM_ENABLE_ZLIB=OFF \
              -DLLVM_ENABLE_ZSTD=FORCE_ON \
              -DLLVM_ENABLE_LLD=OFF \
              -DLLVM_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/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=${{ steps.setup_env.outputs.dependency_build_dir }}/llvm-build \
        MHLO_BUILD_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/mhlo-build \
        DIALECTS_BUILD_DIR=$GITHUB_WORKSPACE/quantum-build \
        RT_BUILD_DIR=$GITHUB_WORKSPACE/runtime-build \
        ENZYME_BUILD_DIR=${{ steps.setup_env.outputs.dependency_build_dir }}/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@v3
      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 Xanadu M2 Mac
    runs-on:
      group: 'Office 24th floor M2 Mac'

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

    - name: Setup Runner Environment
      id: setup_env
      uses: ./.github/workflows/utils/setup_self_hosted_macos_env
      with:
        python_version: ${{ matrix.python_version }}

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

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

    - 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 -v $GITHUB_WORKSPACE/frontend/test/pytest -n auto
        python${{ matrix.python_version }} -m pytest -v $GITHUB_WORKSPACE/frontend/test/pytest --backend="lightning.kokkos" -n auto
        python${{ matrix.python_version }} -m pytest $GITHUB_WORKSPACE/frontend/test/async_tests
        python${{ matrix.python_version }} -m pytest -v $GITHUB_WORKSPACE/frontend/test/pytest --runbraket=LOCAL -n auto