Support control flow in ConsolidateBlocks

qiskit/transpiler/passes/optimization/consolidate_blocks.py

@@ -23,7 +23,11 @@
 from qiskit.extensions import UnitaryGate
 from qiskit.circuit.library.standard_gates import CXGate
 from qiskit.transpiler.basepasses import TransformationPass
+from qiskit.circuit.controlflow import ControlFlowOp
+from qiskit.transpiler.passmanager import PassManager
 from qiskit.transpiler.passes.synthesis import unitary_synthesis
+from .collect_1q_runs import Collect1qRuns
+from .collect_2q_blocks import Collect2qBlocks
 
 
 class ConsolidateBlocks(TransformationPass):
@@ -49,12 +53,16 @@ def __init__(
     ):
         """ConsolidateBlocks initializer.
 
+        If `kak_basis_gate` is not `None` it will be used as the basis gate for KAK decomposition.
+        Otherwise, if `basis_gates` is not `None` a basis gate will be chosen from this list.
+        Otherwise the basis gate will be `CXGate`.
+
         Args:
             kak_basis_gate (Gate): Basis gate for KAK decomposition.
-            force_consolidate (bool): Force block consolidation
+            force_consolidate (bool): Force block consolidation.
             basis_gates (List(str)): Basis gates from which to choose a KAK gate.
             approximation_degree (float): a float between [0.0, 1.0]. Lower approximates more.
-            target (Target): The target object for the compilation target backend
+            target (Target): The target object for the compilation target backend.
         """
         super().__init__()
         self.basis_gates = None
@@ -159,11 +167,32 @@ def run(self, dag):
                         dag.remove_op_node(node)
                 else:
                     dag.replace_block_with_op(run, unitary, {qubit: 0}, cycle_check=False)
+
+        dag = self._handle_control_flow_ops(dag)
+
         # Clear collected blocks and runs as they are no longer valid after consolidation
         if "run_list" in self.property_set:
             del self.property_set["run_list"]
         if "block_list" in self.property_set:
             del self.property_set["block_list"]
+
+        return dag
+
+    def _handle_control_flow_ops(self, dag):
+        """
+        This is similar to transpiler/passes/utils/control_flow.py except that the
+        collect blocks is redone for the control flow blocks.
+        """
+
+        pass_manager = PassManager()
+        if "run_list" in self.property_set:
+            pass_manager.append(Collect1qRuns())
+        if "block_list" in self.property_set:
+            pass_manager.append(Collect2qBlocks())
+
+        pass_manager.append(self)
+        for node in dag.op_nodes(ControlFlowOp):
+            node.op = node.op.replace_blocks(pass_manager.run(block) for block in node.op.blocks)
         return dag
 
     def _check_not_in_basis(self, gate_name, qargs, global_index_map):

releasenotes/notes/add-control-flow-to-consolidate-blocks-e013e28007170377.yaml

@@ -0,0 +1,9 @@
+---
+features:
+  - |
+    Enabled performing the :class:`.ConsolidateBlocks` pass inside the
+    blocks of :class:`.ControlFlowOp`. This pass collects several sequences of gates
+    and replaces each sequence with the equivalent numeric unitary gate. This new feature enables
+    applying this pass recursively to the blocks in control flow operations. Note that the meaning
+    of "block" in :class:`.ConsolidateBlocks` is unrelated to that in
+    :class:`.ControlFlowOp`.

test/python/transpiler/test_consolidate_blocks.py

@@ -17,8 +17,8 @@
 import unittest
 import numpy as np
 
-from qiskit.circuit import QuantumCircuit, QuantumRegister
-from qiskit.circuit.library import U2Gate, SwapGate, CXGate
+from qiskit.circuit import QuantumCircuit, QuantumRegister, IfElseOp
+from qiskit.circuit.library import U2Gate, SwapGate, CXGate, CZGate
 from qiskit.extensions import UnitaryGate
 from qiskit.converters import circuit_to_dag
 from qiskit.transpiler.passes import ConsolidateBlocks
@@ -428,6 +428,132 @@ def test_identity_1q_unitary_is_removed(self):
         pm = PassManager([Collect2qBlocks(), Collect1qRuns(), ConsolidateBlocks()])
         self.assertEqual(QuantumCircuit(5), pm.run(qc))
 
