[MLIR] Enzyme-driven gradients

[pengmai]

Context: The current gradient architecture does not support exact computation of hybrid programs with classical postprocessing or multiple QNodes.

Description of the Change: A reworking of Catalyst's gradient architecture to be driven by Enzyme. Quantum functions are split into their purely classical preprocessing and quantum parts as before, but now the differentiation of the end-to-end circuit is done by Enzyme with registered quantum gradients as custom gradients.

Overview of changes

• Implement a splitting transformation that updates the hybrid computation graph such that QNodes are split into two functions, a .preprocess function that contains the classical preprocessing (just like the argmap function) and the .quantum function that contains the actual quantum computation. The .preprocess function ends in a call to the .quantum function, meaning it can replace QNodes and leave the hybrid graph connected.
  • This assumes that, like PennyLane, a QNode must end in one or more measurements and cannot contain postprocessing.
• GradOps are lowered to one or more BackpropOps of the entire hybrid computation (one for each result entry)
• Custom gradients are registered for the .quantum split out QNodes.
• Modify EinsumLinalgGeneric to support both memref args (in addition to tensor args) and dynamic shapes. This is used in the custom gradient of quantum functions.
• The differential lowerings for adjoint and parameter shift differentiation now operate on individual QNodes instead of GradOps.
  • There are also attributes used to connect the quantum gradients (.qgrad, .adjoint) to the custom gradients across both the lower-gradients and convert-gradient-to-llvm passes.
• The frontend verification is modified to allow for non-QNode callees of catalyst.grad functions while also preserving the same verification as before (return value of a QNode with diff_method="adjoint" must be an expval, etc)
• A bugfix within the runtime to prevent an error because the circuit is run twice with the adjoint method.

The important implementation quirks are documented on Notion.

Benefits: This new architecture supports differentiation of circuits with classical postprocessing, hybrid programs with multiple QNodes, and purely classical programs.

[sc-41364]
[sc-41375]
[sc-42856]

[rmoyard]

Amazing job 🥇 Happy to approve the PR but I would like @dime10 to review it as well before merging

[josh146]

This is great work @pengmai! Super excited to have it in 🎉

A quick question: will we also need to update https://docs.pennylane.ai/projects/catalyst/en/latest/dev/quick_start.html#calculating-quantum-gradients?

[pengmai]

Thanks Josh!

A quick question: will we also need to update https://docs.pennylane.ai/projects/catalyst/en/latest/dev/quick_start.html#calculating-quantum-gradients?

This PR should be totally backwards compatible and thus the quick start shouldn't require any changes to work, unless you wanted to highlight the new features.

[dime10]

This is amazing work 💯

A few comments and questions from my side, but nothing major!

Because this PR is so beefy though, I think it might be really helpful if the PR description contained a bullet point list of what was changed in existing code, what was added, and what are some quirks or compromises of the implantation, .... Doesn't need a lot of explanation, but just so we have a quick overview of everything that was undertaken.

[dime10]

Thanks for the overview in the description, very helpful 💯

[dime10]

🧬