Skip to content

Validate MPI reduction of Enzyme gradients#48

Open
krystophny wants to merge 22 commits into
developfrom
ad-14-mpi-enzyme-gradient-reduction
Open

Validate MPI reduction of Enzyme gradients#48
krystophny wants to merge 22 commits into
developfrom
ad-14-mpi-enzyme-gradient-reduction

Conversation

@krystophny

Copy link
Copy Markdown
Member

Depends on #47.
Refs #15.

Summary

  • add an MPI-C shim for Enzyme validation so the test does not depend on compiler-specific Fortran MPI modules
  • add a two-rank Enzyme gradient test where each rank differentiates a local energy contribution and MPI reduces the scalar energy and full gradient outside __enzyme_autodiff
  • register the MPI validation only when MPI C and mpiexec are available

Scope

This validates the AD contract needed for MPI: Enzyme sees only rank-local energy work, and MPI collectives happen outside the differentiated region.

It does not claim full USE_MPI=ON support for the installed LLVM flang stack. On this machine, flang-new cannot configure a Fortran MPI build because the available MPI Fortran modules are not LLVM-flang modules.

Verification

Test fails on main

Before this PR, the MPI Enzyme reduction target is absent on the stacked parent:

$ cmake --build <parent-build> --target gvec_enzyme_energy_gradient_mpi
build target exit=1
ninja: error: unknown target 'gvec_enzyme_energy_gradient_mpi'

The full USE_MPI=ON flang route is also blocked by the local MPI Fortran toolchain, and this PR does not hide that:

$ cmake -S . -B <build-mpi-fortran> -G Ninja -DCMAKE_Fortran_COMPILER=flang-new -DCMAKE_BUILD_TYPE=Release -DGVEC_FIX_HMAP=RZ -DUSE_OPENMP=OFF -DUSE_MPI=ON -DLINK_GVEC_TO_NETCDF=OFF -DUSE_ENZYME=ON
cmake USE_MPI=ON exit=1
-- The Fortran compiler identification is LLVMFlang 22.1.6
-- Found MPI_C: /usr/lib/libmpi.so (found version "3.1")
-- Found MPI_CXX: /usr/lib/libmpi.so (found version "3.1")
-- Could NOT find MPI_Fortran (missing: MPI_Fortran_WORKS)
Could NOT find MPI (missing: MPI_Fortran_FOUND) (found version "3.1")
-- Configuring incomplete, errors occurred!

Test passes after fix

$ cmake -S . -B <build> -G Ninja -DCMAKE_Fortran_COMPILER=flang-new -DCMAKE_BUILD_TYPE=Release -DGVEC_FIX_HMAP=RZ -DUSE_OPENMP=OFF -DLINK_GVEC_TO_NETCDF=OFF -DUSE_ENZYME=ON
-- The Fortran compiler identification is LLVMFlang 22.1.6
-- Configuring done
-- Build files have been written to: <build>

$ cmake --build <build> --target gvec_enzyme_energy_gradient_mpi -j$(nproc)
[1/4] Building C object CMakeFiles/gvec_enzyme_mpi_shim.dir/test-CI/enzyme/gvec_enzyme_mpi_shim.c.o
[2/4] Linking C static library lib/libgvec_enzyme_mpi_shim.a
[3/4] Generating gvec_enzyme_energy_gradient_mpi.enzyme/gvec_enzyme_energy_gradient_mpi.o
[4/4] Linking Fortran executable bin/gvec_enzyme_energy_gradient_mpi

$ ctest --test-dir <build> -R '^gvec_enzyme_energy_gradient_mpi$' --output-on-failure
1/1 Test #4: gvec_enzyme_energy_gradient_mpi ...   Passed
100% tests passed, 0 tests failed out of 1
$ ctest --test-dir <build-all> -R 'gvec_enzyme_(smoke|array_gradient|rz_flat_gradient|energy_gradient_mpi|gvec_module_smoke|ad_fixture_gradient)$' --output-on-failure
1/6 Test #1: gvec_enzyme_smoke .................   Passed
2/6 Test #2: gvec_enzyme_array_gradient ........   Passed
3/6 Test #3: gvec_enzyme_rz_flat_gradient ......   Passed
4/6 Test #4: gvec_enzyme_energy_gradient_mpi ...   Passed
5/6 Test #5: gvec_enzyme_gvec_module_smoke .....   Passed
6/6 Test #6: gvec_enzyme_ad_fixture_gradient ...   Passed
100% tests passed, 0 tests failed out of 6
$ git diff --check
<no output>

$ fo
fo: no Fortran project detected

@krystophny

Copy link
Copy Markdown
Member Author

Added the AD-14 MPI reduction validation slice.

The fix adds a two-rank MPI-C validation target where Enzyme differentiates only a rank-local flat-kernel energy contribution. The scalar energy and full gradient are reduced with MPI outside __enzyme_autodiff, and the reduced result is compared against a serial full-domain Enzyme gradient.

The PR intentionally avoids Fortran MPI modules in the Enzyme path because the local LLVM flang toolchain cannot consume the installed Fortran MPI modules. Full USE_MPI=ON Enzyme support still needs an LLVM-compatible Fortran MPI stack or a larger MPI abstraction change.

@krystophny krystophny force-pushed the ad-14-mpi-enzyme-gradient-reduction branch from 758a3ec to 761c44d Compare July 9, 2026 20:01
@krystophny krystophny force-pushed the ad-14-mpi-enzyme-gradient-reduction branch from 761c44d to 00d866e Compare July 9, 2026 20:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant