LLVM 和 MLIR 的高性能自动微分工具 Enzyme

Enzyme 是一个插件，用于对可静态分析的 LLVM 和 MLIR 代码执行自动微分（AD）。

使用 Enzyme 时，只需在需要求导的函数上调用 __enzyme_autodiff，如下所示。运行 Enzyme 转换 pass 后，对 __enzyme_autodiff 的调用将被其第一个参数的梯度所替换。

double foo(double);

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

Enzyme 效率极高，并且能够在优化后的代码上执行自动微分，因此其性能可以达到甚至超越当前最先进的自动微分工具。

有关安装和使用 Enzyme 的详细信息，请访问我们的官网：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

或者，您也可以通过包管理器安装 Enzyme：

Homebrew

brew install enzyme

Spack

spack install enzyme

Nix

nix-shell -p enzyme

如需参与或有任何问题，请加入我们的 邮件列表。

如果您在学术环境中使用此代码，请引用以下三篇论文（第一篇为 Enzyme 整体介绍，第二篇为 GPU 加速及优化相关工作，第三篇则涉及对其他并行程序（OpenMP、MPI、Julia Tasks 等）的自动微分）：

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

Enzyme 提供了 Julia 绑定 和 Rust 绑定，方便用户在这些语言中使用。

Enzyme 快速上手指南

Enzyme 是一个基于 LLVM 和 MLIR 的高性能自动微分（Automatic Differentiation, AD）插件。它能够对静态可分析的代码进行高效的梯度计算，性能媲美甚至超越现有的主流 AD 工具，特别适用于科学计算和高性能计算场景。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux 或 macOS（Windows 需通过 WSL 或类似环境）。
• 核心依赖：
  • LLVM：需要安装与 Enzyme 版本兼容的 LLVM 库及开发头文件。
  • 构建工具：CMake (推荐 3.10+)，Ninja 或 Make。
  • 编译器：支持 LLVM 后端的主流编译器（如 Clang/GCC）。
• 可选依赖：若需运行测试套件，需安装 lit (LLVM Integrated Tester)。

提示：对于国内开发者，若直接从源码编译遇到网络问题，建议优先使用包管理器安装，或配置国内镜像源加速 LLVM 相关依赖的下载。

安装步骤

您可以选择从源码编译或使用包管理器安装。

方式一：使用包管理器（推荐）

这是最快捷的安装方式，适合大多数用户。

macOS (Homebrew):

brew install enzyme

Linux (Spack):

spack install enzyme

Nix 环境:

nix-shell -p enzyme

方式二：从源码编译

如果您需要特定版本的 LLVM 支持或进行开发调试，请按以下步骤操作：

1. 克隆仓库并进入目录：

   cd /path/to/Enzyme/enzyme
   

2. 创建构建目录并配置 CMake（请替换 /path/to/llvm... 为您实际的 LLVM 路径）：

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

3. 执行编译：

   ninja
   

基本使用

Enzyme 的核心用法是在 C/C++ 代码中调用 __enzyme_autodiff 函数。编译时加载 Enzyme 插件后，该调用会被自动替换为对应函数的梯度计算逻辑。

最简单示例

假设您有一个函数 foo，想要计算其在点 x 处的导数：

double foo(double);

double grad_foo(double x) {
    // __enzyme_autodiff(目标函数，输入变量)
    // 运行时将被替换为 foo 在 x 处的梯度值
    return __enzyme_autodiff(foo, x);
}

使用说明：

1. 编写包含上述代码的 C/C++ 文件。
2. 使用集成了 Enzyme pass 的编译器命令进行编译（具体编译命令取决于您的构建系统如何链接 Enzyme 插件）。
3. 程序运行时，grad_foo 将直接返回 foo 的导数值，无需手动推导公式。

多语言支持：除了 C/C++，Enzyme 官方还提供了 Julia bindings 和 Rust bindings，可在相应语言的生态中直接使用。