Pybind11 as submodule and build MLIR with Python bindings (#957)

NVIDIA · Nov 23, 2023 · 7f96796 · 7f96796
1 parent ec0e366
commit 7f96796
Show file tree

Hide file tree

Showing 46 changed files with 188 additions and 19,049 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -33,8 +33,13 @@ jobs:
       pull_request_number: ${{ steps.pr_info.outputs.pr_number }}
       pull_request_base: ${{ steps.pr_info.outputs.pr_base }}
       cache_base: ${{ steps.pr_info.outputs.pr_base }}
+      llvm_commit: ${{ steps.repo_info.outputs.llvm_commit }}
+      pybind11_commit: ${{ steps.repo_info.outputs.pybind11_commit }}
 
     steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+
       - id: pr_info
         run: |
           pr_number=`echo ${{ github.ref_name }} | grep pull-request/ | (grep -o [0-9]* || true)`
@@ -49,6 +54,11 @@ jobs:
         env:
           GH_TOKEN: ${{ github.token }}
 
+      - id: repo_info
+        run: |
+          echo "llvm_commit=$(git rev-parse @:./tpls/llvm)" >> $GITHUB_OUTPUT
+          echo "pybind11_commit=$(git rev-parse @:./tpls/pybind11)" >> $GITHUB_OUTPUT
+
   devdeps:
     name: Load dependencies
     needs: metadata
@@ -62,6 +72,9 @@ jobs:
       platforms: linux/${{ matrix.platform }}
       dockerfile: build/devdeps.Dockerfile
       toolchain: ${{ matrix.toolchain }}
+      build_args: |
+        llvm_commit=${{ needs.metadata.outputs.llvm_commit }}
+        pybind11_commit=${{ needs.metadata.outputs.pybind11_commit }}
       registry_cache_from: ${{ inputs.cache_base || needs.metadata.outputs.cache_base }}
       # needed only for the cloudposse GitHub action
       matrix_key: ${{ matrix.platform }}-${{ matrix.toolchain }}
@@ -82,6 +95,8 @@ jobs:
         manylinux_image=manylinux_2_28
         arch=${{ (matrix.platform == 'arm64' && 'aarch64') || (matrix.platform == 'amd64' && 'x86_64') || 'any' }}
         distro=rhel8
+        llvm_commit=${{ needs.metadata.outputs.llvm_commit }}
+        pybind11_commit=${{ needs.metadata.outputs.pybind11_commit }}
       registry_cache_from: ${{ inputs.cache_base || needs.metadata.outputs.cache_base }}
       # needed only for the cloudposse GitHub action
       matrix_key: ${{ matrix.platform }}-python

diff --git a/.github/workflows/create_cache_command.yml b/.github/workflows/create_cache_command.yml
@@ -14,7 +14,7 @@ name: Create CI cache
 
 jobs:
   config:
-    name: Configure build
+    name: Determine PR info
     runs-on: ubuntu-latest
 
     outputs:
@@ -45,9 +45,31 @@ jobs:
             Checking out source code from head `${{ steps.pr_info.outputs.source_label }}` (sha: ${{ steps.pr_info.outputs.source_sha }}).
           edit-mode: append
 
+  repoinfo:
+    name: Configure build
+    needs: config
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+
+    outputs:
+      llvm_commit: ${{ steps.repo_info.outputs.llvm_commit }}
+      pybind11_commit: ${{ steps.repo_info.outputs.pybind11_commit }}
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+        with:
+          ref: "${{ format('refs/pull/{0}/merge', needs.config.outputs.pull_request_number) }}"
+
+      - id: repo_info
+        run: |
+          echo "llvm_commit=$(git rev-parse @:./tpls/llvm)" >> $GITHUB_OUTPUT
+          echo "pybind11_commit=$(git rev-parse @:./tpls/pybind11)" >> $GITHUB_OUTPUT
+
   devdeps_caches:
     name: Cache dev dependencies
-    needs: config
+    needs: [config, repoinfo]
     strategy:
       matrix:
         platform: [amd64, arm64]
@@ -58,6 +80,9 @@ jobs:
       platforms: linux/${{ matrix.platform }}
       dockerfile: build/devdeps.Dockerfile
       toolchain: ${{ matrix.toolchain }}
+      build_args: |
+        llvm_commit=${{ needs.repoinfo.outputs.llvm_commit }}
+        pybind11_commit=${{ needs.repoinfo.outputs.pybind11_commit }}
       create_local_cache: true
       registry_cache_from: ${{ needs.config.outputs.target_branch }}
       pull_request_number: ${{ needs.config.outputs.pull_request_number }}
@@ -66,7 +91,7 @@ jobs:
 
   wheeldeps_caches:
     name: Cache wheel dependencies
-    needs: config
+    needs: [config, repoinfo]
     strategy:
       matrix:
         platform: [amd64, arm64]
@@ -80,6 +105,8 @@ jobs:
         manylinux_image=manylinux_2_28
         arch=${{ (matrix.platform == 'arm64' && 'aarch64') || (matrix.platform == 'amd64' && 'x86_64') || 'any' }}
         distro=rhel8
+        llvm_commit=${{ needs.repoinfo.outputs.llvm_commit }}
+        pybind11_commit=${{ needs.repoinfo.outputs.pybind11_commit }}
       create_local_cache: true
       registry_cache_from: ${{ needs.config.outputs.target_branch }}
       pull_request_number: ${{ needs.config.outputs.pull_request_number }}

diff --git a/.github/workflows/deployments.yml b/.github/workflows/deployments.yml
@@ -59,6 +59,8 @@ jobs:
       pull_request_number: ${{ steps.pr_info.outputs.pr_number }}
       pull_request_base: ${{ steps.pr_info.outputs.pr_base }}
       pull_request_commit: ${{ steps.pr_info.outputs.merge_commit }}
+      llvm_commit: ${{ steps.repo_info.outputs.llvm_commit }}
+      pybind11_commit: ${{ steps.repo_info.outputs.pybind11_commit }}
       cache_base: ${{ steps.build_info.outputs.cache_base }}
       update_cache: ${{ steps.build_info.outputs.update_cache }}
       multi_platform: ${{ steps.build_info.outputs.multi_platform }}
@@ -143,6 +145,17 @@ jobs:
           retention-days: 30
           if-no-files-found: error
 
+      - name: Checkout repository
+        uses: actions/checkout@v3
+        with:
+          ref: "${{ steps.pr_info.outputs.merge_commit }}"
+
+      - name: Configure build
+        id: repo_info
+        run: |
+          echo "llvm_commit=$(git rev-parse @:./tpls/llvm)" >> $GITHUB_OUTPUT
+          echo "pybind11_commit=$(git rev-parse @:./tpls/pybind11)" >> $GITHUB_OUTPUT
+
   devdeps:
     name: Build dev dependencies
     needs: metadata
@@ -157,6 +170,9 @@ jobs:
       platforms: ${{ fromJson(needs.metadata.outputs.multi_platform || needs.metadata.outputs.platforms)[format('{0}', matrix.platform)].docker_flag }}
       dockerfile: build/devdeps.Dockerfile
       toolchain: ${{ matrix.toolchain }}
+      build_args: |
+        llvm_commit=${{ needs.metadata.outputs.llvm_commit }}
+        pybind11_commit=${{ needs.metadata.outputs.pybind11_commit }}
       registry_cache_from: ${{ needs.metadata.outputs.cache_base }}
       registry_cache_update: ${{ needs.metadata.outputs.update_cache == 'true' }}
       registry_cache_update_only: ${{ github.event_name == 'workflow_run' }}
@@ -184,6 +200,8 @@ jobs:
       build_args: |
         arch=${{ fromJson(needs.metadata.outputs.platforms)[format('{0}', matrix.platform)].arch }}
         distro=rhel8
+        llvm_commit=${{ needs.metadata.outputs.llvm_commit }}
+        pybind11_commit=${{ needs.metadata.outputs.pybind11_commit }}
       registry_cache_from: ${{ needs.metadata.outputs.cache_base }}
       registry_cache_update: ${{ needs.metadata.outputs.update_cache == 'true' }}
       registry_cache_update_only: ${{ github.event_name == 'workflow_run' }}

diff --git a/.github/workflows/dev_environment.yml b/.github/workflows/dev_environment.yml
@@ -90,7 +90,6 @@ jobs:
       image_id: ${{ steps.build_info.outputs.image_id }}
       image_tags: ${{ steps.metadata.outputs.tags }}
       image_labels: ${{ steps.metadata.outputs.labels }}
-      llvm_commit: ${{ steps.build_info.outputs.llvm_commit }}
 
     # Needed for access to environment variables (like the registry name).
     environment:
@@ -146,7 +145,6 @@ jobs:
           echo "custom_tags=$custom_tags" >> $GITHUB_OUTPUT
           echo "dockerfile=${{ inputs.dockerfile }}" >> $GITHUB_OUTPUT
           echo "owner=${repo_owner,,}" >> $GITHUB_OUTPUT
-          echo "llvm_commit=$(git rev-parse @:./tpls/llvm)" >> $GITHUB_OUTPUT
 
       - name: Extract metadata for Docker image
         id: metadata
@@ -301,7 +299,6 @@ jobs:
           target: ${{ inputs.build_target }}
           build-args: |
             toolchain=${{ inputs.toolchain }}
-            llvm_commit=${{ needs.metadata.outputs.llvm_commit }}
             ${{ inputs.build_args }}
           tags: ${{ needs.metadata.outputs.image_tags }}
           labels: ${{ needs.metadata.outputs.image_labels }}

diff --git a/.gitmodules b/.gitmodules
@@ -8,6 +8,7 @@
 	path = tpls/llvm
 	url = https://github.com/llvm/llvm-project.git
 	shallow = true
+	ignore = dirty
 [submodule "tpls/eigen"]
 	path = tpls/eigen
 	url = https://gitlab.com/libeigen/eigen.git
@@ -26,3 +27,7 @@
 [submodule "tpls/xtensor"]
 	path = tpls/xtensor
 	url = https://github.com/xtensor-stack/xtensor
+[submodule "tpls/pybind11"]
+	path = tpls/pybind11
+	url = https://github.com/pybind/pybind11
+	ignore = dirty
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -93,6 +93,18 @@ if(GIT_FOUND AND EXISTS "${PROJECT_SOURCE_DIR}/.git")
     if(NOT GIT_SUBMOD_RESULT EQUAL "0")
         message(FATAL_ERROR "Cloning git submodules failed with ${GIT_SUBMOD_RESULT}, please checkout submodules manually")
     endif()
+    # Apply patches; check if patch is already applied first, and apply it if not
+    execute_process(COMMAND ${GIT_EXECUTABLE} -C tpls/pybind11/ apply ../customizations/pybind11/pybind.h.diff --reverse --check
+                    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+                    RESULT_VARIABLE GIT_PATCH_RESULT)
+    if(NOT GIT_PATCH_RESULT EQUAL "0")
+      execute_process(COMMAND ${GIT_EXECUTABLE} -C tpls/pybind11/ apply ../customizations/pybind11/pybind.h.diff
+                      WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+                      RESULT_VARIABLE GIT_PATCH_RESULT)
+    endif()
+    if(NOT GIT_PATCH_RESULT EQUAL "0")
+        message(FATAL_ERROR "Applying patch to submodule failed with ${GIT_PATCH_RESULT}, please update patch")
+    endif()
   endif()
 endif()
 

diff --git a/docker/build/cudaq.wheel.Dockerfile b/docker/build/cudaq.wheel.Dockerfile
@@ -28,14 +28,19 @@ ADD "$workspace" "$destination"
 RUN dnf install -y cuda-nvtx-11-8 cuda-profiler-api-11-8 openblas-devel
 
 ARG python_version=3.10
+RUN echo "Building MLIR bindings for python${python_version}" \
+    && python${python_version} -m pip install --no-cache-dir numpy \
+    && export Python3_EXECUTABLE="$(which python${python_version})" \
+    && export CMAKE_EXE_LINKER_FLAGS="$LLVM_BUILD_LINKER_FLAGS" CMAKE_SHARED_LINKER_FLAGS="$LLVM_BUILD_LINKER_FLAGS" \
+    && bash /scripts/build_llvm.sh -s /llvm-project -c Release -v 
+
 RUN echo "Building wheel for python${python_version}." \
     && cd cuda-quantum && python=python${python_version} \
     # Find any external NVQIR simulator assets to be pulled in during wheel packaging.
     && export CUDAQ_EXTERNAL_NVQIR_SIMS=$(bash scripts/find_wheel_assets.sh assets) \
     && export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$(pwd)/assets" \
     && $python -m pip install --no-cache-dir \
-        cmake auditwheel \
-        cuquantum-cu11==23.10.0 \
+        auditwheel cuquantum-cu11==23.10.0 \
     && cuquantum_location=`$python -m pip show cuquantum-cu11 | grep -e 'Location: .*$'` \
     && export CUQUANTUM_INSTALL_PREFIX="${cuquantum_location#Location: }/cuquantum" \
     && ln -s $CUQUANTUM_INSTALL_PREFIX/lib/libcustatevec.so.1 $CUQUANTUM_INSTALL_PREFIX/lib/libcustatevec.so \

diff --git a/docker/build/devdeps.Dockerfile b/docker/build/devdeps.Dockerfile
@@ -26,6 +26,7 @@ FROM ubuntu:22.04 as llvmbuild
 SHELL ["/bin/bash", "-c"]
 
 ARG llvm_commit
+ARG pybind11_commit
 ARG toolchain=llvm
 
 # When a dialogue box would be needed during install, assume default configurations.
@@ -38,14 +39,15 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 # Install prerequisites for building LLVM.
 RUN apt-get update && apt-get install -y --no-install-recommends \
-        ninja-build cmake python3 \
+        ninja-build cmake python3 python3-dev python3-pip \
+    && python3 -m pip install --no-cache-dir numpy \
     && apt-get autoremove -y --purge && apt-get clean && rm -rf /var/lib/apt/lists/* 
 
 # Clone the LLVM source code.
+# Preserve access to the history to be able to cherry pick specific commits.
 RUN apt-get update && apt-get install -y --no-install-recommends git \
-    && mkdir /llvm-project && cd /llvm-project && git init \
-    && git remote add origin https://github.com/llvm/llvm-project \
-    && git fetch origin --depth=1 $llvm_commit && git reset --hard FETCH_HEAD \
+    && git clone --filter=tree:0 https://github.com/llvm/llvm-project /llvm-project \
+    && cd /llvm-project && git checkout $llvm_commit \
     && apt-get autoremove -y --purge && apt-get clean && rm -rf /var/lib/apt/lists/* 
 
 # Build the the LLVM libraries and compiler toolchain needed to build CUDA Quantum;
@@ -64,7 +66,16 @@ RUN apt-get update && apt-get install -y --no-install-recommends git \
 ADD ./scripts/install_toolchain.sh /scripts/install_toolchain.sh
 ADD ./scripts/build_llvm.sh /scripts/build_llvm.sh
 RUN LLVM_INSTALL_PREFIX=/opt/llvm LLVM_SOURCE=/llvm-project \
-        source scripts/install_toolchain.sh -e /opt/llvm/bootstrap -t ${toolchain}
+        source scripts/install_toolchain.sh -e /opt/llvm/bootstrap -t ${toolchain} \
+    && rm -rf /llvm-project/build
+RUN mkdir /pybind11-project && cd /pybind11-project && git init \
+    && git remote add origin https://github.com/pybind/pybind11 \
+    && git fetch origin --depth=1 $pybind11_commit && git reset --hard FETCH_HEAD \
+    && source /opt/llvm/bootstrap/init_command.sh \
+    && mkdir -p /pybind11-project/build && cd /pybind11-project/build \
+    && cmake -G Ninja ../ -DCMAKE_INSTALL_PREFIX=/usr/local/pybind11 \
+    && cmake --build . --target install --config Release \
+    && cd .. && rm -rf /pybind11-project
 RUN source /opt/llvm/bootstrap/init_command.sh && \
     LLVM_INSTALL_PREFIX=/opt/llvm \
         bash /scripts/build_llvm.sh -s /llvm-project -c Release -v \

diff --git a/docker/build/devdeps.manylinux.Dockerfile b/docker/build/devdeps.manylinux.Dockerfile
@@ -27,6 +27,7 @@ FROM quay.io/pypa/${manylinux_image}_${arch}:latest
 
 ARG distro=rhel8
 ARG llvm_commit
+ARG pybind11_commit
 ARG toolchain=gcc11
 
 # When a dialogue box would be needed during install, assume default configurations.
@@ -35,9 +36,9 @@ ARG toolchain=gcc11
 ARG DEBIAN_FRONTEND=noninteractive
 
 # Clone the LLVM source code.
-RUN mkdir /llvm-project && cd /llvm-project && git init \
-    && git remote add origin https://github.com/llvm/llvm-project \
-    && git fetch origin --depth=1 $llvm_commit && git reset --hard FETCH_HEAD
+# Preserve access to the history to be able to cherry pick specific commits.
+RUN git clone --filter=tree:0 https://github.com/llvm/llvm-project /llvm-project \
+    && cd /llvm-project && git checkout $llvm_commit
 
 # Install the C/C++ compiler toolchain with which the LLVM dependencies have
 # been built. CUDA Quantum needs to be built with that same toolchain, and the
@@ -64,13 +65,19 @@ ENV CXX="$LLVM_INSTALL_PREFIX/bootstrap/cxx"
 
 # Build the the LLVM libraries and compiler toolchain needed to build CUDA Quantum
 ADD ./scripts/build_llvm.sh /scripts/build_llvm.sh
+ENV LLVM_BUILD_LINKER_FLAGS="-static-libgcc -static-libstdc++"
 RUN dnf install -y --nobest --setopt=install_weak_deps=False \
         ninja-build cmake \
-    && export CMAKE_EXE_LINKER_FLAGS="-static-libgcc -static-libstdc++" \
-    && export CMAKE_SHARED_LINKER_FLAGS="-static-libgcc -static-libstdc++" \
-    && bash /scripts/build_llvm.sh -s /llvm-project -c Release -v \
-    && dnf remove -y ninja-build cmake && dnf clean all \
-    && rm -rf /llvm-project && rm /scripts/build_llvm.sh
+    && mkdir /pybind11-project && cd /pybind11-project && git init \
+    && git remote add origin https://github.com/pybind/pybind11 \
+    && git fetch origin --depth=1 $pybind11_commit && git reset --hard FETCH_HEAD \
+    && mkdir -p /pybind11-project/build && cd /pybind11-project/build \
+    && cmake -G Ninja ../ -DCMAKE_INSTALL_PREFIX=/usr/local/pybind11 \
+    && cmake --build . --target install --config Release \
+    && cd .. && rm -rf /pybind11-project \
+    && export CMAKE_EXE_LINKER_FLAGS="$LLVM_BUILD_LINKER_FLAGS" CMAKE_SHARED_LINKER_FLAGS="$LLVM_BUILD_LINKER_FLAGS" \
+    && bash /scripts/build_llvm.sh -s /llvm-project -c Release -v 
+    # No clean up, since we need to re-build llvm for each python version to get the bindings.
 
 # Install additional dependencies required to build the CUDA Quantum wheel.
 ADD ./scripts/install_prerequisites.sh /scripts/install_prerequisites.sh

diff --git a/runtime/cudaq/spin_op.h b/runtime/cudaq/spin_op.h
@@ -128,6 +128,7 @@ class spin_op {
                            _const_iter_type>;
     iterator(iterator &&) = default;
 
+    iterator(iterator const &other) : iter(other.iter) {}
     iterator(iter_type i) : iter(i) {}
     ~iterator() {
       for (auto &c : created) {