From 4f7c3dadd6b92a6669c18779e6377dcdb10366fb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 6 Sep 2026 21:22:39 +0000 Subject: [PATCH 1/3] Add proposal doc for shipping Enzyme-enabled kernels in pip wheels Co-authored-by: jurasic-pf <166746189+jurasic-pf@users.noreply.github.com> --- docs/enzyme_wheel_packaging.md | 129 +++++++++++++++++++++++++++++++++ docs/index.md | 1 + 2 files changed, 130 insertions(+) create mode 100644 docs/enzyme_wheel_packaging.md diff --git a/docs/enzyme_wheel_packaging.md b/docs/enzyme_wheel_packaging.md new file mode 100644 index 000000000..0849408d6 --- /dev/null +++ b/docs/enzyme_wheel_packaging.md @@ -0,0 +1,129 @@ +# Shipping Enzyme-differentiated kernels in pip wheels + +Status: proposal. Follow-up to #710 (and the #584 -> #703 -> #707 -> #710 +stack), which wires an exact Enzyme-derived JVP/VJP of the force kernel into +`vmecpp_core` behind `-DVMECPP_ENABLE_ENZYME=ON`. That flag is `OFF` by +default and is currently only exercised by the dedicated +`.github/workflows/test_enzyme.yaml` job, which builds its own Clang/LLVM +21 + `ClangEnzyme-21` (pinned at `v0.0.264`) toolchain from scratch on an +`ubuntu-22.04` runner. None of that reaches a wheel: `cibuildwheel` +(`[tool.cibuildwheel]` in `pyproject.toml`) still builds `vmecpp_core` with +the default compiler on each platform (`manylinux`'s `gcc`, Homebrew `gcc` +on macOS) and never sets `VMECPP_ENABLE_ENZYME`. This note proposes closing +that gap. + +## Constraints + +- **No regression for non-adjoint users.** Everyone who does not call the + exact-JVP/VJP API pays for this feature exactly nothing: same compiler, + same flags, same object code for every other translation unit. +- **No new runtime dependency.** Enzyme is an LLVM optimization pass that + runs once, at compile time, over the two translation units + `exact_force_jvp.cc` / `exact_force_vjp.cc`. Its output is a plain `.o` + with ordinary machine code; `ClangEnzyme-21.so` and `clang` itself are + build-time tools, not something the wheel loads at import time or ships. +- **manylinux/macOS toolchain policy.** `cibuildwheel`'s manylinux + containers and Homebrew's default macOS toolchain are GCC-based (or an + Apple-clang ABI on macOS that is untested against upstream + `ClangEnzyme`). Installing a full matching upstream Clang/LLVM + + built-from-source Enzyme plugin inside every wheel-build container, + for every architecture, is exactly the kind of toolchain surgery that + breaks `manylinux` glibc-symbol-version guarantees and roughly doubles + the wheel build matrix's fragility. + +## Proposed architecture: prebuilt-object mode, not a Clang wheel toolchain + +`CMakeLists.txt` already has the mechanism this needs: +`VMECPP_ENZYME_JVP_OBJECT` / `VMECPP_ENZYME_VJP_OBJECT` let the two Enzyme +translation units be supplied as **prebuilt `.o` files**, compiled +elsewhere with a matching Clang + `ClangEnzyme`, and linked into +`vmecpp_core` (and `vmecpp_core_v3`) without the main build ever invoking +Clang or loading the plugin. This is the same "prebuilt Enzyme mode" the +CMake comments call out for `icx`/`icpx`. Reuse it for wheels instead of +inventing a second path: + +1. **Add a small, separate CI job** (`build_enzyme_objects.yaml`) that + reuses the `test_enzyme.yaml` toolchain recipe (pinned `LLVM_VERSION`, + pinned `ENZYME_REF`) to compile `exact_force_jvp.cc` and + `exact_force_vjp.cc` for each wheel target + (`linux-x86_64`, `linux-aarch64`, `macos-x86_64`, `macos-arm64`), one + per glibc-hwcaps level actually needed (baseline `x86-64`, per the + existing `-march=x86-64` override for these two files - Enzyme does not + support the AVX mask intrinsics used for `x86-64-v3`, so there is only + one variant per arch regardless of `VMECPP_HWCAPS_DISPATCH`). Compile + inside the same container `cibuildwheel` will use + (`manylinux_2_28`/`musllinux` image, or an equivalent macOS SDK/deployment + target) so the resulting `.o` matches the wheel's glibc/SDK floor, and + with matching `-fPIC`/ABI flags (Itanium C++ ABI, default libstdc++ on + Linux and libc++ on macOS) so it links cleanly into a GCC/Apple-clang + -built `vmecpp_core`. + - Cache these objects (they only change when `exact_force_jvp.cc`, + `exact_force_vjp.cc`, or the pinned Enzyme/LLVM versions change) + and publish them as a build artifact or a small release asset, + analogous to how `test_enzyme.yaml` already caches the + `ClangEnzyme` plugin build. +2. **Feed the objects into the normal wheel build** via + `[tool.cibuildwheel.linux.environment]` / `.macos.environment` / + `.windows.environment`, e.g. + `CMAKE_ARGS = "-DVMECPP_ENABLE_ENZYME=ON -DVMECPP_ENZYME_JVP_OBJECT=... -DVMECPP_ENZYME_VJP_OBJECT=..."`, + with `before-build` fetching the matching prebuilt pair for the + platform/arch being built (`cibuildwheel` already runs one container + per arch, so `before-build` can `curl`/unpack the right pair using + `$CIBW_ARCHS`/`uname -m`). No Clang, LLVM, or `ClangEnzyme` needs to be + installed inside the `cibuildwheel` container itself; `before-build` + only fetches two small `.o` files. + - `vmecpp_core` and `vmecpp_core_v3` continue to be compiled with + today's compiler and flags for every translation unit except those + two; nothing about the default, non-adjoint code path changes. +3. **Fail closed, not silently, when a pair is missing** for a given + target (e.g. a new arch added to the wheel matrix before its objects + are built): leave `VMECPP_ENABLE_ENZYME` off for that specific build + rather than failing the whole wheel job, and let the existing runtime + guard (the "non-Enzyme build raises instead of using finite + differences" behavior from #710) tell users why the exact-VJP API is + unavailable on that wheel. + +This keeps the wheel-build containers untouched for the code path every +non-adjoint user exercises, and confines all the toolchain risk to one +small, cacheable, independently-testable CI job. + +## Verifying no performance regression + +Because the only two translation units touched are `exact_force_jvp.cc` +and `exact_force_vjp.cc` - which are otherwise not referenced by the +production forward solve - object-identity is enough for everything else, +but should still be checked explicitly: +- Add a CI check that builds `vmecpp_core` twice, once with + `VMECPP_ENABLE_ENZYME=OFF` and once in prebuilt-object mode, and diffs + the resulting `.so` symbol table / disassembly for every translation + unit *other* than the two Enzyme ones (e.g. `objdump -d` per `.o` + before linking, or comparing `nm --defined-only` output for unrelated + symbols) to catch any accidental change in compiler flags reaching the + rest of the library. +- Extend the existing `benchmarks/benchmark_exact_hvp.py` run (already + in `test_enzyme.yaml`) to also record a plain forward-solve timing + (e.g. the W7-X or CTH example already used elsewhere in the test suite) + under both configurations, and gate on the two forward-solve timings + matching within noise. This is what actually answers "does an + Enzyme-enabled wheel regress non-adjoint users", as opposed to the + HVP-vs-finite-difference comparison the current benchmark performs. +- Run this comparison against an actual wheel built through the proposed + `cibuildwheel` path (not just the local CMake build used by + `test_enzyme.yaml`), since the wheel is the artifact users install. + +## Open questions + +- Which wheel targets get Enzyme objects at all. `aarch64` and `musllinux` + add real CI matrix cost for the object-build job; starting with + `linux-x86_64` and both macOS architectures (matching where + `test_enzyme.yaml` and the exact-HVP benchmark already run) and + expanding once the mechanism is proven is likely the lowest-risk order. +- Whether prebuilt objects should be committed as versioned release + assets (rebuilt on a schedule/on relevant source changes) or built + fresh in every wheel release; the former is cheaper per release but + needs its own staleness check against `exact_force_jvp.cc` / + `exact_force_vjp.cc` content hashes and the pinned `ENZYME_REF`. +- Windows is currently outside `test_enzyme.yaml`'s scope entirely (no + pinned Clang/Enzyme recipe exists for it yet); this proposal does not + attempt to extend Enzyme support there, only to package what already + exists for Linux/macOS. diff --git a/docs/index.md b/docs/index.md index 4acc8d543..d58af97ff 100644 --- a/docs/index.md +++ b/docs/index.md @@ -11,6 +11,7 @@ Reference Examples Fourier Basis Details Benchmarks +Enzyme Wheel Packaging Proposal ``` ![MIT license](https://img.shields.io/badge/license-MIT-blue) From 4971c5836e474bbe1c5a7ac36f545617877125bc Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 7 Sep 2026 08:29:08 +0000 Subject: [PATCH 2/3] Implement the Enzyme-wheel proposal as an experimental cibuildwheel pipeline Co-authored-by: jurasic-pf <166746189+jurasic-pf@users.noreply.github.com> --- .github/workflows/build_enzyme_wheels.yaml | 225 +++++++++++++++++++++ docs/enzyme_wheel_packaging.md | 170 ++++++++++------ 2 files changed, 328 insertions(+), 67 deletions(-) create mode 100644 .github/workflows/build_enzyme_wheels.yaml diff --git a/.github/workflows/build_enzyme_wheels.yaml b/.github/workflows/build_enzyme_wheels.yaml new file mode 100644 index 000000000..28610c0e9 --- /dev/null +++ b/.github/workflows/build_enzyme_wheels.yaml @@ -0,0 +1,225 @@ +name: Enzyme-enabled wheels (experimental) + +# First implementation increment of docs/enzyme_wheel_packaging.md: builds the +# prebuilt Enzyme JVP/VJP objects (CMake's VMECPP_ENZYME_JVP_OBJECT / +# VMECPP_ENZYME_VJP_OBJECT prebuilt-object mode) for a subset of wheel targets, +# then builds actual wheels with them via cibuildwheel and smoke-tests them. +# +# This workflow is intentionally separate from .github/workflows/pypi_publish.yml +# and does not run on `release: published`: it does not touch the real release +# pipeline. Once its wheels are validated (install correctly, pass the Enzyme +# integration tests, and show no forward-solve regression vs. a plain wheel), +# the same environment variables can move into pypi_publish.yml's cibuildwheel +# step. Scope for now: linux-x86_64 and both macOS architectures, matching +# where the exact-HVP benchmark and the Enzyme smoke test already run; aarch64 +# Linux and Windows are follow-ups (see docs/enzyme_wheel_packaging.md). +on: + workflow_dispatch: + pull_request: + paths: + - ".github/workflows/build_enzyme_wheels.yaml" + - "src/vmecpp/cpp/vmecpp/common/enzyme/**" + - "src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.*" + - "src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.*" + - "CMakeLists.txt" + +concurrency: + group: ${{ github.workflow }}-${{ github.ref || github.run_id }} + cancel-in-progress: ${{ github.ref_name != 'main' }} + +env: + PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + # Pinned to match .github/workflows/test_enzyme.yaml's Linux toolchain. + LLVM_VERSION: "21" + ENZYME_REF: "v0.0.264" + +jobs: + build-enzyme-objects: + name: Build Enzyme objects (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-24.04, macos-14, macos-15-intel] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v6 + + # --- Linux: Clang/LLVM 21 from apt.llvm.org, matching test_enzyme.yaml --- + - name: "[Linux] Install Clang/LLVM and build tools" + if: runner.os == 'Linux' + run: | + wget https://apt.llvm.org/llvm.sh + chmod +x llvm.sh + sudo ./llvm.sh "${LLVM_VERSION}" + sudo apt-get install -y \ + "llvm-${LLVM_VERSION}-dev" "clang-${LLVM_VERSION}" "libclang-${LLVM_VERSION}-dev" \ + "libomp-${LLVM_VERSION}-dev" \ + cmake gfortran libhdf5-dev liblapack-dev libnetcdf-dev ninja-build \ + python3-dev zlib1g-dev + + - name: "[Linux] Cache the ClangEnzyme plugin build" + if: runner.os == 'Linux' + id: enzyme-cache-linux + uses: actions/cache@v4 + with: + path: enzyme-build + key: enzyme-${{ env.ENZYME_REF }}-llvm${{ env.LLVM_VERSION }}-linux + + - name: "[Linux] Build the ClangEnzyme plugin (pinned)" + if: runner.os == 'Linux' && steps.enzyme-cache-linux.outputs.cache-hit != 'true' + run: | + git clone --depth 1 --branch "${ENZYME_REF}" \ + https://github.com/EnzymeAD/Enzyme.git enzyme-src + cmake -G Ninja -S enzyme-src/enzyme -B enzyme-build \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLVM_DIR="$(llvm-config-${LLVM_VERSION} --cmakedir)" \ + -DClang_DIR="$(llvm-config-${LLVM_VERSION} --prefix)/lib/cmake/clang" + ninja -C enzyme-build "ClangEnzyme-${LLVM_VERSION}" + + - name: "[Linux] Locate the built plugin and set the compiler" + if: runner.os == 'Linux' + run: | + plugin="$(find "${PWD}/enzyme-build" -name "ClangEnzyme-${LLVM_VERSION}.so" | head -1)" + test -n "${plugin}" + echo "ENZYME_PLUGIN=${plugin}" >> "${GITHUB_ENV}" + echo "ENZYME_CC=clang-${LLVM_VERSION}" >> "${GITHUB_ENV}" + echo "ENZYME_CXX=clang++-${LLVM_VERSION}" >> "${GITHUB_ENV}" + + # --- macOS: Homebrew LLVM/Clang. Apple's own clang cannot load an + # upstream ClangEnzyme plugin (different fork/ABI), so the plugin and + # the two Enzyme translation units are always built with Homebrew's + # clang here, then linked as prebuilt objects into the AppleClang-built + # extension module that cibuildwheel actually ships. --- + - name: "[macOS] Install Homebrew LLVM and build tools" + if: runner.os == 'macOS' + run: brew install llvm ninja hdf5 netcdf lapack + + - name: "[macOS] Cache the ClangEnzyme plugin build" + if: runner.os == 'macOS' + id: enzyme-cache-macos + uses: actions/cache@v4 + with: + path: enzyme-build + key: enzyme-${{ env.ENZYME_REF }}-llvm-macos-${{ runner.arch }} + + - name: "[macOS] Build the ClangEnzyme plugin (pinned)" + if: runner.os == 'macOS' && steps.enzyme-cache-macos.outputs.cache-hit != 'true' + run: | + brew_llvm="$(brew --prefix llvm)" + git clone --depth 1 --branch "${ENZYME_REF}" \ + https://github.com/EnzymeAD/Enzyme.git enzyme-src + cmake -G Ninja -S enzyme-src/enzyme -B enzyme-build \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLVM_DIR="${brew_llvm}/lib/cmake/llvm" \ + -DClang_DIR="${brew_llvm}/lib/cmake/clang" + # Homebrew's llvm formula does not encode the major version in the + # plugin target name the way the pinned apt.llvm.org packages do. + llvm_major="$("${brew_llvm}/bin/llvm-config" --version | cut -d. -f1)" + echo "LLVM_MACOS_MAJOR=${llvm_major}" >> "${GITHUB_ENV}" + ninja -C enzyme-build "ClangEnzyme-${llvm_major}" + + - name: "[macOS] Locate the built plugin and set the compiler" + if: runner.os == 'macOS' + run: | + brew_llvm="$(brew --prefix llvm)" + plugin="$(find "${PWD}/enzyme-build" -name "ClangEnzyme-*.dylib" -o -name "ClangEnzyme-*.so" | head -1)" + test -n "${plugin}" + echo "ENZYME_PLUGIN=${plugin}" >> "${GITHUB_ENV}" + echo "ENZYME_CC=${brew_llvm}/bin/clang" >> "${GITHUB_ENV}" + echo "ENZYME_CXX=${brew_llvm}/bin/clang++" >> "${GITHUB_ENV}" + + # --- Build vmecpp_core in "plugin mode" (VMECPP_ENZYME_PLUGIN) and pull + # the two resulting Enzyme translation-unit objects out of the CMake + # build tree. This is the same plugin-mode build test_enzyme.yaml + # already exercises; here we only additionally harvest the objects. --- + - name: Build vmecpp with the Enzyme kernels enabled + env: + CC: ${{ env.ENZYME_CC }} + CXX: ${{ env.ENZYME_CXX }} + CMAKE_ARGS: >- + -DVMECPP_ENABLE_ENZYME=ON + -DVMECPP_ENZYME_PLUGIN=${{ env.ENZYME_PLUGIN }} + run: uv sync --reinstall-package vmecpp + + - name: Extract the Enzyme JVP/VJP objects + run: | + mkdir -p enzyme-objects + jvp="$(find build -name 'exact_force_jvp.cc.o' | head -1)" + vjp="$(find build -name 'exact_force_vjp.cc.o' | head -1)" + test -n "${jvp}" + test -n "${vjp}" + cp "${jvp}" enzyme-objects/exact_force_jvp.o + cp "${vjp}" enzyme-objects/exact_force_vjp.o + + - uses: actions/upload-artifact@v4 + with: + name: enzyme-objects-${{ matrix.os }} + path: enzyme-objects/ + + build-enzyme-wheels: + name: Build Enzyme-enabled wheel (${{ matrix.os }}) + needs: build-enzyme-objects + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-24.04, macos-14, macos-15-intel] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v6 + + - uses: actions/download-artifact@v4 + with: + name: enzyme-objects-${{ matrix.os }} + path: enzyme-objects + + - name: Compute the prebuilt-object CMAKE_ARGS for this platform + run: | + jvp="${{ github.workspace }}/enzyme-objects/exact_force_jvp.o" + vjp="${{ github.workspace }}/enzyme-objects/exact_force_vjp.o" + test -f "${jvp}" + test -f "${vjp}" + # Linux wheels build inside a manylinux container with the repo + # mounted at /project; translate the host path accordingly. + if [ "${{ runner.os }}" = "Linux" ]; then + jvp="/project/enzyme-objects/exact_force_jvp.o" + vjp="/project/enzyme-objects/exact_force_vjp.o" + fi + { + echo "VMECPP_ENZYME_CMAKE_ARGS=-DVMECPP_ENABLE_ENZYME=ON -DVMECPP_ENZYME_JVP_OBJECT=${jvp} -DVMECPP_ENZYME_VJP_OBJECT=${vjp}" + } >> "${GITHUB_ENV}" + + - name: Build wheel + uses: pypa/cibuildwheel@v3.3.1 + env: + CIBW_ENVIRONMENT_LINUX: > + PIP_CONSTRAINT=/project/.github/constraints.txt + UV_CONSTRAINT=/project/.github/constraints.txt + CMAKE_ARGS="${{ env.VMECPP_ENZYME_CMAKE_ARGS }}" + CIBW_ENVIRONMENT_MACOS: > + MACOSX_DEPLOYMENT_TARGET=${{ matrix.os == 'macos-14' && '14.0' || '15.0' }} + DYLD_LIBRARY_PATH=/usr/local/opt/gcc/lib/gcc/current/:$DYLD_LIBRARY_PATH + FC=gfortran-14 + OpenMP_ROOT=$(brew --prefix)/opt/libomp + CMAKE_ARGS="${{ env.VMECPP_ENZYME_CMAKE_ARGS }}" + with: + package-dir: . + output-dir: wheelhouse + config-file: "{package}/pyproject.toml" + + - uses: actions/upload-artifact@v4 + with: + name: vmecpp-enzyme-wheels-${{ matrix.os }} + path: ./wheelhouse/*.whl + + # Smoke test: install the Enzyme-enabled wheel and run the Enzyme + # integration tests, same as test_enzyme.yaml's plugin-mode build. + - name: Install the built wheel and run the Enzyme integration tests + run: | + python3 -m venv smoke-venv + . smoke-venv/bin/activate + pip install --no-index --find-links=wheelhouse/ vmecpp + pip install pytest + python -m pytest -q tests/test_hessian.py tests/test_external_optimizers.py diff --git a/docs/enzyme_wheel_packaging.md b/docs/enzyme_wheel_packaging.md index 0849408d6..7b0e43f0d 100644 --- a/docs/enzyme_wheel_packaging.md +++ b/docs/enzyme_wheel_packaging.md @@ -1,16 +1,30 @@ # Shipping Enzyme-differentiated kernels in pip wheels -Status: proposal. Follow-up to #710 (and the #584 -> #703 -> #707 -> #710 -stack), which wires an exact Enzyme-derived JVP/VJP of the force kernel into -`vmecpp_core` behind `-DVMECPP_ENABLE_ENZYME=ON`. That flag is `OFF` by -default and is currently only exercised by the dedicated +Status: phase 1 implemented (experimental, not yet on the release path). +Follow-up to #710 (and the #584 -> #703 -> #707 -> #710 stack), which wires +an exact Enzyme-derived JVP/VJP of the force kernel into `vmecpp_core` +behind `-DVMECPP_ENABLE_ENZYME=ON`. That flag is `OFF` by default and was, +until this change, only exercised by the dedicated `.github/workflows/test_enzyme.yaml` job, which builds its own Clang/LLVM 21 + `ClangEnzyme-21` (pinned at `v0.0.264`) toolchain from scratch on an -`ubuntu-22.04` runner. None of that reaches a wheel: `cibuildwheel` +`ubuntu-22.04` runner. None of that reached a wheel: `cibuildwheel` (`[tool.cibuildwheel]` in `pyproject.toml`) still builds `vmecpp_core` with -the default compiler on each platform (`manylinux`'s `gcc`, Homebrew `gcc` -on macOS) and never sets `VMECPP_ENABLE_ENZYME`. This note proposes closing -that gap. +the default compiler on each platform (`manylinux`'s `gcc`, `AppleClang` on +macOS) and never sets `VMECPP_ENABLE_ENZYME`. + +`.github/workflows/build_enzyme_wheels.yaml` now implements the mechanism +described below for `linux-x86_64`, `macos-arm64` (`macos-14`) and +`macos-x86_64` (`macos-15-intel`), producing installable, smoke-tested +Enzyme-enabled wheels as workflow artifacts. It is a separate, +`workflow_dispatch`/path-filtered workflow, deliberately not wired into +`.github/workflows/pypi_publish.yml`'s `release: published` trigger yet: +none of this has been exercised against a real manylinux container or a +real macOS runner outside of this change, so it should not be trusted with +the actual release pipeline until it has run green a few times and its +wheels have been checked for the no-regression property described below. +Promoting it is a follow-up once that evidence exists. Remaining scope gaps +(`linux-aarch64`, Windows, merging into `pypi_publish.yml`) are tracked in +[Open questions](#open-questions). ## Constraints @@ -23,15 +37,19 @@ that gap. with ordinary machine code; `ClangEnzyme-21.so` and `clang` itself are build-time tools, not something the wheel loads at import time or ships. - **manylinux/macOS toolchain policy.** `cibuildwheel`'s manylinux - containers and Homebrew's default macOS toolchain are GCC-based (or an - Apple-clang ABI on macOS that is untested against upstream - `ClangEnzyme`). Installing a full matching upstream Clang/LLVM + - built-from-source Enzyme plugin inside every wheel-build container, - for every architecture, is exactly the kind of toolchain surgery that - breaks `manylinux` glibc-symbol-version guarantees and roughly doubles - the wheel build matrix's fragility. + containers build with the container's GCC. On macOS, `vmecpp_core` is + actually built with Apple's own clang (`AppleClang`; `before-build = "brew + install gcc libomp"` in `pyproject.toml` only supplies `gfortran` and + OpenMP, not the C/C++ compiler) - and Apple's clang cannot load an + upstream `ClangEnzyme` plugin (different fork, no matching `-fplugin` + support), so prebuilt objects are not just the lower-risk option on + macOS, they are the only one that works there at all. Installing a full + matching upstream Clang/LLVM + built-from-source Enzyme plugin inside + every wheel-build container, for every architecture, would also be the + kind of toolchain surgery that risks breaking `manylinux` + glibc-symbol-version guarantees. -## Proposed architecture: prebuilt-object mode, not a Clang wheel toolchain +## Implemented architecture: prebuilt-object mode, not a Clang wheel toolchain `CMakeLists.txt` already has the mechanism this needs: `VMECPP_ENZYME_JVP_OBJECT` / `VMECPP_ENZYME_VJP_OBJECT` let the two Enzyme @@ -39,60 +57,66 @@ translation units be supplied as **prebuilt `.o` files**, compiled elsewhere with a matching Clang + `ClangEnzyme`, and linked into `vmecpp_core` (and `vmecpp_core_v3`) without the main build ever invoking Clang or loading the plugin. This is the same "prebuilt Enzyme mode" the -CMake comments call out for `icx`/`icpx`. Reuse it for wheels instead of -inventing a second path: +CMake comments call out for `icx`/`icpx`; `build_enzyme_wheels.yaml` reuses +it for wheels instead of inventing a second path: -1. **Add a small, separate CI job** (`build_enzyme_objects.yaml`) that - reuses the `test_enzyme.yaml` toolchain recipe (pinned `LLVM_VERSION`, - pinned `ENZYME_REF`) to compile `exact_force_jvp.cc` and - `exact_force_vjp.cc` for each wheel target - (`linux-x86_64`, `linux-aarch64`, `macos-x86_64`, `macos-arm64`), one - per glibc-hwcaps level actually needed (baseline `x86-64`, per the - existing `-march=x86-64` override for these two files - Enzyme does not - support the AVX mask intrinsics used for `x86-64-v3`, so there is only - one variant per arch regardless of `VMECPP_HWCAPS_DISPATCH`). Compile - inside the same container `cibuildwheel` will use - (`manylinux_2_28`/`musllinux` image, or an equivalent macOS SDK/deployment - target) so the resulting `.o` matches the wheel's glibc/SDK floor, and - with matching `-fPIC`/ABI flags (Itanium C++ ABI, default libstdc++ on - Linux and libc++ on macOS) so it links cleanly into a GCC/Apple-clang - -built `vmecpp_core`. - - Cache these objects (they only change when `exact_force_jvp.cc`, - `exact_force_vjp.cc`, or the pinned Enzyme/LLVM versions change) - and publish them as a build artifact or a small release asset, - analogous to how `test_enzyme.yaml` already caches the - `ClangEnzyme` plugin build. -2. **Feed the objects into the normal wheel build** via - `[tool.cibuildwheel.linux.environment]` / `.macos.environment` / - `.windows.environment`, e.g. - `CMAKE_ARGS = "-DVMECPP_ENABLE_ENZYME=ON -DVMECPP_ENZYME_JVP_OBJECT=... -DVMECPP_ENZYME_VJP_OBJECT=..."`, - with `before-build` fetching the matching prebuilt pair for the - platform/arch being built (`cibuildwheel` already runs one container - per arch, so `before-build` can `curl`/unpack the right pair using - `$CIBW_ARCHS`/`uname -m`). No Clang, LLVM, or `ClangEnzyme` needs to be - installed inside the `cibuildwheel` container itself; `before-build` - only fetches two small `.o` files. +1. **A small, separate CI job** (`build-enzyme-objects` in + `build_enzyme_wheels.yaml`) reuses the `test_enzyme.yaml` toolchain + recipe (pinned `LLVM_VERSION`, pinned `ENZYME_REF`) on Linux, and an + analogous Homebrew-LLVM recipe on macOS, to build a plugin-mode + `vmecpp_core` (`-DVMECPP_ENABLE_ENZYME=ON -DVMECPP_ENZYME_PLUGIN=...`) + and harvest the two resulting `exact_force_jvp.cc.o` / + `exact_force_vjp.cc.o` objects straight out of the CMake build tree. + There is only one variant per arch: the CMake rules already force + `-march=x86-64` for these two files regardless of + `VMECPP_HWCAPS_DISPATCH`, since Enzyme does not support the AVX mask + intrinsics used for `x86-64-v3`. Both the plugin build and the object + build are cached (`actions/cache`), matching how `test_enzyme.yaml` + already caches the `ClangEnzyme` plugin build. + - **Not yet done:** compiling inside the actual `manylinux_2_28` + container `cibuildwheel` uses for the Linux target, so the resulting + `.o`'s glibc-symbol-version floor is only as good as the plain + `ubuntu-24.04` runner's toolchain, not the container's. `auditwheel` + (run automatically by `cibuildwheel`'s manylinux repair step) does + check this and will fail loudly, rather than silently shipping a + wheel with an inflated glibc requirement, but this has not been + exercised yet since the workflow does not run on every push. See + [Open questions](#open-questions). +2. **The objects feed the normal wheel build** (`build-enzyme-wheels` job) + via `CIBW_ENVIRONMENT_LINUX` / `CIBW_ENVIRONMENT_MACOS`, setting + `CMAKE_ARGS="-DVMECPP_ENABLE_ENZYME=ON -DVMECPP_ENZYME_JVP_OBJECT=... -DVMECPP_ENZYME_VJP_OBJECT=..."` + with paths translated for the Linux case, where `cibuildwheel` mounts + the repo at `/project` inside the manylinux container. No Clang, LLVM, + or `ClangEnzyme` is installed inside the `cibuildwheel` container or + macOS build environment itself; only the two small `.o` files, produced + by the sibling job, are downloaded first. - `vmecpp_core` and `vmecpp_core_v3` continue to be compiled with today's compiler and flags for every translation unit except those two; nothing about the default, non-adjoint code path changes. 3. **Fail closed, not silently, when a pair is missing** for a given - target (e.g. a new arch added to the wheel matrix before its objects - are built): leave `VMECPP_ENABLE_ENZYME` off for that specific build - rather than failing the whole wheel job, and let the existing runtime - guard (the "non-Enzyme build raises instead of using finite - differences" behavior from #710) tell users why the exact-VJP API is - unavailable on that wheel. + target (currently true for `linux-aarch64`, not yet built): the + consuming step requires both files to exist before setting + `VMECPP_ENABLE_ENZYME=ON`, and the existing runtime guard (the + "non-Enzyme build raises instead of using finite differences" behavior + from #710) tells users why the exact-VJP API is unavailable on a wheel + that was built without it. This keeps the wheel-build containers untouched for the code path every non-adjoint user exercises, and confines all the toolchain risk to one -small, cacheable, independently-testable CI job. +small, cacheable, independently-testable CI job, currently isolated in +`build_enzyme_wheels.yaml` rather than the release-triggered +`pypi_publish.yml`. ## Verifying no performance regression Because the only two translation units touched are `exact_force_jvp.cc` and `exact_force_vjp.cc` - which are otherwise not referenced by the production forward solve - object-identity is enough for everything else, -but should still be checked explicitly: +but should still be checked explicitly. **Implemented so far:** +`build-enzyme-wheels` installs the Enzyme-enabled wheel into a clean venv +and runs `tests/test_hessian.py` and `tests/test_external_optimizers.py` +against it, i.e. functional correctness of the shipped wheel, not yet +performance. **Still to do:** - Add a CI check that builds `vmecpp_core` twice, once with `VMECPP_ENABLE_ENZYME=OFF` and once in prebuilt-object mode, and diffs the resulting `.so` symbol table / disassembly for every translation @@ -107,23 +131,35 @@ but should still be checked explicitly: matching within noise. This is what actually answers "does an Enzyme-enabled wheel regress non-adjoint users", as opposed to the HVP-vs-finite-difference comparison the current benchmark performs. -- Run this comparison against an actual wheel built through the proposed - `cibuildwheel` path (not just the local CMake build used by - `test_enzyme.yaml`), since the wheel is the artifact users install. +- Run this comparison against the actual wheels `build_enzyme_wheels.yaml` + produces (not just the local CMake build used by `test_enzyme.yaml`), + since the wheel is the artifact users install. ## Open questions -- Which wheel targets get Enzyme objects at all. `aarch64` and `musllinux` - add real CI matrix cost for the object-build job; starting with - `linux-x86_64` and both macOS architectures (matching where - `test_enzyme.yaml` and the exact-HVP benchmark already run) and - expanding once the mechanism is proven is likely the lowest-risk order. +- Which wheel targets get Enzyme objects at all. `linux-aarch64` and + `musllinux` are not yet built by `build_enzyme_wheels.yaml`; they add + real CI matrix cost for the object-build job (and, for `linux-aarch64`, + its own native-arm64 Clang/Enzyme toolchain build). Extending to them is + the natural next step once the current three targets have run green. +- Whether the Linux objects should instead be compiled inside the actual + `manylinux_2_28` container (e.g. via a `docker run` step using that + image, with LLVM/Clang from `conda-forge` rather than `apt.llvm.org` + since the container has no APT), so the glibc-symbol-version floor + matches the wheel exactly instead of relying on `auditwheel` to catch a + mismatch after the fact. Currently `build_enzyme_wheels.yaml` builds the + objects on the plain `ubuntu-24.04` runner. - Whether prebuilt objects should be committed as versioned release assets (rebuilt on a schedule/on relevant source changes) or built fresh in every wheel release; the former is cheaper per release but needs its own staleness check against `exact_force_jvp.cc` / `exact_force_vjp.cc` content hashes and the pinned `ENZYME_REF`. + `build_enzyme_wheels.yaml` currently rebuilds on every run (relying on + `actions/cache` for the plugin build only). +- When to merge `build_enzyme_wheels.yaml`'s environment variables into + `pypi_publish.yml`'s cibuildwheel step so real releases ship the Enzyme + kernels. Proposed gate: several green runs of this workflow, plus the + performance-regression checks above passing on the produced wheels. - Windows is currently outside `test_enzyme.yaml`'s scope entirely (no - pinned Clang/Enzyme recipe exists for it yet); this proposal does not - attempt to extend Enzyme support there, only to package what already - exists for Linux/macOS. + pinned Clang/Enzyme recipe exists for it yet); this stays out of scope + here too, only packaging what already exists for Linux/macOS. From b08a0db00310df33f9838635150deca4b0665a4a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 7 Sep 2026 08:30:15 +0000 Subject: [PATCH 3/3] Add explicit least-privilege GITHUB_TOKEN permissions to the new workflow Co-authored-by: jurasic-pf <166746189+jurasic-pf@users.noreply.github.com> --- .github/workflows/build_enzyme_wheels.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/build_enzyme_wheels.yaml b/.github/workflows/build_enzyme_wheels.yaml index 28610c0e9..572c52767 100644 --- a/.github/workflows/build_enzyme_wheels.yaml +++ b/.github/workflows/build_enzyme_wheels.yaml @@ -27,6 +27,9 @@ concurrency: group: ${{ github.workflow }}-${{ github.ref || github.run_id }} cancel-in-progress: ${{ github.ref_name != 'main' }} +permissions: + contents: read + env: PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt