Implement a decomposition of PPRs into PPMs

doc/releases/changelog-dev.md

@@ -2,6 +2,23 @@
 
 <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), :math:`\exp(-iP_{\{x, y, z\}} \theta)`, into
+  Pauli product measurements (PPMs). Non-Clifford PPR (:math:`\theta = \tfrac{\pi}{8}`) requires
+  the consumption of a magic state, while Clifford PPR (:math:`\theta = \tfrac{\pi}{4}`) 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 PPR (:math:`\theta = \tfrac{\pi}{2}`)).
+
+  Or via the Catalyst CLI as two separate passes:
+  * `decompose_clifford_ppr`: Decompose Clifford PPR (:math:`\theta = \tfrac{\pi}{4}`)
+  rotations into PPMs.
+  * `decompose_non_clifford_ppr`: Decompose non-Cliford PPR (:math:`\theta = \tfrac{\pi}{8}`)
+  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
@@ -47,12 +64,12 @@
   deallocation primitive has been split into deallocation and a separate device release primitive.
   [(#1720)](https://github.com/PennyLaneAI/catalyst/pull/1720)
 
-  - `qunitary_p` is now `unitary_p` (unchanged)
-  - `qmeasure_p` is now `measure_p` (unchanged)
-  - `qdevice_p` is now `device_init_p` (unchanged)
-  - `qdealloc_p` no longer releases the device, thus it can be used at any point of a quantum
+  * `qunitary_p` is now `unitary_p` (unchanged)
+  * `qmeasure_p` is now `measure_p` (unchanged)
+  * `qdevice_p` is now `device_init_p` (unchanged)
+  * `qdealloc_p` no longer releases the device, thus it can be used at any point of a quantum
      execution scope
-  - `device_release_p` is a new primitive that must be used to mark the end of a quantum execution
+  * `device_release_p` is a new primitive that must be used to mark the end of a quantum execution
      scope, which will release the quantum device
 
 * Catalyst has removed the `experimental_capture` keyword from the `qjit` decorator in favour of
@@ -91,13 +108,13 @@
   [(#1729)](https://github.com/PennyLaneAI/catalyst/pull/1729)
 
   Several internal changes were made for this update.
-    - LAPACK kernels are updated to adhere to the new JAX lowering rules for external functions.
+  * LAPACK kernels are updated to adhere to the new JAX lowering rules for external functions.
     [(#1685)](https://github.com/PennyLaneAI/catalyst/pull/1685)
 
-    - The trace stack is removed and replaced with a tracing context manager.
+  * The trace stack is removed and replaced with a tracing context manager.
     [(#1662)](https://github.com/PennyLaneAI/catalyst/pull/1662)
 
-    - A new `debug_info` argument is added to `Jaxpr`, the `make_jaxpr`
+  * A new `debug_info` argument is added to `Jaxpr`, the `make_jaxpr`
     functions, and `jax.extend.linear_util.wrap_init`.
     [(#1670)](https://github.com/PennyLaneAI/catalyst/pull/1670)
     [(#1671)](https://github.com/PennyLaneAI/catalyst/pull/1671)

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

@@ -578,9 +578,94 @@ def circuit():
         . . .
 
     """
-
     if qnode is None:
         return functools.partial(merge_ppr_ppm, max_pauli_size=max_pauli_size)
 
     merge_ppr_ppm_pass = {"merge_ppr_ppm": {"max-pauli-size": max_pauli_size}}
     return PassPipelineWrapper(qnode, merge_ppr_ppm_pass)
+
+
+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)
+    , :math:`\exp(-iP\theta)`, 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 (:math:`\theta = \tfrac{\pi}{8}`) are decomposed first, and then Clifford PPRs
+    (:math:`\theta = \tfrac{\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,
+            "avoid-y-measure": avoid_y_measure,
+        },
+        "decompose_clifford_ppr": {"avoid-y-measure": avoid_y_measure},
+    }
+
+    if qnode is None:
+        return functools.partial(
+            ppr_to_ppm, decompose_method=decompose_method, avoid_y_measure=avoid_y_measure
+        )
+
+    return PassPipelineWrapper(qnode, passes)

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

@@ -19,7 +19,14 @@
 import pytest
 
 from catalyst import measure, pipeline, qjit
-from catalyst.passes import cancel_inverses, commute_ppr, merge_ppr_ppm, merge_rotations, to_ppr
+from catalyst.passes import (
+    cancel_inverses,
+    commute_ppr,
+    merge_ppr_ppm,
+    merge_rotations,
+    ppr_to_ppm,
+    to_ppr,
+)
 
 # pylint: disable=missing-function-docstring
 
@@ -180,9 +187,7 @@ def test_chained_apply_passes_workflow(x: float):
 
 
 def test_convert_clifford_to_ppr():
-    """
-    Test convert_clifford_to_ppr
-    """
+
     pipe = [("pipe", ["enforce-runtime-invariants-pipeline"])]
 
     @qjit(pipelines=pipe, target="mlir")
@@ -204,10 +209,129 @@ def f():
     assert "qec.ppr" in optimized_ir
 
 
+def 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
+
+
+def 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
+
+
+def 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 "quantum.alloc_qb" in optimized_ir
+    assert "qec.fabricate  magic" in optimized_ir
+    assert "qec.select.ppm" in optimized_ir
+    assert 'qec.ppr ["X"]' in optimized_ir
+
+
+def test_ppr_to_ppm_inject_magic_state():
+
+    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
+
+
 def test_commute_ppr_and_merge_ppr_ppm_with_max_pauli_size():
-    """
-    Test commute_ppr
-    """
+
     pipe = [("pipe", ["enforce-runtime-invariants-pipeline"])]
 
     @qjit(pipelines=pipe, target="mlir")

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.
 //===----------------------------------------------------------------------===//