Support control flow in Unroller transpiler pass (#7768)

* create two tests and 1st modification of unroller

* if_else test, parameter test

* black

* linting

* change shallow copy of control flow ops to not copy body

* add special copy

* add `replace_blocks` method

* minor commit

* linting

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

qiskit/circuit/controlflow/control_flow.py

@@ -13,7 +13,7 @@
 "Container to encapsulate all control flow operations."
 
 from abc import ABC, abstractmethod
-from typing import Tuple
+from typing import Tuple, Iterable
 
 from qiskit.circuit import QuantumCircuit, Instruction
 
@@ -29,3 +29,14 @@ class ControlFlowOp(Instruction, ABC):
         parameter to be resolved at run time.
         """
         pass
+
+    @abstractmethod
+    def replace_blocks(self, blocks: Iterable[QuantumCircuit]) -> "ControlFlowOp":
+        """Replace blocks and return new instruction.
+        Args:
+            blocks: Tuple of QuantumCircuits to replace in instruction.
+
+        Returns:
+            New ControlFlowOp with replaced blocks.
+        """
+        pass

qiskit/circuit/controlflow/for_loop.py

@@ -117,6 +117,10 @@ class ForLoopOp(ControlFlowOp):
     def blocks(self):
         return (self._params[2],)
 
+    def replace_blocks(self, blocks):
+        (body,) = blocks
+        return ForLoopOp(self.params[0], self.params[1], body, label=self.label)
+
 
 class ForLoopContext:
     """A context manager for building up ``for`` loops onto circuits in a natural order, without

qiskit/circuit/controlflow/if_else.py

@@ -14,6 +14,7 @@
 
 
 from typing import Optional, Tuple, Union, Iterable, Set
+import itertools
 
 from qiskit.circuit import ClassicalRegister, Clbit, QuantumCircuit
 from qiskit.circuit.instructionset import InstructionSet
@@ -136,6 +137,22 @@ class IfElseOp(ControlFlowOp):
         else:
             return (self.params[0], self.params[1])
 
+    def replace_blocks(self, blocks: Iterable[QuantumCircuit]) -> "IfElseOp":
+        """Replace blocks and return new instruction.
+
+        Args:
+            blocks: Iterable of circuits for "if" and "else" condition. If there is no "else"
+                circuit it may be set to None or ommited.
+
+        Returns:
+            New IfElseOp with replaced blocks.
+        """
+
+        true_body, false_body = (
+            ablock for ablock, _ in itertools.zip_longest(blocks, range(2), fillvalue=None)
+        )
+        return IfElseOp(self.condition, true_body, false_body=false_body, label=self.label)
+
     def c_if(self, classical, val):
         raise NotImplementedError(
             "IfElseOp cannot be classically controlled through Instruction.c_if. "

qiskit/circuit/controlflow/while_loop.py

@@ -95,6 +95,12 @@ class WhileLoopOp(ControlFlowOp):
     def blocks(self):
         return (self._params[0],)
 
+    def replace_blocks(self, blocks):
+        (body,) = blocks
+        if not isinstance(body, QuantumCircuit):
+            raise CircuitError("WhileLoopOp expects a single QuantumCircuit when setting blocks")
+        return WhileLoopOp(self.condition, body, label=self.label)
+
     def c_if(self, classical, val):
         raise NotImplementedError(
             "WhileLoopOp cannot be classically controlled through Instruction.c_if. "

qiskit/transpiler/passes/basis/unroller.py

@@ -14,8 +14,9 @@
 
 from qiskit.transpiler.basepasses import TransformationPass
 from qiskit.exceptions import QiskitError
-from qiskit.circuit import ControlledGate
+from qiskit.circuit import ControlledGate, ControlFlowOp
 from qiskit.converters.circuit_to_dag import circuit_to_dag
+from qiskit.converters.dag_to_circuit import dag_to_circuit
 
 
 class Unroller(TransformationPass):
@@ -68,7 +69,15 @@ class Unroller(TransformationPass):
                 else:
                     continue
 
-            # TODO: allow choosing other possible decompositions
+            if isinstance(node.op, ControlFlowOp):
+                unrolled_blocks = []
+                for block in node.op.blocks:
+                    dag_block = circuit_to_dag(block)
+                    unrolled_dag_block = self.run(dag_block)
+                    unrolled_circ_block = dag_to_circuit(unrolled_dag_block)
+                    unrolled_blocks.append(unrolled_circ_block)
+                node.op = node.op.replace_blocks(unrolled_blocks)
+                continue
             try:
                 phase = node.op.definition.global_phase
                 rule = node.op.definition.data

test/python/transpiler/test_unroller.py

@@ -22,7 +22,7 @@ from qiskit.converters import circuit_to_dag, dag_to_circuit
 from qiskit.quantum_info import Operator
 from qiskit.test import QiskitTestCase
 from qiskit.exceptions import QiskitError
-from qiskit.circuit import Parameter
+from qiskit.circuit import Parameter, Qubit, Clbit
 from qiskit.circuit.library import U1Gate, U2Gate, U3Gate, CU1Gate, CU3Gate
 
 
@@ -297,6 +297,121 @@ class TestUnroller(QiskitTestCase):
 
         self.assertEqual(Operator(qc), Operator(qcd))
 
+    def test_if_simple(self):
+        """Test a simple if statement unrolls correctly."""
+        qubits = [Qubit(), Qubit()]
+        clbits = [Clbit(), Clbit()]
+
+        qc = QuantumCircuit(qubits, clbits)
+        qc.h(0)
+        qc.measure(0, 0)
+        with qc.if_test((clbits[0], 0)):
+            qc.x(0)
+        qc.h(0)
+        qc.measure(0, 1)
+        with qc.if_test((clbits[1], 0)):
+            qc.h(1)
+            qc.cx(1, 0)
+        dag = circuit_to_dag(qc)
+        unrolled_dag = Unroller(["u", "cx"]).run(dag)
+
+        expected = QuantumCircuit(qubits, clbits)
+        expected.u(pi / 2, 0, pi, 0)
+        expected.measure(0, 0)
+        with expected.if_test((clbits[0], 0)):
+            expected.u(pi, 0, pi, 0)
+        expected.u(pi / 2, 0, pi, 0)
+        expected.measure(0, 1)
+        with expected.if_test((clbits[1], 0)):
+            expected.u(pi / 2, 0, pi, 1)
+            expected.cx(1, 0)
+        expected_dag = circuit_to_dag(expected)
+        self.assertEqual(unrolled_dag, expected_dag)
+
+    def test_if_else_simple(self):
+        """Test a simple if-else statement unrolls correctly."""
+        qubits = [Qubit(), Qubit()]
+        clbits = [Clbit(), Clbit()]
+
+        qc = QuantumCircuit(qubits, clbits)
+        qc.h(0)
+        qc.measure(0, 0)
+        with qc.if_test((clbits[0], 0)) as else_:
+            qc.x(0)
+        with else_:
+            qc.z(0)
+        qc.h(0)
+        qc.measure(0, 1)
+        with qc.if_test((clbits[1], 0)) as else_:
+            qc.h(1)
+            qc.cx(1, 0)
+        with else_:
+            qc.h(0)
+            qc.cx(0, 1)
+        dag = circuit_to_dag(qc)
+        unrolled_dag = Unroller(["u", "cx"]).run(dag)
+
+        expected = QuantumCircuit(qubits, clbits)
+        expected.u(pi / 2, 0, pi, 0)
+        expected.measure(0, 0)
+        with expected.if_test((clbits[0], 0)) as else_:
+            expected.u(pi, 0, pi, 0)
+        with else_:
+            expected.u(0, 0, pi, 0)
+        expected.u(pi / 2, 0, pi, 0)
+        expected.measure(0, 1)
+        with expected.if_test((clbits[1], 0)) as else_:
+            expected.u(pi / 2, 0, pi, 1)
+            expected.cx(1, 0)
+        with else_:
+            expected.u(pi / 2, 0, pi, 0)
+            expected.cx(0, 1)
+        expected_dag = circuit_to_dag(expected)
+        self.assertEqual(unrolled_dag, expected_dag)
+
+    def test_nested_control_flow(self):
+        """Test unrolling nested control flow blocks."""
+        qr = QuantumRegister(2)
+        cr1 = ClassicalRegister(1)
+        cr2 = ClassicalRegister(1)
+        cr3 = ClassicalRegister(1)
+        qc = QuantumCircuit(qr, cr1, cr2, cr3)
+        with qc.for_loop(range(3)):
+            with qc.while_loop((cr1, 0)):
+                qc.x(0)
+            with qc.while_loop((cr2, 0)):
+                qc.y(0)
+            with qc.while_loop((cr3, 0)):
+                qc.z(0)
+        dag = circuit_to_dag(qc)
+        unrolled_dag = Unroller(["u", "cx"]).run(dag)
+
+        expected = QuantumCircuit(qr, cr1, cr2, cr3)
+        with expected.for_loop(range(3)):
+            with expected.while_loop((cr1, 0)):
+                expected.u(pi, 0, pi, 0)
+            with expected.while_loop((cr2, 0)):
+                expected.u(pi, pi / 2, pi / 2, 0)
+            with expected.while_loop((cr3, 0)):
+                expected.u(0, 0, pi, 0)
+        expected_dag = circuit_to_dag(expected)
+        self.assertEqual(unrolled_dag, expected_dag)
+
+    def test_parameterized_angle(self):
+        """Test unrolling with parameterized angle"""
+        qc = QuantumCircuit(1)
+        index = Parameter("index")
+        with qc.for_loop((0, 0.5 * pi), index) as param:
+            qc.rx(param, 0)
+        dag = circuit_to_dag(qc)
+        unrolled_dag = Unroller(["u", "cx"]).run(dag)
+
+        expected = QuantumCircuit(1)
+        with expected.for_loop((0, 0.5 * pi), index) as param:
+            expected.u(param, -pi / 2, pi / 2, 0)
+        expected_dag = circuit_to_dag(expected)
+        self.assertEqual(unrolled_dag, expected_dag)
+
 
 class TestUnrollAllInstructions(QiskitTestCase):
     """Test unrolling a circuit containing all standard instructions."""