enzymead / oxide-enzyme Goto Github PK

View Code? Open in Web Editor NEW

101.0 10.0 4.0 134 KB

Enzyme integration into Rust. Experimental, do not use.

License: Apache License 2.0

Rust 100.00%

rust auto-differentiation ad enzyme gradient derivative

oxide-enzyme's People

Contributors

Stargazers

Watchers

Forkers

simonsan strasdat anirudhacharya elite-sheep

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

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.

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/

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!();
}

*/

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

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

Recommend Projects

React

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

🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
Typescript

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

An Open Source Machine Learning Framework for Everyone
Django

The Web framework for perfectionists with deadlines.
Laravel

A PHP framework for web artisans
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.
Visualization

Some thing interesting about visualization, use data art
Game

Some thing interesting about game, make everyone happy.

Recommend Org

Facebook

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

Open source projects and samples from Microsoft.
Google

Google ❤️ Open Source for everyone.
Alibaba

Alibaba Open Source for everyone
D3

Data-Driven Documents codes.
Tencent

China tencent open source team.