c6bf3a01ef
Autodiff batching
Enzyme supports batching, which is especially known from the ML side when training neural networks.
There we would normally have a training loop, where in each iteration we would pass in some data (e.g. an image), and a target vector. Based on how close we are with our prediction we compute our loss, and then use backpropagation to compute the gradients and update our weights.
That's quite inefficient, so what you normally do is passing in a batch of 8/16/.. images and targets, and compute the gradients for those all at once, allowing better optimizations.
Enzyme supports batching in two ways, the first one (which I implemented here) just accepts a Batch size,
and then each Dual/Duplicated argument has not one, but N shadow arguments. So instead of
```rs
for i in 0..100 {
df(x[i], y[i], 1234);
}
```
You can now do
```rs
for i in 0..100.step_by(4) {
df(x[i+0],x[i+1],x[i+2],x[i+3], y[i+0], y[i+1], y[i+2], y[i+3], 1234);
}
```
which will give the same results, but allows better compiler optimizations. See the testcase for details.
There is a second variant, where we can mark certain arguments and instead of having to pass in N shadow arguments, Enzyme assumes that the argument is N times longer. I.e. instead of accepting 4 slices with 12 floats each, we would accept one slice with 48 floats. I'll implement this over the next days.
I will also add more tests for both modes.
For any one preferring some more interactive explanation, here's a video of Tim's llvm dev talk, where he presents his work. https://www.youtube.com/watch?v=edvaLAL5RqU
I'll also add some other docs to the dev guide and user docs in another PR.
r? ghost
Tracking:
- https://github.com/rust-lang/rust/issues/124509
- https://github.com/rust-lang/rust/issues/135283
|
||
|---|---|---|
| .github | ||
| compiler | ||
| library | ||
| LICENSES | ||
| src | ||
| tests | ||
| .clang-format | ||
| .editorconfig | ||
| .git-blame-ignore-revs | ||
| .gitattributes | ||
| .gitignore | ||
| .gitmodules | ||
| .ignore | ||
| .mailmap | ||
| bootstrap.example.toml | ||
| Cargo.lock | ||
| Cargo.toml | ||
| CODE_OF_CONDUCT.md | ||
| configure | ||
| CONTRIBUTING.md | ||
| COPYRIGHT | ||
| INSTALL.md | ||
| LICENSE-APACHE | ||
| license-metadata.json | ||
| LICENSE-MIT | ||
| README.md | ||
| RELEASES.md | ||
| REUSE.toml | ||
| rust-bors.toml | ||
| rustfmt.toml | ||
| triagebot.toml | ||
| x | ||
| x.ps1 | ||
| x.py | ||
This is the main source code repository for Rust. It contains the compiler, standard library, and documentation.
Why Rust?
-
Performance: Fast and memory-efficient, suitable for critical services, embedded devices, and easily integrated with other languages.
-
Reliability: Our rich type system and ownership model ensure memory and thread safety, reducing bugs at compile-time.
-
Productivity: Comprehensive documentation, a compiler committed to providing great diagnostics, and advanced tooling including package manager and build tool (Cargo), auto-formatter (rustfmt), linter (Clippy) and editor support (rust-analyzer).
Quick Start
Read "Installation" from The Book.
Installing from Source
If you really want to install from source (though this is not recommended), see INSTALL.md.
Getting Help
See https://www.rust-lang.org/community for a list of chat platforms and forums.
Contributing
See CONTRIBUTING.md.
License
Rust is primarily distributed under the terms of both the MIT license and the Apache License (Version 2.0), with portions covered by various BSD-like licenses.
See LICENSE-APACHE, LICENSE-MIT, and COPYRIGHT for details.
Trademark
The Rust Foundation owns and protects the Rust and Cargo trademarks and logos (the "Rust Trademarks").
If you want to use these names or brands, please read the Rust language trademark policy.
Third-party logos may be subject to third-party copyrights and trademarks. See Licenses for details.