LLVM 和 MLIR 的酶高性能自动微分器
酶是一个插件,可对静态可分析的 LLVM 和 MLIR 进行自动微分(AD)。
通过对要微分的函数调用 __enzyme_autodiff
,可以使用酶,如下所示。运行酶的转换过程会用其第一个参数的梯度替换对 __enzyme_autodiff
的调用。
double foo(double);
double grad_foo(double x) {
return __enzyme_autodiff(foo, x);
}
酶效率极高,其对优化代码进行 AD 的能力使得酶能够达到或超过最先进 AD 工具的性能。
有关安装和使用酶的详细信息,请访问我们的网站:https://enzyme.mit.edu。
下面是一个简短的例子,说明如何安装酶:
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
或者,使用软件包管理器安装酶:
brew install enzyme
spack install enzyme
要参与或有任何问题,请加入我们的邮件列表。
如果在学术环境中使用此代码,请引用以下三篇论文(第一篇是关于酶整体,第二篇是关于 GPU+优化,第三篇是关于所有其他并行程序(OpenMP、MPI、Julia 任务等)的 AD):
@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}
}