Add merge rotation pass

doc/releases/changelog-dev.md

@@ -117,6 +117,43 @@
   Available MLIR passes are now documented and available within the
   [catalyst.passes module documentation](https://docs.pennylane.ai/projects/catalyst/en/stable/code/__init__.html#module-catalyst.passes).
 
+* A peephole merge rotations pass is now available in MLIR. It can be added to `catalyst.passes.pipeline`, or the 
+  Python function `catalyst.passes.merge_rotations` can be directly called on a `QNode`.
+  [(#1162)](https://github.com/PennyLaneAI/catalyst/pull/1162)
+
+  Using the pipeline, one can run:
+
+  ```python
+  from catalys.passes import pipeline
+
+  my_passes = {
+    "merge_rotations": {}
+  }
+
+  @qjit(circuit_transform_pipeline=my_passes)
+  @qml.qnode(qml.device("lightning.qubit", wires=1))
+  def g(x: float):
+      qml.RX(x, wires=0)
+      qml.RX(x, wires=0)
+      qml.Hadamard(wires=0)
+      return qml.expval(qml.PauliZ(0))
+  ```
+
+  Using the python function, one can run:
+
+  ```python
+  from catalys.passes import merge_rotations
+
+  @qjit
+  @merge_rotations
+  @qml.qnode(qml.device("lightning.qubit", wires=1))
+  def g(x: float):
+      qml.RX(x, wires=0)
+      qml.RX(x, wires=0)
+      qml.Hadamard(wires=0)
+      return qml.expval(qml.PauliZ(0))
+  ```
+
 * Catalyst Autograph now supports updating a single index or a slice of JAX arrays using Python's
   array assignment operator syntax.
   [(#769)](https://github.com/PennyLaneAI/catalyst/pull/769)

frontend/catalyst/passes.py

@@ -310,6 +310,87 @@ def wrapper(*args, **kwrags):
     return fn_clone
 
 
+def merge_rotations(fn=None):
+    """
+    Specify that the ``-merge-rotations`` MLIR compiler pass
+    for merging roations (peephole) will be applied.
+
+    The full list of supported gates are as follows:
+
+    :class:`qml.RX <pennylane.RX>`,
+    :class:`qml.CRX <pennylane.CRX>`,
+    :class:`qml.RY <pennylane.RY>`,
+    :class:`qml.CRY <pennylane.CRY>`,
+    :class:`qml.RZ <pennylane.RZ>`,
+    :class:`qml.CRZ <pennylane.CRZ>`,
+    :class:`qml.PhaseShift <pennylane.PhaseShift>`,
+    :class:`qml.ControlledPhaseShift <pennylane.ControlledPhaseShift>`,
+    :class:`qml.Rot <pennylane.Rot>`,
+    :class:`qml.CRot <pennylane.CRot>`,
+    :class:`qml.MultiRZ <pennylane.MultiRZ>`.
+
+
+    .. note::
+
+        Unlike PennyLane :doc:`circuit transformations <introduction/compiling_circuits>`,
+        the QNode itself will not be changed or transformed by applying these
+        decorators.
+
+        As a result, circuit inspection tools such as :func:`~.draw` will continue
+        to display the circuit as written in Python.
+
+    Args:
+        fn (QNode): the QNode to apply the cancel inverses compiler pass to
+
+    Returns:
+        ~.QNode:
+
+    **Example**
+
+    In this example the three :class:`qml.RX <pennylane.RX>` will be merged in a single
+    one with the sum of angles as parameter.
+
+    .. code-block:: python
+
+        from catalyst.debug import get_compilation_stage
+        from catalyst.passes import merge_rotations
+
+        dev = qml.device("lightning.qubit", wires=1)
+
+        @qjit(keep_intermediate=True)
+        @merge_rotations
+        @qml.qnode(dev)
+        def circuit(x: float):
+            qml.RX(x, wires=0)
+            qml.RX(0.1, wires=0)
+            qml.RX(x**2, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+    >>> circuit(0.54)
+    Array(0.5965506257017892, dtype=float64)
+    """
+    if not isinstance(fn, qml.QNode):
+        raise TypeError(f"A QNode is expected, got the classical function {fn}")
+
+    funcname = fn.__name__
+    wrapped_qnode_function = fn.func
+    uniquer = str(_rename_to_unique())
+
+    def wrapper(*args, **kwrags):
+        if EvaluationContext.is_tracing():
+            apply_registered_pass_p.bind(
+                pass_name="merge-rotations",
+                options=f"func-name={funcname}" + "_merge_rotations" + uniquer,
+            )
+        return wrapped_qnode_function(*args, **kwrags)
+
+    fn_clone = copy.copy(fn)
+    fn_clone.func = wrapper
+    fn_clone.__name__ = funcname + "_merge_rotations" + uniquer
+
+    return fn_clone
+
+
 ## IMPL and helpers ##
 # pylint: disable=missing-function-docstring
 class _PipelineNameUniquer:
@@ -332,7 +413,7 @@ def _rename_to_unique():
 
 
 def _API_name_to_pass_name():
-    return {"cancel_inverses": "remove-chained-self-inverse", "merge_rotations": "merge-rotation"}
+    return {"cancel_inverses": "remove-chained-self-inverse", "merge_rotations": "merge-rotations"}
 
 
 def _inject_transform_named_sequence():

frontend/test/lit/test_peephole_optimizations.py

@@ -31,7 +31,7 @@
 
 from catalyst import qjit
 from catalyst.debug import get_compilation_stage
-from catalyst.passes import cancel_inverses, pipeline
+from catalyst.passes import cancel_inverses, merge_rotations, pipeline
 
 
 def flush_peephole_opted_mlir_to_iostream(QJIT):
@@ -86,7 +86,7 @@ def test_pipeline_lowering():
     """
     my_pipeline = {
         "cancel_inverses": {},
-        "merge_rotations": {"my-option": "aloha"},
+        "merge_rotations": {},
     }
 
     @qjit(keep_intermediate=True)
@@ -104,14 +104,14 @@ def test_pipeline_lowering_workflow(x):
     # CHECK:   pass_name=remove-chained-self-inverse
     # CHECK: ]
     # CHECK: _:AbstractTransformMod() = apply_registered_pass[
-    # CHECK:   options=func-name=test_pipeline_lowering_workflow_transformed0 my-option=aloha
-    # CHECK:   pass_name=merge-rotation
+    # CHECK:   options=func-name=test_pipeline_lowering_workflow_transformed0
+    # CHECK:   pass_name=merge-rotations
     # CHECK: ]
     print_jaxpr(test_pipeline_lowering_workflow, 1.2)
 
     # CHECK: transform.named_sequence @__transform_main
     # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=test_pipeline_lowering_workflow_transformed0"}
-    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotation" to {{%.+}} {options = "func-name=test_pipeline_lowering_workflow_transformed0 my-option=aloha"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=test_pipeline_lowering_workflow_transformed0"}
     # CHECK-NEXT: transform.yield
     print_mlir(test_pipeline_lowering_workflow, 1.2)
 
@@ -160,13 +160,13 @@ def test_pipeline_lowering_keep_original_workflow(x):
     # CHECK: ]
     # CHECK: _:AbstractTransformMod() = apply_registered_pass[
     # CHECK:   options=func-name=f_transformed0
-    # CHECK:   pass_name=merge-rotation
+    # CHECK:   pass_name=merge-rotations
     # CHECK: ]
     print_jaxpr(test_pipeline_lowering_keep_original_workflow, 1.2)
 
     # CHECK: transform.named_sequence @__transform_main
     # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=f_transformed0"}
-    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotation" to {{%.+}} {options = "func-name=f_transformed0"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=f_transformed0"}
     # CHECK-NEXT: transform.yield
     print_mlir(test_pipeline_lowering_keep_original_workflow, 1.2)
 
@@ -223,23 +223,23 @@ def h(x):
     # CHECK: ]
     # CHECK: _:AbstractTransformMod() = apply_registered_pass[
     # CHECK:   options=func-name=g_transformed0
-    # CHECK:   pass_name=merge-rotation
+    # CHECK:   pass_name=merge-rotations
     # CHECK: ]
     # CHECK: _:AbstractTransformMod() = apply_registered_pass[
     # CHECK:   options=func-name=h_transformed1
     # CHECK:   pass_name=remove-chained-self-inverse
     # CHECK: ]
     # CHECK: _:AbstractTransformMod() = apply_registered_pass[
     # CHECK:   options=func-name=h_transformed1
-    # CHECK:   pass_name=merge-rotation
+    # CHECK:   pass_name=merge-rotations
     # CHECK: ]
     print_jaxpr(global_wf)
 
     # CHECK: transform.named_sequence @__transform_main
     # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=g_transformed0"}
-    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotation" to {{%.+}} {options = "func-name=g_transformed0"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=g_transformed0"}
     # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=h_transformed1"}
-    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotation" to {{%.+}} {options = "func-name=h_transformed1"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=h_transformed1"}
     # CHECK-NEXT: transform.yield
     print_mlir(global_wf)
 
@@ -301,20 +301,20 @@ def h(x):
     # CHECK: ]
     # CHECK: _:AbstractTransformMod() = apply_registered_pass[
     # CHECK:   options=func-name=g_transformed1
-    # CHECK:   pass_name=merge-rotation
+    # CHECK:   pass_name=merge-rotations
     # CHECK: ]
     # CHECK: _:AbstractTransformMod() = apply_registered_pass[
     # CHECK:   options=func-name=h_transformed0
     # CHECK-NOT:   pass_name=remove-chained-self-inverse
-    # CHECK:   pass_name=merge-rotation
+    # CHECK:   pass_name=merge-rotations
     # CHECK: ]
     print_jaxpr(global_wf)
 
     # CHECK: transform.named_sequence @__transform_main
     # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=g_transformed1"}
-    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotation" to {{%.+}} {options = "func-name=g_transformed1"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=g_transformed1"}
     # CHECK-NOT: {{%.+}} = transform.apply_registered_pass "remove-chained-self-inverse" to {{%.+}} {options = "func-name=h_transformed0"}
-    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotation" to {{%.+}} {options = "func-name=h_transformed0"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=h_transformed0"}
     # CHECK-NEXT: transform.yield
     print_mlir(global_wf)
 
@@ -563,3 +563,70 @@ def test_cancel_inverses_keep_original_workflow2():
 
 
 test_cancel_inverses_keep_original()
+
+
+#
+# merge_rotations
+#
+
+
+def test_merge_rotations_tracing_and_lowering():
+    """
+    Test merge_rotations during tracing and lowering
+    """
+
+    @qjit
+    def test_merge_rotations_tracing_and_lowering_workflow(xx: float):
+
+        @merge_rotations
+        @qml.qnode(qml.device("lightning.qubit", wires=1))
+        def f(x: float):
+            qml.RX(x, wires=0)
+            qml.RX(x, wires=0)
+            qml.Hadamard(wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        @merge_rotations
+        @qml.qnode(qml.device("lightning.qubit", wires=1))
+        def g(x: float):
+            qml.RX(x, wires=0)
+            qml.RX(x, wires=0)
+            qml.Hadamard(wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        @qml.qnode(qml.device("lightning.qubit", wires=1))
+        def h(x: float):
+            qml.RX(x, wires=0)
+            qml.RX(x, wires=0)
+            qml.Hadamard(wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        _f = f(xx)
+        _g = g(xx)
+        _h = h(xx)
+        return _f, _g, _h
+
+    # CHECK: transform_named_sequence
+    # CHECK: _:AbstractTransformMod() = apply_registered_pass[
+    # CHECK:   options=func-name=f_merge_rotations0
+    # CHECK:   pass_name=merge-rotations
+    # CHECK: ]
+    # CHECK: _:AbstractTransformMod() = apply_registered_pass[
+    # CHECK:   options=func-name=g_merge_rotations1
+    # CHECK:   pass_name=merge-rotations
+    # CHECK: ]
+    # CHECK-NOT: _:AbstractTransformMod() = apply_registered_pass[
+    # CHECK-NOT:   options=func-name=h_merge_rotations
+    # CHECK-NOT:   pass_name=merge-rotations
+    print_jaxpr(test_merge_rotations_tracing_and_lowering_workflow, 1.1)
+
+    # CHECK: module @test_merge_rotations_tracing_and_lowering_workflow
+    # CHECK: transform.named_sequence @__transform_main
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=f_merge_rotations0"}
+    # CHECK-NEXT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=g_merge_rotations1"}
+    # CHECK-NOT: {{%.+}} = transform.apply_registered_pass "merge-rotations" to {{%.+}} {options = "func-name=h_merge_rotations"}
+    # CHECK-NEXT: transform.yield
+    print_mlir(test_merge_rotations_tracing_and_lowering_workflow, 1.1)
+
+
+test_merge_rotations_tracing_and_lowering()