Hacker News Board Hacker News, browsed like a classic message board
Live data from Hacker News

Enzyme – High-performance automatic differentiation of LLVM

enzyme.mit.edu

Post reply on HN
Print view View original

21–30 of 51 posts

Re: Enzyme – High-performance automatic differentiation of LLVM

#21

Post by wsmoses »

cornstalks wrote: post #20
jstrong wrote: post #19

Earlier quoted context omitted.

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.

There was a thread about Enzyme + Rust on Rust's internals forum: https://internals.rust-lang.org/t/automatic-differentiation-...

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 and 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!

Re: Enzyme – High-performance automatic differentiation of LLVM

#22

Post by mhh__ »

PartiallyTyped wrote: post #9
ksr wrote: post #4

> 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?

It constructs an analytical gradient from the code. The reason is that you can compute the gradient directly. This can enable optimizations such as avoiding caching big matrices because you don't need to keep track of states/trace the graph, or you can compute the 2nd, 3rd, 4th... and so on derivatives because you have an analytical gradient. For example in an affine function, the gradient of the bias/intercept is th…

That last line doesn't make much sense to my eye at least.

Re: Enzyme – High-performance automatic differentiation of LLVM

#24

Post by cl3misch »

PartiallyTyped wrote: post #9
ksr wrote: post #4

> 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?

It constructs an analytical gradient from the code. The reason is that you can compute the gradient directly. This can enable optimizations such as avoiding caching big matrices because you don't need to keep track of states/trace the graph, or you can compute the 2nd, 3rd, 4th... and so on derivatives because you have an analytical gradient. For example in an affine function, the gradient of the bias/intercept is th…

Isn't the input of the layer fundamentally a part of the gradient computation? So even in this case (inspecting LLVM code) the computation still needs to look at the input.

Re: Enzyme – High-performance automatic differentiation of LLVM

#25

Post by Agingcoder »

wsmoses wrote: post #14

Hi all, another author here and happy to answer any questions! 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 c…

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)

Re: Enzyme – High-performance automatic differentiation of LLVM

#26

Post by cl3misch »

6d65 wrote: post #5
ksr wrote: post #4

> 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?

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 mathematic…

I don't see how static AD removes the need to store the network state. Is this a fundamental property of static AD?

Also, your statement sounds like pyTorch/TF are doing AD numerically, which is not the case. They build the analytical gradient from the traced computation graph.

Re: Enzyme – High-performance automatic differentiation of LLVM

#27

Post by wsmoses »

mhh__ wrote: post #22
PartiallyTyped wrote: post #9

Earlier quoted context omitted.

It constructs an analytical gradient from the code. The reason is that you can compute the gradient directly. This can enable optimizations such as avoiding caching big matrices because you don't need to keep track of states/trace the graph, or you can compute the 2nd, 3rd, 4th... and so on derivatives because you have an analytical gradient. For example in an affine function, the gradient of the bias/intercept is th…

That last line doesn't make much sense to my eye at least.

I think in essence what PartiallyTyped is trying to say is that one potential optimization opportunity in whole-program AD is that you can avoid having to cache the original inputs of the program if you know that derivative computation won't need it (e.g. its only used in a sum and not a product or something whose derivative depends on the value). Some ML frameworks must cache all of the inputs to an operation since they don't know whether it will be necessary for the reverse pass of an operation. You could go even further and decide to cache a different & smaller set of intermediate values that still lets you compute the gradient.

Beyond cache reduction, in our paper we demonstrate a lot of interesting ways that combining AD with a compiler can create potential speed-up. For example, we are often able to dead-code eliminate part of the original forward-pass code since it's not needed to compute the gradient.

We also have a cool example in the paper showing an asymptotic [O(N^2) => O(N)] speedup on a code for normalizing a vector because doing AD in the compiler lets Enzyme run after optimization (and in that example benefit from loop invariant code motion in a way that tools that aren't in the compiler cannot do).

Re: Enzyme – High-performance automatic differentiation of LLVM

#28

Post by wsmoses »

cl3misch wrote: post #24
PartiallyTyped wrote: post #9

Earlier quoted context omitted.

It constructs an analytical gradient from the code. The reason is that you can compute the gradient directly. This can enable optimizations such as avoiding caching big matrices because you don't need to keep track of states/trace the graph, or you can compute the 2nd, 3rd, 4th... and so on derivatives because you have an analytical gradient. For example in an affine function, the gradient of the bias/intercept is th…

Isn't the input of the layer fundamentally a part of the gradient computation? So even in this case (inspecting LLVM code) the computation still needs to look at the input.

You don't always need the input to compute the gradient. For example the gradient of a sum function doesn't require the original input, it just sets all of the derivative(input)'s to 1.

Re: Enzyme – High-performance automatic differentiation of LLVM

#29

Post by wsmoses »

Agingcoder wrote: post #25
wsmoses wrote: post #14

Hi all, another author here and happy to answer any questions! 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 c…

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)

Enzyme needs to be able to access the IR of any potentially active functions (calls that it deduced could impact the gradient) to be able to differentiate them.

If all of the code you care about is in one compilation unit, you're immediately good to go.

Multiple compilation units can be handled in a couple of ways, depending on how much energy you want to set it up (and we're working on making this easier).

The easiest way is to compile with Link-Time Optimization (LTO) and have Enzyme run during LTO, which ensures it has access to bitcode for all potentially differentiated functions.

The slightly more difficult approach is to have Enzyme ahead-of-time rather than lazily emit derivatives for any functions you may call in an active way (and incidentally this is where Enzyme's rather aggressive activity analysis is super useful). Leveraging Enzyme's support for custom derivatives in which an LLVM function declaration can have metadata that marks its derivative function, Enzyme can then be told to use the "custom" derivatives it generated while compiling other compilation units. This obviously requires more setup so I'm usually lazy and use LTO, but this can definitely be made easier as a workflow.

Re: Enzyme – High-performance automatic differentiation of LLVM

#30

Post by wsmoses »

cl3misch wrote: post #26
6d65 wrote: post #5

Earlier quoted context omitted.

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 mathematic…

I don't see how static AD removes the need to store the network state. Is this a fundamental property of static AD? Also, your statement sounds like pyTorch/TF are doing AD numerically, which is not the case. They build the analytical gradient from the traced computation graph.

Reverse mode AD can always get into situations where it needs to store original values (i.e. network state).

One advantage, however, of doing a more whole-program approach to AD rather than individual operators is that one might be able to avoid caching values unnecessarily. For example if an input isn't modified (and still exists) by the time the value is needed in the reverse pass, you don't need to cache it but can simply use the original input without a copy.

And yes PyTorch/TF tend to perform a (limited) form of AD as well, rather than do numerical differentiation (though I do think there may be an option for numerical?)

I wouldn't really position a tool like Enzyme as a competitor to PyTorch/TF (they may have some better domain-specific knowledge after all), but rather a really nice complement. Enzyme can take derivatives of arbitrary functions, in any LLVM-based language rather than the DSL of operators supported by PyTorch/TF. In fact, we built a plugin for PyTorch/TF that uses Enzyme to import custom foreign code as a differentiable layer!

Post reply on HN