catalyst icon indicating copy to clipboard operation
catalyst copied to clipboard

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

Open paul0403 opened this issue 1 year ago • 6 comments

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.

  1. 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.

  1. 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.

  2. 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.

  3. 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.

paul0403 avatar Oct 20 '24 23:10 paul0403