diff --git a/qiskit/transpiler/passes/routing/sabre_swap.py b/qiskit/transpiler/passes/routing/sabre_swap.py index ff95133f891c..6dcbf3c9823b 100644 --- a/qiskit/transpiler/passes/routing/sabre_swap.py +++ b/qiskit/transpiler/passes/routing/sabre_swap.py @@ -205,21 +205,29 @@ def run(self, dag): original_layout = layout.copy() dag_list = [] - for node in dag.topological_op_nodes(): - if len(node.qargs) == 2: - dag_list.append( - ( - node._node_id, - self._bit_indices[node.qargs[0]], - self._bit_indices[node.qargs[1]], - ) - ) + layers = dag.multigraph_layers() + replay_list = [] + front_layer = None + for layer in layers: + for node in layer: + if isinstance(node, DAGOpNode): + replay_list.append(node) + if len(node.qargs) == 2: + dag_list.append( + ( + node._node_id, + self._bit_indices[node.qargs[0]], + self._bit_indices[node.qargs[1]], + ) + ) + if front_layer is None: + front_layer = dag_list sabre_dag = SabreDAG(len(dag.qubits), dag_list) # A decay factor for each qubit used to heuristically penalize recently # used qubits (to encourage parallelism). qubits_decay = QubitsDecay(len(dag.qubits)) - swap_map = build_swap_map( + swap_map, gate_order = build_swap_map( sabre_dag, self._neighbor_table, self.dist_matrix, @@ -233,28 +241,51 @@ def run(self, dag): output_layout = Layout({dag.qubits[k]: v for (k, v) in layout_mapping}) self.property_set["final_layout"] = output_layout if not self.fake_run: + queued_gates = [] + processed_nodes = set() + count = 0 # Reconstruct circuit by iterating over layers to preserve # relative positioning of 1q gates with SWAPs. - layers = dag.multigraph_layers() - next(layers) - for layer in layers: - for node in layer: - if not isinstance(node, DAGOpNode): + for node in replay_list: + if len(node.qargs) == 2: + if node._node_id != gate_order[count]: + queued_gates.append(node) continue - if node._node_id in swap_map: - for swap in swap_map[node._node_id]: - swap_qargs = [canonical_register[swap[0]], canonical_register[swap[1]]] + self._process_swaps( + swap_map, node, mapped_dag, original_layout, canonical_register + ) + self._apply_gate(mapped_dag, node, original_layout, canonical_register) + processed_nodes.add(node._node_id) + count += 1 + if queued_gates: + for gate_node in queued_gates: + if gate_node._node_id in processed_nodes: + continue + self._process_swaps( + swap_map, gate_node, mapped_dag, original_layout, canonical_register + ) self._apply_gate( - mapped_dag, - DAGOpNode(op=SwapGate(), qargs=swap_qargs), - original_layout, - canonical_register, + mapped_dag, gate_node, original_layout, canonical_register ) - original_layout.swap_logical(*swap) + count += 1 + queued_gates.clear() + else: self._apply_gate(mapped_dag, node, original_layout, canonical_register) return mapped_dag return dag + def _process_swaps(self, swap_map, node, mapped_dag, current_layout, canonical_register): + if node._node_id in swap_map: + for swap in swap_map[node._node_id]: + swap_qargs = [canonical_register[swap[0]], canonical_register[swap[1]]] + self._apply_gate( + mapped_dag, + DAGOpNode(op=SwapGate(), qargs=swap_qargs), + current_layout, + canonical_register, + ) + current_layout.swap_logical(*swap) + def _apply_gate(self, mapped_dag, node, current_layout, canonical_register): new_node = self._transform_gate_for_layout(node, current_layout, canonical_register) if self.fake_run: diff --git a/src/sabre_swap/mod.rs b/src/sabre_swap/mod.rs index f301c44d8c19..195d670a5ac6 100644 --- a/src/sabre_swap/mod.rs +++ b/src/sabre_swap/mod.rs @@ -23,6 +23,7 @@ use std::cmp::Ordering; use hashbrown::{HashMap, HashSet}; use ndarray::prelude::*; +use numpy::IntoPyArray; use numpy::PyReadonlyArray2; use pyo3::prelude::*; use pyo3::wrap_pyfunction; @@ -144,6 +145,7 @@ fn cmap_from_neighor_table(neighbor_table: &NeighborTable) -> DiGraph<(), ()> { /// swaps that should be added before that op. #[pyfunction] pub fn build_swap_map( + py: Python, dag: &SabreDAG, neighbor_table: &NeighborTable, distance_matrix: PyReadonlyArray2, @@ -151,7 +153,8 @@ pub fn build_swap_map( heuristic: &Heuristic, rng: &mut SabreRng, layout: &mut NLayout, -) -> PyResult { +) -> PyResult<(SwapMap, PyObject)> { + let mut gate_order: Vec = Vec::with_capacity(dag.dag.node_count()); let run_in_parallel = getenv_use_multiple_threads(); let mut out_map: HashMap> = HashMap::new(); let mut front_layer: Vec = dag.first_layer.clone(); @@ -259,6 +262,8 @@ pub fn build_swap_map( } if !execute_gate_list.is_empty() { for node in execute_gate_list { + let node_weight = dag.dag.node_weight(node).unwrap(); + gate_order.push(node_weight[0]); let out_swaps: Vec<[usize; 2]> = ops_since_progress.drain(..).collect(); out_map.insert(dag.dag.node_weight(node).unwrap()[0], out_swaps); for edge in dag.dag.edges(node) { @@ -308,7 +313,7 @@ pub fn build_swap_map( } ops_since_progress.push(best_swap); } - Ok(SwapMap { map: out_map }) + Ok((SwapMap { map: out_map }, gate_order.into_pyarray(py).into())) } /// Run the sabre heuristic scoring