Exact autodiff Hessian-vector product for the VMEC force

.github/actions/run-benchmarks/action.yml

@@ -53,6 +53,14 @@ runs:
         export OMP_WAIT_POLICY=passive
         pytest benchmarks/test_benchmarks.py --benchmark-json=benchmark_results.json
 
+    # github-action-benchmark's "pytest" tool always plots pytest-benchmark's
+    # `ops` field as iter/sec, but the C++ suite emits seconds.
+    # We convert both to the same format, so plotting is the same.
+    - name: Convert Python benchmark result to seconds
+      shell: bash
+      run: |
+        python3 benchmarks/to_seconds_json.py pytest benchmark_results.json benchmark_results_seconds.json
+
     - name: Store / compare Python benchmark result
       # Skip the result upload/compare for fork PRs: their GITHUB_TOKEN is
       # read-only, so comment-on-alert/auto-push hit 'Resource not accessible
@@ -61,8 +69,8 @@ runs:
       if: ${{ inputs.auto-push == 'true' || github.event.pull_request.head.repo.full_name == github.repository }}
       uses: benchmark-action/github-action-benchmark@v1.21.0
       with:
-        tool: "pytest"
-        output-file-path: benchmark_results.json
+        tool: "customSmallerIsBetter"
+        output-file-path: benchmark_results_seconds.json
         gh-pages-branch: benchmark-runs
         benchmark-data-dir-path: benchmarks
         alert-threshold: "150%"
@@ -76,7 +84,7 @@ runs:
     # Builds the four benchmark cc_binary targets, runs each with
     # --benchmark_format=json, and merges their output into a single
     # cpp_benchmark_results.json (schema: {"context": ..., "benchmarks": [...]})
-    # for the store step below. Entries with "error_occurred" (e.g. an
+    # for the conversion step below. Entries with "error_occurred" (e.g. an
     # FFT-path benchmark skipped because no FFTX codelet exists for that
     # resolution) are dropped: they carry no valid timing to track.
 
@@ -142,6 +150,14 @@ runs:
               f"({skipped} skipped/errored).")
         PY
 
+    # github-action-benchmark's "pytest" tool always plots pytest-benchmark's
+    # `ops` field as iter/sec, but the C++ suite emits seconds.
+    # We convert both to the same format, so plotting is the same.
+    - name: Convert C++ benchmark result to seconds
+      shell: bash
+      run: |
+        python3 benchmarks/to_seconds_json.py googlecpp cpp_benchmark_results.json cpp_benchmark_results_seconds.json
+
     - name: Store / compare C++ benchmark result
       # Same fork-PR skip rationale as the Python store step above.
       #
@@ -160,8 +176,8 @@ runs:
       uses: benchmark-action/github-action-benchmark@v1.21.0
       with:
         name: "C++ Microbenchmarks"
-        tool: "googlecpp"
-        output-file-path: cpp_benchmark_results.json
+        tool: "customSmallerIsBetter"
+        output-file-path: cpp_benchmark_results_seconds.json
         gh-pages-branch: benchmark-runs
         benchmark-data-dir-path: benchmarks
         skip-fetch-gh-pages: true

.github/workflows/full_validation.yaml

