add control flow support to UnitarySynthesis pass (#8565)

* add control flow support to UnitarySynthesis pass

* Refactor control-flow handling to avoid recomputing setup

The `UnitarySynthesis` runner does rather non-trivial setup of its
internal plugin class before beginning the run.  We don't want to
re-initialise the plugin on each run, nor pay the time penalty for the
setup.  Instead, we refactor the actual "run" loop into its own
function, and use that only for the recursion.

Co-authored-by: Jake Lishman <[email protected]>
Co-authored-by: Kevin Krsulich <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

qiskit/transpiler/passes/synthesis/unitary_synthesis.py

@@ -25,6 +25,7 @@
 from qiskit.quantum_info.synthesis import one_qubit_decompose
 from qiskit.quantum_info.synthesis.xx_decompose import XXDecomposer
 from qiskit.quantum_info.synthesis.two_qubit_decompose import TwoQubitBasisDecomposer
+from qiskit.circuit import ControlFlowOp
 from qiskit.circuit.parameter import Parameter
 from qiskit.circuit.library.standard_gates import (
     iSwapGate,
@@ -35,6 +36,7 @@
     ECRGate,
 )
 from qiskit.transpiler.passes.synthesis import plugin
+from qiskit.transpiler.passes.utils import control_flow
 from qiskit.providers.models import BackendProperties
 
 
@@ -258,7 +260,6 @@ def run(self, dag: DAGCircuit) -> DAGCircuit:
             plugin_method = DefaultUnitarySynthesis()
         plugin_kwargs = {"config": self._plugin_config}
         _gate_lengths = _gate_errors = None
-        dag_bit_indices = {}
 
         if self.method == "default":
             # If the method is the default, we only need to evaluate one set of keyword arguments.
@@ -278,8 +279,6 @@ def run(self, dag: DAGCircuit) -> DAGCircuit:
         for method, kwargs in method_list:
             if method.supports_basis_gates:
                 kwargs["basis_gates"] = self._basis_gates
-            if method.supports_coupling_map:
-                dag_bit_indices = dag_bit_indices or {bit: i for i, bit in enumerate(dag.qubits)}
             if method.supports_natural_direction:
                 kwargs["natural_direction"] = self._natural_direction
             if method.supports_pulse_optimize:
@@ -306,6 +305,29 @@ def run(self, dag: DAGCircuit) -> DAGCircuit:
         if self.method == "default":
             # pylint: disable=attribute-defined-outside-init
             plugin_method._approximation_degree = self._approximation_degree
+        return self._run_main_loop(
+            dag, plugin_method, plugin_kwargs, default_method, default_kwargs
+        )
+
+    def _run_main_loop(self, dag, plugin_method, plugin_kwargs, default_method, default_kwargs):
+        """Inner loop for the optimizer, after all DAG-idependent set-up has been completed."""
+
+        def _recurse(dag):
+            # This isn't quite a trivially recursive call because we need to close over the
+            # arguments to the function.  The loop is sufficiently long that it's cleaner to do it
+            # in a separate method rather than define a helper closure within `self.run`.
+            return self._run_main_loop(
+                dag, plugin_method, plugin_kwargs, default_method, default_kwargs
+            )
+
+        for node in dag.op_nodes(ControlFlowOp):
+            node.op = control_flow.map_blocks(_recurse, node.op)
+
+        dag_bit_indices = (
+            {bit: i for i, bit in enumerate(dag.qubits)}
+            if plugin_method.supports_coupling_map or default_method.supports_coupling_map
+            else {}
+        )
 
         for node in dag.named_nodes(*self._synth_gates):
             if self._min_qubits is not None and len(node.qargs) < self._min_qubits:

test/python/transpiler/test_unitary_synthesis.py

@@ -26,9 +26,9 @@
 from qiskit.test import QiskitTestCase
 from qiskit.providers.fake_provider import FakeVigo, FakeMumbaiFractionalCX
 from qiskit.providers.fake_provider.fake_backend_v2 import FakeBackendV2, FakeBackend5QV2
-from qiskit.circuit import QuantumCircuit, QuantumRegister
+from qiskit.circuit import QuantumCircuit, QuantumRegister, ClassicalRegister
 from qiskit.circuit.library import QuantumVolume
-from qiskit.converters import circuit_to_dag
+from qiskit.converters import circuit_to_dag, dag_to_circuit
 from qiskit.transpiler.passes import UnitarySynthesis
 from qiskit.quantum_info.operators import Operator
 from qiskit.quantum_info.random import random_unitary
@@ -724,6 +724,48 @@ def test_reverse_direction(self, opt_level):
         for instr in tqc.get_instructions("ecr"):
             self.assertEqual((0, 1), (tqc_index[instr.qubits[0]], tqc_index[instr.qubits[1]]))
 
+    def test_if_simple(self):
+        """Test a simple if statement."""
+        basis_gates = {"u", "cx"}
+        qr = QuantumRegister(2)
+        cr = ClassicalRegister(2)
+
+        qc_uni = QuantumCircuit(2)
+        qc_uni.h(0)
+        qc_uni.cx(0, 1)
+        qc_uni_mat = Operator(qc_uni)
+
+        qc_true_body = QuantumCircuit(2)
+        qc_true_body.unitary(qc_uni_mat, [0, 1])
+
+        qc = QuantumCircuit(qr, cr)
+        qc.if_test((cr, 1), qc_true_body, [0, 1], [0, 1])
+        dag = circuit_to_dag(qc)
+        cdag = UnitarySynthesis(basis_gates=basis_gates).run(dag)
+        cqc = dag_to_circuit(cdag)
+        cbody = cqc.data[0].operation.params[0]
+        self.assertEqual(cbody.count_ops().keys(), basis_gates)
+        self.assertEqual(qc_uni_mat, Operator(cbody))
+
+    def test_nested_control_flow(self):
+        """Test unrolling nested control flow blocks."""
+        qr = QuantumRegister(2)
+        cr = ClassicalRegister(1)
+        qc_uni1 = QuantumCircuit(2)
+        qc_uni1.swap(0, 1)
+        qc_uni1_mat = Operator(qc_uni1)
+
+        qc = QuantumCircuit(qr, cr)
+        with qc.for_loop(range(3)):
+            with qc.while_loop((cr, 0)):
+                qc.unitary(qc_uni1_mat, [0, 1])
+        dag = circuit_to_dag(qc)
+        cdag = UnitarySynthesis(basis_gates=["u", "cx"]).run(dag)
+        cqc = dag_to_circuit(cdag)
+        cbody = cqc.data[0].operation.params[2].data[0].operation.params[0]
+        self.assertEqual(cbody.count_ops().keys(), {"u", "cx"})
+        self.assertEqual(qc_uni1_mat, Operator(cbody))
+
 
 if __name__ == "__main__":
     unittest.main()