+    def test_descent_into_control_flow(self):
+        """Test consolidation in blocks when control flow op is the same as at top level."""
+
+        def circuit_of_test_gates():
+            qc = QuantumCircuit(2, 1)
+            qc.cx(0, 1)
+            qc.cx(1, 0)
+            return qc
+
+        def do_consolidation(qc):
+            pass_manager = PassManager()
+            pass_manager.append(Collect2qBlocks())
+            pass_manager.append(ConsolidateBlocks(force_consolidate=True))
+            return pass_manager.run(qc)
+
+        result_top = do_consolidation(circuit_of_test_gates())
+
+        qc_control_flow = QuantumCircuit(2, 1)
+        ifop = IfElseOp((qc_control_flow.clbits[0], False), circuit_of_test_gates(), None)
+        qc_control_flow.append(ifop, qc_control_flow.qubits, qc_control_flow.clbits)
+
+        result_block = do_consolidation(qc_control_flow)
+        gate_top = result_top[0].operation
+        gate_block = result_block[0].operation.blocks[0][0].operation
+        np.testing.assert_allclose(gate_top, gate_block)
+
+    def test_not_crossing_between_control_flow_block_and_parent(self):
+        """Test that consolidation does not occur across the boundary between control flow
+        blocks and the parent circuit."""
+        qc = QuantumCircuit(2, 1)
+        qc.cx(0, 1)
+        qc_true = QuantumCircuit(2, 1)
+        qc_false = QuantumCircuit(2, 1)
+        qc_true.cx(0, 1)
+        qc_false.cz(0, 1)
+        ifop = IfElseOp((qc.clbits[0], True), qc_true, qc_false)
+        qc.append(ifop, qc.qubits, qc.clbits)
+
+        pass_manager = PassManager()
+        pass_manager.append(Collect2qBlocks())
+        pass_manager.append(ConsolidateBlocks(force_consolidate=True))
+        qc_out = pass_manager.run(qc)
+
+        self.assertIsInstance(qc_out[0].operation, UnitaryGate)
+        np.testing.assert_allclose(CXGate(), qc_out[0].operation)
+        op_true = qc_out[1].operation.blocks[0][0].operation
+        op_false = qc_out[1].operation.blocks[1][0].operation
+        np.testing.assert_allclose(CXGate(), op_true)
+        np.testing.assert_allclose(CZGate(), op_false)
+
+    def test_not_crossing_between_control_flow_ops(self):
+        """Test that consolidation does not occur between control flow ops."""
+        qc = QuantumCircuit(2, 1)
+        qc_true = QuantumCircuit(2, 1)
+        qc_false = QuantumCircuit(2, 1)
+        qc_true.cx(0, 1)
+        qc_false.cz(0, 1)
+        ifop1 = IfElseOp((qc.clbits[0], True), qc_true, qc_false)
+        qc.append(ifop1, qc.qubits, qc.clbits)
+        ifop2 = IfElseOp((qc.clbits[0], True), qc_true, qc_false)
+        qc.append(ifop2, qc.qubits, qc.clbits)
+
+        pass_manager = PassManager()
+        pass_manager.append(Collect2qBlocks())
+        pass_manager.append(ConsolidateBlocks(force_consolidate=True))
+        qc_out = pass_manager.run(qc)
+
+        op_true1 = qc_out[0].operation.blocks[0][0].operation
+        op_false1 = qc_out[0].operation.blocks[1][0].operation
+        op_true2 = qc_out[1].operation.blocks[0][0].operation
+        op_false2 = qc_out[1].operation.blocks[1][0].operation
+        np.testing.assert_allclose(CXGate(), op_true1)
+        np.testing.assert_allclose(CZGate(), op_false1)
+        np.testing.assert_allclose(CXGate(), op_true2)
+        np.testing.assert_allclose(CZGate(), op_false2)
+
+    def test_inverted_order(self):
+        """Test that the `ConsolidateBlocks` pass creates matrices that are correct under the
+        application of qubit binding from the outer circuit to the inner block."""
+        body = QuantumCircuit(2, 1)
+        body.h(0)
+        body.cx(0, 1)
+
+        id_op = Operator(np.eye(4))
+        bell = Operator(body)
+
+        qc = QuantumCircuit(2, 1)
+        # The first two 'if' blocks here represent exactly the same operation as each other on the
+        # outer bits, because in the second, the bit-order of the block is reversed, but so is the
+        # order of the bits in the outer circuit that they're bound to, which makes them the same.
+        # The second two 'if' blocks also represnt the same operation as each other, but the 'first
+        # two' and 'second two' pairs represent qubit-flipped operations.
+        qc.if_test((0, False), body.copy(), qc.qubits, qc.clbits)
+        qc.if_test((0, False), body.reverse_bits(), reversed(qc.qubits), qc.clbits)
+        qc.if_test((0, False), body.copy(), reversed(qc.qubits), qc.clbits)
+        qc.if_test((0, False), body.reverse_bits(), qc.qubits, qc.clbits)
+
+        # The first two operations represent Bell-state creation on _outer_ qubits (0, 1), the
+        # second two represent the same creation, but on outer qubits (1, 0).
+        expected = [
+            id_op.compose(bell, qargs=(0, 1)),
+            id_op.compose(bell, qargs=(0, 1)),
+            id_op.compose(bell, qargs=(1, 0)),
+            id_op.compose(bell, qargs=(1, 0)),
+        ]
+
+        actual = []
+        pm = PassManager([Collect2qBlocks(), ConsolidateBlocks(force_consolidate=True)])
+        for instruction in pm.run(qc).data:
+            # For each instruction, the `UnitaryGate` that's been created will always have been made
+            # (as an implementation detail of `DAGCircuit.collect_2q_runs` as of commit e5950661) to
+            # apply to _inner_ qubits (0, 1).  We need to map that back to the _outer_ qubits that
+            # it applies to compare.
+            body = instruction.operation.blocks[0]
+            wire_map = {
+                inner: qc.find_bit(outer).index
+                for inner, outer in zip(body.qubits, instruction.qubits)
+            }
+            actual.append(
+                id_op.compose(
+                    Operator(body.data[0].operation),
+                    qargs=[wire_map[q] for q in body.data[0].qubits],
+                )
+            )
+        self.assertEqual(expected, actual)
+
 
 if __name__ == "__main__":
     unittest.main()