@@ -0,0 +1,70 @@
+name: Full V&V against reference VMEC
+
+# Runs the full Verification & Validation suite from proximafusion/vmecpp-validation
+# (all ~219 input configurations against reference Fortran VMEC2000), using the
+# vmecpp version built from this commit rather than the version pinned on PyPI/GitHub.
+# This is a lot more exhaustive (and slower) than the "short" validation that
+# vmecpp-validation itself runs on its own PRs, so it is not run on every push here.
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: ${{ github.ref_name != 'main' }}
+
+jobs:
+  full-validation:
+    name: Run full VMEC++ validation
+    runs-on: ubuntu-22.04
+    # The full parameter scan runs ~219 configurations, each computing a reference
+    # wout with Fortran VMEC2000 in Docker plus one with VMEC++: budget generously.
+    timeout-minutes: 360
+    steps:
+      - name: Check out VMEC++
+        uses: actions/checkout@v4
+        with:
+          path: vmecpp
+          lfs: true
+
+      - name: Check out vmecpp-validation
+        uses: actions/checkout@v4
+        with:
+          repository: proximafusion/vmecpp-validation
+          path: vmecpp-validation
+          lfs: true
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Install required system packages for Ubuntu
+        run: |
+          sudo apt-get update && sudo apt-get install -y \
+            build-essential cmake libnetcdf-dev liblapack-dev liblapacke-dev libopenmpi-dev \
+            libomp-dev libeigen3-dev nlohmann-json3-dev libhdf5-dev
+
+      - name: Install vmecpp-validation's Python requirements
+        run: |
+          cd vmecpp-validation
+          # Install everything except vmecpp itself: we build and install vmecpp
+          # from this commit below instead of the version pinned in requirements.txt.
+          grep -v '^vmecpp@' requirements.txt > requirements.no-vmecpp.txt
+          python -m pip install -r requirements.no-vmecpp.txt
+
+      - name: Install VMEC++ from this commit
+        run: python -m pip install ./vmecpp
+
+      - name: Run full validation
+        working-directory: vmecpp-validation
+        run: python validate_vmec.py
+
+      - name: Upload V&V results
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: vnvresults
+          path: vmecpp-validation/vnvresults_*/
+          retention-days: 30

.github/workflows/test_enzyme.yaml

@@ -32,7 +32,8 @@ jobs:
           sudo ./llvm.sh "${LLVM_VERSION}"
           sudo apt-get install -y \
             "llvm-${LLVM_VERSION}-dev" "clang-${LLVM_VERSION}" "libclang-${LLVM_VERSION}-dev" \
-            cmake ninja-build
+            "libomp-${LLVM_VERSION}-dev" \
+            cmake gfortran libhdf5-dev liblapack-dev libnetcdf-dev ninja-build python3-dev
 
       - name: Cache the ClangEnzyme plugin build
         id: enzyme-cache
@@ -75,3 +76,29 @@ jobs:
             --copt="-fplugin=${ENZYME_PLUGIN}" \
             --test_output=all \
             //vmecpp/common/enzyme:enzyme_smoke_test
+
+      - uses: astral-sh/setup-uv@v6
+
+      - name: Run Enzyme force-derivative integration tests
+        env:
+          CC: clang-${{ env.LLVM_VERSION }}
+          CXX: clang++-${{ env.LLVM_VERSION }}
+          CMAKE_ARGS: >-
+            -DVMECPP_ENABLE_ENZYME=ON
+            -DVMECPP_ENZYME_PLUGIN=${{ env.ENZYME_PLUGIN }}
+        run: |
+          uv sync --reinstall-package vmecpp
+          uv run --with pytest pytest -q \
+            tests/test_hessian.py \
+            tests/test_hessian_transpose.py \
+            tests/test_external_optimizers.py
+
+      - name: Run exact-HVP benchmark
+        run: |
+          uv run python benchmarks/benchmark_exact_hvp.py | tee exact_hvp_benchmark.txt
+          {
+            echo "### Exact-HVP benchmark"
+            echo '```text'
+            cat exact_hvp_benchmark.txt
+            echo '```'
+          } >> "${GITHUB_STEP_SUMMARY}"

CMakeLists.txt

@@ -185,26 +185,65 @@ install(TARGETS indata2json DESTINATION vmecpp/cpp/third_party/indata2json/)
 #   -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")
+  if(NOT CMAKE_CXX_COMPILER_ID MATCHES "Clang|IntelLLVM")
     message(FATAL_ERROR
-      "VMECPP_ENABLE_ENZYME requires a Clang compiler (got "
-      "${CMAKE_CXX_COMPILER_ID}); Enzyme attaches as a Clang plugin.")
+      "VMECPP_ENABLE_ENZYME requires a Clang-based compiler (got "
+      "${CMAKE_CXX_COMPILER_ID}); Enzyme attaches as a Clang plugin. "
+      "Intel oneAPI icx/icpx (IntelLLVM) is Clang-based and supported.")
   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")
