Enzyme.jl

LLVM 高性能自动微分酶

这是一个包含Enzyme的 Julia 绑定的包。这仍在开发中,欢迎提交错误报告和讨论!

Enzyme 是一个插件,用于对可静态分析的 LLVM 进行自动微分(AD)。它高度高效,并且能够对优化代码进行自动微分,使 Enzyme 的性能能够达到或超过最先进的自动微分工具。

Enzyme.jl 可以通常的方式安装 Julia 包:

] add Enzyme

可以通过在要微分的函数上调用 autodiff 来使用 Enzyme.jl,如下所示:

using Enzyme, Test

f1(x) = x*x
# 返回一个活跃返回的元组,在这种情况下只是 (2.0,)
@test first(autodiff(Reverse, f1, Active(1.0))[1]) ≈ 2.0

更多详细信息,请参阅软件包文档。

有关直接安装和使用 Enzyme (不通过 Julia) 的更多信息,请访问我们的网站: https://enzyme.mit.edu。

如果您有任何问题或想参与,请加入我们的邮件列表。

如果在学术环境中使用此代码,请引用以下两篇论文(第一篇是针对 Enzyme 整体,第二篇是针对 GPU+优化):

@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}
}