An extended RZZ pass for unbounded parameters

qiskit_ibm_runtime/transpiler/passes/basis/fold_rzz_angle.py

@@ -12,14 +12,16 @@
 
 """Pass to wrap Rzz gate angle in calibrated range of 0-pi/2."""
 
-from typing import Tuple
+from typing import Tuple, Optional, Union, List
 from math import pi
+from operator import mod
 
 from qiskit.converters import dag_to_circuit, circuit_to_dag
-from qiskit.circuit.library.standard_gates import RZZGate, RZGate, XGate, GlobalPhaseGate
+from qiskit.circuit.library.standard_gates import RZZGate, RZGate, XGate, GlobalPhaseGate, RXGate
 from qiskit.circuit.parameterexpression import ParameterExpression
 from qiskit.circuit import Qubit, ControlFlowOp
 from qiskit.dagcircuit import DAGCircuit
+from qiskit.transpiler import Target
 from qiskit.transpiler.basepasses import TransformationPass
 
 import numpy as np
@@ -42,15 +44,21 @@ class FoldRzzAngle(TransformationPass):
 
     This pass allows the Qiskit users to naively use the Rzz gates
     with angle of arbitrary real numbers.
-
-    .. note::
-        This pass doesn't transform the circuit when the
-        Rzz gate angle is an unbound parameter.
-        In this case, the user must assign a gate angle before
-        transpilation, or be responsible for choosing parameters
-        from the calibrated range of [0, pi/2].
     """
 
+    def __init__(self, target: Optional[Union[Target, List[str]]] = None):
+        """
+        Args:
+            target - either a target or only a list of basis gates, either way it can be checked
+            if an instruction is supported using the `in` operator, for example `"rx" in target`.
+            If None then we assume that there is no limit on the gates in the transpiled circuit.
+        """
+        super().__init__()
+        if target is None:
+            self._target = ["rz", "x", "rx", "rzz"]
+        else:
+            self._target = target
+
     def run(self, dag: DAGCircuit) -> DAGCircuit:
         self._run_inner(dag)
         return dag
@@ -85,41 +93,140 @@ def _run_inner(self, dag: DAGCircuit) -> bool:
                 continue
 
             angle = node.op.params[0]
-            if isinstance(angle, ParameterExpression) or 0 <= angle <= pi / 2:
+
+            if not isinstance(angle, ParameterExpression) and 0 <= angle <= pi / 2:
                 # Angle is an unbound parameter or a calibrated value.
                 continue
 
             # Modify circuit around Rzz gate to address non-ISA angles.
-            modified = True
-            wrap_angle = np.angle(np.exp(1j * angle))
-            if 0 <= wrap_angle <= pi / 2:
-                # In the first quadrant.
-                replace = self._quad1(wrap_angle, node.qargs)
-            elif pi / 2 < wrap_angle <= pi:
-                # In the second quadrant.
-                replace = self._quad2(wrap_angle, node.qargs)
-            elif -pi <= wrap_angle <= -pi / 2:
-                # In the third quadrant.
-                replace = self._quad3(wrap_angle, node.qargs)
-            elif -pi / 2 < wrap_angle < 0:
-                # In the forth quadrant.
-                replace = self._quad4(wrap_angle, node.qargs)
+            if isinstance(angle, ParameterExpression):
+                replace = self._unbounded_parameter(angle, node.qargs)
             else:
-                raise RuntimeError("Unreacheable.")
-            if pi < angle % (4 * pi) < 3 * pi:
-                replace.apply_operation_back(GlobalPhaseGate(pi))
-            dag.substitute_node_with_dag(node, replace)
+                replace = self._numeric_parameter(angle, node.qargs)
+
+            if replace is not None:
+                dag.substitute_node_with_dag(node, replace)
+                modified = True
+
         return modified
 
-    @staticmethod
-    def _quad1(angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
+    # The next functions are required because sympy doesn't convert Boolean values to integers.
+    # symengine maybe does but I failed to find it in its documentation.
+    def gt_op(self, exp1: ParameterExpression, exp2: ParameterExpression) -> ParameterExpression:
+        """Return an expression which, after substitution, will be equal to 1 if `exp1` is
+        greater than `exp2`, and 0 otherwise"""
+        tmp = (exp1 - exp2).sign()
+
+        # We want to return 0 if tmp is -1 or 0, and 1 otherwise
+        return tmp * tmp * (tmp + 1) / 2
+
+    def gteq_op(self, exp1: ParameterExpression, exp2: ParameterExpression) -> ParameterExpression:
+        """Return an expression which, after substitution, will be equal to 1 if `exp1` is
+        greater or equal than `exp2`, and 0 otherwise"""
+        tmp = (exp1 - exp2).sign()
+
+        # We want to return 1 if tmp is 1 or 0, and 0 otherwise
+        return (1 - tmp * tmp) + tmp * tmp * (tmp + 1) / 2
+
+    def and_op(self, exp1: ParameterExpression, exp2: ParameterExpression) -> ParameterExpression:
+        """Return an expression which, after substitution, will be equal to 1 if `exp1` and `exp2`
+        are both 1, and 0 if at least one of then is 0"""
+        return exp1 * exp2
+
+    def between(
+        self, exp: ParameterExpression, lower: ParameterExpression, upper: ParameterExpression
+    ) -> ParameterExpression:
+        """Return an expression which, after substitution, will be equal to 1 if `exp1` is
+        greater or equal than `lower` and smaller than `upper`, and 0 otherwise"""
+        return self.and_op(self.gteq_op(exp, lower), self.gt_op(upper, exp))
+
+    def _unbounded_parameter(
+        self, angle: ParameterExpression, qubits: Tuple[Qubit, ...]
+    ) -> DAGCircuit:
+        if "rz" not in self._target or "rx" not in self._target or "rzz" not in self._target:
+            return None
+
+        wrap_angle = (angle + pi)._apply_operation(mod, 2 * pi) - pi
+        pi_phase = self.between(angle._apply_operation(mod, 4 * pi), pi, 3 * pi)
+
+        quad1 = self.between(wrap_angle, 0, pi / 2)
+        quad2 = self.between(wrap_angle, pi / 2, pi)
+        quad3 = self.between(wrap_angle, -pi, -pi / 2)
+        quad4 = self.between(wrap_angle, -pi / 2, 0)
+
+        global_phase = quad2 * (-pi / 2) + quad3 * (-pi / 2) + quad4 * pi + pi_phase * pi
+        rz_angle = quad2 * pi + quad3 * pi
+        rx_angle = quad2 * pi + quad4 * pi
+        rzz_angle = (
+            quad1 * wrap_angle
+            + quad2 * (pi - wrap_angle)
+            + quad3 * (pi + wrap_angle)
+            + quad4 * (-wrap_angle)
+        )
+
+        new_dag = DAGCircuit()
+        new_dag.add_qubits(qubits=qubits)
+        new_dag.apply_operation_back(GlobalPhaseGate(global_phase))
+        new_dag.apply_operation_back(
+            RZGate(rz_angle),
+            qargs=(qubits[0],),
+            check=False,
+        )
+        new_dag.apply_operation_back(
+            RZGate(rz_angle),
+            qargs=(qubits[1],),
+            check=False,
+        )
+        new_dag.apply_operation_back(
+            RXGate(rx_angle),
+            qargs=(qubits[0],),
+            check=False,
+        )
+        new_dag.apply_operation_back(
+            RZZGate(rzz_angle),
+            qargs=qubits,
+            check=False,
+        )
+        new_dag.apply_operation_back(
+            RXGate(rx_angle),
+            qargs=(qubits[0],),
+            check=False,
+        )
+
+        return new_dag
+
+    def _numeric_parameter(self, angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
+        wrap_angle = np.angle(np.exp(1j * angle))
+        if 0 <= wrap_angle <= pi / 2:
+            # In the first quadrant.
+            replace = self._quad1(wrap_angle, qubits)
+        elif pi / 2 < wrap_angle <= pi:
+            # In the second quadrant.
+            replace = self._quad2(wrap_angle, qubits)
+        elif -pi <= wrap_angle <= -pi / 2:
+            # In the third quadrant.
+            replace = self._quad3(wrap_angle, qubits)
+        elif -pi / 2 < wrap_angle < 0:
+            # In the forth quadrant.
+            replace = self._quad4(wrap_angle, qubits)
+        else:
+            raise RuntimeError("Unreacheable.")
+        if pi < angle % (4 * pi) < 3 * pi:
+            replace.apply_operation_back(GlobalPhaseGate(pi))
+
+        return replace
+
+    def _quad1(self, angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
         """Handle angle between [0, pi/2].
 
         Circuit is not transformed - the Rzz gate is calibrated for the angle.
 
         Returns:
             A new dag with the same Rzz gate.
         """
+        if "rzz" not in self._target:
+            return None
+
         new_dag = DAGCircuit()
         new_dag.add_qubits(qubits=qubits)
         new_dag.apply_operation_back(
@@ -129,21 +236,23 @@ def _quad1(angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
         )
         return new_dag
 
-    @staticmethod
-    def _quad2(angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
+    def _quad2(self, angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
         """Handle angle between (pi/2, pi].
 
         Circuit is transformed into the following form:
 
-                ┌───────┐┌───┐            ┌───┐
+                 ┌───────┐┌───┐            ┌───┐
             q_0: ┤ Rz(π) ├┤ X ├─■──────────┤ X ├
-                ├───────┤└───┘ │ZZ(π - θ) └───┘
+                 ├───────┤└───┘ │ZZ(π - θ) └───┘
             q_1: ┤ Rz(π) ├──────■───────────────
-                └───────┘
+                 └───────┘
 
         Returns:
             New dag to replace Rzz gate.
         """
+        if "rz" not in self._target or "x" not in self._target or "rzz" not in self._target:
+            return None
+
         new_dag = DAGCircuit()
         new_dag.add_qubits(qubits=qubits)
         new_dag.apply_operation_back(GlobalPhaseGate(pi / 2))
@@ -176,21 +285,23 @@ def _quad2(angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
             )
         return new_dag
 
-    @staticmethod
-    def _quad3(angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
+    def _quad3(self, angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
         """Handle angle between [-pi, -pi/2].
 
         Circuit is transformed into following form:
 
-                ┌───────┐
+                 ┌───────┐
             q_0: ┤ Rz(π) ├─■───────────────
-                ├───────┤ │ZZ(π - Abs(θ))
+                 ├───────┤ │ZZ(π - Abs(θ))
             q_1: ┤ Rz(π) ├─■───────────────
-                └───────┘
+                 └───────┘
 
         Returns:
             New dag to replace Rzz gate.
         """
+        if "rz" not in self._target or "rzz" not in self._target:
+            return None
+
         new_dag = DAGCircuit()
         new_dag.add_qubits(qubits=qubits)
         new_dag.apply_operation_back(GlobalPhaseGate(-pi / 2))
@@ -212,20 +323,22 @@ def _quad3(angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
             )
         return new_dag
 
-    @staticmethod
-    def _quad4(angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
+    def _quad4(self, angle: float, qubits: Tuple[Qubit, ...]) -> DAGCircuit:
         """Handle angle between (-pi/2, 0).
 
         Circuit is transformed into following form:
 
-                ┌───┐             ┌───┐
+                 ┌───┐             ┌───┐
             q_0: ┤ X ├─■───────────┤ X ├
-                └───┘ │ZZ(Abs(θ)) └───┘
+                 └───┘ │ZZ(Abs(θ)) └───┘
             q_1: ──────■────────────────
 
         Returns:
             New dag to replace Rzz gate.
         """
+        if "x" not in self._target or "rzz" not in self._target:
+            return None
+
         new_dag = DAGCircuit()
         new_dag.add_qubits(qubits=qubits)
         new_dag.apply_operation_back(

qiskit_ibm_runtime/transpiler/plugin.py

@@ -133,5 +133,5 @@ def pass_manager(
             pre_passes.append(ConvertIdToDelay(instruction_durations))
         if "rzz" in target:
             # Apply this pass after SU4 is translated.
-            post_passes.append(FoldRzzAngle())
+            post_passes.append(FoldRzzAngle(target))
         return PassManager(pre_passes) + translator_pm + PassManager(post_passes)

test/unit/transpiler/passes/basis/test_fold_rzz_angle.py

@@ -18,12 +18,14 @@
 
 from qiskit.circuit import QuantumCircuit
 from qiskit.circuit.parameter import Parameter
+from qiskit.circuit.library import RZZGate
+from qiskit.transpiler.target import InstructionProperties
 from qiskit.transpiler.passmanager import PassManager
 from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
 from qiskit.quantum_info import Operator
 
 from qiskit_ibm_runtime.transpiler.passes.basis import FoldRzzAngle
-from qiskit_ibm_runtime.fake_provider import FakeFractionalBackend
+from qiskit_ibm_runtime.fake_provider import FakeFractionalBackend, FakeSherbrooke
 from .....ibm_test_case import IBMTestCase
 
 
@@ -63,13 +65,42 @@ def test_folding_rzz_angles(self, angle):
                 self.assertGreaterEqual(fold_angle, 0.0)
                 self.assertLessEqual(fold_angle, pi / 2)
 
-    def test_folding_rzz_angle_unbound(self):
-        """Test skip folding unbound gate angle."""
+    @named_data(
+        ("pi/2_pos", pi / 2),
+        ("pi/2_neg", -pi / 2),
+        ("pi_pos", pi),
+        ("pi_neg", -pi),
+        ("quad1_no_wrap", 0.1),
+        ("quad2_no_wrap", pi / 2 + 0.1),
+        ("quad3_no_wrap", -pi + 0.1),
+        ("quad4_no_wrap", -0.1),
+        ("quad1_2pi_wrap", 2 * pi + 0.1),
+        ("quad2_2pi_wrap", -3 * pi / 2 + 0.1),
+        ("quad3_2pi_wrap", pi + 0.1),
+        ("quad4_2pi_wrap", 2 * pi - 0.1),
+        ("quad1_12pi_wrap", -12 * pi + 0.1),
+        ("quad2_12pi_wrap", 23 * pi / 2 + 0.1),
+        ("quad3_12pi_wrap", 11 * pi + 0.1),
+        ("quad4_12pi_wrap", -12 * pi - 0.1),
+    )
+    def test_folding_rzz_angle_unbound(self, angle):
+        """Test transformation in the case of an unbounded parameter"""
+        param = Parameter("angle")
+
         qc = QuantumCircuit(2)
-        qc.rzz(Parameter("θ"), 0, 1)
+        qc.rzz(param, 0, 1)
         pm = PassManager([FoldRzzAngle()])
         isa = pm.run(qc)
-        self.assertEqual(qc, isa)
+
+        qc.assign_parameters({param: angle}, inplace=True)
+        isa.assign_parameters({param: angle}, inplace=True)
+
+        self.assertEqual(Operator.from_circuit(qc), Operator.from_circuit(isa))
+        for inst_data in isa.data:
+            if inst_data.operation.name == "rzz":
+                fold_angle = inst_data.operation.params[0]
+                self.assertGreaterEqual(fold_angle, 0.0)
+                self.assertLessEqual(fold_angle, pi / 2)
 
     def test_controlflow(self):
         """Test non-ISA Rzz gates inside/outside a control flow branch."""
@@ -115,3 +146,23 @@ def test_fractional_plugin(self):
         self.assertEqual(isa_circ.data[0].operation.name, "global_phase")
         self.assertEqual(isa_circ.data[1].operation.name, "rzz")
         self.assertTrue(np.isclose(isa_circ.data[1].operation.params[0], 7 - 2 * pi))
+
+    def test_unsupported_instructions_skip(self):
+        """Verify that the pass does not output gates that are not in the basis gates"""
+        backend = FakeSherbrooke()
+        backend.target.add_instruction(RZZGate(Parameter("θ")), {(0, 1): InstructionProperties()})
+
+        p = Parameter("p")
+        circ = QuantumCircuit(2)
+        circ.rzz(p, 0, 1)
+
+        pm = generate_preset_pass_manager(
+            optimization_level=0,
+            backend=FakeSherbrooke(),
+            translation_method="ibm_fractional",
+            basis_gates=backend.target.operation_names,
+        )
+        isa_circ = pm.run(circ)
+
+        self.assertEqual(isa_circ.data[0].operation.name, "rzz")
+        self.assertTrue(isinstance(isa_circ.data[0].operation.params[0], Parameter))