[MLIR] Split registerCoreDialectAutodiffInterfaces into its own TU; trim pass deps (#2988)
* Split registerCoreDialectAutodiffInterfaces into its own TU; trim pass deps
Two dialect-agnostic, non-Fortran registration changes, factored out so they can
land independently of the FIR/HLFIR autodiff work that motivated them.
* Move registerCoreDialectAutodiffInterfaces (the aggregate that force-links
every core-dialect autodiff model, Linalg/NVVM included) out of
CoreDialectsAutoDiffImplementations.cpp into a dedicated translation unit,
CoreDialectsAutoDiffRegistration.cpp. That old TU also defines the
memory-identity handler the FIR/HLFIR models depend on, so anything linking
those models used to drag in the whole aggregate. Isolating it lets a
consumer that only wants a subset of dialects (e.g. a lean `flang -fc1`
plugin that must not pull in Linalg symbols flang does not export) link the
individual register*DialectAutoDiffInterface entry points and never reference
the aggregate TU.
* Drop the stale linalg and tensor entries from the `enzyme` / `enzyme-wrap`
pass dependentDialects. Neither pass creates linalg ops; tensor is only
needed for vector-mode concat and every in-tree tool registers it itself.
These were the last undefined dialect symbols that blocked dlopen'ing the
plugin against a host that does not export Linalg/Tensor TypeIDs.
No behavior change for the existing tools (enzymemlir-opt still calls the
aggregate). Full MLIR suite green except the 2 pre-existing LLVM-24 failures
(complex_create_im, affine_parallel_mincut).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BkvJByk8nTti2GiF4RudNG
* Address #2988 review: keep tensor pass dep; drop NOTE comment
- Restore tensor::TensorDialect in the `enzyme` pass dependentDialects
(Passes.td) and the DifferentiatePass registry.insert override
(EnzymeMLIRPass.cpp). The pass batches scalar ops into tensors by
default (ArithConstantOpBatchInterface), so tensor must be loaded.
Only the stale linalg entry is dropped.
- Remove the explanatory NOTE comment left where
registerCoreDialectAutodiffInterfaces used to live.
- Add trailing newline to CoreDialectsAutoDiffRegistration.cpp.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* apply clang-format
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.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.