[MLIR] Integrated BackpropOp lowering with Enzyme

[pengmai]

Context: The existing Enzyme integration has issues that prevent Enzyme from compiling our code and producing the correct results.

Description of the Change:

• Convert the argmap function to be in destination-passing style. This function is given directly to Enzyme instead of the wrapper, which is necessary to avoid Enzyme's type analysis failing.
• Unpack the MemRef descriptors when passing to Enzyme rather than wrapping the descriptor structs to avoid issues where Enzyme was incorrectly computing zero gradients. This lets us tell Enzyme that allocated pointers are constant while aligned pointers are not.

Benefits: The integration tests around hybrid gradients produce the correct results using Enzyme. Enzyme is computing exact classical gradients via backpropagation instead of estimating them with finite differences.

Note that although the hybrid gradients are producing correct results, the complete hybrid gradient work is still considered WIP. This PR primarily focuses on the lowering of the BackpropOp.

[erick-xanadu]

Great work! Just a high level comment regarding memory management functions.

[rmoyard]

Great work @pengmai 💯 Some questions and comments, happy to approve after that!

[rmoyard]

Maybe @erick-xanadu and @dime10 know more about this, can it impact negatively the compilation time for other programs (no gradients), is it negligible in general?

[erick-xanadu]

It is possible, but if it is necessary for correctness, then it is necessary. What we should do first is measure compilation/execution time on a set of representative programs to have certainty of how compiler options affect compilation time and execution time. But that is a work in progress...

[rmoyard]

Thanks @erick-xanadu !

[dime10]

The optimization is necessary for correctness? 😮

[erick-xanadu]

Maybe correctness is the wrong word, but I believe it is necessary for feasibility (it might even improve compile time of gradients when compared to compile time of gradients without -O2) ? @pengmai can correct me if I'm wrong. One example that I keep in mind from previous conversation is the -mem2reg optimization. Instead of tracking the indirection due to pointers, pointers get converted to registers. Same with inlining. -O2 might inline some functions that will make the compilation (and execution) of gradients perform better. However, the cost is that we do need to optimize it first.

I believe a lot of Enzyme presentations start with a diagram showing "optimize first" -> "enzyme" -> "optimize again".

[pengmai]

Right, Enzyme's main research thesis is to improve performance by differentiating optimized code, so getting good performance is dependent on first optimizing the code before running Enzyme.

Separately, we see some issues where Enzyme fails to generate code for our tests if we don't optimize the code first. These are likely bugs or unsupported cases within Enzyme, so a workaround is running this optimization first.

[rmoyard]

[sc-41510]

[codecov]

Codecov Report

Merging #193 (1859a53) into enzyme_integration (856080d) will decrease coverage by 0.85%.
The diff coverage is 100.00%.

@@                  Coverage Diff                   @@
##           enzyme_integration     #193      +/-   ##
======================================================
- Coverage               99.95%   99.10%   -0.85%     
======================================================
  Files                      21       37      +16     
  Lines                    4096     6462    +2366     
  Branches                    0      333     +333     
======================================================
+ Hits                     4094     6404    +2310     
- Misses                      2       33      +31     
- Partials                    0       25      +25     

Impacted Files	Coverage Δ
frontend/catalyst/compiler.py	100.00% <100.00%> (ø)

... and 15 files with indirect coverage changes

[rmoyard]

Thanks @pengmai, great PR 🥇

[erick-xanadu]

Awesome!