Support control flow in ConsolidateBlocks (Qiskit#10355)

* Support control flow in ConsolidateBlocks

* Add release note for support control flow in ConsolidateBlocks

* Move imports from inside function to top level

* Make construction of ConsolidateBlocks for control flow ops more efficient

* Do IfElseOp test without builder interface

Before, we used the builder interface for an ifelse op. Some details of
the circuit built, in particular the mapping of wires is not deterministic.
We could have used canonicalize_control_flow. But instead we construct the
IfElseOp manually. This removes the complexity of the builder interface from
this test.

* Linting

* Avoid cyclic import

* Try to fix cyclic import (in lint tool only)

* Reuse top-level consolidation pass and add tests

* Before, we created an new ConsolidationPass when descending into control flow
  blocks. With this commit, we use the existing pass.

* Add some tests for cases where consolidation should not happen.

* Remove argument `decomposer` from constructor of ConsolidateBlocks

This was added in a previous commit in the series of commits for this
PR. The code has been redesigned so that this argument is no longer
necessary.

* Remove cruft accidentally left

In the first version of the associated PR, we created a new pass
when descending into control flow blocks. These lines were included to
support that construction and are no longer needed.

* Move function-level import to module level

Remove unused import in tests

* Write loop more concisely

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

Co-authored-by: Jake Lishman <[email protected]>

* Use assertion in tests with better diagnostics

* Remove reference to decomposer from docstring to ConsolidateBlocks

The previous doc string was a bit imprecise. It also referred to a decomposer which although
implied, is not meant to be accessible by the user.

* Use more informative tests for ConsolidateBlocks with control flow

* Simplify test in test_consolidate_blocks and factor

* Gates used in testing ConsolidateBlocks with control flow ops were copied
from another test in the file. They complicate the test, and removing them does
not weaken the test at all.

* Factor some code within a test into a function

* Factor more code in test

* Use clbit in circuit as test bit when appending IfElse

This has no effect on the test. But, previously, we used a clbit taken
from an unrelated circuit. This might give the impression that there is
some significance to this unexpected choice.

* In test, use bits in circuit as specifiers when appending gate to that circuit

Previously we used qubits and clbits from an unrelated circuit to specify which
bits to apply an IfElse to when appending. Here, we use bits from the same circuit that we are
appending the gate to. This does not change the test. But again, one might look for
significance in the previous unusual choice.

Qiskit considers two registers of the same length and the same name to be the
same register. The previous behavior depended on this choice.

* In a test, don't use qubits from unrelated circuit when constructing circuit

This has no effect on the test. As in the previous commits, the choice was
unusual and might imply some significance. The more natural choice is to
create new qubits by specifying the number.

* Factor code in test to simplify test

* Simplify remaining control flow op tests for ConsolidateBlocks

These simplifications are similar to those made for the first test.

* Run black

* Update test/python/transpiler/test_consolidate_blocks.py

Co-authored-by: Jake Lishman <[email protected]>

---------

Co-authored-by: Jake Lishman <[email protected]>

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()