Repository navigation
Conversation
Add native Fortran support to Enzyme-JAX through fortran_call, enabling seamless execution, JIT compilation, and automatic differentiation (forward-mode JVP, reverse-mode VJP/grad) of Fortran routines directly within JAX. Key changes: - In-process Flang compilation: Compile Fortran source directly to LLVM IR using LLVM Flang frontend (Fortran::frontend::CompilerInstance / CompilerInvocation) and FIR pass pipeline (createMLIRToLLVMPassPipeline), with in-memory standard intrinsic module headers (iso_c_binding, iso_fortran_env, ieee_arithmetic, etc.). - ABI and calling convention mapping: Analyze FIR signatures to automatically handle: * Raw pointers (void*) for standard dummy arrays/references. * Pass-by-value (VALUE attribute) scalar parameters. * Standard ISO Fortran descriptors (CFI_cdesc_t) for assumed-shape arrays, allocatables, and pointer dummy arguments. * Reverse multi-dimensional array bounds and byte strides to map row-major JAX arrays to Fortran column-major layout seamlessly with zero-copy. * Subroutines (multiple outputs) and scalar-returning functions. - Static Flang runtime integration: Bundle and link static @llvm-project//flang-rt:runtime via LLJIT StaticLibraryDefinitionGenerator so Fortran runtime symbols (I/O, descriptors, intrinsics) resolve during kernel execution. - Python API: Export fortran_call in enzyme_ad.jax, backed by Language::FORTRAN (lang=3). Enhance JVP/VJP rules for multi-output primitives and non-differentiable integers/booleans. - Comprehensive test suite: Unit tests in test/fortran_test.py covering scalar operations, vector/matrix math, assumed-shape descriptors, subroutines, functions, non-differentiable inputs, JIT compilation, and forward/reverse automatic differentiation.
… for libc++ Fix CI failures: - Reformat primitives.py and test/fortran_test.py with black==26.10.1 to pass Python-Black lint workflow. - Add patches/llvm_flang_features.patch and wire into workspace.bzl to include <unordered_map>, <string>, and <array> in flang/include/flang/Support/Fortran-features.h when building with libc++ on macOS / Darwin.
Guard Microsoft CRT-specific Bessel functions (_j0, _j1, _jn, _y0, _y1, _yn) with #if defined(_MSC_VER) instead of #ifdef _WIN32 in flang/lib/Evaluate/intrinsics-library.cpp, fixing compilation errors when building for Windows with MinGW toolchains.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3821 +/- ##
==========================================
+ Coverage 29.15% 30.79% +1.64%
==========================================
Files 249 249
Lines 54687 55898 +1211
==========================================
+ Hits 15942 17214 +1272
+ Misses 38745 38684 -61 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…ned long long on MinGW
joewallwork
left a comment
joewallwork
left a comment
There was a problem hiding this comment.
This seems like it will be very handy, thanks @wsmoses. It would've been nice to have a heads-up that you were going to do this, since we had already been working on it, but thanks for making me aware of this PR.
I can't claim to understand the internals of your implementation, since I'm not familiar with Enzyme-JAX. However, I've shared some feedback on the Fortran code and tests.
It would be good to follow this up by adding some docs and an example to the Fortran Guide on the main Enzyme website.
|
|
||
| class FortranCallTest(absltest.TestCase): | ||
| def test_explicit_shape_and_layout_f32_f64(self): | ||
| for dtype, kind in ((jnp.float32, "4"), (jnp.float64, "8")): |
There was a problem hiding this comment.
Could we not use a map {jnp.float32: "4", jnp.float64: "8"}, since this is a direct correspondence?
There was a problem hiding this comment.
Better still, how about adding use, intrinsic :: iso_fortran_env, only: {kind} and using the map {jnp.float32: "real32", jnp.float64: "real64"}?
| module my_fortran_mod | ||
| implicit none | ||
| contains | ||
| subroutine poly_2d(y, x) | ||
| real(8), intent(out) :: y(:, :) | ||
| real(8), intent(in) :: x(:, :) |
There was a problem hiding this comment.
Again, how about
| module my_fortran_mod | |
| implicit none | |
| contains | |
| subroutine poly_2d(y, x) | |
| real(8), intent(out) :: y(:, :) | |
| real(8), intent(in) :: x(:, :) | |
| module my_fortran_mod | |
| use, intrinsic :: iso_fortran_env, only: real64 | |
| implicit none | |
| contains | |
| subroutine poly_2d(y, x) | |
| real(real64), intent(out) :: y(:, :) | |
| real(real64), intent(in) :: x(:, :) |
There was a problem hiding this comment.
Ah okay, I see that you have a different test for iso_fortran_env further down.
| x = jnp.array([[1.0, -2.0, 0.5], [3.0, 1.5, -1.0]], dtype=jnp.float64) | ||
| out_shape = [jax.core.ShapedArray(x.shape, jnp.float64)] | ||
|
|
||
| for proc_name in ("my_fortran_mod::poly_2d", "poly_2d"): |
There was a problem hiding this comment.
What is the loop doing here? Testing that you can access the subroutine directly or from the module?
| real(8) function weighted_sum(x, y) result(res) | ||
| real(8), intent(in) :: x(:), y(:) |
There was a problem hiding this comment.
As above, although I'm not sure whether you would be able to have a use statement here, given you have an orphan function. Ideally, this would be put inside a module, although I suppose Fortran codes in the wild don't necessarily follow best practices so we should test handling such procedures.
There was a problem hiding this comment.
Ah okay, I see that you have a different test for iso_fortran_env further down.
| use iso_c_binding, only: c_int, c_double | ||
| implicit none | ||
| contains | ||
| subroutine axpy_c(y, a, x, n) bind(C, name="custom_axpy_c") |
There was a problem hiding this comment.
Where is the C file for the custom_axpy_c implementation?
| complex_source = """ | ||
| subroutine cscale(y, z, alpha) | ||
| complex(8), intent(out) :: y(2) | ||
| complex(8), intent(in) :: z(2), alpha | ||
| y = alpha * z | ||
| end subroutine cscale | ||
| """ | ||
| z = jnp.array([1.0 + 2.0j, -3.0 + 0.5j], dtype=jnp.complex128) | ||
| alpha = jnp.array(0.5 - 1.5j, dtype=jnp.complex128) | ||
| (cy,) = jax.jit( | ||
| lambda z_, a_: fortran_call( | ||
| z_, | ||
| a_, | ||
| out_shapes=[jax.core.ShapedArray((2,), jnp.complex128)], | ||
| source=complex_source, | ||
| fn="cscale", | ||
| ) | ||
| )(z, alpha) | ||
| np.testing.assert_allclose(cy, alpha * z, rtol=1e-12) |
There was a problem hiding this comment.
Why isn't this a separate test? Or at least handled by a parameter
| )(z, alpha) | ||
| np.testing.assert_allclose(cy, alpha * z, rtol=1e-12) | ||
|
|
||
| def test_fixed_form_and_error_reporting(self): |
There was a problem hiding this comment.
Why aren't these broken down into separate tests?
Ah sorry, this is entirely AI generated. Per our call Monday where you mentioned wanting to look into jax<->Fortran, I mentioned that this repo was the right place to connect things [having already the connection and most of the infra from C++]. Please feel free to completely take this PR, I just figure some starter code [that builds and hopefully does something reasonable] would help get you bootstrapped on |
No worries. This repo does seem like the right place to do this stuff so thanks for the help getting something set up. I'll chat with @Mikolaj-A-Kowalski in our standup tomorrow about how we want to proceed. |
Overview
This PR adds native Fortran support to Enzyme-JAX via a new
fortran_callAPI, enabling direct execution, JIT compilation, and automatic differentiation (forward-mode JVP, reverse-mode VJP/grad) of Fortran routines in JAX.Key Changes
In-Process Flang Compilation (
flang_compile.h,flang_compile.cc):Fortran::frontend::CompilerInstance,CompilerInvocation).createMLIRToLLVMPassPipeline).iso_c_binding,iso_fortran_env,ieee_arithmetic, etc.) from an in-memory virtual filesystem._QPfootofoo_) to match LLVM IR export definitions.Fortran ABI & Calling Convention Adapter (
clang_compile.cc):VALUE.void*) for standard dummy arrays/scalars.CFI_cdesc_tdescriptors for assumed-shape arrays, allocatables, and pointer dummy arguments.enzyme_constso Enzyme AD focuses strictly on differentiable floating-point tensors.Static Flang Runtime Linking (
enzyme_call.cc,BUILD):@llvm-project//flang-rt:runtimeintoenzyme_callviaflang_rt_archive.inc.StaticLibraryDefinitionGeneratorwith the archive so all standard Fortran runtime routines (I/O, descriptors, math intrinsics) resolve seamlessly inside the JIT execution session.Python Frontend & AD Integration (
primitives.py,__init__.py):fortran_callfromenzyme_ad.jax.enzyme_jvpandenzyme_vjpto handle non-differentiable outputs (float0/ad.Zero) and multiple outputs cleanly.Comprehensive Test Suite (
test/fortran_test.py,test/BUILD):VALUEand reference arguments).jax.jvp) and reverse-mode (jax.grad/jax.vjp) automatic differentiation with numerical finite difference validation.jax.jit.Verification
bazel test --config=public_cache --repo_env=CC=/usr/bin/clang-18 --test_env=JAX_PLATFORMS=cpu --test_output=errors //test:fortran_test //test:test //test:testffi//test:fortran_test: PASSED//test:test: PASSED//test:testffi: PASSEDbuildifierchecked and clean.git clang-formatclean on all touched lines.ord(c) < 128on all files).