diff --git a/.github/workflows/benchmarks.yaml b/.github/workflows/benchmarks.yaml index c2f3fdf4e..54326ec6a 100644 --- a/.github/workflows/benchmarks.yaml +++ b/.github/workflows/benchmarks.yaml @@ -36,7 +36,11 @@ jobs: pytest benchmarks/test_benchmarks.py --benchmark-json=benchmark_results.json - name: Compare benchmark result - if: github.event_name == 'pull_request' + # Skip the result upload/compare for fork PRs: their GITHUB_TOKEN is + # read-only, so comment-on-alert/auto-push hit 'Resource not accessible + # by integration'. The benchmarks still run above; only the write-back + # is skipped for forks. + if: github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name == github.repository uses: benchmark-action/github-action-benchmark@v1.21.0 with: tool: "pytest" diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index a22e0cbdc..0f68ab6e9 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -50,15 +50,33 @@ jobs: run: | # install VMEC++ deps as well as VMEC2000 deps (we need to import VMEC2000 in a test) sudo apt-get update && sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev + - name: Cache the VMEC2000 wheel + if: ${{ matrix.os == 'ubuntu-22.04' && matrix.python-version == '3.10' }} + uses: actions/cache@v4 + with: + path: /tmp/vmec2000-wheel + # Keyed on the pinned VMEC2000 source commit: rebuild only when it changes. + key: vmec2000-728af8bd6c79-cp310-ubuntu22.04 - name: Also install VMEC2000 (only on Ubuntu 22.04) if: ${{ matrix.os == 'ubuntu-22.04' && matrix.python-version == '3.10' }} run: | - # mpi4py is needed for VMEC2000 - sudo apt-get install -y libopenmpi-dev + sudo apt-get install -y libopenmpi-dev libnetcdff-dev libscalapack-mpi-dev \ + libopenblas-dev ninja-build python -m pip install mpi4py - # custom wheel for VMEC2000, needed for some VMEC++/VMEC2000 compatibility tests - # NOTE: this wheel is only guaranteed to work on Ubuntu 22.04 - python -m pip install vmec@https://anaconda.org/eguiraud-pf/vmec/0.0.6/download/vmec-0.0.6-cp310-cp310-linux_x86_64.whl + if ! ls /tmp/vmec2000-wheel/vmec-*.whl >/dev/null 2>&1; then + # Build with the SYSTEM cmake + ninja; do NOT pip-install the + # cmake/ninja wheels (they shadow the system cmake that + # scikit-build-core's editable vmecpp rebuild records). + python -m pip install numpy scikit-build f90wrap setuptools wheel + git clone https://github.com/hiddenSymmetries/VMEC2000.git /tmp/VMEC2000 + git -C /tmp/VMEC2000 checkout 728af8bd6c796b36a0aa85fe298e507791e57c6e + cp /tmp/VMEC2000/cmake/machines/ubuntu.json /tmp/VMEC2000/cmake_config_file.json + LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu \ + python -m pip wheel /tmp/VMEC2000 --no-build-isolation -w /tmp/vmec2000-wheel + fi + python -m pip install /tmp/vmec2000-wheel/vmec-*.whl + # fail loudly here if the binding is still broken, not in the test step + python -c "import vmec; print('VMEC2000 import OK')" - name: Install package run: | # on Ubuntu we would not need this, but on MacOS we need to point CMake to gfortran-14 and gcc-14 diff --git a/CMakeLists.txt b/CMakeLists.txt index 9a65418c6..e6034bef0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -66,10 +66,9 @@ find_package(LAPACK REQUIRED) FetchContent_Declare( abseil-cpp GIT_REPOSITORY https://github.com/abseil/abseil-cpp.git - # 20260107.1 LTS: required for Clang >= 21, where the 2024 pin fails to - # compile (absl::Nonnull SFINAE in absl/strings/ascii.cc). Clang is the - # compiler used for the Enzyme autodiff build. - GIT_TAG 20260107.1 + # 20260107.1 LTS: older abseil fails to compile under Clang >= 21 (the + # Enzyme build) on absl::Nonnull SFINAE in absl/strings/ascii.cc. + GIT_TAG 255c84dadd029fd8ad25c5efb5933e47beaa00c7 GIT_SHALLOW TRUE ) FetchContent_Declare( @@ -175,3 +174,34 @@ endif() install(TARGETS _vmecpp LIBRARY DESTINATION vmecpp/cpp/.) install(TARGETS indata2json DESTINATION vmecpp/cpp/third_party/indata2json/) + +# Optional Enzyme automatic-differentiation targets. +# Enzyme (https://enzyme.mit.edu) differentiates LLVM IR via a Clang plugin, so +# these targets need a Clang frontend and the matching ClangEnzyme plugin. OFF +# by default; the production build and all existing targets are unaffected. +# Enable with: +# -DVMECPP_ENABLE_ENZYME=ON -DVMECPP_ENZYME_PLUGIN=/path/to/ClangEnzyme-NN.so +option(VMECPP_ENABLE_ENZYME "Build Enzyme autodiff targets" OFF) +if(VMECPP_ENABLE_ENZYME) + if(NOT CMAKE_CXX_COMPILER_ID MATCHES "Clang") + message(FATAL_ERROR + "VMECPP_ENABLE_ENZYME requires a Clang compiler (got " + "${CMAKE_CXX_COMPILER_ID}); Enzyme attaches as a Clang plugin.") + endif() + set(VMECPP_ENZYME_PLUGIN "" CACHE FILEPATH "Path to ClangEnzyme-NN.so") + if(NOT VMECPP_ENZYME_PLUGIN OR NOT EXISTS "${VMECPP_ENZYME_PLUGIN}") + message(FATAL_ERROR + "VMECPP_ENABLE_ENZYME=ON requires " + "-DVMECPP_ENZYME_PLUGIN=/path/to/ClangEnzyme-NN.so") + endif() + message(STATUS "Enzyme plugin: ${VMECPP_ENZYME_PLUGIN}") + enable_testing() + add_executable(enzyme_smoke_test + ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/common/enzyme/enzyme_smoke_test.cc) + # Enzyme runs as an optimization-time pass; it needs optimizations enabled and + # the plugin attached. -O2 guards against a Debug (-O0) configuration where the + # AD pass would otherwise not run. + target_compile_options(enzyme_smoke_test PRIVATE + -O2 -fplugin=${VMECPP_ENZYME_PLUGIN}) + add_test(NAME enzyme_smoke COMMAND enzyme_smoke_test) +endif() diff --git a/src/vmecpp/cpp/vmecpp/common/enzyme/enzyme.h b/src/vmecpp/cpp/vmecpp/common/enzyme/enzyme.h new file mode 100644 index 000000000..6010b055c --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/common/enzyme/enzyme.h @@ -0,0 +1,52 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT +#ifndef VMECPP_COMMON_ENZYME_ENZYME_H_ +#define VMECPP_COMMON_ENZYME_ENZYME_H_ + +// Thin declarations for the Enzyme automatic-differentiation intrinsics. +// +// Enzyme (https://enzyme.mit.edu) differentiates at the LLVM-IR level via a +// Clang plugin (ClangEnzyme-NN.so). These intrinsics are resolved by that +// plugin; without it the symbols do not link, so any translation unit that +// calls them must be compiled with -fplugin=. The CMake option +// VMECPP_ENABLE_ENZYME wires that flag and is OFF by default. +// +// Differentiation activity is selected per argument with the marker globals +// below: pass `enzyme_dup, primal_ptr, shadow_ptr` for an active buffer (the +// gradient accumulates into shadow_ptr) and `enzyme_const, value` for an input +// held fixed. Enzyme matches these markers by symbol name. +// +// Constraint that shapes how differentiable kernels here are written: Enzyme's +// allocation analysis does not track Eigen's aligned allocator, so a heap +// temporary from an Eigen expression (e.g. a dynamic-size `A * x`) crossing the +// differentiated call aborts with "freeing without malloc". Differentiable +// kernels therefore operate on caller-owned buffers via Eigen::Map and avoid +// allocating expression temporaries on the differentiated path. + +// The marker globals and the __enzyme_* intrinsic names are part of the Enzyme +// ABI: the plugin matches them by exact symbol name, so they cannot be const or +// renamed to satisfy the in-tree naming/identifier lint rules. + +// NOLINTBEGIN(cppcoreguidelines-avoid-non-const-global-variables) +extern int enzyme_dup; +extern int enzyme_const; +extern int enzyme_dupnoneed; +extern int enzyme_out; +// NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) + +// Reverse mode: returns nothing useful here; gradients land in the shadow +// buffers passed alongside each `enzyme_dup` argument. +template +void __enzyme_autodiff( // NOLINT(bugprone-reserved-identifier,readability-identifier-naming) + void*, Args...); + +// Forward mode: propagates the seed in the shadow argument to the directional +// derivative of the result. +template +Return +__enzyme_fwddiff( // NOLINT(bugprone-reserved-identifier,readability-identifier-naming) + void*, Args...); + +#endif // VMECPP_COMMON_ENZYME_ENZYME_H_ diff --git a/src/vmecpp/cpp/vmecpp/common/enzyme/enzyme_smoke_test.cc b/src/vmecpp/cpp/vmecpp/common/enzyme/enzyme_smoke_test.cc new file mode 100644 index 000000000..5da2fac17 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/common/enzyme/enzyme_smoke_test.cc @@ -0,0 +1,97 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT + +// Toolchain smoke test for the Enzyme autodiff build. +// +// It checks the two properties the differentiable VMEC++ kernels rely on: +// 1. The ClangEnzyme plugin is attached and resolves the autodiff intrinsics. +// 2. Enzyme differentiates a scalar objective expressed over Eigen::Map'd +// caller buffers (the pattern all differentiable kernels here use). +// +// The objective is f(x) = sum_i 0.5 * w_i * x_i^2 + c_i * x_i, with the closed +// form gradient df/dx_i = w_i * x_i + c_i. Reverse- and forward-mode Enzyme +// gradients are checked against that closed form and against central finite +// differences. Exit code 0 on success, 1 on any mismatch. + +#include +#include +#include + +#include "vmecpp/common/enzyme/enzyme.h" + +namespace { + +using Eigen::VectorXd; + +// Quadratic objective over caller-owned buffers. No allocating Eigen +// expression temporaries cross the differentiated boundary. +__attribute__((noinline)) double Objective(const double* x, const double* w, + const double* c, int n) { + Eigen::Map xv(x, n); + Eigen::Map wv(w, n); + Eigen::Map cv(c, n); + double sum = 0.0; + for (int i = 0; i < n; ++i) { + sum += 0.5 * wv[i] * xv[i] * xv[i] + cv[i] * xv[i]; + } + return sum; +} + +double MaxAbsDiff(const VectorXd& a, const VectorXd& b) { + return (a - b).cwiseAbs().maxCoeff(); +} + +} // namespace + +int main() { + const int n = 8; + VectorXd x = VectorXd::LinSpaced(n, 1.0, n); + VectorXd w = VectorXd::LinSpaced(n, 2.0, 2.0 + 0.5 * (n - 1)); + VectorXd c = VectorXd::Constant(n, 0.25); + + VectorXd analytic(n); + for (int i = 0; i < n; ++i) analytic[i] = w[i] * x[i] + c[i]; + + // Reverse mode: gradient accumulates into the shadow buffer. + VectorXd g_rev = VectorXd::Zero(n); + __enzyme_autodiff((void*)Objective, enzyme_dup, x.data(), g_rev.data(), + enzyme_const, w.data(), enzyme_const, c.data(), + enzyme_const, n); + + // Forward mode: one directional derivative per coordinate seed. + VectorXd g_fwd = VectorXd::Zero(n); + for (int j = 0; j < n; ++j) { + VectorXd seed = VectorXd::Zero(n); + seed[j] = 1.0; + g_fwd[j] = __enzyme_fwddiff( + (void*)Objective, enzyme_dup, x.data(), seed.data(), enzyme_const, + w.data(), enzyme_const, c.data(), enzyme_const, n); + } + + // Central finite differences. + VectorXd g_fd = VectorXd::Zero(n); + const double h = 1e-6; + for (int j = 0; j < n; ++j) { + VectorXd xp = x, xm = x; + xp[j] += h; + xm[j] -= h; + g_fd[j] = (Objective(xp.data(), w.data(), c.data(), n) - + Objective(xm.data(), w.data(), c.data(), n)) / + (2.0 * h); + } + + const double err_rev = MaxAbsDiff(g_rev, analytic); + const double err_fwd = MaxAbsDiff(g_fwd, analytic); + const double err_fd = MaxAbsDiff(g_rev, g_fd); + + std::printf("enzyme smoke test (n=%d)\n", n); + std::printf(" max|reverse - analytic| = %.3e\n", err_rev); + std::printf(" max|forward - analytic| = %.3e\n", err_fwd); + std::printf(" max|reverse - finite-diff| = %.3e\n", err_fd); + + const bool ok = err_rev < 1e-10 && err_fwd < 1e-10 && err_fd < 1e-5; + std::printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} diff --git a/tests/test_simsopt_compat.py b/tests/test_simsopt_compat.py index 609dffea2..1443d8cbc 100644 --- a/tests/test_simsopt_compat.py +++ b/tests/test_simsopt_compat.py @@ -303,6 +303,11 @@ def test_ensure_vmec2000_input_from_vmecpp_input(): if varname[1:-1] == "axis_": # these are called differently in VMEC2000, e.g. raxis_c -> raxis_cc varname_vmec2000 = f"{varname[:-1]}c{varname[-1]}" + if not hasattr(vmec2000.indata, varname_vmec2000): + # vmecpp-only field (e.g. free_boundary_method) with no counterpart + # in the legacy VMEC2000 INDATA namelist; not part of the common + # subset under test. + continue vmec2000_var = getattr(vmec2000.indata, varname_vmec2000) if vmecpp_var is None: