16 ms·

Enzyme – High-performance automatic differentiation of LLVM

160 points by albertzeyer 6y ago

bdauvergne 6y ago
Funny, I worked on Tapenade (one of the compared automatic differentiation software). I'm happy that it still reaches 60% of the performance of something written directly inside an optimizing compiler.
- thechao 6y ago
  The first time my co-advisor said "you're going to take the derivative of that code by the end of the day" I felt like I'd taken a wrong turn into crazy-town. I enjoyed reading about & understanding Tapenade — such a great piece of code!
- LockAndLol 6y ago
  What is up with these app names? Enzyme? Tapenade?
  - wsmoses 6y ago
    Enzyme is named such as it's a tool that "synthesizes derivatives" and also as a pun referencing Zygote (another AD tool) since Enzyme operates at a lower level (LLVM rather than Julia).
  - bdauvergne 6y ago
    Tapenade is just a traditional recipe from south-east of France, where the research team developing Tapenane is based (they are at INRIA Sophia-Antipolis).
KenoFischer 6y ago
Also see the Julia package that makes it acessible with a high level interface and probably one of the easier ways to play with it: https://github.com/wsmoses/Enzyme.jl https://github.com/wsmoses/Enzyme.jl.
ksr 6y ago
> The Enzyme project is a tool for performing reverse-mode automatic differentiation (AD) of statically-analyzable LLVM IR. This allows developers to use Enzyme to automatically create gradients of their source code without much additional work. Can someone please explain applications of creating gradients of my source code?
- 6d65 6y ago
  AFAIK, It's mainly used for implementing gradient descent, which is used for training neural networks. Frameworks like pytorch, tensorflow, probably used back propagation to calculate the gradient of a multidimensional function. But in involves tracing, and storing the network state during the forward pass. Static automatic differentiation should be faster and should look a lot like differentiation is done mathematically rather than numerically. Of course there are more applications to AD in scientific computing.
  - oscargrouch 6y ago
    I think Swift is also going into this direction of backing in directly into the compiler and provide it as a higher level language construction. https://github.com/apple/swift/blob/main/docs/DifferentiableProgramming.md https://github.com/apple/swift/blob/main/docs/Differentiable... Which leads to "Swift for Tensorflow" that unlike other languages like Java, Go or Python is not just about bindings to the C++ tensorflow library.
    - 6d65 6y ago
      You're right about Swift for Tensorflow. It is an interesting development. I have installed swift on Linux just to play with it. Was dissapointed to see that Swift for Tensorflow was a fork of the Swift compiler, and haven't followed whether it's being refreshed, and whether there are any plans to merge it back. Ideally an LLVM tool should allow languages that compile to LLVM(I'm mostly interested in Rust and Julia) to leverage it without dramatically changing their compilers. It would be interesting to see something like JAX in Rust, exposing the AD functionality in the Standard Library, paired with a high performance SIMD/GPU array type. Things could get very interesting.
vchuravy 6y ago
One of the authors here, happy to answer any questions! (Particularly about the Julia integration)
- jstrong 6y ago
  do you have any sense for how this would integrate with rust? as someone who isn't familiar with how it works, it's not clear whether that would be as easy as normal C ffi interop or more involved.
  - cornstalks 6y ago
    There was a thread about Enzyme + Rust on Rust's internals forum: https://internals.rust-lang.org/t/automatic-differentiation-differential-programming-via-llvm/13188 https://internals.rust-lang.org/t/automatic-differentiation-...
    - wsmoses 6y ago
      Oh man that was a fun hack to write. Basically we demonstrated an easy-to-setup AD on rust by leveraging link-time optimization (LTO) as a way to make sure Enzyme's generate derivatives "optimization pass" was run. We're currently working with the Rust ML infrastructure group to make a nice integration of Enzyme into Rust (e.g. nice type-checking, safety, etc). If you're interested in helping, you should join the Rust ML meetings and/or Enzyme weekly meetings and check out https://github.com/rust-ml/Meetings https://github.com/rust-ml/Meetings and https://github.com/tiberiusferreira/oxide-enzyme/tree/c-api https://github.com/tiberiusferreira/oxide-enzyme/tree/c-api . There's a bunch of interesting optimizations and nicer UX for interoperability we want to add so more manpower is really helpful! The most interesting thing from the Rust standpoint is that ideally we'd want Enzyme to be loaded into the Rust compiler as a plugin (much like it is for Julia, Clang for C/C++, etc) -- but Rust doesn't support the option for that yet. This means we can either help push for plugins/custom codegen in Rust, make script-based compilation tools within rustc [I don't remember the specific name but someone who is more of a Rust expert I'm sure can chime in], or do the sketchy LTO approach above [not always desirable as it requires running LTO]. Alternatively Enzyme can just become part of LLVM mainline so everyone can use it without a plugin :P We're not quite there yet but we're in the process of becoming a formal LLVM incubator project!
wsmoses 6y ago
Hi all, another author here and happy to answer any questions! Some more relevant links for the curious Github: https://github.com/wsmoses/Enzyme https://github.com/wsmoses/Enzyme Paper: https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf 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!
- turnersr 6y ago
  Hello, Thank you for sharing and releasing usable code! Do you know if this would work for GPU based applications? Tensorflow models that are trained on a GPU, for example?
  - wsmoses 6y ago
    For GPU's, there's a couple of different things that you might want to do. 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).
- Agingcoder 6y ago
  Very interesting. I especially like the second item. What happens if the function you want to differentiate calls multiple other functions, in multiple other compilation units? (I haven't read the paper yet but definitely will)

160 points by albertzeyer 6y ago

Funny, I worked on Tapenade (one of the compared automatic differentiation software). I'm happy that it still reaches 60% of the performance of something written directly inside an optimizing compiler.

thechao 6y ago
The first time my co-advisor said "you're going to take the derivative of that code by the end of the day" I felt like I'd taken a wrong turn into crazy-town. I enjoyed reading about & understanding Tapenade — such a great piece of code!
LockAndLol 6y ago
What is up with these app names? Enzyme? Tapenade?
- wsmoses 6y ago
  Enzyme is named such as it's a tool that "synthesizes derivatives" and also as a pun referencing Zygote (another AD tool) since Enzyme operates at a lower level (LLVM rather than Julia).
- bdauvergne 6y ago
  Tapenade is just a traditional recipe from south-east of France, where the research team developing Tapenane is based (they are at INRIA Sophia-Antipolis).

KenoFischer 6y ago

Also see the Julia package that makes it acessible with a high level interface and probably one of the easier ways to play with it: https://github.com/wsmoses/Enzyme.jl https://github.com/wsmoses/Enzyme.jl.

ksr 6y ago

> The Enzyme project is a tool for performing reverse-mode automatic differentiation (AD) of statically-analyzable LLVM IR. This allows developers to use Enzyme to automatically create gradients of their source code without much additional work. Can someone please explain applications of creating gradients of my source code?

6d65 6y ago
AFAIK, It's mainly used for implementing gradient descent, which is used for training neural networks. Frameworks like pytorch, tensorflow, probably used back propagation to calculate the gradient of a multidimensional function. But in involves tracing, and storing the network state during the forward pass. Static automatic differentiation should be faster and should look a lot like differentiation is done mathematically rather than numerically. Of course there are more applications to AD in scientific computing.
- oscargrouch 6y ago
  I think Swift is also going into this direction of backing in directly into the compiler and provide it as a higher level language construction. https://github.com/apple/swift/blob/main/docs/DifferentiableProgramming.md https://github.com/apple/swift/blob/main/docs/Differentiable... Which leads to "Swift for Tensorflow" that unlike other languages like Java, Go or Python is not just about bindings to the C++ tensorflow library.
  - 6d65 6y ago
    You're right about Swift for Tensorflow. It is an interesting development. I have installed swift on Linux just to play with it. Was dissapointed to see that Swift for Tensorflow was a fork of the Swift compiler, and haven't followed whether it's being refreshed, and whether there are any plans to merge it back. Ideally an LLVM tool should allow languages that compile to LLVM(I'm mostly interested in Rust and Julia) to leverage it without dramatically changing their compilers. It would be interesting to see something like JAX in Rust, exposing the AD functionality in the Standard Library, paired with a high performance SIMD/GPU array type. Things could get very interesting.

vchuravy 6y ago

One of the authors here, happy to answer any questions! (Particularly about the Julia integration)

jstrong 6y ago
do you have any sense for how this would integrate with rust? as someone who isn't familiar with how it works, it's not clear whether that would be as easy as normal C ffi interop or more involved.
- cornstalks 6y ago
  There was a thread about Enzyme + Rust on Rust's internals forum: https://internals.rust-lang.org/t/automatic-differentiation-differential-programming-via-llvm/13188 https://internals.rust-lang.org/t/automatic-differentiation-...
  - wsmoses 6y ago
    Oh man that was a fun hack to write. Basically we demonstrated an easy-to-setup AD on rust by leveraging link-time optimization (LTO) as a way to make sure Enzyme's generate derivatives "optimization pass" was run. We're currently working with the Rust ML infrastructure group to make a nice integration of Enzyme into Rust (e.g. nice type-checking, safety, etc). If you're interested in helping, you should join the Rust ML meetings and/or Enzyme weekly meetings and check out https://github.com/rust-ml/Meetings https://github.com/rust-ml/Meetings and https://github.com/tiberiusferreira/oxide-enzyme/tree/c-api https://github.com/tiberiusferreira/oxide-enzyme/tree/c-api . There's a bunch of interesting optimizations and nicer UX for interoperability we want to add so more manpower is really helpful! The most interesting thing from the Rust standpoint is that ideally we'd want Enzyme to be loaded into the Rust compiler as a plugin (much like it is for Julia, Clang for C/C++, etc) -- but Rust doesn't support the option for that yet. This means we can either help push for plugins/custom codegen in Rust, make script-based compilation tools within rustc [I don't remember the specific name but someone who is more of a Rust expert I'm sure can chime in], or do the sketchy LTO approach above [not always desirable as it requires running LTO]. Alternatively Enzyme can just become part of LLVM mainline so everyone can use it without a plugin :P We're not quite there yet but we're in the process of becoming a formal LLVM incubator project!

wsmoses 6y ago

Hi all, another author here and happy to answer any questions! Some more relevant links for the curious Github: https://github.com/wsmoses/Enzyme https://github.com/wsmoses/Enzyme Paper: https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf 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!

turnersr 6y ago
Hello, Thank you for sharing and releasing usable code! Do you know if this would work for GPU based applications? Tensorflow models that are trained on a GPU, for example?
- wsmoses 6y ago
  For GPU's, there's a couple of different things that you might want to do. 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).
Agingcoder 6y ago
Very interesting. I especially like the second item. What happens if the function you want to differentiate calls multiple other functions, in multiple other compilation units? (I haven't read the paper yet but definitely will)