Skip to content

Commit 559c909

Browse files
committed
Reuse LLVM build tree for MLIR python bindings
The wheel container previously re-ran build_llvm.sh per Python version with python-bindings added, which reconfigured from scratch and recompiled the full LLVM/MLIR/clang/lld tree (~5,800 objects, ~65 min). The manylinux base already builds and retains /llvm-project/build without bindings. Add a focused script that reuses that tree, flips MLIR_ENABLE_BINDINGS_PYTHON on, and runs only the install-MLIRPythonModules and install-mlir-python-sources targets. nanobind is still rebuilt per Python since it ABI-binds to the interpreter. Signed-off-by: mdzurick <mitch_dz@hotmail.com>
1 parent 7744cc7 commit 559c909

3 files changed

Lines changed: 89 additions & 6 deletions

File tree

docker/build/devdeps.manylinux.Dockerfile

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,14 +83,17 @@ RUN curl -L https://github.com/Kitware/CMake/releases/download/v3.28.4/cmake-3.2
8383

8484
# Build the the LLVM libraries and compiler toolchain needed to build CUDA-Q.
8585
ADD ./scripts/build_llvm.sh /scripts/build_llvm.sh
86+
ADD ./scripts/build_mlir_python_bindings.sh /scripts/build_mlir_python_bindings.sh
8687
ADD ./cmake/caches/LLVM.cmake /cmake/caches/LLVM.cmake
8788
ADD ./tpls/customizations/llvm/ /tpls/customizations/llvm/
8889
RUN LLVM_PROJECTS='clang;lld;mlir' LLVM_SOURCE=/llvm-project \
8990
LLVM_CMAKE_CACHE=/cmake/caches/LLVM.cmake \
9091
LLVM_CMAKE_PATCHES=/tpls/customizations/llvm \
9192
bash /scripts/build_llvm.sh -c Release -v
92-
# No clean up of the build or source directory,
93-
# since we need to re-build llvm for each python version to get the bindings.
93+
# The build directory at /llvm-project/build is intentionally retained:
94+
# build_mlir_python_bindings.sh reuses it in the wheel container to add
95+
# the python-binding targets per Python version without recompiling
96+
# the rest of the LLVM/MLIR/clang/lld tree.
9497

9598
# Install CUDA
9699

docker/release/cudaq.wheel.Dockerfile

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,11 +51,11 @@ RUN --mount=from=ccache-data,target=/tmp/ccache-import,rw \
5151
fi
5252
RUN echo "Building MLIR bindings for python${python_version}" && \
5353
CCACHE_DISABLE=1 python${python_version} -m pip install --no-cache-dir numpy "nanobind>=2.9.0" && \
54-
rm -rf "$LLVM_INSTALL_PREFIX/src" "$LLVM_INSTALL_PREFIX/python_packages" && \
54+
rm -rf "$LLVM_INSTALL_PREFIX/python_packages" && \
5555
Python3_EXECUTABLE="$(which python${python_version})" \
56-
LLVM_PROJECTS='clang;lld;mlir;python-bindings' \
57-
LLVM_CMAKE_CACHE=/cmake/caches/LLVM.cmake LLVM_SOURCE=/llvm-project \
58-
bash /scripts/build_llvm.sh -c Release -v
56+
LLVM_SOURCE=/llvm-project \
57+
LLVM_INSTALL_PREFIX="$LLVM_INSTALL_PREFIX" \
58+
bash /scripts/build_mlir_python_bindings.sh
5959