+  # The Enzyme pass is an LLVM-version-specific plugin. icx/icpx (IntelLLVM)
+  # cannot load a plugin built against a different upstream LLVM. Enzyme's output
+  # is an ordinary object file, though, so for such compilers the single Enzyme
+  # TU (exact_force_jvp.cc) can be precompiled with a matching ClangEnzyme and
+  # linked in via VMECPP_ENZYME_OBJECT, while the rest is built with icx.
+  set(VMECPP_ENZYME_OBJECT "" CACHE FILEPATH
+    "Prebuilt exact_force_jvp object (compiled elsewhere with a matching ClangEnzyme)")
+  if(VMECPP_ENZYME_OBJECT)
+    if(NOT EXISTS "${VMECPP_ENZYME_OBJECT}")
+      message(FATAL_ERROR "VMECPP_ENZYME_OBJECT not found: ${VMECPP_ENZYME_OBJECT}")
+    endif()
+    message(STATUS "Enzyme: linking prebuilt object ${VMECPP_ENZYME_OBJECT}")
+    set_source_files_properties("${VMECPP_ENZYME_OBJECT}" PROPERTIES
+      EXTERNAL_OBJECT TRUE GENERATED TRUE)
+    target_sources(vmecpp_core PRIVATE "${VMECPP_ENZYME_OBJECT}")
+    target_compile_definitions(vmecpp_core PUBLIC VMECPP_ENABLE_ENZYME)
+  else()
+    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 "
+        "(or -DVMECPP_ENZYME_OBJECT=/path/to/exact_force_jvp.o for icx)")
+    endif()
+    message(STATUS "Enzyme plugin: ${VMECPP_ENZYME_PLUGIN}")
+    enable_testing()
+    # Exact autodiff (forward and reverse) of a real VMEC nonlinear kernel: the
+    # half-grid Jacobian. Enzyme runs as an optimization-time pass, so it needs
+    # -O2 and the plugin attached.
+    add_executable(jacobian_kernel_autodiff_test
+      ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/common/enzyme/jacobian_kernel_autodiff_test.cc)
+    target_compile_options(jacobian_kernel_autodiff_test PRIVATE
+      -O2 -fplugin=${VMECPP_ENZYME_PLUGIN})
+    add_test(NAME jacobian_kernel_autodiff COMMAND jacobian_kernel_autodiff_test)
+    # Exact Hessian of the composed local force map (all six force-chain kernels):
+    # forward/reverse Jacobian vs finite differences, and JVP cost vs FD-HVP.
+    add_executable(local_force_hessian_test
+      ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/common/enzyme/local_force_hessian_test.cc
+      ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc)
+    target_compile_options(local_force_hessian_test PRIVATE
+      -O2 -fplugin=${VMECPP_ENZYME_PLUGIN})
+    add_test(NAME local_force_hessian COMMAND local_force_hessian_test)
+
+    # Wire the exact force Hessian-vector product into the core library and the
+    # Python extension. The Enzyme translation units are compiled with the plugin;
+    # the rest of the core stays normally compiled.
+    target_sources(vmecpp_core PRIVATE
+      ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc
+      ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.cc)
+    set_source_files_properties(
+      ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc
+      ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.cc
+      PROPERTIES COMPILE_OPTIONS "-O2;-fplugin=${VMECPP_ENZYME_PLUGIN}")
+    target_compile_definitions(vmecpp_core PUBLIC VMECPP_ENABLE_ENZYME)
   endif()
-  message(STATUS "Enzyme plugin: ${VMECPP_ENZYME_PLUGIN}")
-  enable_testing()
-  # Self-contained over flat buffers; checks Jv and J^T u against finite
-  # differences and against each other (adjoint identity). Enzyme runs as an
-  # optimization-time pass, so it needs -O2 and the plugin attached. -O2 also
-  # guards against a Debug (-O0) configuration where the AD pass would not run.
-  add_executable(jacobian_kernel_autodiff_test
-    ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/common/enzyme/jacobian_kernel_autodiff_test.cc)
-  target_compile_options(jacobian_kernel_autodiff_test PRIVATE
-    -O2 -fplugin=${VMECPP_ENZYME_PLUGIN})
-  add_test(NAME jacobian_kernel_autodiff COMMAND jacobian_kernel_autodiff_test)
 endif()

