Git Product home page Git Product logo

oxide-enzyme's Introduction

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.

oxide-enzyme's People

Contributors

bytesnake avatar tgymnich avatar zusez4 avatar

Stargazers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

Watchers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

oxide-enzyme's Issues

ODE example failing

Using

 20     let fnc_j_0 = FncInfo::new("J_0", "unsafe_d_J0",
 21                              vec![DFT_DUP_ARG, DFT_CONSTANT, DFT_CONSTANT, DFT_CONSTANT,],
 22                              Some((DFT_OUT_DIFF, true)));

one the following test function:

  1 pub fn J_0(t: &mut Theta, y0: State, E: SigmaMat, input_true: FVec) -> f64 {
  2     let mut stepper = Dop853::new(t.clone(), 0.0, 100.0, 1e-3, y0, 1e-4, 1e-4);
  3     let res = stepper.integrate();
  4 
  5     if let Err(e) = res {
  6         panic!("An error occured: {}", e);
  7     }
  8     let rand_out = stepper.y_out().to_vec();
  9     let tmp = rand_out.iter().sum::<State>() / (rand_out.len() as f64);
 10     let input = f(tmp);
 11 
 12     let inner: FVec = input - input_true;
 13     let res: SMatrix<f64,1,1> = 0.5 * inner.transpose() * E * inner;
 14 
 15     res[0]
 16 }

