Autodiff support [feature request]

[adam-urbanczyk]

It would be great if lpython (lfortran too BTW) supported automatic differentiation. Ideally on the intermediate level, if not possible, for specific backends, e.g

• for LLVM: using enzyme
• for C++ using clad or simply integrating an autodiff library in the backend (e.g. https://github.com/autodiff/autodiff )

[certik]

Yes, this has been requested many times: https://github.com/lfortran/lfortran/issues/204 (and the original issue https://gitlab.com/lfortran/lfortran/-/issues/97).

I am not focusing on it right now, first we need to get to beta, but if others are interested, they can work on it.

[adam-urbanczyk]

Thanks for the quick reply and background links! Let's say I'd like to experiment with using the c++ autodiff library, would you then have any pointers on how to achieve this? E.g. is there and extension API for defining custom intrinsics or pragmas that could be expanded into something in the backend?

[adam-urbanczyk]

Not really an extension API, but I was able to find this

https://github.com/lcompilers/lpython/blob/f84afdd19d8b53273e3a2dbfc95a19f12c741504/src/libasr/pass/intrinsic_function_registry.h#L18

and this

https://github.com/lcompilers/lpython/blob/f84afdd19d8b53273e3a2dbfc95a19f12c741504/src/libasr/codegen/asr_to_c_cpp.h#L2891

No support for pragmas though, or did I miss something?

[certik]

We don't have support for pragmas yet. For LFortran at least we'll have to add it for sure.

If you are interested in implementing it, I am happy to meet over video and we can figure out a good design. It would make sense to me to add it as an ASR->ASR pass, and it can be triggered by some function or decorator in the frontend Python language. Or if enzyme can cover all use cases, then we can use it. To use the C++ solution, our C++ generator would generate the code using the autodiff's API. One would represent things in ASR by adding some nodes, either by the IntrinsicFunction or even just new nodes / types. I don't understand the use cases well enough to have an opinion yet. But if you are a user of autodiff, then you can help us with this.

[wsmoses]

Lead Enzyme dev here, would be happy to help you all set things up if desired.

Depending on what path you want to go down, you can either use Enzyme as a modulepass (which replaces calls to __enzyme_autodiff / __enzyme_fwddiff / etc ) with calls to a corresponding implementation generated in the pass, or as an API which takes an LLVM Function, and relevant information, and spits out a new LLVM function to be called.

[certik]

Thanks @wsmoses. I would like to integrate with your Enzyme. My apologies that we didn't manage to do that yet. I personally do not use autodiff, so I was waiting for somebody who actually needs this to help and ensure we do it properly.

[adam-urbanczyk]

@certik thanks, I'll try to hack something together using the C++ backend. Are there any resources on extending/developing lpython?

[wsmoses]

@adam-urbanczyk if you want to use Enzyme's C++ frontend (you just add an extra -fplugin=/path/to/ClangEnzyme.so), I'd be happy to help. That might actually be pretty nice since then the same autodiff features would be possible in either a C++, LLVM, etc backend lowering.

A colleague of ours recently added some more C++ Enzyme examples if that would be helpful: https://enzyme.mit.edu/getting_started/Examples/

[certik]

Thanks a lot @wsmoses for the offer. I don't have time right now to lead this effort, but I am happy to meet with anyone who wants to pursue this. @adam-urbanczyk do you want to go ahead and try this? The best resource to develop LPython/LFortran is at https://docs.lfortran.org/en/contributing/ and this video: https://www.youtube.com/watch?v=yuYsyM08bss.