README.md

@@ -19,6 +19,7 @@
 
 [![CI](https://github.com/proximafusion/vmecpp/actions/workflows/tests.yaml/badge.svg)](https://github.com/proximafusion/vmecpp/actions/workflows/tests.yaml)
 [![C++ core tests](https://github.com/proximafusion/vmecpp/actions/workflows/test_bazel.yaml/badge.svg)](https://github.com/proximafusion/vmecpp/actions/workflows/test_bazel.yaml)
+[![Full V&V against reference VMEC](https://github.com/proximafusion/vmecpp/actions/workflows/full_validation.yaml/badge.svg)](https://github.com/proximafusion/vmecpp/actions/workflows/full_validation.yaml)
 [![Publish wheels to PyPI](https://github.com/proximafusion/vmecpp/actions/workflows/pypi_publish.yml/badge.svg)](https://github.com/proximafusion/vmecpp/actions/workflows/pypi_publish.yml)
 
 VMEC++ is a Python-friendly, from-scratch reimplementation in C++ of the Variational Moments Equilibrium Code (VMEC),
@@ -315,6 +316,10 @@ The single-thread runtimes as well as the contents of the "wout" file produced b
 The full validation test can be found at https://github.com/proximafusion/vmecpp-validation, including a set of sensible input configurations,
 parameter scan values and tolerances that make the comparison pass. See that repo for more information.
 
+This full validation (~219 input configurations) is run against every commit to `main` and can also be triggered
+on demand from the [Full V&V against reference VMEC](https://github.com/proximafusion/vmecpp/actions/workflows/full_validation.yaml) workflow
+(click "Run workflow"). It builds `vmecpp` from the corresponding commit rather than using the version pinned by `vmecpp-validation`.
+
 ## Differences with respect to PARVMEC/VMEC2000
 
 VMEC++:

benchmarks/benchmark_exact_hvp.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH
+# <info@proximafusion.com>
+#
+# SPDX-License-Identifier: MIT
+"""Benchmark exact autodiff Hessian-vector products as an internal solver.
+
+Compare preconditioned JFNK, finite-difference Newton-HVP, and exact-autodiff
+Newton-HVP on the same VMEC++ residual. Report force evaluations, iterations,
+wall time, residual norm, and energy error. Requires an Enzyme-enabled build.
+
+Run: python benchmarks/benchmark_exact_hvp.py
+"""
+
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "examples"))
+import external_optimizers as eo  # type: ignore  # noqa: E402
+
+DATA = ROOT / "examples" / "data"
+CASES = [
+    ("solovev", DATA / "solovev.json", 11),
+    ("cth_like", DATA / "cth_like_fixed_bdy.json", 11),
+]
+
+SOLVERS = [
+    ("precond JFNK", eo.solve_newton_krylov_preconditioned),
+    ("Newton FD-HVP + M^-1", eo.solve_newton_hvp),
+    ("Newton exact-HVP + M^-1", eo.solve_newton_exact_hvp),
+]
+
+
+def main():
+    for case, path, ns in CASES:
+        print(f"\n=== {case} (ns={ns}) ===")
+        _, wref = eo.reference_equilibrium(path, ns)
+        print(f"native solve: W = {wref:.8e}")
+        print(
+            f"{'optimizer':28s} {'F-evals':>8s} {'iters':>6s} {'time[s]':>8s} "
+            f"{'||F||':>10s} {'dW':>10s}"
+        )
+        for label, solver in SOLVERS:
+            result = solver(path, ns)[1]
+            print(
+                f"{label:28s} {result.force_evals:8d} {result.outer_iters:6d} "
+                f"{result.seconds:8.2f} {result.residual_norm:10.1e} "
+                f"{abs(result.energy - wref):10.1e}"
+            )
+
+
+if __name__ == "__main__":
+    main()