Implement a decomposition of PPRs into PPMs

doc/releases/changelog-dev.md

@@ -2,6 +2,20 @@
 
 <h3>New features since last release</h3>
 
+* A new compilation pass called :func:`~.passes.ppr_to_ppm` has been added to Catalyst to decompose Pauli product rotations (PPRs) into
+  Pauli product measurements (PPMs). Non-Clifford (pi/8) rotations require the consumption of a
+  magic state, while Clifford rotations will not. The methods are implemented as described in
+  [arXiv:1808.02892](https://arxiv.org/abs/1808.02892v3).
+  [(#1664)](https://github.com/PennyLaneAI/catalyst/pull/1664)
+
+  The new compilation pass can be accessed in the frontend from the :mod:`~.passes` module:
+  * :func:`~.passes.ppr_to_ppm`: Decomposes all PPRs into PPMs (except Pauli rotations).
+
+  Or via the Catalyst CLI as two separate passes:
+  * `decompose_clifford_ppr`: Decompose pi/4 rotations into PPMs.
+  * `decompose_non_clifford_ppr`: Decompose pi/8 rotations into PPMs using a magic state.
+
+
 <h3>Improvements 🛠</h3>
 
 * The behaviour of measurement processes executed on `null.qubit` with QJIT is now more in line with

frontend/catalyst/passes/__init__.py

@@ -39,6 +39,7 @@
     ions_decomposition,
     merge_ppr_ppm,
     merge_rotations,
+    ppr_to_ppm,
     to_ppr,
 )
 from catalyst.passes.pass_api import Pass, PassPlugin, apply_pass, apply_pass_plugin
@@ -54,4 +55,5 @@
     "apply_pass",
     "apply_pass_plugin",
     "merge_ppr_ppm",
+    "ppr_to_ppm",
 )

frontend/catalyst/passes/builtin_passes.py

@@ -499,3 +499,83 @@ def circuit():
 
     """
     return PassPipelineWrapper(qnode, "merge_ppr_ppm")
+
+
+def ppr_to_ppm(qnode=None, decompose_method="auto-corrected", avoid_y_measure=False):
+    R"""Specify that the MLIR compiler passes for decomposing Pauli Product rotations (PPR)
+    into Pauli Pauli measurements (PPM) will be applied.
+
+    This pass is used to decompose both non-Clifford and Clifford PPRs into PPMs. The non-Clifford
+    PPRs (pi/8) are decomposed first, and then Clifford PPRs (pi/4) are decomposed. Non-Clifford
+    decomposition can be performed in one of two ways: "clifford-corrected" or "auto-corrected",
+    by default the latter is used. Both methods are based on
+    `A Game of Surface Codes <https://arxiv.org/abs/1808.02892>`, figures 7 and 17(b) respectively.
+
+    Args:
+        qnode (QNode, optional): QNode to apply the pass to. If None, returns a decorator.
+        decompose_method (str, optional): The method to use for decomposing non-Clifford PPRs.
+            Options are "auto-corrected" and "clifford-corrected". Defaults to "auto-corrected".
+            "auto-corrected" uses an additional measurement for correction.
+            "clifford-corrected" uses a Clifford rotation for correction.
+        avoid_y_measure (bool): Rather than performing a Pauli-Y measurement for Clifford rotations
+            (sometimes more costly), a :math:`Y` state (:math:`Y\vert 0 \rangle`) is used instead
+            (requires :math:`Y` state preparation). Defaults to ``False``.
+
+    Returns:
+        ~.QNode or callable: Returns decorated QNode if qnode is provided,
+            otherwise returns a decorator.
+
+    **Example**
+
+    This example shows the sequence of passes that will be applied. The last pass
+    will convert the non-Clifford PPR into Pauli Product Measurements.
+
+    .. code-block:: python
+
+        import pennylane as qml
+        from catalyst import qjit, measure
+        from catalyst.passes import to_ppr, commute_ppr, merge_ppr_ppm, ppr_to_ppm
+
+        pipeline = [("pipe", ["enforce-runtime-invariants-pipeline"])]
+
+        @qjit(pipelines=pipeline, target="mlir")
+        @to_ppr
+        @commute_ppr
+        @merge_ppr_ppm
+        @ppr_to_ppm
+        @qml.qnode(qml.device("null.qubit", wires=2))
+        def circuit():
+            qml.H(0)
+            qml.T(0)
+            qml.CNOT([0, 1])
+            return measure(0), measure(1)
+
+        print(circuit.mlir_opt)
+
+    Example MLIR Representation:
+
+    .. code-block:: mlir
+
+        . . .
+        %5 = qec.fabricate  zero : !quantum.bit
+        %6 = qec.fabricate  magic : !quantum.bit
+        %mres, %out_qubits:2 = qec.ppm ["X", "Z"] %1, %6 : !quantum.bit, !quantum.bit
+        %mres_0, %out_qubits_1:2 = qec.ppm ["Z", "Y"] %5, %out_qubits#1 : !quantum.bit, !quantum.bit
+        %mres_2, %out_qubits_3 = qec.ppm ["X"] %out_qubits_1#1 : !quantum.bit
+        %mres_4, %out_qubits_5 = qec.select.ppm(%mres, ["X"], ["Z"]) %out_qubits_1#0 : !quantum.bit
+        %7 = arith.xori %mres_0, %mres_2 : i1
+        %8 = qec.ppr ["X"](2) %out_qubits#0 cond(%7) : !quantum.bit
+        . . .
+
+    """
+    passes = {
+        "decompose_non_clifford_ppr": {"decompose-method": decompose_method},
+        "decompose_clifford_ppr": {"avoid-y-measure": avoid_y_measure},
+    }
+
+    def decorator(qnode_func):
+        return PassPipelineWrapper(qnode_func, passes)
+
+    if qnode is None:
+        return decorator
+    return decorator(qnode)

frontend/catalyst/passes/pass_api.py

@@ -377,4 +377,5 @@ def _API_name_to_pass_name():
         "to_ppr": "to_ppr",
         "commute_ppr": "commute_ppr",
         "merge_ppr_ppm": "merge_ppr_ppm",
+        "ppr_to_ppm": "ppr_to_ppm",
     }

frontend/test/pytest/test_peephole_optimizations.py

@@ -18,8 +18,15 @@
 import pennylane as qml
 import pytest
 
-from catalyst import pipeline, qjit
-from catalyst.passes import cancel_inverses, merge_rotations, to_ppr
+from catalyst import measure, pipeline, qjit
+from catalyst.passes import (
+    cancel_inverses,
+    commute_ppr,
+    merge_ppr_ppm,
+    merge_rotations,
+    ppr_to_ppm,
+    to_ppr,
+)
 
 # pylint: disable=missing-function-docstring
 
@@ -189,8 +196,6 @@ def test_chained_apply_passes_workflow(x: float):
     assert "merge-rotations" in test_chained_apply_passes_workflow.mlir
 
 
-test_chained_passes()
-
 #
 # to_ppr
 #
@@ -221,7 +226,153 @@ def f():
     assert "qec.ppr" in optimized_ir
 
 
-test_convert_clifford_to_ppr()
+#
+# commute_ppr
+#
+
+
+def test_commute_ppr():
+    """
+    Test commute_ppr
+    """
+    pipe = [("pipe", ["enforce-runtime-invariants-pipeline"])]
+
+    @qjit(pipelines=pipe, target="mlir")
+    def test_commute_ppr_workflow():
+
+        @to_ppr
+        @commute_ppr
+        @qml.qnode(qml.device("lightning.qubit", wires=2))
+        def f():
+            qml.H(0)
+            qml.S(1)
+            qml.T(0)
+            qml.CNOT([0, 1])
+            return measure(0), measure(1)
+
+        return f()
+
+    assert 'transform.apply_registered_pass "commute_ppr"' in test_commute_ppr_workflow.mlir
+    optimized_ir = test_commute_ppr_workflow.mlir_opt
+    assert 'transform.apply_registered_pass "commute_ppr"' not in optimized_ir
+    assert "qec.ppr" in optimized_ir
+    assert "qec.ppm" in optimized_ir
+
+
+#
+# merge_ppr_ppm
+#
+
+
+def test_merge_ppr_ppm():
+    """
+    Test merge_ppr_ppm
+    """
+    pipe = [("pipe", ["enforce-runtime-invariants-pipeline"])]
+
+    @qjit(pipelines=pipe, target="mlir")
+    def test_merge_ppr_ppm_workflow():
+
+        @to_ppr
+        @merge_ppr_ppm
+        @qml.qnode(qml.device("lightning.qubit", wires=2))
+        def f():
+            qml.H(0)
+            qml.S(1)
+            qml.CNOT([0, 1])
+            return measure(0), measure(1)
+
+        return f()
+
+    assert 'transform.apply_registered_pass "merge_ppr_ppm"' in test_merge_ppr_ppm_workflow.mlir
+    optimized_ir = test_merge_ppr_ppm_workflow.mlir_opt
+    assert 'transform.apply_registered_pass "merge_ppr_ppm"' not in optimized_ir
+    assert 'qec.ppm ["Z", "X"]' in optimized_ir
+    assert 'qec.ppm ["X"]' in optimized_ir
+
+
+#
+# ppr_to_ppm
+#
+
+
+def test_ppr_to_ppm():
+    """
+    Test ppr_to_ppm
+    """
+    pipe = [("pipe", ["enforce-runtime-invariants-pipeline"])]
+
+    @qjit(pipelines=pipe, target="mlir")
+    def test_ppr_to_ppm_workflow():
+
+        @to_ppr
+        @ppr_to_ppm
+        @qml.qnode(qml.device("lightning.qubit", wires=2))
+        def f():
+            qml.H(0)
+            qml.S(1)
+            qml.T(0)
+            qml.CNOT([0, 1])
+            return measure(0), measure(1)
+
+        return f()
+
+    assert (
+        'transform.apply_registered_pass "decompose_non_clifford_ppr"'
+        in test_ppr_to_ppm_workflow.mlir
+    )
+    assert (
+        'transform.apply_registered_pass "decompose_clifford_ppr"' in test_ppr_to_ppm_workflow.mlir
+    )
+    optimized_ir = test_ppr_to_ppm_workflow.mlir_opt
+    assert 'transform.apply_registered_pass "decompose_non_clifford_ppr"' not in optimized_ir
+    assert 'transform.apply_registered_pass "decompose_clifford_ppr"' not in optimized_ir
+    assert "qec.fabricate  zero" in optimized_ir
+    assert "qec.fabricate  magic" in optimized_ir
+    assert "qec.select.ppm" in optimized_ir
+    assert 'qec.ppr ["X"]' in optimized_ir
+
+
+#
+# ppr_to_ppm with clifford-corrected
+#
+
+
+def test_ppr_to_ppm_inject_magic_state():
+    """
+    Test ppr_to_ppm
+    """
+    pipe = [("pipe", ["enforce-runtime-invariants-pipeline"])]
+
+    @qjit(pipelines=pipe, target="mlir")
+    def test_ppr_to_ppm_workflow():
+
+        @to_ppr
+        @ppr_to_ppm(decompose_method="clifford-corrected", avoid_y_measure=True)
+        @qml.qnode(qml.device("lightning.qubit", wires=2))
+        def f():
+            qml.H(0)
+            qml.S(1)
+            qml.T(0)
+            qml.CNOT([0, 1])
+            return measure(0), measure(1)
+
+        return f()
+
+    assert (
+        'transform.apply_registered_pass "decompose_non_clifford_ppr"'
+        in test_ppr_to_ppm_workflow.mlir
+    )
+    assert (
+        'transform.apply_registered_pass "decompose_clifford_ppr"' in test_ppr_to_ppm_workflow.mlir
+    )
+    optimized_ir = test_ppr_to_ppm_workflow.mlir_opt
+    assert 'transform.apply_registered_pass "decompose_non_clifford_ppr"' not in optimized_ir
+    assert 'transform.apply_registered_pass "decompose_clifford_ppr"' not in optimized_ir
+    assert "qec.fabricate  magic" in optimized_ir
+    assert "qec.fabricate  plus_i" in optimized_ir
+    assert 'qec.ppm ["X", "Z"]' in optimized_ir
+    assert 'qec.ppr ["X"]' in optimized_ir
 
 
 if __name__ == "__main__":

mlir/include/QEC/IR/CMakeLists.txt

@@ -2,3 +2,10 @@ add_mlir_dialect(QECDialect qec)
 add_mlir_interface(QECOpInterfaces)
 add_mlir_doc(QECDialect QECDialect QEC/ -gen-dialect-doc -gen-op-doc)
 add_mlir_doc(QECOpInterfaces QECOpInterfaces QEC/ -gen-op-interface-docs)
+
+set(LLVM_TARGET_DEFINITIONS QECDialect.td)
+mlir_tablegen(QECEnums.h.inc -gen-enum-decls)
+mlir_tablegen(QECEnums.cpp.inc -gen-enum-defs)
+mlir_tablegen(QECAttributes.h.inc -gen-attrdef-decls)
+mlir_tablegen(QECAttributes.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRQECEnumsIncGen)

mlir/include/QEC/IR/QECDialect.h

@@ -34,6 +34,19 @@
 #define GET_TYPEDEF_CLASSES
 #include "QEC/IR/QECDialectTypes.h.inc"
 
+//===----------------------------------------------------------------------===//
+// QEC enum definitions.
+//===----------------------------------------------------------------------===//
+
+#include "QEC/IR/QECEnums.h.inc"
+
+//===----------------------------------------------------------------------===//
+// QEC attribute declarations.
+//===----------------------------------------------------------------------===//
+
+#define GET_ATTRDEF_CLASSES
+#include "QEC/IR/QECAttributes.h.inc"
+
 //===----------------------------------------------------------------------===//
 // QEC ops declarations.
 //===----------------------------------------------------------------------===//