[MLIR] Add StoreLikeInterface to decouple activity analysis from stores (#2978) * [MLIR] Add ActiveStoreOpInterface to decouple activity analysis from stores Activity analysis hard-codes which ops are stores (dyn_cast<LLVM::StoreOp> / memref::StoreOp) when reasoning about stored-value / pointer activity. That prevents out-of-tree dialects from participating and duplicates logic per dialect. Add enzyme::ActiveStoreOpInterface exposing (getStoredValue, getStoredPointer), attach it to memref.store and llvm.store, and use it at the potential-active- store site that already handled both dialects -- now a single dialect-agnostic branch. Any dialect (e.g. an out-of-tree fir.store / hlfir.assign) can opt in by attaching the interface. Behavior-preserving: llvm.store/memref.store return the same value/pointer through the interface. The MLIR test suite is unchanged (the two ReverseMode failures are pre-existing on this LLVM build, independent of this change). The remaining LLVM-specific store sites (which walk llvm.alloca / llvm.load origins) are intentionally left as-is. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [MLIR] Rename ActiveStoreOpInterface to StoreLikeInterface Address review: rename the interface and clarify getStoredPointer's description to note it returns the base pointer before any in-op offsets. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [MLIR] Finish dialect-agnostic activity analysis for stores/mutable refs Extends the StoreLikeInterface generalization to the remaining store sites in ActivityAnalysis and gates the pointer-like type check on the type interface, so out-of-tree dialects (e.g. FIR/HLFIR) participate fully in activity analysis without hard-coded dyn_casts. These are the Fortran-independent pieces of EnzymeAD/Enzyme#2969. - isConstantValue: besides LLVM ptr / memref, treat any type whose AutoDiffTypeInterface reports isMutable() (e.g. !fir.ref) as a reference that carries active memory. Behavior-preserving in-tree: among core dialects only memref and LLVM pointer report isMutable(), and both are already matched by the existing isa<> check. - isOperationInactiveFromOrigin: dialect-agnostic store branch mirroring the LLVM::StoreOp case (inactive iff stored value or pointer is constant). - isValueActivelyStoredOrReturned: dialect-agnostic store branch mirroring the LLVM::StoreOp case. Deliberately excludes the Fortran-coupled / behavior-changing parts of #2969 (FIR/HLFIR models, flang plugin, the CoreDialects registration TU split, and dropping tensor/linalg from the pass dependent dialects). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Cleanup PR a bit * apply clang-format * Update enzyme/Enzyme/MLIR/Interfaces/AutoDiffOpInterface.td Co-authored-by: Paul Berg <naydex.mc+github@gmail.com> * [MLIR] Attach StoreLikeInterface to affine.store Lets activity analysis reason about affine.store's stored-value/pointer activity generically, the same way memref.store and llvm.store now do, instead of requiring a hard-coded dyn_cast. getStoredPointer returns the base memref (the affine map's indices are applied within the op). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [MLIR] Add forward-mode conditional-store tests for memref/llvm The affine.store `@if_then` test in ForwardMode/affine.mlir covers a memref that is unconditionally initialized with a constant and then conditionally overwritten with an active value: forward mode must zero-initialize the shadow before the conditional store, otherwise the not-taken path reads uninitialized shadow memory and returns a garbage tangent instead of 0. No equivalent coverage existed for memref.store or llvm.store. Add memref_if.mlir and llvm_if.mlir as the direct analogues. Both are marked XFAIL: forward-mode differentiation of memref.store / llvm.store does not currently emit the shadow zero-initialization, so the tests document the known bug and will XPASS once it is fixed. Verified with a local enzymemlir-opt build: FileCheck fails on exactly the zero-init CHECK-DAG line and matches everywhere else. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix constant handling * clang format * [MLIR] Zero shadow memory for inactive stores in forward mode Restore the original store-inactivity condition in activity analysis: a store is inactive if *either* the stored value or the pointer is inactive, matching LLVM Enzyme's isInstructionInactiveFromOrigin. That condition alone made the shadow store disappear, because MGradientUtils::visitChild skips any constant operation in forward mode, so memoryIdentityForwardHandler never ran and never emitted the null value it already knows how to produce for a constant operand. LLVM Enzyme has no such gate: its forward-mode visitCommonStore returns early only when the pointer is constant. Mirror that by continuing to visit store-like ops whose pointer is active. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * [MLIR] Address review: generalize the forward-mode skip condition Replace the store-specific storesThroughActivePointer check with the general condition wsmoses suggested: only skip a constant operation in forward mode if it is pure, or if all of its operands are constant. A side-effecting op with an active operand may still need to touch shadow memory, which now also covers llvm.memcpy / llvm.memset rather than only StoreLikeInterface ops. This also makes the stale /*iface.hasNoEffect()*/ note redundant. Retain the "if we are being stored into, not storing this value" comment in isValueActivelyStoredOrReturned. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Paul Berg <naydex.mc+github@gmail.com>
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:
brew install enzyme
spack install enzyme
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.