Feature implementations: Make Variable and Graph generic to support possible symbolic differentiation

[Feiyang472]

Hi @avhz . Thanks for writing and open sourcing this project.

I would like to add symbolic differentiation capabilities to the autodiff subcrate.
There are a couple of rust-autodiff packages out there but none look as actively maintained as your package.
My end goal here is to have a procedural macro for compile-time autodiffed code generation, so that I can have autodiff without a on-heap computation graph.

This PR just extends a bunch of generic typing to the Graph/Vertex/Variable structs to support the possibility of putting expressions to Graphs. I'd be glad to hear your thoughts!

[avhz]

Hi :)

Could you provide an example of how this would be used in contrast to the current implementation ?

Cheers !

[Feiyang472]

For sure!
I envision it looking like this

fn european_vomma(s, k, v, …) {
    partial_diff!(european_price, v, 2)
}

and

[derive(Cumulant)]
struct Heston {…} // gets an impl of Cumulants trait which forms a part of Distribution

But we are not there yet. I stopped the PR here, to have a discussion about current and next directions I am taking, and to stop the PR becoming too long to review.

In this PR I only made the container structs generic. To have symbolic derivatives, we would further support a Variable which puts and assembles necessary derivatives nodes on the graph in a process similar to the current accumulate. the resultant
partial = Variable<'a Expression>
is equivalent to a syntax tree and can be processed into a tokenstream by a proc macro.

Thanks

[avhz]

This would be cool, however I suspect it will achieve what Enyzme does, which is making it's way into nightly soon I think.

Edit: link to rust-lang tracking issue

[Feiyang472]

Absolutely enzyme will be great. I havent used it myself but have been looking forward to it for a while. I would say we are not reinventing the wheel here

1. because enzyme is a much more powerful tool which will take longer to come along and to get into stable.
2. afaik it's llvm-based instead of syntax based, which I suspect means it will be less trivial for it to optimise away obviously unused variables like x d(y*x)/dy doesnt depend on y
3. complex numbers, for characteristic functions: link https://enzyme.mit.edu/julia/dev/faq/#Complex-numbers . finance maths can pretty much treat imaginary unit as a special case and ignore branches and singularities and get away with it 90% of the time whereas enzyme cant afford to do it.

Once enzyme is rolled out, I would be very glad to integrate it where we can. I imagine it would be quite fun to do some typing gymnastics with date/putcall/exercise rule stuff.

[ZuseZ4]

I'd be very curious if you can come up with an example where Enzyme can't eliminate dead code, do you have a concrete one in mind?
Compilers like LLVM have specific passes for that https://llvm.org/doxygen/DCE_8cpp_source.html and Enzyme has ActivityAnalysis, so it will only compute whatever is necessary for those arguments that you marked as active.
That doesn't necessarily mean that they work exactly how they should, Enzyme is a lot younger than LLVM and still has some bugs, but it would be worth reporting them.

Edit: In case that you literally meant your small example above, I've implemented it here (in C, since our rust explorer is just getting updated) https://fwd.gymni.ch/ApN1pK We have

double square(double x, double y) { return x * y; }

double dsquare(double x, double y) {
  // d(y*x)/dy, so y is active (implicit default), x is const.
  return x * __enzyme_autodiff((void *)square, enzyme_const, x, y);
}

If you add -S emit-llvm than you can see what LLVM+Enzyme generate for dsquare, which is

define dso_local double @dsquare(double noundef %0, double %1) local_unnamed_addr #0 !dbg !30 {
  call void @llvm.dbg.value(metadata double %0, metadata !31, metadata !DIExpression()), !dbg !32
  call void @llvm.dbg.value(metadata double poison, metadata !33, metadata !DIExpression()), !dbg !32
  %3 = fmul double %0, %0, !dbg !34
  ret double %3, !dbg !35
}

so it returns x*x and ignores y, which is your desired results, IIUC?
The other reasons to not use it are ofc. totally valid, it would be very brave to build a whole project on Rust-enzyme in the current state.

[Feiyang472]

Hi @ZuseZ4 thanks for your comment, your work on integrating enzyme will be a game-changer for me.
In terms of deadcode, I just wanted some ergonomics. In here https://fwd.gymni.ch/rQS5m7 dsquare can be made not to take y as argument by just passing y as nan into enzyme. But I didnt know if at rust level I could extract the fact that y can be passed as NaN into enzyme, or there is some LLVM metadata specifier I missed from the docs. However now that I think about it I am not even sure if this kind of ergonomics is really desirable. At the time of writing I just thought "should gamma have put call type as an argument"
I will pull together an example from heston characteristic functions, hopefully later this week.

[ZuseZ4]

Ah, I think I see what you mean. The problem here is that whether it can be NaN will (in more complex cases) at some point be optimization dependent, e.g. in Debug mode your LLVM-IR will be much more complex due to the lack of optimizations, so Enzyme might also generate code where y is used (even if it cancels itself out). Also in the Rust frontend, you can't ask the backend like LLVM if something will be optimized out, that would require something like Reflection which Rust doesn't have right now.

What you can do though is asking cargo to emit the IR of your functions and scan either the debug info (poison for y), or whether an argument is noundef. It's a bit hacky, but you can probably get something through a build script to work, that's the closest to reflection what we have.

All in all you can do the same in Rust as in C++ in this case. Also, I assume this is just to have a nice design. If an argument ends up being unused (like here) it won't be moved/copied around, so you as a user/dev don't have to try to optimize the function argument away, LLVM already does that for you for free.