Differentiable Fortran with LFortran and Enzyme | Hacker News Reader

Differentiable Fortran with LFortran and Enzyme

docs.pasteurlabs.ai

57 pointsdionhaefner17 comments

Author here — I work on Tesseract at Pasteur Labs, and I wrote this up because the "what if this was possible" was bugging me for way too long :)

I was surprised by how well this worked, the LFortran + Enzyme stack seems to be a very clean way to get gradients through Fortran code via LLVM IR transformations. Pretty cool to see a 220-line Fortran heat solver turn into ~6,900-line reverse pass automatically if I dare say so.

Would be awesome to see this applied to a real scientific codebase, and I hope that the demo is enough to convince people that it’s worth trying.

Very interesting. Does LFortran have the same internal array layout as the standard C runtime ?

A shared layout and a shared calling convention would be very nice.

Sorry about my naive question. Haven't touched Fortran directly in three decades I think.

EDIT: thanks for your reply. For some reason it has been flagged dead. So am responding here. You can mail dang hn at ycombinator dot co m about the flagging. He is very nice.

I would also like to know this. Fortran itself is column-major, so I would guess the internal layout isn't same for multi-dimensional arrays when compared to row-major C? I'm not sure how LFortran represents arrays internally though.

LFortran internally uses column-major, so interchanging data with C should be done carefully for multi-dimensional arrays. If row-major representation is highly needed feature, We can introduce a flag to do that. I'm not totally sure about that but it's doable under some conditions for sure.

We can internally do both row-major and column-major arrays. For Fortran we enable the column-major flag by default, but the door is open to have both at the same time.

Lots of scientific code in Fortran has sparse arrays, so a NxN array that only has values on 5 diagonals will store that as 5xN array to save memory allowing you to run a larger problem.

That's a very orthogonal issue.

Sparse arrays are supported on C libraries too. I have done my time with CSC and CSR even inside Python that called out to C libraries.

Not naive! The answer is yes and no.

Array layout: yes, for the simple cases this post relies on. In the demo the work arrays are literally calloc'd in a C wrapper and handed straight to the Fortran routine, and it just works. (Column-major vs. C's row-major is still on you to keep straight, of course.) This wouldn't be so easy for fancier Fortran array features — allocatables, assumed-shape (:) dummies, pointer arrays — which in general get a descriptor struct (bounds, strides, etc.) rather than a bare pointer. Flang uses those descriptors much more aggressively, but I don't know how that would manifest at the C-Fortran bridge. Would be interesting to repeat the experiment with Flang (if at all possible) and comparing the pain.

Calling convention: there are well established contracts, much older than the layout story. Fortran passes everything by reference, so a double precision :: x argument is an x* at the ABI level — which is why every argument in the C wrapper is a pointer (&n_, &k0_, ...). Combine that with C-compatible name mangling (a trailing underscore historically, controllable via LFortran/bind(C)) and Bob's your uncle. Nothing here is new, the pipeline is really just leaning on that decades-old interop contract and then letting Enzyme differentiate across the boundary.

Very interesting stuff. How would I get GPU offload working? I have a rather complex scientific code I'm working on with JAX. Most of it can be expressed well with JAX's programming model, but the last 10% really sucks. It's still worth it so I don't have to mess around with offload onto whatever XPU flavor of the week. But going to C++ would really make my life easier, as long as I could use e.g. Kokkos.

You can write the cuda kernels and directly differentiate them using Enzyme. Last time I checked KOKKOS was still WIP: https://github.com/EnzymeAD/Enzyme/issues?q=kokkos

No idea lol. I assume it’s possible since both Enzyme and GPU programming are pervasive in Julia. Let us know if you end up trying.

When you say you 'wrote this up', you mean you had an AI write (at least) chunks of it.

Indeed, just like I let my compiler write (at least) chunks of my AD logic. Not great when the tool becomes a leaky abstraction, but overall net positive don't you think?

Disagree; it irks me to read AI slop, and writing is meant to be persuasive and engaging to human readers.

Happy to take the blame for the lack of persuasian and engagement then :) Thanks for the feedback, although I’d like to believe there’s a way to have that cake and eat it too.

Really nice work. I love Enzyme, and used it in my project about differentiable atomic descriptors. Idea was that I can quickly gobble up existing C++ and fortran codes alike for atomic descriptors and create a encompassing library what differentiate against hyper-params as well! But at time Enzyme was very early ~0.0.50 version or so. In our observations also Enzyme was fast enough that performance wise it matched the analytical gradients (when embedded inside entire pipeline) ![libdescriptor](https://libdescriptor.readthedocs.io/en/latest/_images/E_F_C...) .

Great work, I am glad LFortran worked well for this work. We are very close to beta quality now, many codes just work, pretty much all Fortran features are implemented now, but some codes still don't compile due to various small bugs, so we've been fixing them all in the last couple months.

Reply on news.ycombinator.com