add controlflow handling to Layout2qDistance transpiler pass

qiskit/transpiler/passes/layout/layout_2q_distance.py

@@ -18,7 +18,11 @@
 Therefore, 0 is a perfect layout selection.
 """
 
+import math
 from qiskit.transpiler.basepasses import AnalysisPass
+from qiskit.converters import circuit_to_dag
+from qiskit.circuit.controlflow import IfElseOp, WhileLoopOp, ForLoopOp
+from qiskit.circuit.controlflow import ControlFlowOp
 
 
 class Layout2qDistance(AnalysisPass):
@@ -30,16 +34,19 @@ class Layout2qDistance(AnalysisPass):
     No CX direction is considered.
     """
 
-    def __init__(self, coupling_map, property_name="layout_score"):
+    def __init__(self, coupling_map, property_name="layout_score", weight_loops=True):
         """Layout2qDistance initializer.
 
         Args:
             coupling_map (CouplingMap): Directed graph represented a coupling map.
             property_name (str): The property name to save the score. Default: layout_score
+            weight_loops (bool): Whether to weight 2q distance of loops by number of loops. If
+               while loop exists the returned distance will be NaN.
         """
         super().__init__()
         self.coupling_map = coupling_map
         self.property_name = property_name
+        self.weight_loops = weight_loops
 
     def run(self, dag):
         """
@@ -67,5 +74,31 @@ def run(self, dag):
             physical_q1 = virtual_physical_map[gate.qargs[1]]
 
             sum_distance += dist_matrix[physical_q0, physical_q1] - 1
+        for node in dag.op_nodes(op=ControlFlowOp):
+            sum_distance += self._control_flow_2q_distance(node.op)
 
         self.property_set[self.property_name] = sum_distance
+
+    def _control_flow_2q_distance(self, cfop):
+        distance = 0
+        if self.weight_loops:
+            if isinstance(cfop, IfElseOp):
+                for block in cfop.blocks:
+                    dag_block = circuit_to_dag(block)
+                    self.run(dag_block)
+                    distance += self.property_set[self.property_name]
+            elif isinstance(cfop, ForLoopOp):
+                index_set, _, block = cfop.params
+                dag_block = circuit_to_dag(block)
+                self.run(dag_block)
+                distance += len(index_set) * self.property_set[self.property_name]
+            elif isinstance(cfop, WhileLoopOp):
+                body = cfop.blocks[0]
+                if body.num_nonlocal_gates():
+                    # Indeterminate number of loops so indeterminate 2q distance
+                    distance = math.nan
+        else:
+            for block in cfop.blocks:
+                dag_block = circuit_to_dag(block)
+                self.run(dag_block)
+        return distance

test/python/transpiler/test_layout_score.py

@@ -13,8 +13,8 @@
 """Test the Layout Score pass"""
 
 import unittest
-
-from qiskit import QuantumRegister, QuantumCircuit
+import math
+from qiskit import QuantumRegister, QuantumCircuit, ClassicalRegister
 from qiskit.transpiler.passes import Layout2qDistance
 from qiskit.transpiler import CouplingMap, Layout
 from qiskit.converters import circuit_to_dag
@@ -101,6 +101,64 @@ def test_swap_mapped_false(self):
 
         self.assertEqual(pass_.property_set["layout_score"], 1)
 
+    def test_control_flow_if_else(self):
+        """Circuit with control flow if else blocks"""
+        qr = QuantumRegister(4)
+        cr = ClassicalRegister(4)
+        circuit = QuantumCircuit(qr, cr)
+        true_body = QuantumCircuit(qr, cr)
+        false_body = QuantumCircuit(qr, cr)
+        true_body.cx(0, 2)
+        false_body.cx(0, 2)
+        false_body.cx(0, 3)
+        circuit.h(qr)
+        circuit.cx(0, 2)
+        circuit.if_else((cr, 0), true_body, false_body, qr, cr)
+        coupling = CouplingMap([[0, 1], [1, 2], [2, 3]])
+        layout = Layout().generate_trivial_layout(qr)
+        dag = circuit_to_dag(circuit)
+        pass_ = Layout2qDistance(coupling)
+        pass_.property_set["layout"] = layout
+        pass_.run(dag)
+        self.assertEqual(pass_.property_set["layout_score"], 5)
+
+    def test_control_flow_for_loop(self):
+        """Circuit with control flow for loop"""
+        nloops = 5
+        qr = QuantumRegister(4)
+        cr = ClassicalRegister(4)
+        circuit = QuantumCircuit(qr, cr)
+        body = QuantumCircuit(qr, cr)
+        body.cx(0, 3)
+        circuit.h(qr)
+        circuit.cx(0, 2)
+        circuit.for_loop(range(nloops), None, body, qr, cr)
+        coupling = CouplingMap([[0, 1], [1, 2], [2, 3]])
+        layout = Layout().generate_trivial_layout(qr)
+        dag = circuit_to_dag(circuit)
+        pass_ = Layout2qDistance(coupling)
+        pass_.property_set["layout"] = layout
+        pass_.run(dag)
+        self.assertEqual(pass_.property_set["layout_score"], 11)
+
+    def test_control_flow_while_loop(self):
+        """Circuit with control flow while loop"""
+        qr = QuantumRegister(4)
+        cr = ClassicalRegister(4)
+        circuit = QuantumCircuit(qr, cr)
+        body = QuantumCircuit(qr, cr)
+        body.cx(0, 3)
+        circuit.h(qr)
+        circuit.cx(0, 2)
+        circuit.while_loop((cr, 0), body, qr, cr)
+        coupling = CouplingMap([[0, 1], [1, 2], [2, 3]])
+        layout = Layout().generate_trivial_layout(qr)
+        dag = circuit_to_dag(circuit)
+        pass_ = Layout2qDistance(coupling)
+        pass_.property_set["layout"] = layout
+        pass_.run(dag)
+        self.assertTrue(math.isnan(pass_.property_set["layout_score"]))
+
 
 if __name__ == "__main__":
     unittest.main()