Skip to content

Commit 3089d70

Browse files
committed
Improve Documentation
1 parent 59bbfa7 commit 3089d70

10 files changed

+71
-22
lines changed

src/Debugging.md

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
# How to Debug AD?
2+
3+

src/SUMMARY.md

+8-5
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,15 @@
44
- [What is Autodiff?](./chapter_1.md)
55
- [Motivation](./motivation.md)
66
- [Prior Art](./prior_art.md)
7+
- [Current Limitations](./limitations.md)
8+
- [Safety](./limitations/safety.md)
9+
- [Runtime Performance](./limitations/runtime.md)
10+
- [Compile Times](./limitations/comptime.md)
11+
- [Higher Order Derivatives](./limitations/higher.md)
12+
- [How to Debug](./Debugging.md)
13+
# Reference Guide
14+
- [Other Enzyme frontends](./other_Frontends.md)
715
- [User facing design](./user_design.md)
816
- [rustc internal design](./rustc_design.md)
9-
- [Other Enzyme frontends](./other_Frontends.md)
10-
# Reference Guide
11-
- [Forward Mode](./fwd.md)
12-
- [Reverse Mode](./rev.md)
13-
- [Current Limitations](./limitations.md)
1417

1518
- [Acknowledgments](./acknowledgments.md)

src/acknowledgments.md

+2
Original file line numberDiff line numberDiff line change
@@ -1 +1,3 @@
11
# Acknowledgments
2+
3+

src/fwd.md

+4-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
# Forward Mode
22

3-
In Forward mode we are only allowed to mark input arguments
4-
The return value of forward mode with a Duplicated return is a tuple containing as the first value the primal return value and as the second value the derivative.
3+
In Forward mode we are only allowed to mark input arguments with Dual or Const.
4+
The return value of forward mode with a Dual return is a tuple containing as the first value the primal return value and as the second value the derivative.
5+
6+
In forward mode Dual(x, 0.0) is equivalent to Const(x), except that we can perform more optimizations for Const.
57

6-
In forward mode Duplicated(x, 0.0) is equivalent to Const(x), except that we can perform more optimizations for Const.
78

src/limitations.md

+1-14
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,10 @@
11
# Current limitations
22

33
1) Enzyme currently does only support freestanding functions. We added some support for `self`, but don't link pieces together correctly in all cases. `self` does not exist on llvm-ir level, so it's just a matter of fixing our macro and should be easy to solve.
4-
5-
2) Soundness: Enzyme currently does assume that the user passes shadow arguments (`dx`, `dy`, ...) of appropriate size. That's a current research project of Manuel, so we hope to check at least basic DST (vectors, enums) till end of November. If we remember the backprop function from above, there is no way for the type system to guarantee that `dweights` is at least as large as the `weights` vector. Adding length checks for vectors and making sure that in case of enums primal and shadow are of the same variant should get us a large step towards soundness. Once implemented, we can still evaluate how many more checks we can insert automatically and where we want to fall back to unsafe. We also consider to allow the user to (unsafely) implement a safety check for his own types which we would then insert. Concretely, here we would add the following check at the top of backprop (above the code generated by enzyme) `assert!(dweights.len() >= weights.len())`
6-
```rust
7-
fn backprop(images: &[f32], weights: &[f32], dweights: &mut [f32]) { ... }
8-
```
9-
10-
3) Computing higher order derivatives (hessians, ...) can be done with Enzyme by differentiating functions that compute lower order derivatives. [This example](https://github.com/EnzymeAD/rust/blob/master/library/autodiff/examples/hessian_sin.rs) requires that rustc first uses Enzyme to fill the implementation of the `jac` function, before it uses Enzyme to fill the implementation of `hess`, by differentiatng `jac`. This is currently not guaranteed. It should be comparably easy to fix. Enzyme also considers adding helper function to directly compute common higher order derivatives.
11-
12-
4+
135
4) Parallelism: Enzyme currently does not handle Rust parallelism (rayon). Enzyme does (partly) support various parallel paradigms: OpenMP, MPI, CUDA, Rocm, Julia tasks. Enzyme only does need to support the lowest level of parallelism for each language, so adding support for Rust is not hard, but also not a high priority.
146

157

16-
5) Compile Times: Enzyme can create the TypeTrees it requires for each variable based on its LLVM-IR reads/writes/dereferences/usages. This is slow for large types and programms. A (real-world) examples which updates a vector of length 50k element by element, line by line e.g. takes 6hrs to compile with C++ Enzyme due to TypeTree creation, while JAX only takes 1hr. Due to leveraging rustc information (see next point) Rust-Enzyme manages to compile the code in 5 minutes. Compiling the C++ code without AD however only takes ~1 second.
17-
18-
6) Rust ABI optimizations: In order to improve the compile times mentioned above we create Enzyme TypeTrees based on the rustc type knowledge. These trees unfortunately are not correct anymore once Rust ABI optimization take place. E.g. `fn rosenbrock(x: &[f64; 2]) -> f64 {...}` will get lowered into an LLVM-IR function comparable to `fn rosenbrock(f64, f64) -> f64`. A TypeTree therefore needs to be updated on each applied optimization. For now, we just block optimizations on the outermost functions since they tend to have a small performance effect, so we want to focus on other parts first.
19-
20-
7) FAT-LTO requirement: Rust-Enzyme currently requires fat-lto when AutoDiff is used. We technically only need lto if the function being differentiated calls functions in other compilation units. Other solutions are possible but this is the most simple one. Since the compile time overhead of lto is small compared to the compile time overhead of differentiating larger functions this is not a priority.
218