6060
# Build wheel using unified wheel build script
6161
RUN cd /cuda-quantum && \
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
#!/bin/bash
2+
3+
# ============================================================================ #
4+
# Copyright (c) 2022 - 2026 NVIDIA Corporation & Affiliates. #
5+
# All rights reserved. #
6+
# #
7+
# This source code and the accompanying materials are made available under #
8+
# the terms of the Apache License 2.0 which accompanies this distribution. #
9+
# ============================================================================ #
10+
11+
# Builds the MLIR Python bindings against an existing LLVM build tree and
12+
# installs them into LLVM_INSTALL_PREFIX. Requires the manylinux devdeps
13+
# base image, where build_llvm.sh has already configured and built the
14+
# LLVM/MLIR libraries (without python-bindings) at LLVM_SOURCE/build.
15+
#
16+
# This avoids reconfiguring/rebuilding the full LLVM tree once per Python
17+
# version: only the binding objects (linked against the active interpreter
18+
# ABI) and a small amount of tablegen output get produced.
19+
#
20+
# Required environment:
21+
# LLVM_SOURCE path to the llvm-project source (default: /llvm-project)
22+
# LLVM_INSTALL_PREFIX install destination (default: /usr/local/llvm)
23+
# Python3_EXECUTABLE interpreter to bind against
24+
# NANOBIND_INSTALL_PREFIX nanobind install dir (default: /usr/local/nanobind)
25+
#
26+
# Usage:
27+
# Python3_EXECUTABLE=$(which python3.11) bash scripts/build_mlir_python_bindings.sh
28+
29+
set -e
30+
31+
LLVM_SOURCE=${LLVM_SOURCE:-/llvm-project}
32+
LLVM_INSTALL_PREFIX=${LLVM_INSTALL_PREFIX:-/usr/local/llvm}
33+
NANOBIND_INSTALL_PREFIX=${NANOBIND_INSTALL_PREFIX:-/usr/local/nanobind}
34+
LLVM_BUILD_FOLDER=${LLVM_BUILD_FOLDER:-build}
35+
36+
if [ -z "$Python3_EXECUTABLE" ]; then
37+
echo "Python3_EXECUTABLE must be set."
38+
exit 1
39+
fi
40+
41+
llvm_build_dir="$LLVM_SOURCE/$LLVM_BUILD_FOLDER"
42+
if [ ! -f "$llvm_build_dir/CMakeCache.txt" ]; then
43+
echo "Expected pre-configured LLVM build at $llvm_build_dir."
44+
echo "This script must run on top of the manylinux devdeps base image."
45+
exit 1
46+
fi
47+
48+
this_file_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
49+
50+
# nanobind links against the active interpreter ABI; build it if missing.
51+
if [ ! -d "$NANOBIND_INSTALL_PREFIX" ] || [ -z "$(ls -A "$NANOBIND_INSTALL_PREFIX"/* 2>/dev/null)" ]; then
52+
echo "Building nanobind..."
53+
cd "$this_file_dir/.." && repo_root=$(git rev-parse --show-toplevel) && cd "$repo_root"
54+
git submodule update --init --recursive --recommend-shallow --single-branch tpls/nanobind
55+
mkdir -p tpls/nanobind/build && cd tpls/nanobind/build
56+
cmake -G Ninja ../ \
57+
-DCMAKE_INSTALL_PREFIX="$NANOBIND_INSTALL_PREFIX" \
58+
-DPython3_EXECUTABLE="$Python3_EXECUTABLE" \
59+
-DNB_TEST=False
60+
cmake --build . --target install --config Release
61+
fi
62+
63+
cd "$llvm_build_dir"
64+
65+
# Reconfigure the existing build tree: enable MLIR python bindings and
66+
# point at the active interpreter. Other cache entries (compiler flags,
67+
# enabled projects, ccache launcher, build type) are preserved.
68+
echo "Reconfiguring LLVM build for python bindings (Python: $Python3_EXECUTABLE)..."
69+
cmake . \
70+
-DMLIR_ENABLE_BINDINGS_PYTHON=ON \
71+
-DPython3_EXECUTABLE="$Python3_EXECUTABLE" \
72+
-Dnanobind_DIR="$NANOBIND_INSTALL_PREFIX/nanobind/cmake"
73+
74+
# Build and install only the python-binding components. Other targets in
75+
# the existing tree are left untouched and remain installed at
76+
# LLVM_INSTALL_PREFIX from the base image.
77+
echo "Building and installing MLIR python bindings..."
78+
ninja install-MLIRPythonModules install-mlir-python-sources
79+
80+
echo "Installed MLIR python bindings into $LLVM_INSTALL_PREFIX."

0 commit comments

Comments
 (0)