Fixing BlockCollapser with Clbits

qiskit/dagcircuit/collect_blocks.py

@@ -14,7 +14,8 @@
 """Various ways to divide a DAG into blocks of nodes, to split blocks of nodes
 into smaller sub-blocks, and to consolidate blocks."""
 
-from qiskit.circuit import QuantumCircuit, CircuitInstruction
+from qiskit.circuit import QuantumCircuit, CircuitInstruction, Clbit
+from qiskit.circuit.controlflow.condition import condition_bits
 from . import DAGOpNode, DAGCircuit, DAGDependency
 from .exceptions import DAGCircuitError
 
@@ -254,22 +255,37 @@ def collapse_to_operation(self, blocks, collapse_fn):
         then uses collapse_fn to collapse this circuit into a single operation.
         """
         global_index_map = {wire: idx for idx, wire in enumerate(self.dag.qubits)}
+        global_index_map.update({wire: idx for idx, wire in enumerate(self.dag.clbits)})
+
         for block in blocks:
-            # Find the set of qubits used in this block (which might be much smaller than
-            # the set of all qubits).
+            # Find the sets of qubits/clbits used in this block (which might be much smaller
+            # than the set of all qubits/clbits).
             cur_qubits = set()
+            cur_clbits = set()
+
             for node in block:
                 cur_qubits.update(node.qargs)
+                cur_clbits.update(_get_node_cargs(node))
 
-            # For reproducibility, order these qubits compatibly with the global order.
+            # For reproducibility, order these qubits/clbits compatibly with the global order.
             sorted_qubits = sorted(cur_qubits, key=lambda x: global_index_map[x])
+            sorted_clbits = sorted(cur_clbits, key=lambda x: global_index_map[x])
 
             # Construct a quantum circuit from the nodes in the block, remapping the qubits.
             wire_pos_map = {qb: ix for ix, qb in enumerate(sorted_qubits)}
-            qc = QuantumCircuit(len(cur_qubits))
+            wire_pos_map.update({qb: ix for ix, qb in enumerate(sorted_clbits)})
+
+            qc = QuantumCircuit(sorted_qubits, sorted_clbits)
             for node in block:
-                remapped_qubits = [wire_pos_map[qarg] for qarg in node.qargs]
-                qc.append(CircuitInstruction(node.op, remapped_qubits, node.cargs))
+                instructions = qc.append(CircuitInstruction(node.op, node.qargs, node.cargs))
+                cond = getattr(node.op, "condition", None)
+                if cond is not None:
+                    if not isinstance(cond[0], Clbit):
+                        raise DAGCircuitError(
+                            "Conditional gates must be conditioned on individual Clbits, "
+                            "not ClassicalRegisters to be able to collapse."
+                        )
+                    instructions.c_if(*cond)
 
             # Collapse this quantum circuit into an operation.
             op = collapse_fn(qc)
@@ -278,3 +294,14 @@ def collapse_to_operation(self, blocks, collapse_fn):
             # (the function replace_block_with_op is implemented both in DAGCircuit and DAGDependency).
             self.dag.replace_block_with_op(block, op, wire_pos_map, cycle_check=False)
         return self.dag
+
+
+def _get_node_cargs(node):
+    """Get cargs for possibly conditional node"""
+    # Fix because DAG circuits hide conditional node cargs in
+    # node.op.condition
+    cargs = set(node.cargs)
+    cond = getattr(node.op, "condition", None)
+    if cond is not None:
+        cargs.update(condition_bits(cond))
+    return cargs

qiskit/dagcircuit/dagcircuit.py

@@ -30,6 +30,7 @@
 import rustworkx as rx
 
 from qiskit.circuit import ControlFlowOp, ForLoopOp, IfElseOp, WhileLoopOp
+from qiskit.circuit.controlflow.condition import condition_bits
 from qiskit.circuit.exceptions import CircuitError
 from qiskit.circuit.quantumregister import QuantumRegister, Qubit
 from qiskit.circuit.classicalregister import ClassicalRegister, Clbit
@@ -1134,8 +1135,10 @@ def replace_block_with_op(self, node_block, op, wire_pos_map, cycle_check=True):
 
         for nd in node_block:
             block_qargs |= set(nd.qargs)
-            if isinstance(nd, DAGOpNode) and getattr(nd.op, "condition", None):
-                block_cargs |= set(nd.cargs)
+            block_cargs |= set(nd.cargs)
+            cond = getattr(nd.op, "condition", None)
+            if cond is not None:
+                block_cargs.update(condition_bits(cond))
 
         # Create replacement node
         new_node = DAGOpNode(

qiskit/dagcircuit/dagdependency.py

@@ -19,6 +19,7 @@
 
 import rustworkx as rx
 
+from qiskit.circuit.controlflow.condition import condition_bits
 from qiskit.circuit.quantumregister import QuantumRegister, Qubit
 from qiskit.circuit.classicalregister import ClassicalRegister, Clbit
 from qiskit.dagcircuit.exceptions import DAGDependencyError
@@ -394,11 +395,8 @@ def _create_op_node(self, operation, qargs, cargs):
                 #   (1) cindices_list are specific to template optimization and should not be computed
                 #       in this place.
                 #   (2) Template optimization pass needs currently does not handle general conditions.
-                if isinstance(operation.condition[0], Clbit):
-                    condition_bits = [operation.condition[0]]
-                else:
-                    condition_bits = operation.condition[0]
-                cindices_list = [self.clbits.index(clbit) for clbit in condition_bits]
+                cond_bits = condition_bits(operation.condition)
+                cindices_list = [self.clbits.index(clbit) for clbit in cond_bits]
             else:
                 cindices_list = []
         else:
@@ -655,8 +653,10 @@ def replace_block_with_op(self, node_block, op, wire_pos_map, cycle_check=True):
 
         for nd in node_block:
             block_qargs |= set(nd.qargs)
-            if nd.op.condition:
-                block_cargs |= set(nd.cargs)
+            block_cargs |= set(nd.cargs)
+            cond = getattr(nd.op, "condition", None)
+            if cond is not None:
+                block_cargs.update(condition_bits(cond))
 
         # Create replacement node
         new_node = self._create_op_node(

releasenotes/notes/fix-collapse-with-clbits-e14766353303d442.yaml

@@ -0,0 +1,9 @@
+---
+fixes:
+  - |
+    Fixed a bug in :class:`~BlockCollapser` where classical bits were ignored when collapsing
+    a block of nodes.
+  - |
+    Fixed a bug in :meth:`~qiskit.dagcircuit.DAGCircuit.replace_block_with_op` and
+    :meth:`~qiskit.dagcircuit.DAGDependency.replace_block_with_op`
+    that led to ignoring classical bits.

test/python/dagcircuit/test_collect_blocks.py

@@ -15,10 +15,17 @@
 
 import unittest
 
-from qiskit.converters import circuit_to_dag, circuit_to_dagdependency
+from qiskit import QuantumRegister, ClassicalRegister
+from qiskit.converters import (
+    circuit_to_dag,
+    circuit_to_dagdependency,
+    circuit_to_instruction,
+    dag_to_circuit,
+    dagdependency_to_circuit,
+)
 from qiskit.test import QiskitTestCase
-from qiskit.circuit import QuantumCircuit
-from qiskit.dagcircuit.collect_blocks import BlockCollector, BlockSplitter
+from qiskit.circuit import QuantumCircuit, Measure, Clbit
+from qiskit.dagcircuit.collect_blocks import BlockCollector, BlockSplitter, BlockCollapser
 
 
 class TestCollectBlocks(QiskitTestCase):
@@ -449,6 +456,239 @@ def test_do_not_split_blocks(self):
         )
         self.assertEqual(len(blocks), 1)
 
+    def test_collect_blocks_with_cargs(self):
+        """Test collecting and collapsing blocks with classical bits appearing as cargs."""
+
+        qc = QuantumCircuit(3)
+        qc.h(0)
+        qc.h(1)
+        qc.h(2)
+        qc.measure_all()
+
+        dag = circuit_to_dag(qc)
+
+        # Collect all measure instructions
+        blocks = BlockCollector(dag).collect_all_matching_blocks(
+            lambda node: isinstance(node.op, Measure), split_blocks=False, min_block_size=1
+        )
+
+        # We should have a single block consisting of 3 measures
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(len(blocks[0]), 3)
+        self.assertEqual(blocks[0][0].op, Measure())
+        self.assertEqual(blocks[0][1].op, Measure())
+        self.assertEqual(blocks[0][2].op, Measure())
+
+        def _collapse_fn(circuit):
+            op = circuit_to_instruction(circuit)
+            op.name = "COLLAPSED"
+            return op
+
+        # Collapse block with measures into a single "COLLAPSED" block
+        dag = BlockCollapser(dag).collapse_to_operation(blocks, _collapse_fn)
+        collapsed_qc = dag_to_circuit(dag)
+
+        self.assertEqual(len(collapsed_qc.data), 5)
+        self.assertEqual(collapsed_qc.data[0].operation.name, "h")
+        self.assertEqual(collapsed_qc.data[1].operation.name, "h")
+        self.assertEqual(collapsed_qc.data[2].operation.name, "h")
+        self.assertEqual(collapsed_qc.data[3].operation.name, "barrier")
+        self.assertEqual(collapsed_qc.data[4].operation.name, "COLLAPSED")
+        self.assertEqual(collapsed_qc.data[4].operation.definition.num_qubits, 3)
+        self.assertEqual(collapsed_qc.data[4].operation.definition.num_clbits, 3)
+
+    def test_collect_blocks_with_cargs_dagdependency(self):
+        """Test collecting and collapsing blocks with classical bits appearing as cargs,
+        using DAGDependency."""
+
+        qc = QuantumCircuit(3)
+        qc.h(0)
+        qc.h(1)
+        qc.h(2)
+        qc.measure_all()
+
+        dag = circuit_to_dagdependency(qc)
+
+        # Collect all measure instructions
+        blocks = BlockCollector(dag).collect_all_matching_blocks(
+            lambda node: isinstance(node.op, Measure), split_blocks=False, min_block_size=1
+        )
+
+        # We should have a single block consisting of 3 measures
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(len(blocks[0]), 3)
+        self.assertEqual(blocks[0][0].op, Measure())
+        self.assertEqual(blocks[0][1].op, Measure())
+        self.assertEqual(blocks[0][2].op, Measure())
+
+        def _collapse_fn(circuit):
+            op = circuit_to_instruction(circuit)
+            op.name = "COLLAPSED"
+            return op
+
+        # Collapse block with measures into a single "COLLAPSED" block
+        dag = BlockCollapser(dag).collapse_to_operation(blocks, _collapse_fn)
+        collapsed_qc = dagdependency_to_circuit(dag)
+
+        self.assertEqual(len(collapsed_qc.data), 5)
+        self.assertEqual(collapsed_qc.data[0].operation.name, "h")
+        self.assertEqual(collapsed_qc.data[1].operation.name, "h")
+        self.assertEqual(collapsed_qc.data[2].operation.name, "h")
+        self.assertEqual(collapsed_qc.data[3].operation.name, "barrier")
+        self.assertEqual(collapsed_qc.data[4].operation.name, "COLLAPSED")
+        self.assertEqual(collapsed_qc.data[4].operation.definition.num_qubits, 3)
+        self.assertEqual(collapsed_qc.data[4].operation.definition.num_clbits, 3)
+
+    def test_collect_blocks_with_clbits(self):
+        """Test collecting and collapsing blocks with classical bits appearing under
+        condition."""
+
+        qc = QuantumCircuit(4, 3)
+        qc.cx(0, 1).c_if(0, 1)
+        qc.cx(2, 3)
+        qc.cx(1, 2)
+        qc.cx(0, 1)
+        qc.cx(2, 3).c_if(1, 0)
+
+        dag = circuit_to_dag(qc)
+
+        # Collect all cx gates (including the conditional ones)
+        blocks = BlockCollector(dag).collect_all_matching_blocks(
+            lambda node: node.op.name == "cx", split_blocks=False, min_block_size=1
+        )
+
+        # We should have a single block consisting of all CX nodes
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(len(blocks[0]), 5)
+
+        def _collapse_fn(circuit):
+            op = circuit_to_instruction(circuit)
+            op.name = "COLLAPSED"
+            return op
+
+        # Collapse block with measures into a single "COLLAPSED" block
+        dag = BlockCollapser(dag).collapse_to_operation(blocks, _collapse_fn)
+        collapsed_qc = dag_to_circuit(dag)
+
+        self.assertEqual(len(collapsed_qc.data), 1)
+        self.assertEqual(collapsed_qc.data[0].operation.name, "COLLAPSED")
+        self.assertEqual(collapsed_qc.data[0].operation.definition.num_qubits, 4)
+        self.assertEqual(collapsed_qc.data[0].operation.definition.num_clbits, 2)
+
+    def test_collect_blocks_with_clbits_dagdependency(self):
+        """Test collecting and collapsing blocks with classical bits appearing
+        under conditions, using DAGDependency."""
+
+        qc = QuantumCircuit(4, 3)
+        qc.cx(0, 1).c_if(0, 1)
+        qc.cx(2, 3)
+        qc.cx(1, 2)
+        qc.cx(0, 1)
+        qc.cx(2, 3).c_if(1, 0)
+
+        dag = circuit_to_dagdependency(qc)
+
+        # Collect all cx gates (including the conditional ones)
+        blocks = BlockCollector(dag).collect_all_matching_blocks(
+            lambda node: node.op.name == "cx", split_blocks=False, min_block_size=1
+        )
+
+        # We should have a single block consisting of all CX nodes
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(len(blocks[0]), 5)
+
+        def _collapse_fn(circuit):
+            op = circuit_to_instruction(circuit)
+            op.name = "COLLAPSED"
+            return op
+
+        # Collapse block with measures into a single "COLLAPSED" block
+        dag = BlockCollapser(dag).collapse_to_operation(blocks, _collapse_fn)
+        collapsed_qc = dagdependency_to_circuit(dag)
+
+        self.assertEqual(len(collapsed_qc.data), 1)
+        self.assertEqual(collapsed_qc.data[0].operation.name, "COLLAPSED")
+        self.assertEqual(collapsed_qc.data[0].operation.definition.num_qubits, 4)
+        self.assertEqual(collapsed_qc.data[0].operation.definition.num_clbits, 2)
+
+    def test_collect_blocks_with_clbits2(self):
+        """Test collecting and collapsing blocks with classical bits appearing under
+        condition."""
+
+        qreg = QuantumRegister(4, "qr")
+        creg = ClassicalRegister(3, "cr")
+        cbit = Clbit()
+
+        qc = QuantumCircuit(qreg, creg, [cbit])
+        qc.cx(0, 1).c_if(creg[1], 1)
+        qc.cx(2, 3).c_if(cbit, 0)
+        qc.cx(1, 2)
+        qc.cx(0, 1).c_if(creg[2], 1)
+
+        dag = circuit_to_dag(qc)
+
+        # Collect all cx gates (including the conditional ones)
+        blocks = BlockCollector(dag).collect_all_matching_blocks(
+            lambda node: node.op.name == "cx", split_blocks=False, min_block_size=1
+        )
+
+        # We should have a single block consisting of all CX nodes
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(len(blocks[0]), 4)
+
+        def _collapse_fn(circuit):
+            op = circuit_to_instruction(circuit)
+            op.name = "COLLAPSED"
+            return op
+
+        # Collapse block with measures into a single "COLLAPSED" block
+        dag = BlockCollapser(dag).collapse_to_operation(blocks, _collapse_fn)
+        collapsed_qc = dag_to_circuit(dag)
+
+        self.assertEqual(len(collapsed_qc.data), 1)
+        self.assertEqual(collapsed_qc.data[0].operation.name, "COLLAPSED")
+        self.assertEqual(collapsed_qc.data[0].operation.definition.num_qubits, 4)
+        self.assertEqual(collapsed_qc.data[0].operation.definition.num_clbits, 3)
+
+    def test_collect_blocks_with_clbits2_dagdependency(self):
+        """Test collecting and collapsing blocks with classical bits appearing under
+        condition."""
+
+        qreg = QuantumRegister(4, "qr")
+        creg = ClassicalRegister(3, "cr")
+        cbit = Clbit()
+
+        qc = QuantumCircuit(qreg, creg, [cbit])
+        qc.cx(0, 1).c_if(creg[1], 1)
+        qc.cx(2, 3).c_if(cbit, 0)
+        qc.cx(1, 2)
+        qc.cx(0, 1).c_if(creg[2], 1)
+
+        dag = circuit_to_dag(qc)
+
+        # Collect all cx gates (including the conditional ones)
+        blocks = BlockCollector(dag).collect_all_matching_blocks(
+            lambda node: node.op.name == "cx", split_blocks=False, min_block_size=1
+        )
+
+        # We should have a single block consisting of all CX nodes
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(len(blocks[0]), 4)
+
+        def _collapse_fn(circuit):
+            op = circuit_to_instruction(circuit)
+            op.name = "COLLAPSED"
+            return op
+
+        # Collapse block with measures into a single "COLLAPSED" block
+        dag = BlockCollapser(dag).collapse_to_operation(blocks, _collapse_fn)
+        collapsed_qc = dag_to_circuit(dag)
+
+        self.assertEqual(len(collapsed_qc.data), 1)
+        self.assertEqual(collapsed_qc.data[0].operation.name, "COLLAPSED")
+        self.assertEqual(collapsed_qc.data[0].operation.definition.num_qubits, 4)
+        self.assertEqual(collapsed_qc.data[0].operation.definition.num_clbits, 3)
+
 
 if __name__ == "__main__":
     unittest.main()