Cache only the last-iteration value of loop PHIs used after loop exit

In split mode (augmented forward + reverse), a loop header PHI chosen to be
cached is stored once per iteration of its loop. When the loop is not rotated,
a running reduction (e.g. `s` in `for (...) s += x[i]*x[i]; return sqrt(s);`)
is the header PHI itself and is used after the loop exits. The reverse pass
then only ever reads the entry of the final iteration, yet the tape holds
trip-count + 1 values (per iteration of any enclosing loop). For a mag() call
inside an O(n) loop this made the tape O(n^2) (2 GB at n = 16000).

A use outside the PHI's loop can only observe the value of the PHI's last
execution before the loop was exited (by dominance the use follows the loop
within the same iteration of every enclosing loop, and the loop cannot be
re-entered without a new iteration of the enclosing loop), independent of
the number of exits. So, if no reverse-pass computation inside the loop needs
the PHI, it suffices to keep one slot per iteration of the enclosing loops
that is overwritten on every iteration; the slot holds the exit value.

computeMinCache (augmented pass only) records such PHIs in
cacheOnlyLastIteration: cached (knownRecomputeHeuristic == false), non-pointer
header PHIs, not needed for loop bounds, all of whose in-loop users are
non-terminator, non-memory, non-pointer instructions that neither use the PHI
directly in their adjoint nor are needed in the reverse pass. For these,
cacheForReverse uses the loop preheader as cache scope (a plain tape value if
the loop is outermost). The decision is stored in AugmentedReturn and reused
by the reverse pass and the split forward-mode derivative, so all tape users
agree on the layout. Anything else keeps the previous behaviour. Combined
mode is unaffected (it already uses fixLCSSA's manual_lcssa phi).

On the reproducer (normalize/mag with only SROA before Enzyme, n = 16000):
peak RSS 1956 MB -> 12 MB, gradient/primal time 5.9x -> 1.2x, identical
gradient.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019RNZGgCdK1fbT2WPQpR7Br
6 files changed
tree: 2eefa0bb3373d4b46eb1a6a6447ede89241c7b70
  1. .devcontainer/
  2. .github/
  3. .packaging/
  4. enzyme/
  5. integration/
  6. .gitattributes
  7. .gitignore
  8. CITATION.bib
  9. CONTRIBUTING.md
  10. fortitude.toml
  11. fpm.toml
  12. LICENSE
  13. Readme.md
  14. zizmor.yml
Readme.md

The Enzyme High-Performance Automatic Differentiator of LLVM and MLIR

Enzyme is a plugin that performs automatic differentiation (AD) of statically analyzable LLVM and MLIR.

Enzyme can be used by calling __enzyme_autodiff on a function to be differentiated as shown below. Running the Enzyme transformation pass then replaces the call to __enzyme_autodiff with the gradient of its first argument.

double foo(double);

double grad_foo(double x) {
    return __enzyme_autodiff(foo, x);
}

Enzyme is highly-efficient and its ability to perform AD on optimized code allows Enzyme to meet or exceed the performance of state-of-the-art AD tools.

Detailed information on installing and using Enzyme can be found on our website: https://enzyme.mit.edu.

A short example of how to install Enzyme is below:

cd /path/to/Enzyme/enzyme
mkdir build && cd build
cmake -G Ninja .. -DLLVM_DIR=/path/to/llvm/lib/cmake/llvm -DLLVM_EXTERNAL_LIT=/path/to/lit/lit.py
ninja

Or, install Enzyme using a package manager:

Homebrew

brew install enzyme

Spack

spack install enzyme

Nix

nix-shell -p enzyme

To get involved or if you have questions, please join our mailing list.

If using this code in an academic setting, please cite the following three papers (first for Enzyme as a whole, second for GPU+optimizations, and third for AD of all other parallel programs (OpenMP, MPI, Julia Tasks, etc.)):

@inproceedings{NEURIPS2020_9332c513,
 author = {Moses, William and Churavy, Valentin},
 booktitle = {Advances in Neural Information Processing Systems},
 editor = {H. Larochelle and M. Ranzato and R. Hadsell and M. F. Balcan and H. Lin},
 pages = {12472--12485},
 publisher = {Curran Associates, Inc.},
 title = {Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients},
 url = {https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf},
 volume = {33},
 year = {2020}
}
@inproceedings{10.1145/3458817.3476165,
author = {Moses, William S. and Churavy, Valentin and Paehler, Ludger and H\"{u}ckelheim, Jan and Narayanan, Sri Hari Krishna and Schanen, Michel and Doerfert, Johannes},
title = {Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme},
year = {2021},
isbn = {9781450384421},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3458817.3476165},
doi = {10.1145/3458817.3476165},
booktitle = {Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis},
articleno = {61},
numpages = {16},
keywords = {CUDA, LLVM, ROCm, HPC, AD, GPU, automatic differentiation},
location = {St. Louis, Missouri},
series = {SC '21}
}
@inproceedings{10.5555/3571885.3571964,
author = {Moses, William S. and Narayanan, Sri Hari Krishna and Paehler, Ludger and Churavy, Valentin and Schanen, Michel and H\"{u}ckelheim, Jan and Doerfert, Johannes and Hovland, Paul},
title = {Scalable Automatic Differentiation of Multiple Parallel Paradigms through Compiler Augmentation},
year = {2022},
isbn = {9784665454445},
publisher = {IEEE Press},
booktitle = {Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis},
articleno = {60},
numpages = {18},
keywords = {automatic differentiation, tasks, OpenMP, compiler, Julia, parallel, Enzyme, C++, RAJA, hybrid parallelization, MPI, distributed, LLVM},
location = {Dallas, Texas},
series = {SC '22}
}

Julia bindings, Rust bindings, and Fortran bindings are available for Enzyme.