VJP/JVP support pytree #365
Workflow file for this run
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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: | |
| if: >- | |
| ${{ | |
| github.event_name != 'pull_request' || | |
| github.event.action != 'labeled' || | |
| github.event.label.name == 'ci:build-wheels' | |
| }} | |
| 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 \ | |
| -DOpenMP_ROOT=$(brew --prefix libomp) \ | |
| -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: | | |
| # 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 -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 -v $GITHUB_WORKSPACE/frontend/test/pytest --runbraket=LOCAL -n auto |