enzyme: opt-in Clang/Enzyme build option + AD smoke test

.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"

.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

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()

src/vmecpp/cpp/vmecpp/common/enzyme/enzyme.h

@@ -0,0 +1,52 @@
+// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH
+// <info@proximafusion.com>
+//
+// 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=<ClangEnzyme>. 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 <typename... Args>
+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 <typename Return, typename... Args>
+Return
+__enzyme_fwddiff(  // NOLINT(bugprone-reserved-identifier,readability-identifier-naming)
+    void*, Args...);
+
+#endif  // VMECPP_COMMON_ENZYME_ENZYME_H_

src/vmecpp/cpp/vmecpp/common/enzyme/enzyme_smoke_test.cc

@@ -0,0 +1,97 @@
+// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH
+// <info@proximafusion.com>
+//
+// 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 <Eigen/Dense>
+#include <cmath>
+#include <cstdio>
+
+#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<const VectorXd> xv(x, n);
+  Eigen::Map<const VectorXd> wv(w, n);
+  Eigen::Map<const VectorXd> 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<double>(
+        (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;
+}

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: