From 2d164aaf55d8a683ec6d2502b018557dc213c30b Mon Sep 17 00:00:00 2001 From: Josh Izaac Date: Mon, 29 Jan 2024 14:32:43 -0500 Subject: [PATCH] v0.4.1 bugfix RC (#481) **Context:** v0.4.1 bugfix PR **Description of the Change:** Incorporates the following PRs: #439 #457 #437 #465 #471 --------- Co-authored-by: David Ittah Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com> Co-authored-by: Romain Moyard Co-authored-by: Rashid N H M <95639609+rashidnhm@users.noreply.github.com> Co-authored-by: erick-xanadu <110487834+erick-xanadu@users.noreply.github.com> --- .../workflows/build-wheel-linux-x86_64.yaml | 54 ++++-- .../workflows/build-wheel-macos-arm64.yaml | 177 +++++++++++------- .../workflows/build-wheel-macos-x86_64.yaml | 26 ++- .../setup_self_hosted_macos_env/action.yml | 122 ++++++++++++ Makefile | 4 +- doc/changelog.md | 62 ++++++ doc/dev/installation.rst | 5 +- doc/requirements.txt | 5 + frontend/catalyst/compilation_pipelines.py | 2 +- frontend/catalyst/compiler.py | 16 +- frontend/catalyst/jax_primitives.py | 2 +- frontend/catalyst/jax_tracer.py | 80 ++++++-- frontend/catalyst/pennylane_extensions.py | 8 +- frontend/test/pytest/test_decomposition.py | 28 ++- frontend/test/pytest/test_gradient.py | 36 +++- mlir/CMakeLists.txt | 6 + mlir/Makefile | 28 +-- mlir/python/CMakeLists.txt | 3 + pyproject.toml | 2 +- runtime/Makefile | 12 +- runtime/lib/capi/ExecutionContext.hpp | 16 +- runtime/lib/capi/RuntimeCAPI.cpp | 5 +- setup.py | 25 --- 23 files changed, 548 insertions(+), 176 deletions(-) create mode 100644 .github/workflows/utils/setup_self_hosted_macos_env/action.yml diff --git a/.github/workflows/build-wheel-linux-x86_64.yaml b/.github/workflows/build-wheel-linux-x86_64.yaml index 1c9c236bc0..650e63432d 100644 --- a/.github/workflows/build-wheel-linux-x86_64.yaml +++ b/.github/workflows/build-wheel-linux-x86_64.yaml @@ -119,32 +119,48 @@ jobs: run: | # Reduce wait time for repos not responding cat /etc/yum.conf | sed "s/\[main\]/\[main\]\ntimeout=5/g" > /etc/yum.conf - yum update -y && yum install -y cmake ninja-build + yum update -y && yum install -y cmake ninja-build libzstd-devel - name: Install Dependencies (Python) run: | python${{ matrix.python_version }} -m pip install wheel numpy pybind11 PyYAML + # Required for MHLO and building MLIR with protected symbols. + # (Don't forget to add the build directory to PATH in subsequent steps, so + # other tools can find it, in particular collect2 invoked by gcc.) + - name: Build LLD + 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_TARGETS_TO_BUILD="host" \ + -DLLVM_ENABLE_PROJECTS="lld" + + cmake --build llvm-build --target lld + - name: Build LLVM / MLIR if: steps.cache-llvm-build.outputs.cache-hit != 'true' run: | + export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH 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="lld;mlir" \ + -DLLVM_ENABLE_PROJECTS="mlir" \ -DLLVM_ENABLE_ASSERTIONS=ON \ -DLLVM_INSTALL_UTILS=ON \ -DLLVM_ENABLE_ZLIB=OFF \ - -DLLVM_ENABLE_ZSTD=OFF \ + -DLLVM_ENABLE_ZSTD=FORCE_ON \ -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())") + -DPython3_NumPy_INCLUDE_DIRS=$(python${{ matrix.python_version }} -c "import numpy as np; print(np.get_include())") \ + -DCMAKE_CXX_VISIBILITY_PRESET=protected \ + -DLLVM_ENABLE_LLD=ON # 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 lld check-mlir + LIT_FILTER_OUT="Bytecode" cmake --build llvm-build --target check-mlir - name: Build MHLO Dialect if: steps.cache-mhlo-build.outputs.cache-hit != 'true' @@ -155,19 +171,23 @@ jobs: -DCMAKE_BUILD_TYPE=Release \ -DLLVM_ENABLE_ASSERTIONS=ON \ -DMLIR_DIR=$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir \ - -DLLVM_ENABLE_LLD=ON \ -DLLVM_ENABLE_ZLIB=OFF \ - -DLLVM_ENABLE_ZSTD=OFF + -DLLVM_ENABLE_ZSTD=FORCE_ON \ + -DCMAKE_CXX_VISIBILITY_PRESET=protected \ + -DLLVM_ENABLE_LLD=ON cmake --build mhlo-build --target check-mlir-hlo - name: Build Enzyme if: steps.cache-enzyme-build.outputs.cache-hit != 'true' run: | + export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH 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 + -DENZYME_STATIC_LIB=ON \ + -DCMAKE_CXX_VISIBILITY_PRESET=protected \ + -DLLVM_ENABLE_LLD=ON cmake --build enzyme-build --target EnzymeStatic-18 @@ -193,11 +213,11 @@ jobs: run: | # Reduce wait time for repos not responding cat /etc/yum.conf | sed "s/\[main\]/\[main\]\ntimeout=5/g" > /etc/yum.conf - yum update -y && yum install -y cmake ninja-build openmpi-devel + yum update -y && yum install -y cmake ninja-build openmpi-devel libzstd-devel - name: Install Dependencies (Python) run: | - python${{ matrix.python_version }} -m pip install wheel numpy pybind11 PyYAML scipy + python${{ matrix.python_version }} -m pip install wheel numpy pybind11 PyYAML - name: Get Cached LLVM Source id: cache-llvm-source @@ -263,13 +283,15 @@ jobs: -DKokkos_ENABLE_OPENMP=ON \ -DENABLE_WARNINGS=OFF \ -DENABLE_OPENQASM=ON \ - -DENABLE_OPENMP=ON + -DENABLE_OPENMP=OFF \ + -DLQ_ENABLE_KERNEL_OMP=OFF cmake --build runtime-build --target rt_capi rtd_lightning rtd_openqasm # Build Quantum and Gradient Dialects - name: Build MLIR Dialects run: | + export PATH=$GITHUB_WORKSPACE/llvm-build/bin:$PATH cmake -S mlir -B quantum-build -G Ninja \ -DCMAKE_BUILD_TYPE=Release \ -DLLVM_ENABLE_ASSERTIONS=ON \ @@ -282,7 +304,8 @@ jobs: -DEnzyme_DIR=$GITHUB_WORKSPACE/enzyme-build \ -DENZYME_SRC_DIR=$GITHUB_WORKSPACE/mlir/Enzyme \ -DLLVM_ENABLE_ZLIB=OFF \ - -DLLVM_ENABLE_ZSTD=OFF + -DLLVM_ENABLE_ZSTD=FORCE_ON \ + -DLLVM_ENABLE_LLD=ON cmake --build quantum-build --target check-dialects compiler_driver @@ -296,11 +319,16 @@ jobs: ENZYME_BUILD_DIR=$GITHUB_WORKSPACE/enzyme-build \ make wheel + - name: Repair wheel using auditwheel + run: | + # exclude libopenblas as we rely on the openblas/lapack library shipped by scipy + auditwheel repair dist/*.whl -w ./wheel --no-update-tags --exclude libopenblasp-r0-23e5df77.3.21.dev.so + - name: Upload Wheel Artifact uses: actions/upload-artifact@v3 with: name: catalyst-manylinux2014_x86_64-wheel-py-${{ matrix.python_version }}.zip - path: dist/ + path: wheel/ retention-days: 14 test-wheels: diff --git a/.github/workflows/build-wheel-macos-arm64.yaml b/.github/workflows/build-wheel-macos-arm64.yaml index 0c34a08efd..1550cc7a8b 100644 --- a/.github/workflows/build-wheel-macos-arm64.yaml +++ b/.github/workflows/build-wheel-macos-arm64.yaml @@ -6,7 +6,7 @@ on: workflow_dispatch: env: - MACOSX_DEPLOYMENT_TARGET: 13.0 + MACOSX_DEPLOYMENT_TARGET: 14.0 concurrency: group: ${{ github.workflow }}-${{ github.ref }} @@ -23,23 +23,28 @@ jobs: 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 + 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 }} + # Cache external project sources - name: Cache LLVM Source id: cache-llvm-source uses: actions/cache@v3 with: - path: mlir/llvm-project + path: ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/llvm-project key: Linux-llvm-${{ needs.constants.outputs.llvm_version }}-generic-source enableCrossOsArchive: True @@ -47,7 +52,7 @@ jobs: id: cache-mhlo-source uses: actions/cache@v3 with: - path: mlir/mlir-hlo + path: ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/mlir-hlo key: Linux-mhlo-${{ needs.constants.outputs.mhlo_version }}-generic-source enableCrossOsArchive: True @@ -55,7 +60,7 @@ jobs: id: cache-enzyme-source uses: actions/cache@v3 with: - path: mlir/Enzyme + path: ${{ steps.setup_env.outputs.dependency_build_dir }}/mlir/Enzyme key: Linux-enzyme-${{ needs.constants.outputs.enzyme_version }}-generic-source enableCrossOsArchive: True @@ -83,40 +88,64 @@ jobs: 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: llvm-build + 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: mhlo-build + 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: enzyme-build + 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: Set up Python via brew - run: | - brew install python@${{ matrix.python_version }} - - name: Install Dependencies (Python) - run: | - python${{ matrix.python_version }} -m pip install wheel numpy pybind11 PyYAML cmake ninja + 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 mlir/llvm-project/llvm -B llvm-build -G Ninja \ + 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" \ @@ -124,70 +153,74 @@ jobs: -DLLVM_ENABLE_ASSERTIONS=ON \ -DLLVM_INSTALL_UTILS=ON \ -DLLVM_ENABLE_ZLIB=OFF \ - -DLLVM_ENABLE_ZSTD=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())") + -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 + 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=$GITHUB_WORKSPACE/llvm-build/bin:$PATH - cmake -S mlir/mlir-hlo -B mhlo-build -G Ninja \ + 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=$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir \ + -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=OFF + -DLLVM_ENABLE_ZSTD=FORCE_ON \ + -DCMAKE_CXX_VISIBILITY_PRESET=hidden - cmake --build mhlo-build --target check-mlir-hlo + 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 mlir/Enzyme/enzyme -B enzyme-build -G Ninja \ + 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=$GITHUB_WORKSPACE/llvm-build/lib/cmake/llvm \ - -DENZYME_STATIC_LIB=ON + -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 enzyme-build --target EnzymeStatic-18 + 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: - 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 }} + runs-on: + group: 'Office 24th floor M2 Mac' steps: - name: Checkout Catalyst repo uses: actions/checkout@v3 - - name: Set up Python via brew - run: | - brew install python@${{ matrix.python_version }} + - 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 scipy + 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: mlir/llvm-project + 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 @@ -196,7 +229,7 @@ jobs: id: cache-llvm-build uses: actions/cache@v3 with: - path: llvm-build + 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 @@ -204,7 +237,7 @@ jobs: id: cache-mhlo-source uses: actions/cache@v3 with: - path: mlir/mlir-hlo + 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 @@ -213,7 +246,7 @@ jobs: id: cache-mhlo-build uses: actions/cache@v3 with: - path: mhlo-build + 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 @@ -221,7 +254,7 @@ jobs: id: cache-enzyme-source uses: actions/cache@v3 with: - path: mlir/Enzyme + 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 @@ -229,7 +262,7 @@ jobs: id: cache-enzyme-build uses: actions/cache@v3 with: - path: enzyme-build + 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 @@ -244,16 +277,21 @@ jobs: -DENABLE_LIGHTNING_KOKKOS=ON \ -DLIGHTNING_GIT_TAG="v0.34.0" \ -DKokkos_ENABLE_SERIAL=ON \ - -DKokkos_ENABLE_OPENMP=OFF \ + -DKokkos_ENABLE_OPENMP=ON \ -DKokkos_ENABLE_COMPLEX_ALIGN=OFF \ + -DOpenMP_ROOT=$(brew --prefix libomp) \ -DENABLE_WARNINGS=OFF \ -DENABLE_OPENQASM=ON \ - -DENABLE_OPENMP=OFF + -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 @@ -268,32 +306,38 @@ jobs: -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 \ + -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=OFF \ - -DLLVM_ENABLE_LLD=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=$GITHUB_WORKSPACE/llvm-build \ - MHLO_BUILD_DIR=$GITHUB_WORKSPACE/mhlo-build \ + 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=$GITHUB_WORKSPACE/enzyme-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: dist/ + path: wheel/ retention-days: 14 test-wheels: @@ -301,28 +345,29 @@ jobs: 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 }} + 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: Set up Python via brew - run: | - brew install python@${{ matrix.python_version }} - - name: Install Python dependencies run: | # tensorflow-cpu is not distributed for macOS ARM @@ -340,6 +385,6 @@ jobs: - 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 + 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 diff --git a/.github/workflows/build-wheel-macos-x86_64.yaml b/.github/workflows/build-wheel-macos-x86_64.yaml index 7d29eaa1d8..dad7e3c586 100644 --- a/.github/workflows/build-wheel-macos-x86_64.yaml +++ b/.github/workflows/build-wheel-macos-x86_64.yaml @@ -125,11 +125,12 @@ jobs: -DLLVM_ENABLE_ASSERTIONS=ON \ -DLLVM_INSTALL_UTILS=ON \ -DLLVM_ENABLE_ZLIB=OFF \ - -DLLVM_ENABLE_ZSTD=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())") + -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 @@ -146,7 +147,8 @@ jobs: -DMLIR_DIR=$GITHUB_WORKSPACE/llvm-build/lib/cmake/mlir \ -DLLVM_ENABLE_LLD=OFF \ -DLLVM_ENABLE_ZLIB=OFF \ - -DLLVM_ENABLE_ZSTD=OFF + -DLLVM_ENABLE_ZSTD=FORCE_ON \ + -DCMAKE_CXX_VISIBILITY_PRESET=hidden cmake --build mhlo-build --target check-mlir-hlo @@ -156,7 +158,8 @@ jobs: 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 + -DENZYME_STATIC_LIB=ON \ + -DCMAKE_CXX_VISIBILITY_PRESET=hidden cmake --build enzyme-build --target EnzymeStatic-18 @@ -183,7 +186,7 @@ jobs: - name: Install Dependencies (Python) run: | - python${{ matrix.python_version }} -m pip install wheel numpy pybind11 PyYAML cmake ninja scipy + python${{ matrix.python_version }} -m pip install wheel numpy pybind11 PyYAML cmake ninja delocate - name: Get Cached LLVM Source id: cache-llvm-source @@ -238,6 +241,7 @@ jobs: # Build Catalyst-Runtime - name: Build Catalyst-Runtime run: | + # Segfaults in computing Lightning's adjoint-jacobian when building with OMP cmake -S runtime -B runtime-build -G Ninja \ -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$GITHUB_WORKSPACE/runtime-build/lib \ @@ -251,7 +255,8 @@ jobs: -DKokkos_ENABLE_DEPRECATION_WARNINGS=OFF \ -DENABLE_WARNINGS=OFF \ -DENABLE_OPENQASM=ON \ - -DENABLE_OPENMP=OFF + -DENABLE_OPENMP=OFF \ + -DLQ_ENABLE_KERNEL_OMP=OFF cmake --build runtime-build --target rt_capi rtd_lightning rtd_openqasm @@ -277,7 +282,7 @@ jobs: -DEnzyme_DIR=$GITHUB_WORKSPACE/enzyme-build \ -DENZYME_SRC_DIR=$GITHUB_WORKSPACE/mlir/Enzyme \ -DLLVM_ENABLE_ZLIB=OFF \ - -DLLVM_ENABLE_ZSTD=OFF \ + -DLLVM_ENABLE_ZSTD=FORCE_ON \ -DLLVM_ENABLE_LLD=OFF cmake --build quantum-build --target check-dialects compiler_driver @@ -292,11 +297,16 @@ jobs: 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=x86_64 -w ./wheel -v dist/*.whl --ignore-missing-dependencies -vv + - name: Upload Wheel Artifact uses: actions/upload-artifact@v3 with: name: catalyst-macos_x86_64-wheel-py-${{ matrix.python_version }}.zip - path: dist/ + path: wheel/ retention-days: 14 test-wheels: diff --git a/.github/workflows/utils/setup_self_hosted_macos_env/action.yml b/.github/workflows/utils/setup_self_hosted_macos_env/action.yml new file mode 100644 index 0000000000..786624b88a --- /dev/null +++ b/.github/workflows/utils/setup_self_hosted_macos_env/action.yml @@ -0,0 +1,122 @@ +name: Setup PATH for Self-Hosted MacOS Runner +description: | + This workflow sets up the workflow environment for the Self-Hosted MacOS runner. + It should be called right after actions/checkout. + + +inputs: + python_version: + description: Version of Python being used by the Job + required: true + +outputs: + python_version: + description: 'The version of Python that was setup' + value: ${{ steps.python_venv.outputs.python_version }} + python_path: + description: 'Full absolute path to the python binary' + value: ${{ steps.python_venv.outputs.python_path }} + venv_location: + description: 'The path to the virtualenv created for python' + value: ${{ steps.python_venv.outputs.venv_location }} + venv_site_packages: + description: 'The path to the site-packages folder for the python venv created' + value: ${{ steps.site_packages.outputs.site_packages_dir }} + dependency_build_dir: + description: 'The directory where the LLVM/MLIR/MHLO sources should reside in' + value: ${{ steps.build_dir.outputs.dir_name }} + + +runs: + using: composite + steps: + + # These environment variables ensure multi-threading workloads (ex: pytest-xdist) + # work without any issues. Without them, the tests take hours to complete. + - name: Setup MacOS Env Vars + shell: bash + run: | + # The following two environment variables are set to ensure no threading related commands are blocked/slowed + # by the OS. They may or may not be needed but added to ensure there are no slowdowns. + echo 'no_proxy=*' >> $GITHUB_ENV + echo 'OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES' >> $GITHUB_ENV + + # The following two flags are exported to ensure the correct homebrew installed libomp libraries are used + # during build time. + echo "LDFLAGS=-L/opt/homebrew/opt/libomp/lib $LDFLAGS" >> $GITHUB_ENV + echo "CPPFLAGS=-I/opt/homebrew/opt/libomp/include $CPPFLAGS" >> $GITHUB_ENV + + + # These binaries are added to PATH as there were issues adding them by default on job start + - name: Add additional binaries to PATH + shell: bash + run: | + additional_paths=("/opt/homebrew/bin" "$HOME/.pyenv/shims" "$HOME/.cargo/bin") + for _path in ${additional_paths[@]}; + do + [[ ":$PATH:" != *":$_path:"* ]] && echo "$_path" >> $GITHUB_PATH + done + echo "PYENV_ROOT=$HOME/.pyenv" >> $GITHUB_ENV + + - name: Setup Python ${{ inputs.python_version }} + id: python_setup + shell: bash + env: + input_python_version: ${{ inputs.python_version }} + run: | + # Check if the requested Python version exists on the runner + pyenv versions | grep "$input_python_version" + + if [ $? -ne 0 ]; then + echo "Installing Python $input_python_version" + pyenv install "$input_python_version" + fi + + - name: Setup Python venv + id: python_venv + shell: bash + env: + input_python_version: ${{ inputs.python_version }} + venv_name: py_venv_${{ inputs.python_version }}_${{ github.sha }} + run: | + rm -rf "$venv_name" + + echo "$input_python_version" > "$GITHUB_WORKSPACE/.python-version" + + python -m venv "$venv_name" + + echo "$GITHUB_WORKSPACE/$venv_name/bin" >> $GITHUB_PATH + echo "venv_location=$GITHUB_WORKSPACE/$venv_name" >> $GITHUB_OUTPUT + + PY_PATH="$GITHUB_WORKSPACE/$venv_name/bin/python" + PY_VERSION=$(python --version) + echo "python_path=$PY_PATH" >> $GITHUB_OUTPUT + echo "python_version=$PY_VERSION" >> $GITHUB_OUTPUT + + - name: Get site-packages path + id: site_packages + shell: bash + run: | + echo "site_packages_dir=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')" >> $GITHUB_OUTPUT + + - name: Display Python Setup + shell: bash + run: | + set -x + which python + which pip + python --version + pip --version + set +x + + - name: Setup Build Dir + id: build_dir + shell: bash + env: + BUILD_DIR: /tmp/llvm_cache_${{ inputs.python_version }} + run: | + rm -rf "$BUILD_DIR"* + mkdir "$BUILD_DIR" + + echo "dir_name=${BUILD_DIR}" >> $GITHUB_OUTPUT + diff --git a/Makefile b/Makefile index 46f7adf042..dbf23d2e48 100644 --- a/Makefile +++ b/Makefile @@ -222,11 +222,11 @@ endif $(MAKE) -C mlir format $(MAKE) -C runtime format ifdef check - python3 ./bin/format.py --check $(if $(version:-=),--cfversion $(version)) ./frontend/catalyst/utils + $(PYTHON) ./bin/format.py --check $(if $(version:-=),--cfversion $(version)) ./frontend/catalyst/utils black --check --verbose . isort --check --diff . else - python3 ./bin/format.py $(if $(version:-=),--cfversion $(version)) ./frontend/catalyst/utils + $(PYTHON) ./bin/format.py $(if $(version:-=),--cfversion $(version)) ./frontend/catalyst/utils black . isort . endif diff --git a/doc/changelog.md b/doc/changelog.md index a09290de36..96fa68fea0 100644 --- a/doc/changelog.md +++ b/doc/changelog.md @@ -1,3 +1,65 @@ +# Release 0.4.1 + +

Improvements

+ +* Catalyst wheels are now packaged with OpenMP and ZStd, which avoids installing additional + requirements separately in order to use pre-packaged Catalyst binaries. + [(#457)](https://github.com/PennyLaneAI/catalyst/pull/457) + [(#478)](https://github.com/PennyLaneAI/catalyst/pull/478) + + Note that OpenMP support for the `lightning.kokkos` backend has been disabled on macOS x86_64, due + to memory issues in the computation of Lightning's adjoint-jacobian in the presence of multiple + OMP threads. + +

Bug fixes

+ +* Resolve an infinite recursion in the decomposition of the `Controlled` + operator whenever computing a Unitary matrix for the operator fails. + [(#468)](https://github.com/PennyLaneAI/catalyst/pull/468) + +* Resolve a failure to generate gradient code for specific input circuits. + [(#439)](https://github.com/PennyLaneAI/catalyst/pull/439) + + In this case, [`jnp.mod`](https://github.com/PennyLaneAI/catalyst/issues/437) + was used to compute wire values in a for loop, which prevented the gradient + architecture from fully separating quantum and classical code. The following + program is now supported: + ```py + @qjit + @grad + @qml.qnode(dev) + def f(x): + def cnot_loop(j): + qml.CNOT(wires=[j, jnp.mod((j + 1), 4)]) + + for_loop(0, 4, 1)(cnot_loop)() + + return qml.expval(qml.PauliZ(0)) + ``` + +* Resolve unpredictable behaviour when importing libraries that share Catalyst's LLVM dependency + (e.g. TensorFlow). In some cases, both packages exporting the same symbols from their shared + libraries can lead to process crashes and other unpredictable behaviour, since the wrong functions + can be called if both libraries are loaded in the current process. + The fix involves building shared libraries with hidden (macOS) or protected (linux) symbol + visibility by default, exporting only what is necessary. + [(#465)](https://github.com/PennyLaneAI/catalyst/pull/465) + +* Resolve a failure to find the SciPy OpenBLAS library when running Catalyst, + due to a different SciPy version being used to build Catalyst than to run it. + [(#471)](https://github.com/PennyLaneAI/catalyst/pull/471) + +* Resolve a memory leak in the runtime stemming from missing calls to device destructors + at the end of programs. + [(#446)](https://github.com/PennyLaneAI/catalyst/pull/446) + +

Contributors

+ +This release contains contributions from (in alphabetical order): + +Ali Asadi, +David Ittah. + # Release 0.4.0

New features

diff --git a/doc/dev/installation.rst b/doc/dev/installation.rst index c57cb0f988..9cdceb4a20 100644 --- a/doc/dev/installation.rst +++ b/doc/dev/installation.rst @@ -79,7 +79,7 @@ installed and available on the path (depending on the platform): - The `clang `_ compiler, `LLD `_ linker (Linux only), `CCache `_ compiler cache (optional, recommended), and - `OpenMP `_ (Linux only). + `OpenMP `_. - The `Ninja `_, `Make `_, and `CMake `_ (v3.20 or greater) build tools. @@ -100,11 +100,12 @@ They can be installed on **Debian/Ubuntu** via: versions of it via ``pip install cmake``. On **macOS**, it is strongly recommended to install the official XCode Command Line Tools -(for ``clang`` & ``make``). The remaining packages can then be installed via ``pip``: +(for ``clang`` & ``make``). The remaining packages can then be installed via ``pip`` and ``brew``: .. code-block:: console pip install cmake ninja + brew install libomp If you install Catalyst on a macOS system with ``ARM`` architecture (e.g. Apple M1/M2), you additionally need to install `Rust `_ and the diff --git a/doc/requirements.txt b/doc/requirements.txt index 693b8a711b..9ebd25e7ec 100644 --- a/doc/requirements.txt +++ b/doc/requirements.txt @@ -1,6 +1,11 @@ docutils==0.16 sphinxcontrib-bibtex sphinx==4.5.0 +sphinxcontrib-applehelp==1.0.4 +sphinxcontrib-devhelp==1.0.2 +sphinxcontrib-qthelp==1.0.3 +sphinxcontrib-htmlhelp==2.0.1 +sphinxcontrib-serializinghtml==1.1.5 ipykernel nbsphinx==0.7 Jinja2==3.1.2 diff --git a/frontend/catalyst/compilation_pipelines.py b/frontend/catalyst/compilation_pipelines.py index 8e50d2e5b6..75fc62e28f 100644 --- a/frontend/catalyst/compilation_pipelines.py +++ b/frontend/catalyst/compilation_pipelines.py @@ -582,7 +582,7 @@ def get_mlir(self, *args): inject_functions(mlir_module, ctx) self._jaxpr = jaxpr canonicalizer_options = deepcopy(self.compile_options) - canonicalizer_options.pipelines = [("pipeline", ["canonicalize"])] + canonicalizer_options.pipelines = [("0_canonicalize", ["canonicalize"])] canonicalizer_options.lower_to_llvm = False canonicalizer = Compiler(canonicalizer_options) _, self._mlir, _ = canonicalizer.run(mlir_module, self.workspace) diff --git a/frontend/catalyst/compiler.py b/frontend/catalyst/compiler.py index 0ed172aff1..45ce7ff7fb 100644 --- a/frontend/catalyst/compiler.py +++ b/frontend/catalyst/compiler.py @@ -14,6 +14,7 @@ """This module contains functions for lowering, compiling, and linking MLIR/LLVM representations. """ +import glob import importlib import os import pathlib @@ -120,6 +121,7 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt "canonicalize", "scatter-lowering", "hlo-custom-call-lowering", + "cse", ], ) @@ -267,20 +269,29 @@ def get_default_flags(): f"-L{DEFAULT_CUSTOM_CALLS_LIB_PATH}", ] - # Discover the LAPACK library provided by scipy & add it to the rpath. - package_name = "scipy" + # Discover the LAPACK library provided by scipy & add link against it. + # Doing this here ensures we will always have the correct library name. if platform.system() == "Linux": file_path_within_package = "../scipy.libs/" + file_extension = ".so" elif platform.system() == "Darwin": # pragma: nocover file_path_within_package = ".dylibs/" + file_extension = ".dylib" + package_name = "scipy" scipy_package = importlib.util.find_spec(package_name) package_directory = path.dirname(scipy_package.origin) scipy_lib_path = path.join(package_directory, file_path_within_package) + file_prefix = "libopenblas" + search_pattern = path.join(scipy_lib_path, f"{file_prefix}*{file_extension}") + openblas_so_file = glob.glob(search_pattern)[0] + openblas_lib_name = path.basename(openblas_so_file)[3 : -len(file_extension)] + lib_path_flags += [ f"-Wl,-rpath,{scipy_lib_path}", + f"-L{scipy_lib_path}", ] system_flags = [] @@ -297,6 +308,7 @@ def get_default_flags(): "-lrt_capi", "-lpthread", "-lmlir_c_runner_utils", # required for memref.copy + f"-l{openblas_lib_name}", # required for custom_calls lib "-lcustom_calls", "-lmlir_async_runtime", ] diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index c8ef501752..6d973126e8 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -1538,7 +1538,7 @@ def _scalar_abstractify(t): # pylint: disable=protected-access if t in {int, float, complex, bool} or isinstance(t, jax._src.numpy.lax_numpy._ScalarMeta): return core.ShapedArray([], dtype=t, weak_type=True) - raise TypeError(f"Cannot convert given type {t} scalar ShapedArray.") + raise TypeError(f"Cannot convert given type {t} to scalar ShapedArray.") # pylint: disable=protected-access diff --git a/frontend/catalyst/jax_tracer.py b/frontend/catalyst/jax_tracer.py index 43925fa07f..1ffe0a4e63 100644 --- a/frontend/catalyst/jax_tracer.py +++ b/frontend/catalyst/jax_tracer.py @@ -30,6 +30,7 @@ from catalyst.jax_primitives import ( AbstractQreg, compbasis_p, + cond_p, counts_p, expval_p, func_p, @@ -44,6 +45,7 @@ qextract_p, qinsert_p, qinst_p, + qmeasure_p, qunitary_p, sample_p, state_p, @@ -705,6 +707,31 @@ def trace_post_processing(ctx, trace, post_processing: Callable, pp_args): return closed_jaxpr, out_type, out_tree_promise() +def reset_qubit(qreg_in, w): + """Perform a qubit reset on a single wire. Suitable for use during late-stage tracing, + as JAX primitives are used directly. These operations will not appear on tape.""" + + def flip(qreg): + """Flip a qubit.""" + qbit = qextract_p.bind(qreg, w) + qbit2 = qinst_p.bind(qbit, op="PauliX", qubits_len=1)[0] + return qinsert_p.bind(qreg, w, qbit2) + + def dont_flip(qreg): + """Identity function.""" + return qreg + + qbit = qextract_p.bind(qreg_in, w) + m, qbit2 = qmeasure_p.bind(qbit) + qreg_mid = qinsert_p.bind(qreg_in, w, qbit2) + + jaxpr_true = jax.make_jaxpr(flip)(qreg_mid) + jaxpr_false = jax.make_jaxpr(dont_flip)(qreg_mid) + qreg_out = cond_p.bind(m, qreg_mid, branch_jaxprs=[jaxpr_true, jaxpr_false])[0] + + return qreg_out + + def trace_quantum_function( f: Callable, device: QubitDevice, args, kwargs, qnode=None ) -> Tuple[ClosedJaxpr, Any]: @@ -765,28 +792,32 @@ def is_leaf(obj): # (2) - Quantum tracing transformed_results = [] is_program_transformed = qnode and qnode.transform_program - for tape in tapes: - # If the program is batched, that means that it was transformed. - # If it was transformed, that means that the program might have - # changed the output. See `split_non_commuting` - if is_program_transformed: - # TODO: In the future support arbitrary output from the user function. - output = tape.measurements - _, trees = jax.tree_util.tree_flatten(output, is_leaf=is_leaf) - else: - output = return_values_flat - trees = return_values_tree - - with EvaluationContext.frame_tracing_context(ctx, trace): - qdevice_p.bind( - rtd_lib=device.backend_lib, - rtd_name=device.backend_name, - rtd_kwargs=str(device.backend_kwargs), - ) - qreg_in = qalloc_p.bind(len(device.wires)) + + with EvaluationContext.frame_tracing_context(ctx, trace): + # Set up same device and quantum register for all tapes in the program. + # We just need to ensure the qubits are reset in between each. + qdevice_p.bind( + rtd_lib=device.backend_lib, + rtd_name=device.backend_name, + rtd_kwargs=str(device.backend_kwargs), + ) + qreg_in = qalloc_p.bind(len(device.wires)) + + for i, tape in enumerate(tapes): + # If the program is batched, that means that it was transformed. + # If it was transformed, that means that the program might have + # changed the output. See `split_non_commuting` + if is_program_transformed: + # TODO: In the future support arbitrary output from the user function. + output = tape.measurements + _, trees = jax.tree_util.tree_flatten(output, is_leaf=is_leaf) + else: + output = return_values_flat + trees = return_values_tree + qrp_out = trace_quantum_tape(tape, device, qreg_in, ctx, trace) meas, meas_trees = trace_quantum_measurements(device, qrp_out, output, trees) - qdealloc_p.bind(qrp_out.actualize()) + qreg_out = qrp_out.actualize() meas_tracers = [trace.full_raise(m) for m in meas] meas_results = tree_unflatten(meas_trees, meas_tracers) @@ -801,6 +832,15 @@ def is_leaf(obj): else: transformed_results.append(meas_results) + # Reset the qubits and update the register value for the next tape. + if len(tapes) > 1 and i < len(tapes) - 1: + for w in device.wires: + qreg_out = reset_qubit(qreg_out, w) + qreg_in = qreg_out + + # Deallocate the register before tracing the post-processing. + qdealloc_p.bind(qreg_out) + closed_jaxpr, out_type, out_tree = trace_post_processing( ctx, trace, post_processing, transformed_results ) diff --git a/frontend/catalyst/pennylane_extensions.py b/frontend/catalyst/pennylane_extensions.py index 24b877726f..6bc9275b75 100644 --- a/frontend/catalyst/pennylane_extensions.py +++ b/frontend/catalyst/pennylane_extensions.py @@ -323,7 +323,13 @@ def default_expand_fn(self, circuit, max_expansion=10): decompose_to_qubit_unitary = QJITDevice._get_operations_to_convert_to_matrix(self.config) def _decomp_to_unitary(self, *_args, **_kwargs): - return [qml.QubitUnitary(qml.matrix(self), wires=self.wires)] + try: + mat = self.matrix() + except Exception as e: + raise CompileError( + f"Operation {self} could not be decomposed, it might be unsupported." + ) from e + return [qml.QubitUnitary(mat, wires=self.wires)] # Fallback for controlled gates that won't decompose successfully. # Doing so before rather than after decomposition is generally a trade-off. For low diff --git a/frontend/test/pytest/test_decomposition.py b/frontend/test/pytest/test_decomposition.py index 8395e7519f..4aef376bd2 100644 --- a/frontend/test/pytest/test_decomposition.py +++ b/frontend/test/pytest/test_decomposition.py @@ -16,7 +16,7 @@ import pytest from jax import numpy as jnp -from catalyst import for_loop, measure, qjit +from catalyst import CompileError, ctrl, measure, qjit from catalyst.compiler import get_lib_path # This is used just for internal testing @@ -77,3 +77,29 @@ def mid_circuit(x: float): return measure(wires=1) assert mid_circuit(param) == expected + + +class TestControlledDecomposition: + """Test behaviour around the decomposition of the `Controlled` class.""" + + def test_no_matrix(self, backend): + """Test that controlling an operation without a matrix method raises an error.""" + dev = qml.device(backend, wires=4) + + class OpWithNoMatrix(qml.operation.Operation): + num_wires = qml.operation.AnyWires + + def matrix(self): + raise NotImplementedError() + + @qml.qnode(dev) + def f(): + ctrl(OpWithNoMatrix(wires=[0, 1]), control=[2, 3]) + return qml.probs() + + with pytest.raises(CompileError, match="could not be decomposed, it might be unsupported"): + qjit(f, target="jaxpr") + + +if __name__ == "__main__": + pytest.main(["-x", __file__]) diff --git a/frontend/test/pytest/test_gradient.py b/frontend/test/pytest/test_gradient.py index c344f705ad..266f436ca7 100644 --- a/frontend/test/pytest/test_gradient.py +++ b/frontend/test/pytest/test_gradient.py @@ -21,9 +21,11 @@ from jax import numpy as jnp import catalyst.utils.calculate_grad_shape as infer -from catalyst import CompileError, cond, for_loop, grad, jacobian, qjit +from catalyst import cond, for_loop, grad, jacobian, qjit from catalyst.pennylane_extensions import DifferentiableCompileError +# pylint: disable=too-many-lines + class TestGradShape: """Unit tests for the calculate_grad_shape module.""" @@ -965,5 +967,37 @@ def compiled(x: float, y: float): assert actual_entry == pytest.approx(expected_entry) +def test_loop_with_dyn_wires(backend): + """Test the gradient on a function with a loop and modular wire arithmetic.""" + num_wires = 4 + dev = qml.device(backend, wires=num_wires) + + @qml.qnode(dev) + def cat(phi): + @for_loop(0, 3, 1) + def loop(i): + qml.RY(phi, wires=jnp.mod(i, num_wires)) + + loop() + + return qml.expval(qml.prod(*[qml.PauliZ(i) for i in range(num_wires)])) + + @qml.qnode(dev) + def pl(phi): + @for_loop(0, 3, 1) + def loop(i): + qml.RY(phi, wires=i % num_wires) + + loop() + + return qml.expval(qml.prod(*[qml.PauliZ(i) for i in range(num_wires)])) + + arg = 0.75 + result = qjit(grad(cat))(arg) + expected = qml.grad(pl, argnum=0)(arg) + + assert np.allclose(result, expected) + + if __name__ == "__main__": pytest.main(["-x", __file__]) diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt index e1368cfbad..124cd18c95 100644 --- a/mlir/CMakeLists.txt +++ b/mlir/CMakeLists.txt @@ -78,6 +78,12 @@ include_directories(${LLVM_INCLUDE_DIRS} link_directories(${LLVM_BUILD_LIBRARY_DIR}) add_definitions(${LLVM_DEFINITIONS}) +if(APPLE) + set(CMAKE_CXX_VISIBILITY_PRESET hidden) +else() + set(CMAKE_CXX_VISIBILITY_PRESET protected) +endif() + add_subdirectory(include) add_subdirectory(lib) add_subdirectory(tools) diff --git a/mlir/Makefile b/mlir/Makefile index 03a0a07140..b1ef841b9d 100644 --- a/mlir/Makefile +++ b/mlir/Makefile @@ -11,20 +11,19 @@ MHLO_BUILD_DIR?=$(MK_DIR)/mlir-hlo/bazel-build ENZYME_BUILD_DIR?=$(MK_DIR)/Enzyme/build RT_BUILD_DIR?=$(MK_DIR)/../runtime/build ENABLE_ASAN?=OFF +BUILD_TYPE?=Release ifeq ($(shell uname), Darwin) DEFAULT_ENABLE_LLD := OFF -DEFAULT_ENABLE_ZLIB := OFF -DEFAULT_ENABLE_ZSTD := OFF +SYMBOL_VISIBILITY := hidden else DEFAULT_ENABLE_LLD := ON -DEFAULT_ENABLE_ZLIB := ON -DEFAULT_ENABLE_ZSTD := ON +SYMBOL_VISIBILITY := protected endif ENABLE_LLD?=$(DEFAULT_ENABLE_LLD) -ENABLE_ZLIB?=$(DEFAULT_ENABLE_LLD) -ENABLE_ZSTD?=$(DEFAULT_ENABLE_LLD) +ENABLE_ZLIB?=OFF +ENABLE_ZSTD?=ON ifeq ($(ENABLE_ASAN), ON) USE_SANITIZER_NAMES="Address" @@ -54,7 +53,7 @@ all: llvm mhlo enzyme dialects llvm: @echo "build LLVM and MLIR enabling Python bindings" cmake -G Ninja -S llvm-project/llvm -B $(LLVM_BUILD_DIR) \ - -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \ -DLLVM_BUILD_EXAMPLES=OFF \ -DLLVM_TARGETS_TO_BUILD="host" \ -DLLVM_ENABLE_PROJECTS="mlir" \ @@ -69,7 +68,8 @@ llvm: -DLLVM_USE_SANITIZER=$(USE_SANITIZER_NAMES) \ -DLLVM_ENABLE_LLD=$(ENABLE_LLD) \ -DLLVM_ENABLE_ZLIB=$(ENABLE_ZLIB) \ - -DLLVM_ENABLE_ZSTD=$(ENABLE_ZSTD) + -DLLVM_ENABLE_ZSTD=$(ENABLE_ZSTD) \ + -DCMAKE_CXX_VISIBILITY_PRESET=$(SYMBOL_VISIBILITY) # TODO: when updating LLVM, test to see if mlir/unittests/Bytecode/BytecodeTest.cpp:55 is passing # and remove filter @@ -79,7 +79,7 @@ llvm: mhlo: @echo "build MLIR-HLO" cmake -G Ninja -S mlir-hlo -B $(MHLO_BUILD_DIR) \ - -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \ -DLLVM_ENABLE_ASSERTIONS=ON \ -DMLIR_DIR=$(LLVM_BUILD_DIR)/lib/cmake/mlir \ -DCMAKE_C_COMPILER=$(C_COMPILER) \ @@ -89,7 +89,8 @@ mhlo: -DCMAKE_EXE_LINKER_FLAGS=$(USE_SANITIZER_FLAGS) \ -DLLVM_ENABLE_LLD=$(ENABLE_LLD) \ -DLLVM_ENABLE_ZLIB=$(ENABLE_ZLIB) \ - -DLLVM_ENABLE_ZSTD=$(ENABLE_ZSTD) + -DLLVM_ENABLE_ZSTD=$(ENABLE_ZSTD) \ + -DCMAKE_CXX_VISIBILITY_PRESET=$(SYMBOL_VISIBILITY) cmake --build $(MHLO_BUILD_DIR) --target check-mlir-hlo @@ -98,13 +99,14 @@ enzyme: @echo "build enzyme" cmake -G Ninja -S Enzyme/enzyme -B $(ENZYME_BUILD_DIR) \ -DENZYME_STATIC_LIB=ON \ - -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \ -DLLVM_DIR=$(LLVM_BUILD_DIR)/lib/cmake/llvm \ -DCMAKE_C_COMPILER=$(C_COMPILER) \ -DCMAKE_CXX_COMPILER=$(CXX_COMPILER) \ -DCMAKE_C_COMPILER_LAUNCHER=$(COMPILER_LAUNCHER) \ -DCMAKE_CXX_COMPILER_LAUNCHER=$(COMPILER_LAUNCHER) \ - -DCMAKE_EXE_LINKER_FLAGS=$(USE_SANITIZER_FLAGS) + -DCMAKE_EXE_LINKER_FLAGS=$(USE_SANITIZER_FLAGS) \ + -DCMAKE_CXX_VISIBILITY_PRESET=$(SYMBOL_VISIBILITY) cmake --build $(ENZYME_BUILD_DIR) --target EnzymeStatic-18 @@ -113,7 +115,7 @@ dialects: @echo "build quantum-lsp compiler_driver and dialects" cmake -G Ninja -S . -B $(DIALECTS_BUILD_DIR) \ - -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \ -DLLVM_ENABLE_ASSERTIONS=ON \ -DQUANTUM_ENABLE_BINDINGS_PYTHON=ON \ -DPython3_EXECUTABLE=$$(which $(PYTHON)) \ diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index e6af50a1cf..0cef7c6358 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -47,6 +47,9 @@ declare_mlir_python_extension(QuantumPythonSources.Extension ADD_TO_PARENT QuantumPythonSources SOURCES QuantumExtension.cpp + PRIVATE_LINK_LIBS + LLVMSupport + MLIRIR EMBED_CAPI_LINK_LIBS QuantumCAPI ) diff --git a/pyproject.toml b/pyproject.toml index 6b30f2990a..20bc95cab6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ extend_skip_glob = [ ] [build-system] -requires = ["setuptools>=62", "wheel", "pybind11>=2.7.0", "numpy>=1.22", "scipy"] +requires = ["setuptools>=62", "wheel", "pybind11>=2.7.0", "numpy>=1.22"] build-backend = "setuptools.build_meta" [tool.pytest.ini_options] diff --git a/runtime/Makefile b/runtime/Makefile index 4c52dea67d..dfa6f0c254 100644 --- a/runtime/Makefile +++ b/runtime/Makefile @@ -17,12 +17,6 @@ ENABLE_ASAN?=OFF LIGHTNING_GIT_TAG_VALUE?="v0.34.0" NPROC?=$(shell python3 -c "import os; print(os.cpu_count())") -ifeq ($(shell uname), Darwin) - DEFAULT_ENABLE_OPENMP := OFF -else - DEFAULT_ENABLE_OPENMP := ON -endif - BUILD_TARGETS := rt_capi TEST_TARGETS := "" @@ -39,7 +33,8 @@ ifeq ($(ENABLE_OPENQASM), ON) TEST_TARGETS += runner_tests_openqasm endif -ENABLE_OPENMP?=$(DEFAULT_ENABLE_OPENMP) +LIGHTNING_ENABLE_OPENMP?=OFF +KOKKOS_ENABLE_OPENMP?=ON coverage: CODE_COVERAGE=ON coverage: BUILD_TYPE=Debug @@ -76,7 +71,8 @@ configure: -DENABLE_LIGHTNING=$(ENABLE_LIGHTNING) \ -DENABLE_LIGHTNING_KOKKOS=$(ENABLE_LIGHTNING_KOKKOS) \ -DENABLE_OPENQASM=$(ENABLE_OPENQASM) \ - -DENABLE_OPENMP=$(ENABLE_OPENMP) \ + -DENABLE_OPENMP=$(LIGHTNING_ENABLE_OPENMP) \ + -DKokkos_ENABLE_OPENMP=$(KOKKOS_ENABLE_OPENMP) \ -DENABLE_CODE_COVERAGE=$(CODE_COVERAGE) \ -DRUNTIME_ENABLE_WARNINGS=$(ENABLE_WARNINGS) \ -DENABLE_WARNINGS=OFF \ diff --git a/runtime/lib/capi/ExecutionContext.hpp b/runtime/lib/capi/ExecutionContext.hpp index c6281620b6..2bd8067cb6 100644 --- a/runtime/lib/capi/ExecutionContext.hpp +++ b/runtime/lib/capi/ExecutionContext.hpp @@ -208,7 +208,7 @@ class RTDevice { std::string rtd_kwargs; std::unique_ptr rtd_dylib{nullptr}; - QuantumDevice *rtd_qdevice{nullptr}; + std::unique_ptr rtd_qdevice{nullptr}; RTDeviceStatus status{RTDeviceStatus::Inactive}; @@ -255,7 +255,7 @@ class RTDevice { _pl2runtime_device_info(rtd_lib, rtd_name); } - ~RTDevice() { rtd_dylib.reset(nullptr); } + ~RTDevice() = default; auto operator==(const RTDevice &other) const -> bool { @@ -263,7 +263,7 @@ class RTDevice { this->rtd_kwargs == other.rtd_kwargs; } - [[nodiscard]] auto getQuantumDevicePtr() -> QuantumDevice * + [[nodiscard]] auto getQuantumDevicePtr() -> const std::unique_ptr & { if (rtd_qdevice) { return rtd_qdevice; @@ -272,9 +272,9 @@ class RTDevice { rtd_dylib = std::make_unique(rtd_lib); std::string factory_name{rtd_name + "Factory"}; void *f_ptr = rtd_dylib->getSymbol(factory_name); - rtd_qdevice = + rtd_qdevice = std::unique_ptr( f_ptr ? reinterpret_cast(f_ptr)(rtd_kwargs.c_str()) - : nullptr; + : nullptr); return rtd_qdevice; } @@ -316,11 +316,7 @@ class ExecutionContext final { memory_man_ptr = std::make_unique(); } - ~ExecutionContext() - { - memory_man_ptr.reset(nullptr); - py_guard.reset(nullptr); - } + ~ExecutionContext() = default; void setDeviceRecorderStatus(bool status) noexcept { initial_tape_recorder_status = status; } diff --git a/runtime/lib/capi/RuntimeCAPI.cpp b/runtime/lib/capi/RuntimeCAPI.cpp index c6d85a3a88..ce215ad703 100644 --- a/runtime/lib/capi/RuntimeCAPI.cpp +++ b/runtime/lib/capi/RuntimeCAPI.cpp @@ -64,7 +64,10 @@ thread_local static RTDevice *RTD_PTR = nullptr; /** * @brief get the active device. */ -auto getQuantumDevicePtr() -> QuantumDevice * { return RTD_PTR->getQuantumDevicePtr(); } +auto getQuantumDevicePtr() -> const std::unique_ptr & +{ + return RTD_PTR->getQuantumDevicePtr(); +} /** * @brief Inactivate the active device instance. diff --git a/setup.py b/setup.py index cd54d44140..d2f05d1b19 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,6 @@ # limitations under the License. import glob -import importlib.util import platform import subprocess from distutils import sysconfig @@ -115,25 +114,11 @@ def run(self): ) -package_name = "scipy" - -scipy_package = importlib.util.find_spec(package_name) -package_directory = path.dirname(scipy_package.origin) - # Compile the library of custom calls in the frontend if system_platform == "Linux": - file_path_within_package_linux = "../scipy.libs/" - scipy_lib_path = path.join(package_directory, file_path_within_package_linux) - file_prefix = "libopenblasp" - file_extension = ".so" - search_pattern = path.join(scipy_lib_path, f"{file_prefix}*{file_extension}") - openblas_so_file = glob.glob(search_pattern)[0] - openblas_lib_name = path.basename(openblas_so_file)[3 : -len(file_extension)] custom_calls_extension = Extension( "catalyst.utils.libcustom_calls", sources=["frontend/catalyst/utils/libcustom_calls.cpp"], - libraries=[openblas_lib_name], - library_dirs=[scipy_lib_path], ) cmdclass = {"build_ext": CustomBuildExtLinux} @@ -141,19 +126,9 @@ def run(self): variables = sysconfig.get_config_vars() # Here we need to switch the deault to MacOs dynamic lib variables["LDSHARED"] = variables["LDSHARED"].replace("-bundle", "-dynamiclib") - - file_path_within_package_macos = ".dylibs/" - scipy_lib_path = path.join(package_directory, file_path_within_package_macos) - file_prefix = "libopenblas" - file_extension = ".dylib" - search_pattern = path.join(scipy_lib_path, f"{file_prefix}*{file_extension}") - openblas_dylib_file = glob.glob(search_pattern)[0] - openblas_lib_name = path.basename(openblas_dylib_file)[3 : -len(file_extension)] custom_calls_extension = Extension( "catalyst.utils.libcustom_calls", sources=["frontend/catalyst/utils/libcustom_calls.cpp"], - libraries=[openblas_lib_name], - library_dirs=[scipy_lib_path], ) cmdclass = {"build_ext": CustomBuildExtMacos}