[ad-gpu-01] prove Enzyme differentiation for one CUDA physics kernel

[krystophny]

Stage: CUDA autodiff / Enzyme toolchain
Source language: C++20 and CUDA C++, built with Clang CUDA for the Enzyme target
Manuals to read first: PR #561 (VMECPP_USE_CUDA execution path), PR #585 (exact CPU Enzyme QS adjoint stack), docs/cuda_acceleration.md, src/vmecpp/cpp/vmecpp/common/enzyme/enzyme.h, src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/local_force_composition.h, Enzyme README (https://github.com/EnzymeAD/Enzyme), Enzyme GPU paper (https://c.wsmoses.com/papers/EnzymeGPU.pdf), Clang CUDA docs (https://llvm.org/docs/CompileCudaWithLLVM.html)
Depends on: #561, #585

Goal

Add the first opt-in proof that a VMEC++ CUDA physics kernel can be differentiated by Enzyme on the GPU. The landed PR should compile one small CUDA device-kernel JVP or VJP target with ClangEnzyme, run it on device buffers, and validate the derivative against a trusted CPU or finite-difference oracle. It should not try to differentiate the full CUDA equilibrium loop.

Files to edit

• CMakeLists.txt: add VMECPP_ENABLE_CUDA_ENZYME, default OFF; require VMECPP_USE_CUDA=ON, Clang as the CUDA compiler, and VMECPP_ENZYME_PLUGIN.
• src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/CMakeLists.txt: add the opt-in CUDA-Enzyme test target only when VMECPP_ENABLE_CUDA_ENZYME=ON.
• src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/jacobian_kernel.h: factor the per-point half-grid Jacobian arithmetic into a small __host__ __device__ helper without changing the existing CPU loop contract.
• src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/cuda_enzyme_jacobian_jvp.cu: NEW; launch a device-side Enzyme JVP for the half-grid Jacobian helper.
• src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/cuda_enzyme_jacobian_jvp_test.cc: NEW; deterministic GPU test, skipped or not built unless CUDA-Enzyme is enabled.

Syntax or behavior to implement

Add a build-only experimental switch:

-DVMECPP_USE_CUDA=ON \
-DVMECPP_ENABLE_CUDA_ENZYME=ON \
-DVMECPP_ENZYME_PLUGIN="${VMECPP_ENZYME_PLUGIN}" \
-DCMAKE_CXX_COMPILER=clang++-21 \
-DCMAKE_CUDA_COMPILER=clang++-21

The switch must not affect the default CPU build, the default Enzyme CPU build, or the default CUDA build from #561. When enabled, it builds a CUDA-Enzyme test target that differentiates one pure device function. The first target should be the half-grid Jacobian because the CPU branch already treats it as a small, allocation-free Enzyme kernel and because it avoids cuFFT, CUDA graphs, persistent CudaToroidalState, free-boundary NESTOR bridges, and solver control flow.

Behavioral contract:

• Inputs, tangents, outputs, and output tangents live in CUDA device buffers.
• The differentiated function is visible to Enzyme as CUDA device LLVM IR.
• The test compares the GPU-Enzyme tangent for r12, ru12, zu12, rs, zs, and tau against central finite differences of the same helper or against the existing CPU Enzyme Jacobian test oracle.
• The test uses small deterministic dimensions, for example nZnT = 8, three full-grid surfaces, and two half-grid surfaces.
• The test uses tolerances tight enough to catch missing activity markers, wrong shadow strides, and host/device layout drift.

Design notes

PR #561 is a primal CUDA execution backend. It adds an opt-in VMECPP_USE_CUDA path, device-resident iteration state, cuFFT transforms, CUDA graphs, custom reductions, batched configuration axes, device convergence and timestep logic, and host/device bridges for free-boundary NESTOR work. That path is not an Enzyme target today. The only Enzyme content on that branch is the existing CPU smoke-test infrastructure.

The current Enzyme stack is narrower by design. It differentiates flat, allocation-free CPU kernels such as the local force-density composition and QS harmonic maps, then wraps them with analytic linear transforms and adjoint identities. It does not differentiate the OpenMP phase-barrier solver or the whole nonlinear VMEC solve.

Enzyme can differentiate GPU kernels when the kernel computation is visible in LLVM device IR. That makes a Clang CUDA route plausible. It does not make the full CUDA backend automatically differentiable. The following #561 pieces need explicit derivative boundaries or follow-up issues:

• cuFFT and cuBLAS calls: treat transforms as linear maps with hand-coded forward and transpose wrappers.
• CUDA graph capture/replay: keep outside differentiated kernels until there is a proven derivative path.
• Host/device copies and pinned host caches: keep orchestration outside the differentiated region.
• Persistent CudaToroidalState: add separate tangent or adjoint buffers only after one local kernel works.
• Reductions and atomics: define derivative accumulation order and numerical drift policy explicitly.
• Convergence, restart, timestep, and active-configuration masks: keep fixed for local derivative tests.
• Free-boundary NESTOR calls: not included until fixed-boundary GPU local derivatives are validated.

A good sequence after this first PR is:

