Add merge rotation patterns for qml.Rot and qml.CRot

[paul0403]

Context

In quantum circuits, two neighbouring rotation gates about the same rotation axis on the same wire can be merged into one, i.e. $R_x(a) R_x(b) = R_x(a+b)$.

In #1162, a merge rotation pass was implemented in Catalyst. The pass included qml.Rot/qml.CRot, but then it was realized that the angles in these two functions don't simply add when being merged, because y and z rotations do not commute. These two functions were removed from merge rotation pass in #1206 .

However, qml.Rot can be merged; it's just that the new angles of the merged gate are more complicated: https://docs.pennylane.ai/en/stable/code/api/pennylane.transforms.merge_rotations.html#details-on-rot

Goal

We should include qml.Rot and qml.CRot in merge rotation pass, with the full angle formulas, instead of just removing them from the pattern.

Requirements:

• In Catalyst mlir, two neighbouring "Rot" gates, and two neighbouring "CRot" gates, should be merged into a single one, with the new angles appropriately given by the formulas linked to above.
• The pass must not change the functionality (aka returned state/probabilities/...) of the circuit. This should be automatic if the formulas above are followed.
• For CRot gates, since now two wires are taken in, merging should only happen if the two wires on both gates are in the same order. Note that we have a verifier analysis that checks this, and is already used by the existing merge rotation pass.

Technical details

1. Catalyst has its own mlir quantum dialect. The operation definitions are in QuantumOps.td. The operation of particular interest to us is the CustomOp, which corresponds to a general named gate, with the name being stored as a string.

2. To see what kind of input mlir the pass will receive, see merge rotation mlir test for an example. In particular, the "Rot" and "CRot" tests right now checks that no merging happens; they should be updated to check that the correct angles are produced and then used by the merged gate.

3. We currently already have merge rotation patterns for all other rotation gates. The pattern simply adds the two rotation angles, because the simple rotation gates can be easily merged like that. You can add the pattern for qml.Rot and qml.CRot to this same file.

You should not need to modify the actual merge rotation pass containing the pattern applicator, though you can if you wish.

4. To compile your pattern/pass, run make dialects from the top-level catalyst directory. This will build all passes in Catalyst and run all filecheck tests in mlir/test.

You can also invoke individual passes on an arbitrary input mlir file with catalyst/mlir/build/bin/quantum-opt, and use it like the standard mlir opt. At the top of each mlir test file, you can see how each pass is invoked.

You can use quantum-opt --help to see all available passes.

5. You might want to see what Catalyst mlir a frontend PennyLane python program would produce. To lower frontend PennyLane python code to mlir, you can use the keep_intermediate option in qjit. See here.

6. To apply the merge rotation pass from frontend PennyLane, you can use the catalyst.passes.merge_rotations decorator. See here for the documentation on a similar optimization, cancel_inverses,alongside how to use keep_intermediate, and here for merge_rotation.

7. Various mathematical operations, like sqrt and arctan, are in the mlir math dialect

Installation help

Complete instructions to install Catalyst from source can be found here. Note that due to the size of the llvm-project it can take a while (~3 hrs on a personal laptop) to compile.

[Mohxen]

Hi @paul0403,

In tackling the merging of qml.Rot and qml.CRot gates, I’m considering an approach that goes beyond simple angle addition due to the non-commutativity of Y and Z rotations. Specifically, based on your document, I’m using the Z-Y-Z Euler angle decomposition to calculate the combined rotations, treating each gate as a rotation matrix.

The steps involve:

Converting individual qml.Rot or qml.CRot angles into matrices.
Multiplying these matrices to get a composite rotation.
Converting the resulting matrix back to Euler angles, ensuring alignment with Z-Y-Z decomposition.

This approach aims to preserve the correct sequence and dependencies of rotations. Could you advise if this method is consistent with the intended use in Catalyst, or if adjustments are recommended?

[paul0403]

Hi @Mohxen , I'll write up a more detailed response tomorrow. For now let me link you to the explicit formulas for the result rotation gate's angles in PennyLane's parametrization

[paul0403]

Hi @Mohxen, I would advise against using the matrix approach for the following reasons:

1. While the encoding between a qml.Rot(a, b, c) and the rotation matrix is fixed, when converting a rotation matrix back to the angles, the logarithms of complex numbers will be involved, which might not be straightforward in mlir/cpp backend and require some extra work.
2. Matrix multiplication (MM) is an expensive process. While the merge rotation pass occurs at compile time, since Catalyst is a JIT compiler, the compile time is actually felt by the user (as the runtime of their first run). If possible, we'd like to avoid matrix multiplication if there is a scalar formula for the purpose. Also on this point, the product rotation matrix carries four terms, but we really just need one of these four terms to compute the new rotation angles, which means some computations are unnecessary in the MM approach.
3. In general, we do not want to associate gates to matrices in Catalyst. If you take a look at the quantum.customOp definition (point 1 in technical details above), you will see that the gates are characterized by name, not by matrices. This is because Catalyst is designed to be backend agnostic. For real quantum hardware, they don't use the notion of "gate = matrices" and perform the computations through MM. They just, well, do the computation.
4. Finally, if you feel like the explicit scalar angle formulas are too long, rest assured that since this is an assessment, we fully understand that you might have other commitments at the moment. There is no need to get everything perfect and production ready! Just do the best you can and we'll take it from there 😉 (Though the rotation matrices and the following extraction from production matrix to new angles are also quite complicated, so my gut feeling is the IR would probably be just as long in the matrix approach).

Don't hesitate to let me know if you need any help!

[Mohxen]

Hi @paul0403,

Thank you so much for the guidance and detailed points! I’m working to make this contribution usable within Catalyst, beyond just completing my assessment.

Using the scalar formulas from PennyLane’s documentation, I’ve implemented the approach in merge_rotation.cpp with MLIR’s Arith and Math libraries, and I’ve successfully passed make all and make test. Currently, I’m focusing on MergeRotationsTest.mlir, although I’m not entirely sure which specific checks we need, given that we already have existing tests for rot and crot.

Additionally, I’m setting up a frontend test in test_peephole_optimizations.py to observe the merged rotations' output, and I’m considering whether a new test is needed explicitly for this merge behavior.

[paul0403]

@Mohxen Thanks for working on this. Can you open a PR? It would be easier to review.

You can fork catalyst repo, push to your fork, and make a PR to merge your branch into catalyst main branch. See, for example, what I did in my interview #458

[Mohxen]

Hi @paul0403,

Thank you for the feedback! I've followed your instructions and opened a pull request as you asked. The PR for my changes is available at #1270. This PR includes the merge rotation patterns for qml.Rot and qml.CRot with full-angle formulas, as discussed, and I have added tests to verify the merging behavior.

Could you please review my PR and provide any feedback? I’m happy to make any edits based on your suggestions.

Thank you again for your guidance!