Skip to content

Commit

Permalink
add higher order, update other pieces
Browse files Browse the repository at this point in the history
  • Loading branch information
ZuseZ4 committed Apr 7, 2024
1 parent 344fd1f commit be6be52
Show file tree
Hide file tree
Showing 5 changed files with 42 additions and 21 deletions.
3 changes: 1 addition & 2 deletions src/SUMMARY.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,10 @@
- [Usage](./usage/usage.md)
- [Forward Mode](./usage/fwd.md)
- [Reverse Mode](./usage/rev.md)
- [Higher Order Derivatives](./usage/higher.md)
- [Current Limitations](./limitations.md)
- [Safety](./limitations/safety.md)
- [Runtime Performance](./limitations/runtime.md)
- [Compile Times](./limitations/comptime.md)
- [Higher Order Derivatives](./limitations/higher.md)
- [How to Debug](./Debugging.md)
# Reference Guide
- [Other Enzyme frontends](./other_Frontends.md)
Expand Down
4 changes: 2 additions & 2 deletions src/limitations/comptime.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,13 @@ Enzyme's TA will create TypeTrees based on usage patterns in the code.
Due to a suboptimal datastructure this process scales very poorly.
Transfer the code (~1200 Lines of C++) to a better suited trie should remove most of this overhead, please reach out if you can help.
For the meantime, we do initialize TypeTrees for outermost function (those to which you apply '#[autodiff(...)]` based on the Rust types.
We also annotate a very small set of recognized LLVM functions with TypeInformation, but this work should be expanded.
In some real-worl applications (50k LoC), this improved the compile times by over 1000x - reducing them from hours to single minutes.

## Duplicated Optimizations
The key reason for Enzyme offering often excellent performance is that Enzyme does differentiate already optimized LLVM-IR.
However, we also (have to) run LLVM's optimization pipeline after differentiating, to make sure that the code which Enzyme generates is optimized properly.
This is currently done approximately, but in certain cases some code will be optimized too often, while other code is not getting optimized enough. Tuning this could allow both compile time and runtime improvements.

As a result you should have excellent runtime performance (please fill an issue if not), but at a compile time cost.

## FAT-LTO
The usage of '#[autodiff(...)]' currently requires compiling your project with fat-lto.
Expand Down
9 changes: 0 additions & 9 deletions src/limitations/higher.md

This file was deleted.

8 changes: 0 additions & 8 deletions src/limitations/runtime.md

This file was deleted.

39 changes: 39 additions & 0 deletions src/usage/higher.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
# Higher Order Derivatives

Computing higher order derivatives like hessians can be done with Enzyme by differentiating functions that compute lower order derivatives.
This requires that functions are differentiated in the right order, which we currently don't handle.
As a workaround, we introduce two new AD modes `ForwardFirst` and ReverseFirst` that will be differentiated (and optimized)
before we differentiate the default `Forward` and `Reverse` mode invocations. An example is given below.


```rust
// A direct translation of
// https://enzyme.mit.edu/index.fcgi/julia/stable/generated/autodiff/#Forward-over-reverse

#[autodiff(ddf, Forward, Dual, Dual, Dual, Dual)]
fn df2(x: &[f32;2], dx: &mut [f32;2], out: &mut [f32;1], dout: &mut [f32;1]) {
df(x, dx, out, dout);
}

#[autodiff(df, ReverseFirst, Duplicated, Duplicated)]
fn f(x: &[f32;2], y: &mut [f32;1]) {
y[0] = x[0] * x[0] + x[1] * x[0]
}

#[test]
fn main() {
let mut y = [0.0];
let x = [2.0, 2.0];

let mut dy = [0.0];
let mut dx = [1.0, 0.0];

let mut bx = [0.0, 0.0];
let mut by = [1.0];
let mut dbx = [0.0, 0.0];
let mut dby = [0.0];

ddf(&x, &mut bx, &mut dx, &mut dbx,
&mut y, &mut by, &mut dy, &mut dby);
}
```

0 comments on commit be6be52

Please sign in to comment.