1. Prove one CUDA-Enzyme local kernel, covered here.
2. Share more per-point physics helpers between the CPU flat kernels and CUDA kernels so the primal formula is not maintained twice.
3. Add device tangent buffers to the CUDA state for fixed-boundary local JVPs.
4. Add hand-coded adjoints for the spectral transforms, matching the CPU exact-HVP chain T^T J_g T.
5. Build a GPU local-force JVP/VJP API for fixed-boundary, symmetric inputs.
6. Wire that API into an opt-in exact HVP or QS-gradient path and benchmark against the CPU Enzyme path.
7. File separate issues for batched mode, free-boundary mode, and any reverse-mode full-solve adjoint.

OpenMP target and OpenACC should not be part of this path. LLVM documents OpenACC support as still under development for Flang and Clang (https://openmp.llvm.org/openacc/Overview.html). The existing VMEC++ OpenMP code is a host synchronization model, not a GPU kernel layout. The CUDA backend is already the project-local GPU execution design.

Scaffold

#if defined(__CUDACC__)
#define VMECPP_HD __host__ __device__
#else
#define VMECPP_HD
#endif

VMECPP_HD void HalfGridJacobianPoint(const JacobianPointInput& x,
                                     JacobianPointOutput* y);

__global__ void k_jacobian_jvp(const JacobianPointInput* x,
                               const JacobianPointInput* dx,
                               JacobianPointOutput* y,
                               JacobianPointOutput* dy) {
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  __enzyme_fwddiff<void>((void*)HalfGridJacobianPoint,
                         enzyme_dup, x + i, dx + i,
                         enzyme_dup, y + i, dy + i);
}

The actual implementation can use a structure-of-arrays layout instead of JacobianPointInput if that better matches ComputeHalfGridJacobian. Keep the differentiated call pure: no cudaMemcpy, no cuFFT, no graph capture, no CudaToroidalState, no host-side Eigen object inside the differentiated device function.

Positive fixtures to add

• cuda_enzyme_jacobian_jvp_test: deterministic small-shape device test. Seed all geometry arrays with nonzero smooth values, seed one tangent direction, launch the GPU-Enzyme JVP, and compare every output tangent to a central-difference oracle with relative error below 1e-8 for regular entries and an absolute floor for near-zero entries.
• cuda_enzyme_jacobian_jvp_test second case: seed only one input block at a time (r1e, r1o, z1e, z1o, rue, ruo, zue, zuo) and verify that inactive input tangents produce zero output tangents where expected.

Negative fixtures to add

• Configure with VMECPP_ENABLE_CUDA_ENZYME=ON and VMECPP_USE_CUDA=OFF: CMake must fail with a clear message.
• Configure with VMECPP_ENABLE_CUDA_ENZYME=ON and a non-Clang CUDA compiler: CMake must fail with a clear message.
• Configure with VMECPP_ENABLE_CUDA_ENZYME=ON and a missing VMECPP_ENZYME_PLUGIN: CMake must fail with the same plugin-path guidance used by the CPU Enzyme path.

Makefile target

No Makefile target. Add a CTest target named cuda_enzyme_jacobian_jvp and a build target named cuda_enzyme_jacobian_jvp_test.

Success criteria

cmake -S . -B build-cuda-enzyme -G Ninja \
  -DVMECPP_USE_CUDA=ON \
  -DVMECPP_ENABLE_CUDA_ENZYME=ON \
  -DVMECPP_ENZYME_PLUGIN="${VMECPP_ENZYME_PLUGIN}" \
  -DCMAKE_CXX_COMPILER=clang++-21 \
  -DCMAKE_CUDA_COMPILER=clang++-21
cmake --build build-cuda-enzyme --target cuda_enzyme_jacobian_jvp_test -j"$(nproc)"
ctest --test-dir build-cuda-enzyme -R cuda_enzyme_jacobian_jvp --output-on-failure

Non-goals

• Do not differentiate the full CUDA iteration body.
• Do not add OpenMP target or OpenACC support.
• Do not differentiate cuFFT, cuBLAS, CUDA graph capture, host/device transfers, convergence gates, timestep control, restart logic, or NESTOR in this PR.
• Do not change the default VMECPP_USE_CUDA behavior or performance path.
• Do not add a second copy of the Jacobian formula. Share or factor the existing kernel arithmetic.
• Do not require CUDA-Enzyme in CI unless a GPU runner with the matching ClangEnzyme plugin is available.

Verification

export VMECPP_ENZYME_PLUGIN="$(find /tmp/enzyme-build -name 'ClangEnzyme-21.so' | head -1)"
test -n "${VMECPP_ENZYME_PLUGIN}"
cmake -S . -B build-cuda-enzyme -G Ninja -DVMECPP_USE_CUDA=ON -DVMECPP_ENABLE_CUDA_ENZYME=ON -DVMECPP_ENZYME_PLUGIN="${VMECPP_ENZYME_PLUGIN}" -DCMAKE_CXX_COMPILER=clang++-21 -DCMAKE_CUDA_COMPILER=clang++-21
cmake --build build-cuda-enzyme --target cuda_enzyme_jacobian_jvp_test -j"$(nproc)"
ctest --test-dir build-cuda-enzyme -R cuda_enzyme_jacobian_jvp --output-on-failure