integrating Enzyme Automatic Differentiation

[mppf]

It's interesting to consider if Enzyme can be integrated into Chapel. It supports automatic differentiation and that's useful for machine learning.

This issue serves as a place for discussion on this topic.

[mppf]

If anybody wants to try this, I think a first next step is to see if you can follow https://enzyme.mit.edu/getting_started/UsingEnzyme/ . It might be worth trying to follow along first with the C example they show and then try to reproduce it with a Chapel program.

Note that there are 2 ways to get the LLVM IR out of the Chapel compiler (so you can try doing things like running LLVM's opt tool as the Enzyme documentation shows)

1. You can use something like chpl myprogram.chpl --llvm-print-ir myfunction to see a dump of LLVM IR for the function myfunction
2. You can compile with --savec=tmp as in chpl --savec=tmp myprogram.chpl. This will create tmp/chpl__module-nopt.bc and tmp/chpl__module-opt1.bc files that can be processed by opt or dissasembled to a textual form with llvm-dis. Note that these .bc files will contain lots of code that is not relevant to your input program, because Chapel is currently a whole-program compiler. But, they should work with tools like opt nonetheless.

Either way, make sure that any LLVM tools you are using (opt, llvm-dis, etc) match the LLVM version being used by Chapel. You can see the LLVM version being used by Chapel with printchplenv --all (which prints out CHPL_LLVM_VERSION).

[lucaferranti]

Got a first prototype working this is my source file todiff.chpl pretty much identical to the C source code in the enzyme doc, except that the __enzyme_autodiff is wrapped in an extern block

use CTypes;

extern {
  extern double __enzyme_autodiff(double (*fun)(double), double);
}

proc square(x: real) {
  return x * x;
}

proc dsquare(x: real): real {
  return __enzyme_autodiff(c_ptrTo(square), x);
}

for i in 1..4 {
  writeln("x = ", i, " f(x) = ", square(i), " f'(x) = ", dsquare(i));
}

The compilation process is then

# Chapel compilation phase
chpl --driver-compilation-phase --driver-tmp-dir tmp todiff.chpl

# Enzyme pass
opt tmp/chpl__module-opt1.bc -load-pass-plugin=/home/lferrant/Enzyme/enzyme/build/Enzyme/LLVMEnzyme-14.so -passes=enzyme -o tmp/chpl__module-opt1.ll -S

# Optimize enzyme code
opt tmp/chpl__module-opt1.ll -O2 -o tmp/chpl__module-opt1.ll -S

# Convert to bc
llvm-as tmp/chpl__module-opt1.ll -o tmp/chpl__module-opt1.bc

# Finish Chapel compilation
chpl --driver-makebinary-phase --driver-tmp-dir tmp todiff.chpl

Quite a lot of work to get a proper binding and nice user interface, but this is promising.

[lucaferranti]

some initial experiments: https://github.com/lucaferranti/chapel-enzyme

it's easy for simple functions (working with numbers), challenge will be getting it to work with more complex data structures and chapel peculiarities, plus of course a good binding APi to make it user-friendly(ish).

[mppf]

that's great @lucaferranti! BTW, if you want to look at integrating enzyme into the compilation process in a more user-friendly way, I can see these steps:

• adjust the compiler build to include enzyme somehow. You could do something hacky at first. Eventually we will need to answer these open questions: is enzyme optional? If it is, how do users enable it? Does that mean it's reflected in printchplenv ?
• adjust the compiler's LLVM pass pipeline to run enzyme. You will need to know if it is a Module pass or a Function pass to do this but you might be able to find an example to start with. You can register it in or around runModuleOptPipeline in compiler/llvm/clangUtil.cpp. Open question: does this pass need to be toggle-able with a command-line option? Or does configuring a build with enzyme mean that it is used? If we want to toggle it, you would also add a flag to compiler/main/driver.cpp.

These open questions will probably need their own issues to support discussion.

There are also things to figure out about what we want it to look like to use it, from the point of view of the source code.

it's easy for simple functions (working with numbers), challenge will be getting it to work with more complex data structures and chapel peculiarities, plus of course a good binding APi to make it user-friendly(ish).

I'm not really familiar with how enzyme is typically used, but my expectation is that working for simple functions adds significant value & so we should seek to make it easy to use for that purpose / try to get something merge-able with what you have working already.

[lucaferranti]

Thanks for the pointers @mppf !

(incidentally I filled this form to chat about enzyme integration, although I udnerstand it's out of scope strictly speaking).

There are indeed a lot of open questions, which also depend on design choices. E.g. what API we want, we can go pretty far I think by monkeypatching the C/C++ one into an extern block (that's what I did so far). This has not been too bad so far for hte simple examples in the docs, indeed the chapel translation is a nice 1-1 correspondence of the C one. I think however that as soon as we go beyond proc f(x: real): real the amount of needed glue code would grow annoying and a higher level interface woud be very much needed. The higher-level functionality could generate the glue C-code or maybe operate directly on LLVM under the hood.

I think in general enzyme should be optional, that is one should be able to use Chapel even if they don't care about AD (plenty of happy users so far). This makes me wonder whether it should be integrated into the main repo itself or be an external plugin / mason library that one can add to their project if they need AD.

Some ideas

• One approach could be to do something similar to LAPACK. I.e. functionalities API defined in the core language, but when one compiles, needs to pass the path to enzyme. The compiler could check for this flag and do the compilation ot include the enzyme passes if present.

• One could simply document, if you want to use enzyme you need to isntall this and compile this way. We could have a utility function that generates the Makefile, similar to what one needs to do to call Chapel from C.

• One approach could be to have the enzyme binding as mason library (enzyme itself is avaliable from spack, which makes things a bit easier). This could keep it separate from the core language, but might limit how intrusive with the compiler we can be.

To me, the next logical step would be to get it to work with arrays first, at least for simple restricted cases. (I guess a function with a small array counts as simple for you?), after that I think we would have a much clearer picture of what is needed for integration and what we want the source code to look like.

Current issue with arrays (I can push a better documented minimum non-working example to the repo): I was bitten by https://github.com/chapel-lang/chapel/issues/25075, after some work-around, enzyme was complaining that it could not find the function to differentiate. My first uneducated guess would that some rewriting happens during compilation / optimization and between creating the c_fn_ptr and calling __autodiff that pointer does not point to the function anymore, but need to look closer

[lucaferranti]

btw, this talk from EnzymeCon reporting progress in integrating enzyme with rust was very interesting. I think some of the challenges and lessons learnt from there could be useful here too.

[mppf]

To me, the next logical step would be to get it to work with arrays first, at least for simple restricted cases. (I guess a function with a small array counts as simple for you?), after that I think we would have a much clearer picture of what is needed for integration and what we want the source code to look like.

I don't think the array case is simple. I was more thinking that there's enough value in proc f(x: real): real to try to proceed at making that usable. Looking at arrays first to try to understand how that will impact the integration & user interface seems reasonable as well.