Fix control-flow routing in StochasticSwap (#8880) (#8888)

* Fix control-flow routing in StochasticSwap

The `StochasticSwap` pass has some fairly complex hand-offs between
different parts of its API, including in recursive calls and the regular
hand-off between Python and Rust.  In the course of adding the
control-flow support, some of these became muddled, and the mapping
between different virtual/physical/integer representations got mixed up,
resulting in invalid swaps being output in the final circuit.

This commit simplifies much of the internal mapping, removing many
superfluous `DAGCircuit` creations and compositions.  This also removes
instances where two layouts were "chained"; this was not well typed (the
output of a "virtual -> physical" mapping can't be the input for another
"virtual -> physical" mapping), and in general was being used to "undo"
some of compositions that were about to be applied.

This fixes a tacit assumption in the original code that the initial
layout was a trivial layout in the hand-off between Rust and Python.
This worked until the recursive call added the `initial_layout` option,
making this assumption invalid. Previously, virtual qubit bit instances
were converted to integers (to allow them to be passed to Rust) using
their indices into the original DAG, but the integer outputs were then
converted _back_ using the `initial_layout`.  In the old form, this
worked anyway, but wasn't logically correct and consequently broke when
the assumptions about `initial_layout` changed.

For the recursive calls, we now ensure that the inner passes are
essentially created with the same internal structure as the outer pass;
the passed in `DAGCircuit` uses the same bit instances and same meaning
of the virtual qubits as the outer circuit, and the `initial_layout`
ensures that the inner passes start with at the same layout as the outer
pass.  This makes the inner passes more like a logical continuation of
the current operation, rather than a completely separate entity that
needs to have its virtual qubits remapped.

The changes to the tests are twofold:

- move the `CheckMap` calls earlier and apply them directly to the
  `StochasticSwap` output rather than the expected circuit, to improve
  the quality of failure error messages

- use the same physical qubits inside the expected control-flow blocks;
  the new simpler form of doing the circuit rewriting internally in the
  pass ensures that the same bit objects are used all the way through
  the control-flow stack now, rather than creating new instances.

* Add tests for stochastic swap valid output

This commit adds full path transpile() tests for running with stochastic
swap that validates a full path transpilation outputs a valid physical
circuit. These tests are purposefully high level to provide some
validation that stochastic swap is not creating invalid output by
inserting incorrect swaps. It's not meant as a test of valid unitary
equivalent output of the full transpilation.

Co-authored-by: Jake Lishman <jake.lishman@ibm.com>

Co-authored-by: Matthew Treinish <mtreinish@kortar.org>
(cherry picked from commit b3cf64f)

Co-authored-by: Jake Lishman <jake.lishman@ibm.com>
Co-authored-by: Matthew Treinish <mtreinish@kortar.org>

Author: 3 people

qiskit/transpiler/passes/routing/stochastic_swap.py

@@ -26,7 +26,7 @@
 from qiskit.circuit import IfElseOp, WhileLoopOp, ForLoopOp, ControlFlowOp
 from qiskit._accelerate import stochastic_swap as stochastic_swap_rs
 
-from .utils import combine_permutations, get_swap_map_dag
+from .utils import get_swap_map_dag
 
 logger = logging.getLogger(__name__)
 
@@ -70,7 +70,8 @@ def __init__(self, coupling_map, trials=20, seed=None, fake_run=False, initial_l
         self.fake_run = fake_run
         self.qregs = None
         self.initial_layout = initial_layout
-        self._qubit_indices = None
+        self._qubit_to_int = None
+        self._int_to_qubit = None
 
     def run(self, dag):
         """Run the StochasticSwap pass on `dag`.
@@ -97,7 +98,10 @@ def run(self, dag):
         canonical_register = dag.qregs["q"]
         if self.initial_layout is None:
             self.initial_layout = Layout.generate_trivial_layout(canonical_register)
-        self._qubit_indices = {bit: idx for idx, bit in enumerate(dag.qubits)}
+        # Qubit indices are used to assign an integer to each virtual qubit during the routing: it's
+        # a mapping of {virtual: virtual}, for converting between Python and Rust forms.
+        self._qubit_to_int = {bit: idx for idx, bit in enumerate(dag.qubits)}
+        self._int_to_qubit = tuple(dag.qubits)
 
         self.qregs = dag.qregs
         logger.debug("StochasticSwap rng seeded with seed=%s", self.seed)
@@ -174,18 +178,18 @@ def _layer_permutation(self, layer_partition, layout, qubit_subset, coupling, tr
 
         cdist2 = coupling._dist_matrix**2
         int_qubit_subset = np.fromiter(
-            (self._qubit_indices[bit] for bit in qubit_subset),
+            (self._qubit_to_int[bit] for bit in qubit_subset),
             dtype=np.uintp,
             count=len(qubit_subset),
         )
 
         int_gates = np.fromiter(
-            (self._qubit_indices[bit] for gate in gates for bit in gate),
+            (self._qubit_to_int[bit] for gate in gates for bit in gate),
             dtype=np.uintp,
             count=2 * len(gates),
         )
 
-        layout_mapping = {self._qubit_indices[k]: v for k, v in layout.get_virtual_bits().items()}
+        layout_mapping = {self._qubit_to_int[k]: v for k, v in layout.get_virtual_bits().items()}
         int_layout = stochastic_swap_rs.NLayout(layout_mapping, num_qubits, coupling.size())
 
         trial_circuit = DAGCircuit()  # SWAP circuit for slice of swaps in this trial
@@ -204,16 +208,15 @@ def _layer_permutation(self, layer_partition, layout, qubit_subset, coupling, tr
             edges,
             seed=self.seed,
         )
-        # If we have no best circuit for this layer, all of the
-        # trials have failed
+        # If we have no best circuit for this layer, all of the trials have failed
         if best_layout is None:
             logger.debug("layer_permutation: failed!")
             return False, None, None, None
 
         edges = best_edges.edges()
         for idx in range(len(edges) // 2):
-            swap_src = self.initial_layout._p2v[edges[2 * idx]]
-            swap_tgt = self.initial_layout._p2v[edges[2 * idx + 1]]
+            swap_src = self._int_to_qubit[edges[2 * idx]]
+            swap_tgt = self._int_to_qubit[edges[2 * idx + 1]]
             trial_circuit.apply_operation_back(SwapGate(), [swap_src, swap_tgt], [])
         best_circuit = trial_circuit
 
@@ -234,24 +237,17 @@ def _layer_update(self, dag, layer, best_layout, best_depth, best_circuit):
             best_depth (int): depth returned from _layer_permutation
             best_circuit (DAGCircuit): swap circuit returned from _layer_permutation
         """
-        layout = best_layout
-        logger.debug("layer_update: layout = %s", layout)
+        logger.debug("layer_update: layout = %s", best_layout)
         logger.debug("layer_update: self.initial_layout = %s", self.initial_layout)
 
         # Output any swaps
         if best_depth > 0:
             logger.debug("layer_update: there are swaps in this layer, depth %d", best_depth)
-            dag.compose(best_circuit)
+            dag.compose(best_circuit, qubits={bit: bit for bit in best_circuit.qubits})
         else:
             logger.debug("layer_update: there are no swaps in this layer")
         # Output this layer
-        layer_circuit = layer["graph"]
-        initial_v2p = self.initial_layout.get_virtual_bits()
-        new_v2p = layout.get_virtual_bits()
-        initial_order = [initial_v2p[qubit] for qubit in dag.qubits]
-        new_order = [new_v2p[qubit] for qubit in dag.qubits]
-        order = combine_permutations(initial_order, new_order)
-        dag.compose(layer_circuit, qubits=order)
+        dag.compose(layer["graph"], qubits=best_layout.reorder_bits(dag.qubits))
 
     def _mapper(self, circuit_graph, coupling_graph, trials=20):
         """Map a DAGCircuit onto a CouplingMap using swap gates.
@@ -376,29 +372,18 @@ def _controlflow_layer_update(self, dagcircuit_output, layer_dag, current_layout
         """
         cf_opnode = layer_dag.op_nodes()[0]
         if isinstance(cf_opnode.op, IfElseOp):
-            updated_ctrl_op, cf_layout, idle_qubits = self._route_control_flow_multiblock(
+            new_op, new_qargs, new_layout = self._route_control_flow_multiblock(
                 cf_opnode, current_layout, root_dag
             )
         elif isinstance(cf_opnode.op, (ForLoopOp, WhileLoopOp)):
-            updated_ctrl_op, cf_layout, idle_qubits = self._route_control_flow_looping(
+            new_op, new_qargs, new_layout = self._route_control_flow_looping(
                 cf_opnode, current_layout, root_dag
             )
         else:
             raise TranspilerError(f"unsupported control flow operation: {cf_opnode}")
-        if self.fake_run:
-            return cf_layout
-
-        cf_layer_dag = DAGCircuit()
-        cf_qubits = [qubit for qubit in root_dag.qubits if qubit not in idle_qubits]
-        qreg = QuantumRegister(len(cf_qubits), "q")
-        cf_layer_dag.add_qreg(qreg)
-        for creg in layer_dag.cregs.values():
-            cf_layer_dag.add_creg(creg)
-        cf_layer_dag.apply_operation_back(updated_ctrl_op, cf_layer_dag.qubits, cf_opnode.cargs)
-        target_qubits = [qubit for qubit in dagcircuit_output.qubits if qubit not in idle_qubits]
-        order = current_layout.reorder_bits(target_qubits)
-        dagcircuit_output.compose(cf_layer_dag, qubits=order)
-        return cf_layout
+        if not self.fake_run:
+            dagcircuit_output.apply_operation_back(new_op, new_qargs, cf_opnode.cargs)
+        return new_layout
 
     def _new_seed(self):
         """Get a seed for a new RNG instance."""
@@ -432,19 +417,19 @@ def _route_control_flow_multiblock(self, node, current_layout, root_dag):
 
         Returns:
             ControlFlowOp: routed control flow operation.
-            final_layout (Layout): layout after instruction.
-            list(Qubit): list of idle qubits in controlflow layer.
+            List[Qubit]: the new physical-qubit arguments that the output `ControlFlowOp` should be
+                applied to. This might be wider than the input node if internal routing was needed.
+            Layout: the new layout after the control-flow operation is applied.
         """
         # For each block, expand it up be the full width of the containing DAG so we can be certain
         # that it is routable, then route it within that.  When we recombine later, we'll reduce all
         # these blocks down to remove any qubits that are idle.
         block_dags = []
         block_layouts = []
-        order = [self._qubit_indices[bit] for bit in node.qargs]
         for block in node.op.blocks:
             inner_pass = self._recursive_pass(current_layout)
             full_dag_block = root_dag.copy_empty_like()
-            full_dag_block.compose(circuit_to_dag(block), qubits=order)
+            full_dag_block.compose(circuit_to_dag(block), qubits=node.qargs)
             block_dags.append(inner_pass.run(full_dag_block))
             block_layouts.append(inner_pass.property_set["final_layout"].copy())
 
@@ -455,42 +440,28 @@ def _route_control_flow_multiblock(self, node, current_layout, root_dag):
         deepest_index = np.argmax([block.depth(recurse=True) for block in block_dags])
         final_layout = block_layouts[deepest_index]
         if self.fake_run:
-            return None, final_layout, None
-        p2v = current_layout.get_physical_bits()
+            return None, None, final_layout
         idle_qubits = set(root_dag.qubits)
         for i, updated_dag_block in enumerate(block_dags):
             if i != deepest_index:
-                swap_circuit, swap_qubits = get_swap_map_dag(
+                swap_dag, swap_qubits = get_swap_map_dag(
                     root_dag,
                     self.coupling_map,
                     block_layouts[i],
                     final_layout,
                     seed=self._new_seed(),
                 )
-                if swap_circuit.depth():
-                    virtual_swap_dag = updated_dag_block.copy_empty_like()
-                    order = [p2v[virtual_swap_dag.qubits.index(qubit)] for qubit in swap_qubits]
-                    virtual_swap_dag.compose(swap_circuit, qubits=order)
-                    updated_dag_block.compose(virtual_swap_dag)
+                if swap_dag.depth():
+                    updated_dag_block.compose(swap_dag, qubits=swap_qubits)
             idle_qubits &= set(updated_dag_block.idle_wires())
 
         # Now for each block, expand it to be full width over all active wires (all blocks of a
         # control-flow operation need to have equal input wires), and convert it to circuit form.
         block_circuits = []
         for i, updated_dag_block in enumerate(block_dags):
             updated_dag_block.remove_qubits(*idle_qubits)
-            new_dag_block = DAGCircuit()
-            new_num_qubits = updated_dag_block.num_qubits()
-            qreg = QuantumRegister(new_num_qubits, "q")
-            new_dag_block.add_qreg(qreg)
-            for creg in updated_dag_block.cregs.values():
-                new_dag_block.add_creg(creg)
-            for inner_node in updated_dag_block.op_nodes():
-                new_qargs = [qreg[updated_dag_block.qubits.index(bit)] for bit in inner_node.qargs]
-                new_dag_block.apply_operation_back(inner_node.op, new_qargs, inner_node.cargs)
-            block_circuits.append(dag_to_circuit(new_dag_block))
-
-        return node.op.replace_blocks(block_circuits), final_layout, idle_qubits
+            block_circuits.append(dag_to_circuit(updated_dag_block))
+        return node.op.replace_blocks(block_circuits), block_circuits[0].qubits, final_layout
 
     def _route_control_flow_looping(self, node, current_layout, root_dag):
         """Route a control-flow operation that represents a loop, such as :class:`.ForOpLoop` or
@@ -505,45 +476,36 @@ def _route_control_flow_looping(self, node, current_layout, root_dag):
 
         Returns:
             ControlFlowOp: routed control flow operation.
-            Layout: layout after instruction (this will be the same as the input layout).
-            list(Qubit): list of idle qubits in controlflow layer.
+            List[Qubit]: the new physical-qubit arguments that the output `ControlFlowOp` should be
+                applied to. This might be wider than the input node if internal routing was needed.
+            Layout: the new layout after the control-flow operation is applied.
         """
         if self.fake_run:
-            return None, current_layout, None
+            return None, None, current_layout
         # Temporarily expand to full width, and route within that.
         inner_pass = self._recursive_pass(current_layout)
-        order = [self._qubit_indices[bit] for bit in node.qargs]
         full_dag_block = root_dag.copy_empty_like()
-        full_dag_block.compose(circuit_to_dag(node.op.blocks[0]), qubits=order)
+        full_dag_block.compose(circuit_to_dag(node.op.blocks[0]), qubits=node.qargs)
         updated_dag_block = inner_pass.run(full_dag_block)
 
         # Ensure that the layout at the end of the block is returned to being the layout at the
         # start of the block again, so the loop works.
-        swap_circuit, swap_qubits = get_swap_map_dag(
+        swap_dag, swap_qubits = get_swap_map_dag(
             root_dag,
             self.coupling_map,
             inner_pass.property_set["final_layout"],
             current_layout,
             seed=self._new_seed(),
         )
-        if swap_circuit.depth():
-            p2v = current_layout.get_physical_bits()
-            virtual_swap_dag = updated_dag_block.copy_empty_like()
-            order = [p2v[virtual_swap_dag.qubits.index(qubit)] for qubit in swap_qubits]
-            virtual_swap_dag.compose(swap_circuit, qubits=order)
-            updated_dag_block.compose(virtual_swap_dag)
+        if swap_dag.depth():
+            updated_dag_block.compose(swap_dag, qubits=swap_qubits)
 
         # Contract the routed block back down to only operate on the qubits that it actually needs.
         idle_qubits = set(root_dag.qubits) & set(updated_dag_block.idle_wires())
         updated_dag_block.remove_qubits(*idle_qubits)
-        new_dag_block = DAGCircuit()
-        new_num_qubits = updated_dag_block.num_qubits()
-        qreg = QuantumRegister(new_num_qubits, "q")
-        new_dag_block.add_qreg(qreg)
-        for creg in updated_dag_block.cregs.values():
-            new_dag_block.add_creg(creg)
-        for inner_node in updated_dag_block.op_nodes():
-            new_qargs = [qreg[updated_dag_block.qubits.index(bit)] for bit in inner_node.qargs]
-            new_dag_block.apply_operation_back(inner_node.op, new_qargs, inner_node.cargs)
-        updated_circ_block = dag_to_circuit(new_dag_block)
-        return node.op.replace_blocks([updated_circ_block]), current_layout, idle_qubits
+        updated_circ_block = dag_to_circuit(updated_dag_block)
+        return (
+            node.op.replace_blocks([updated_circ_block]),
+            updated_dag_block.qubits,
+            current_layout,
+        )

qiskit/transpiler/passes/routing/utils.py

@@ -12,49 +12,24 @@
 
 """Utility functions for routing"""
 
-from qiskit.transpiler import CouplingMap
 from qiskit.transpiler.exceptions import TranspilerError
 from .algorithms import ApproximateTokenSwapper
 
 
-def combine_permutations(*permutations):
-    """
-    Chain a series of permutations.
-
-    Args:
-        *permutations (list(int)): permutations to combine
-
-    Returns:
-        list: combined permutation
-    """
-    order = permutations[0]
-    for this_order in permutations[1:]:
-        order = [order[i] for i in this_order]
-    return order
-
-
 def get_swap_map_dag(dag, coupling_map, from_layout, to_layout, seed, trials=4):
-    """Get the circuit of swaps to go from from_layout to to_layout, and the qubit ordering of the
-    qubits in that circuit."""
-
+    """Get the circuit of swaps to go from from_layout to to_layout, and the physical qubits
+    (integers) that the swap circuit should be applied on."""
     if len(dag.qregs) != 1 or dag.qregs.get("q", None) is None:
         raise TranspilerError("layout transformation runs on physical circuits only")
-
     if len(dag.qubits) > len(coupling_map.physical_qubits):
         raise TranspilerError("The layout does not match the amount of qubits in the DAG")
-
-    if coupling_map:
-        graph = coupling_map.graph.to_undirected()
-    else:
-        coupling_map = CouplingMap.from_full(len(to_layout))
-        graph = coupling_map.graph.to_undirected()
-
-    token_swapper = ApproximateTokenSwapper(graph, seed)
+    token_swapper = ApproximateTokenSwapper(coupling_map.graph.to_undirected(), seed)
     # Find the permutation between the initial physical qubits and final physical qubits.
     permutation = {
-        pqubit: to_layout.get_virtual_bits()[vqubit]
-        for vqubit, pqubit in from_layout.get_virtual_bits().items()
+        pqubit: to_layout[vqubit] for vqubit, pqubit in from_layout.get_virtual_bits().items()
     }
-    permutation_circ = token_swapper.permutation_circuit(permutation, trials)
-    permutation_qubits = [dag.qubits[i] for i in sorted(permutation_circ.inputmap.keys())]
-    return permutation_circ.circuit, permutation_qubits
+    # The mapping produced here maps physical qubit indices of the outer dag to the bits used to
+    # represent them in the inner map.  For later composing, we actually want the opposite map.
+    swap_circuit, phys_to_circuit_qubits = token_swapper.permutation_circuit(permutation, trials)
+    circuit_to_phys = {inner: outer for outer, inner in phys_to_circuit_qubits.items()}
+    return swap_circuit, [circuit_to_phys[bit] for bit in swap_circuit.qubits]