will lead to Enzyme failing with the following message:

   %_2.i.i6.i.i = inttoptr i64 %new_layout.1 to i8*, !dbg !5582: {[-1]:Integer}, intvals: {}
    %20 = tail call i8* @__rust_alloc(i64 %new_layout.0, i64 %new_layout.1) #34, !dbg !5587: {}, intvals: {}
    %21 = tail call i8* @__rust_realloc(i8* nonnull %7, i64 %9, i64 %new_layout.1, i64 %new_layout.0) #34, !dbg !5601: {}, intvals: {}
  i8 48: {[-1]:Integer}, intvals: {48,}
  i8 127: {[-1]:Integer}, intvals: {127,}
  i8 49: {[-1]:Integer}, intvals: {49,}
  i8 53: {[-1]:Integer}, intvals: {53,}
  i8 40: {[-1]:Integer}, intvals: {40,}
  i8 32: {[-1]:Integer}, intvals: {32,}
  i8 56: {[-1]:Integer}, intvals: {56,}
  i8 46: {[-1]:Integer}, intvals: {46,}
  i8 47: {[-1]:Integer}, intvals: {47,}
  i8 45: {[-1]:Integer}, intvals: {45,}
  i8 99: {[-1]:Integer}, intvals: {99,}
  i8 114: {[-1]:Integer}, intvals: {114,}
  i8 109: {[-1]:Integer}, intvals: {109,}
  i8 101: {[-1]:Integer}, intvals: {101,}
  i8 111: {[-1]:Integer}, intvals: {111,}
  i8 116: {[-1]:Integer}, intvals: {116,}
  i8 117: {[-1]:Integer}, intvals: {117,}
  i8 110: {[-1]:Integer}, intvals: {110,}
  i8 108: {[-1]:Integer}, intvals: {108,}
  i8 105: {[-1]:Integer}, intvals: {105,}
  i8 103: {[-1]:Integer}, intvals: {103,}
  i8 100: {[-1]:Integer}, intvals: {100,}
  i8 115: {[-1]:Integer}, intvals: {115,}
  i8 104: {[-1]:Integer}, intvals: {104,}
  i8 122: {[-1]:Integer}, intvals: {122,}
  i8 121: {[-1]:Integer}, intvals: {121,}
  i8 54: {[-1]:Integer}, intvals: {54,}
  i8 98: {[-1]:Integer}, intvals: {98,}
  i8 120: {[-1]:Integer}, intvals: {120,}
  i8 95: {[-1]:Integer}, intvals: {95,}
  i8 52: {[-1]:Integer}, intvals: {52,}
  i8 107: {[-1]:Integer}, intvals: {107,}
  i8 119: {[-1]:Integer}, intvals: {119,}
  i8 118: {[-1]:Integer}, intvals: {118,}
  i8 -127: {[-1]:Integer}, intvals: {-127,}
  i8 113: {[-1]:Integer}, intvals: {113,}
  i8 96: {[-1]:Integer}, intvals: {96,}
  i8 -74: {[-1]:Integer}, intvals: {-74,}
  i8 -9: {[-1]:Integer}, intvals: {-9,}
  </analysis>
  Illegal updateAnalysis prev:{[-1]:Integer} new: {[-1]:Pointer}
  val:   %_2.i.i.i = inttoptr i64 %new_layout.1 to i8*, !dbg !5488 origin=  %.sroa.0.0.i.i.pn = phi i8* [ %_2.i.i.i, %bb3.i.i14 ], [ %13, %bb7.i.i15 ], [ %_2.i.i6.i.i, %bb3.i7.i.i ], [ %20, %bb7.i8.i.i ], [ %21, %bb10.i.i ]
  build-script-build: ../Enzyme/TypeAnalysis/TypeAnalysis.cpp:617: void TypeAnalyzer::updateAnalysis(llvm::Value*, TypeTree, llvm::Value*): Assertion `0 && "Performed illegal updateAnalysis"' failed.

I'm using an Enzyme fork with the Rust debug parser merged, but the failing code doesn't seem to be affected by that:
https://github.com/ZuseZ4/Enzyme/blob/08eb0646ee6f0ae7eca99fcc41f0c54fcc09cbc1/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp#L617

It could be due to the sloppy initialization of some values which (based on my understanding) shouldn't be seriously considered
by Enzyme: https://github.com/rust-ml/oxide-enzyme/blob/a1dfbde391b32cde6a594b9fe87f7e8b8940a9e2/src/enzyme/enzyme_wrapper.rs#L77

I guess this issue about allocating is outdated, correct? https://github.com/wsmoses/Enzyme.jl/blob/05c9b84f9051c5bd8863925e54a2fc4eb11fc784/src/Enzyme.jl#L126

Structure of user-facing macros

Enzyme's integration into Rust should give the user

  • precise control of what are constant, primary and adjoint variables
  • make it easy to mark a function as a candidate for differentiation
  • inline derivatives into existing structs or calls directly in code

There several ways how to structure macros.

fn test(a: f32, b: f32) -> f32 {
    a * b
}

derive_diff!(test, test_first, a: const, b: dup);
derive_diff!(test_first, test_second, a: const, b: dup);

/*
Created by the derive_diff! proc-macro:
fn test_second(a: f32, b: f32, b_dup: f32) -> f32 {
 unreachable!();
}

*/

or

#[diff(
   test_diff: (const, dup) -> active
)]
fn test(a: f32, b: f32) -> f32 {
    a * b
}

or

diff fn test(a: (f32, const), b: (f32, dup)) -> (f32, const) {
}

Add documentation

To manage expectations, we should document what is expected to work already, what we will possibly fix and what is probably
not going to work till we have finished a better Enzyme integration. All issues should be solved in the next iteration.
This overview should be moved into a real documentation..

Not Working, unlikely to be fixed:

  • Differentiating a function having one or more f32 parameters (f64 works and f32 values are allowed inside of the function).
  • Not ffi-safe types
  • Differentiating a function which accepts or uses a dyn trait object anywhere.
  • AD across language barriers
  • CUDA and HIP support.
  • Using oxide-enzyme in a dependency, rather than in your current main project.

Likely to be fixed

  • Only a combined forward+ReverseAD pass is supported at the moment. Enzyme does support ForwardAD and splitting
    the forward and reverse pass of ReverseAD. (Working on it).
  • Differentiating a function which uses parallelism (e.g. rayon) other than OpenMP / MPI.
  • Differentiating a function using BLAS / Lapack routines. If you compile them with debug symbols it will already work.
    If you use a pre-compiled version it might or might not work, Enzyme doesn't cover all functions yet.

Fixed

  • Generating Functions returning one f64 parameter in the return struct.
  • Generating Functions returning three or more parameters in the return struct.

Workarounds:

  • There is no real workaround for the dyn trait issue, since it requires creating a new vTable. You might be lucky if that object
    is not interacting with your active values, but that's not granted.
  • You can cast f32 to f64 values before passing them to the function and cast them back to f32 inside of your function.
  • If you want to differentiate functions calling C/C++/Julia/Fortran/... routines, first compile those other languages to .bc bitcode files. Look up the build script on where to place them. You might run into some issues due to missing symbols when using other languages, let me know about your experiences.
  • Similar to differentiating across language barriers. Technically it will work if you use rust wrappers around them or if you use the clang version which we build in ~/.cache/enzyme/rustc-<version>-src, but that is probably too messy to set up and also requires looking at https://enzyme.mit.edu/getting_started/CUDAGuide/

Add a License

We should agree on a License @bytesnake.
I would prefer MIT / Apache2, as it's the default for the rust side.
However, we should verify that this isn't conflicting with LLVM's / Enzyme's Apache2-with-LLVM-exceptions licsense.

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    ๐Ÿ–– Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. ๐Ÿ“Š๐Ÿ“ˆ๐ŸŽ‰

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google โค๏ธ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.