enzyme: opt-in Clang/Enzyme build option + AD smoke test#6
Draft
krystophny wants to merge 11 commits into
Draft
Conversation
Add VMECPP_ENABLE_ENZYME (OFF by default), which requires a Clang compiler and a ClangEnzyme plugin path and builds a self-contained autodiff smoke test. The test differentiates a scalar objective written over Eigen::Map'd caller buffers and checks reverse- and forward-mode Enzyme gradients against the closed form and central finite differences. enzyme.h documents the intrinsic ABI and the allocation constraint that shapes the differentiable kernels: Enzyme cannot track Eigen's aligned allocator, so differentiable paths use Eigen::Map over caller-owned buffers and avoid heap expression temporaries. With the option off the build is unchanged.
The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing.
The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure.
With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one.
…mit pin Bring this stack branch up to the corrected CI baseline (from proximafusion#583/proximafusion#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag.
Move the Enzyme autodiff smoke test into the bazel test framework, which
owns every other C++ test in this repository, and drop the separate CMake
ctest path that nothing in CI exercised.
- vmecpp/common/enzyme/BUILD.bazel: an `enzyme` header library plus an
`enzyme_smoke_test` cc_test. The test is tagged `manual` so the default
GCC `bazel test //...` skips it (the Enzyme intrinsics only resolve under
Clang with the plugin attached) and never tries to compile it with GCC.
- .bazelrc: a `--config=enzyme` that sets -O2 so the Enzyme optimization
pass fires. Select Clang with CC/CXX and pass the plugin path the way
-DVMECPP_ENZYME_PLUGIN did under CMake:
CC=clang CXX=clang++ bazel test --config=enzyme \
--copt=-fplugin=/path/to/ClangEnzyme-NN.so \
//vmecpp/common/enzyme:enzyme_smoke_test
- CMakeLists.txt: remove the VMECPP_ENABLE_ENZYME option and the ctest
registration it only existed to drive.
Add a GitHub Actions job that gives the Enzyme autodiff smoke test actual CI coverage. It mirrors the EnzymeAD upstream recipe: install Clang/LLVM 21 from apt.llvm.org, build a pinned ClangEnzyme-21 plugin (v0.0.264, the version this stack is developed against) against the installed LLVM and Clang, then run the bazel target under --config=enzyme with the plugin attached. The plugin build is cached on the pinned ref so only the first run pays for it. This is what the enzyme test needed beyond the bazel move: the default GCC test_bazel job skips the manual-tagged target, so without a Clang/Enzyme job nothing exercised it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Add an opt-in Clang/Enzyme build path and a toolchain smoke test:
VMECPP_ENABLE_ENZYMECMake option, off by default. When on, it requiresa Clang compiler and
-DVMECPP_ENZYME_PLUGIN=/path/to/ClangEnzyme-NN.so, andbuilds the
enzyme_smoketest target.common/enzyme/enzyme.h: thin declarations of the Enzyme autodiff intrinsicsand the activity markers, with the allocation constraint that shapes every
differentiable kernel in this stack documented in place.
common/enzyme/enzyme_smoke_test.cc: differentiates a scalar objective overEigen::Map'd caller buffers and checks reverse- and forward-mode gradientsagainst the closed form and central finite differences.
Why
Stacked on #5. Enzyme differentiates LLVM IR through a Clang plugin, so the
differentiable VMEC++ work needs a Clang build and a way to attach the plugin.
This patch adds that switch and a test that fails loudly if the plugin is not
attached or if Enzyme cannot differentiate the
Eigen::Mapbuffer pattern thelater kernels depend on. It adds no production code; with the option off, the
build is byte-for-byte the previous one.
The smoke test also pins down the one hard Enzyme constraint for this codebase:
Enzyme's allocation analysis does not track Eigen's aligned allocator, so a
dynamic-size Eigen heap temporary crossing the differentiated call aborts with
"freeing without malloc". Differentiable kernels therefore operate on
caller-owned buffers via
Eigen::Map. The test exercises exactly that pattern.Verification
Configure and build with Clang 21.1.8 and the ClangEnzyme-21 plugin:
ctest:Smoke test output (reverse and forward both exact, agree with finite
differences):
No regression with the option off. A default configure (GCC, no Enzyme flags)
succeeds and emits no
enzyme_smoketarget: