Enzyme – High-performance automatic differentiation of LLVM

Some more relevant links for the curious

Github: https://github.com/wsmoses/Enzyme

Paper: https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b...

Basically the long story short is that Enzyme has a couple of interesting contributions:

1) Low-level Automatic Differentiation (AD) IS possible and can be high performance

2) By working at LLVM we get cross-language and cross-platform AD

3) Working at the LLVM level actually can give more speedups (since it's able to be performed after optimization)

4) We made a plugin for PyTorch/TF that uses Enzyme to import foreign code into those frameworks with ease!

First of all they suffer from accuracy decay. For example if you were to do the standard f'(x) \approx [f(x+h)-f(x)]/h, you'd subtract two similar numbers and waste many bits of precision. In contrast if you were to generate the derivative function directly like below, you'd end up far more accurate.

double square(double x) { return x * x; }

double d_square(double x) { return __enzyme_autodiff(square, x); }

becomes

double d_square(double x) { return 2 * x; }

Secondly, from a performance perspective numerical differentiation is really slow -- especially for gradient computation. For example you would need to evaluate the function at once per argument in numeric differentiation to get the whole gradient. In contrast, reverse mode AD lets you generate the entire gradient in one call.

In addition to these generic issues, we illustrate in our paper how doing this at a compiler level allows for significant additional optimization (by removing unnecessary code from the forward pass, finding common expressions, etc).

These issues also are amplified as you make higher-order derivatives and so on.

You can use existing tools within LLVM to automatically generate GPU code out of existing code, and this works perfectly fine, even running Enzyme first to synthesize the derivative.

You can also consider taking an existing GPU kernel and then automatically differentiating it. We currently support a limited set of cases for this (certain CUDA instructions, shared memory etc), and are working on expanding as well as doing performance improvements. AD of existing general GPU kernels is interesting [and more challenging] since racey reads in your original code become racey writes in the gradient -- which must have extra care taken to make sure they don't conflict. To my knowledge GPU AD on general programs (e.g. not a specific code) really hasn't been done before, so it's a fun research problem to work on (and if someone knows of existing tools for this please email me at wmoses at mit dot edu).