229
Enzyme does support custom allocators, but Rust-Enzyme does not expose support for it yet. Low priority.
2310

src/limitations/comptime.md

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# Compile Times
2+
3+
Enzyme will often achieve excellent runtime performance, but might increase your compile time by a large factor.
4+
For Rust, we already have made significant improvements and have a list of further improvements planed - please reach out if you have time to help here.
5+
6+
## Type Analysis
7+
Most of the times, Type Analysis (TA) is the reason of large (>5x) compile time increases when using Enzyme.
8+
This poster explains why we need to run Type Analysis in the bottom left part: [Poster Link](https://c.wsmoses.com/posters/Enzyme-llvmdev.pdf).
9+
10+
Enzyme's TA will create TypeTrees based on usage patterns in the code.
11+
Due to a suboptimal datastructure this process scales very poorly.
12+
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.
13+
For the meantime, we do initialize TypeTrees for outermost function (those to which you apply '#[autodiff(...)]` based on the Rust types.
14+
In some real-worl applications (50k LoC), this improved the compile times by over 1000x - reducing them from hours to single minutes.
15+
16+
## Duplicated Optimizations
17+
The key reason for Enzyme offering often excellent performance is that Enzyme does differentiate already optimized LLVM-IR.
18+
However, we also (have to) run LLVM's optimization pipeline after differentiating, to make sure that the code which Enzyme generates is optimized properly.
19+
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.
20+
21+
22+
## FAT-LTO
23+
The usage of '#[autodiff(...)]' currently requires compiling your project with fat-lto.
24+
We technically only need lto if the function being differentiated calls functions in other compilation units.
25+
Therefore other solutions are possible but this is the most simple one to get started.
26+
The compile time overhead of lto is small compared to the current compile time overhead of differentiating larger functions so this limitation is currently not a priority.
27+

src/limitations/higher.md

+9
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
# Higher Order Derivatives
2+
3+
Computing higher order derivatives like hessians can be done with Enzyme by differentiating functions that compute lower order derivatives.
4+
[This example](https://github.com/EnzymeAD/rust/blob/master/library/autodiff/examples/hessian_sin.rs) requires that rustc first uses Enzyme to fill the implementation of the `jac` function, before it uses Enzyme to fill the implementation of `hess`, by differentiatng `jac`.
5+
This is currently not guaranteed and only works by coincidence in some cases.
6+
This should be easy to fix, so please reach out if you would like to contribute and need some help to get started!
7+
8+
Enzyme also considers adding helper function to directly compute common higher order derivatives in the future.
9+

src/limitations/runtime.md

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
# Runtime Performance
2+
3+
While Enzymes performance should already be good in most cases, there are some optimizations left to apply. One is mentioned in the following compile time section.
4+
The other optimization left to apply is re-enabling Rust's ABI optimizations.
5+
The Rust compiler might change how Rust types are represented on a lower level, to allow faster function calls. These optimizations are mainly relevant when you call a small functions many times.
6+
We don't expect this to be the main application of autodiff, where we assume that you will often differentiate math-heavy code that for example calls faer, ndarray, or nalgebra matrix operations.
7+
We therefore disabled this optimization for the outermost function (the one to which one applies '#[autodiff(...)]`, to enable compile time improvements.
8+
However, it would be nice to teach Enzyme about these Rust ABI optimizations so we can have the best of both worlds.

src/limitations/safety.md

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
# Safety and Soundness
2+
3+
Enzyme currently does assume that the user passes shadow arguments (`dx`, `dy`, ...) of appropriate size. That's a current research project of Manuel, so we hope to check at least basic DST (vectors, enums) soon. If we remember the backprop function from above, there is no way for the type system to guarantee that `dweights` is at least as large as the `weights` vector. Adding length checks for vectors and making sure that in case of enums primal and shadow are of the same variant should get us a large step towards soundness. Once implemented, we can still evaluate how many more checks we can insert automatically and where we want to fall back to unsafe. We also consider to allow the user to (unsafely) implement a safety check for his own types which we would then insert. Concretely, here we would add the following check at the top of backprop (above the code generated by enzyme) `assert!(dweights.len() >= weights.len())`
4+
```rust
5+
fn backprop(images: &[f32], weights: &[f32], dweights: &mut [f32]) { ... }
6+
```
7+
8+
This is an ongoing effort. Please use `carg expand <FunctionName>` to get a feeling for which safety checks we currently insert.

src/safety.md

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
# Safety

0 commit comments

Comments
 (0)