impl new modes for higher order ad

[ZuseZ4]

Does currently not work, fails to diff the fwd function, the rev succeeds (forward-over-reverse example from julia docs). Fails inside of Enzyme

    5 rustc: /h/344/drehwald/prog/rust/build/x86_64-unknown-linux-gnu/llvm/include/llvm/IR/Instructions.h:3144: llvm::Value* llvm::ReturnInst::getOperand(unsigned int) const: Assertion `i_nocapture < OperandTraits<ReturnInst>::operands(this) && "getOperand() out of range!"' failed.

Extra info during compiling the example is available through rustc's tracing infra: RUSTC_LOG=rustc_codegen_llvm::llvm=trace cargo +enzyme build

Can you please look into it @wsmoses? I hope there is just some call setting wrong, but it would be nice if it would hit an enzyme assertion instead of an llvm one. Also, I here reuse the enzymelogicref, but the same issue happens if I don't.

[jedbrown]

Are you trying with a matching test? Perhaps an update to https://github.com/EnzymeAD/rustbook/pull/4?

[ZuseZ4]

Did you miss my zulip message @jedbrown ?

[jedbrown]

Oh, I see. https://github.com/EnzymeAD/rustbook/pull/5/files

[jedbrown]

error[E0425]: cannot find value `logic_ref` in this scope
    --> compiler/rustc_codegen_llvm/src/back/write.rs:1135:63
     |
1135 |         let res = enzyme_ad(llmod, llcx, &diag_handler, item, logic_ref);
     |                                                               ^^^^^^^^^
     |
help: a local variable with a similar name exists, consider changing it
     |
1124 |     let logic_ref: EnzymeLogicRef = CreateEnzymeLogic(fnc_opt as u8);
     |         ~~~~~~~~~

[jedbrown]

The error message is confusing because the local variable is _logic_ref, but that underscore doesn't make it into the message. Did you intend to pick just one of the opt or non-opt paths?

[ZuseZ4]

The logic is: If it's base for higher order => force optimization. Good for perf, but also increases chance of enzyme succeeding (since generating simpler code). If anything else (e.g. the forward part of forward-over-reverse, or just a single autodiff call), then it will follow your decision on dbg vs release

[jedbrown]

Thanks, do you know what's wrong with my example now?

diff --git a/samples/tests/neohookean/mod.rs b/samples/tests/neohookean/mod.rs
index 3d3dabb..76270d2 100644
--- a/samples/tests/neohookean/mod.rs
+++ b/samples/tests/neohookean/mod.rs
@@ -241,7 +241,7 @@ impl NH {

 // We can only differentiate free functions, not methods (yet)
 // Helmholtz free energy density
-#[autodiff(d_psi, Reverse, Duplicated, Const, Active)]
+#[autodiff(d_psi, ReverseFirst, Duplicated, Const, Active)]
 fn psi(e: &KM, nh: &NH) -> f64 {
     let mu = nh.mu;
     let lambda = nh.lambda;
@@ -335,7 +335,6 @@ fn check_stress() {
     assert!(diff < 1e-14);
 }

-#[cfg(broken)]
 #[test]
 fn check_dstress() {
     let nh = NH::from_youngs(1.0, 0.3);

And then

$ cargo +enzyme t
[...]
num_fnc_args: 2
input_activity.len(): 2
num_fnc_args: 2
input_activity.len(): 2
Function return type does not match operand type of return inst!
  ret { double } %5
 doubleLLVM ERROR: Broken function found, compilation aborted!
warning: `samples` (test "mod") generated 1 warning
error: could not compile `samples` (test "mod"); 1 warning emitted

[jedbrown]

Or should I try working around this by using an output argument instead of return value?

[ZuseZ4]

If you can get it to work just by changing the return style I'd appreciate it, since I hope for more perf from my tt work. If that doesn't work I can look into fixing it. It's only Rust error, so hopefully not too much work.