Skip to content
This repository has been archived by the owner on Feb 24, 2023. It is now read-only.
/ oxide-enzyme Public archive

Enzyme integration into Rust. Experimental, do not use.

License

Apache-2.0, MIT licenses found

Licenses found

Apache-2.0
LICENSE-APACHE
MIT
LICENSE-MIT
Notifications You must be signed in to change notification settings

EnzymeAD/oxide-enzyme

Repository files navigation

ARCHIVED

This is a package containing a Rust frontend for Enzyme. This is very much a work in progress and bug reports/discussion is greatly appreciated!

Enzyme is a plugin that performs automatic differentiation (AD) of statically analyzable LLVM. It is highly-efficient and its ability perform AD on optimized code allows Enzyme to meet or exceed the performance of state-of-the-art AD tools.

Supported types

  • Scalars
  • Structs, Unions
  • Tuple, Array, Vec
  • Box, Reference, Raw pointer

We are working on adding support for dyn trait objects, slices and enums.
Adding Generics to your types or implementing traits is already working fine.

Setup

First you have to get an adequate rustc/llvm/enzyme build here: enzyme_build.
Afterwards for your convenience you should export this path for LLVM_SYS

$ export LLVM_SYS_130_PREFIX=$HOME/.cache/enzyme/rustc-1.59.0-src/build/x86_64-unknown-linux-gnu/llvm  

and tell Enzyme about your library locations:

$ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.cache/enzyme/Enzyme-0.0.29/enzyme/build/Enzyme:$HOME/.cache/enzyme/rustc-1.59.0-src/build/x86_64-unknown-linux-gnu/llvm/build/lib/  

As an alternative you can also run

$ ninja install  

inside of your enzyme and llvm build directory.

Afterwards you can execute the following lines in oxide-enzyme/example, in order to compile the example.

$ cargo enzyme

You will find your executable in ./target/$TARGET/debug/

Compilation

We generate gradient functions based on LLVM-IR code. Therefore we currently need two compilation runs. The first to generate a llvm-bc file with the LLVM-IR code, the second to process the bc file, generate the gradients, and build the entire crate. You can do that manually using

RUSTFLAGS="--emit=llvm-bc" cargo +enzyme -Z build-std rustc --target x86_64-unknown-linux-gnu -- --emit=llvm-bc -g -C opt-level=3 -Zno-link && RUSTFLAGS="--emit=llvm-bc" cargo +enzyme -Z build-std rustc --target x86_64-unknown-linux-gnu -- --emit=llvm-bc -g -C opt-level=3

We have created a wrapper for this command which you can call with:

cargo enzyme

Please be aware that our wrapper will ignore all additional commands.
This approach won't work on dependencies since cargo doesn't support such a build process.

FAQ

  • Q: How about Windows / Mac?
  • A: WSL might work, the others probably not. Please let us know if you try.

Further Information

More information on installing and using Enzyme directly (not through Rust) can be found on our website: https://enzyme.mit.edu.

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 paper to appear in NeurIPS 2020

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

License

Dual-licensed to be compatible with the Rust project.

Licensed under the Apache License, Version 2.0 http://www.apache.org/licenses/LICENSE-2.0 or the MIT license http://opensource.org/licenses/MIT, at your option. This file may not be copied, modified, or distributed except according to those terms.

About

Enzyme integration into Rust. Experimental, do not use.

Topics

Resources

License

Apache-2.0, MIT licenses found

Licenses found

Apache-2.0
LICENSE-APACHE
MIT
LICENSE-MIT

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Contributors 4

  •  
  •  
  •  
  •  

Languages