Fix performance of Sabre rust<->Python boundary (backport #10652)

crates/accelerate/src/sabre_layout.rs

@@ -12,6 +12,7 @@
 #![allow(clippy::too_many_arguments)]
 
 use ndarray::prelude::*;
+use numpy::IntoPyArray;
 use numpy::PyReadonlyArray2;
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
@@ -24,10 +25,12 @@ use crate::getenv_use_multiple_threads;
 use crate::nlayout::NLayout;
 use crate::sabre_swap::neighbor_table::NeighborTable;
 use crate::sabre_swap::sabre_dag::SabreDAG;
-use crate::sabre_swap::{build_swap_map_inner, Heuristic, SabreResult};
+use crate::sabre_swap::swap_map::SwapMap;
+use crate::sabre_swap::{build_swap_map_inner, Heuristic, NodeBlockResults, SabreResult};
 
 #[pyfunction]
 pub fn sabre_layout_and_routing(
+    py: Python,
     dag: &SabreDAG,
     neighbor_table: &NeighborTable,
     distance_matrix: PyReadonlyArray2<f64>,
@@ -36,7 +39,7 @@ pub fn sabre_layout_and_routing(
     num_swap_trials: usize,
     num_layout_trials: usize,
     seed: Option<u64>,
-) -> ([NLayout; 2], SabreResult) {
+) -> ([NLayout; 2], (SwapMap, PyObject, NodeBlockResults)) {
     let run_in_parallel = getenv_use_multiple_threads();
     let outer_rng = match seed {
         Some(seed) => Pcg64Mcg::seed_from_u64(seed),
@@ -47,7 +50,7 @@ pub fn sabre_layout_and_routing(
         .take(num_layout_trials)
         .collect();
     let dist = distance_matrix.as_array();
-    if run_in_parallel && num_layout_trials > 1 {
+    let res = if run_in_parallel && num_layout_trials > 1 {
         seed_vec
             .into_par_iter()
             .enumerate()
@@ -91,7 +94,15 @@ pub fn sabre_layout_and_routing(
             })
             .min_by_key(|(_, result)| result.map.map.values().map(|x| x.len()).sum::<usize>())
             .unwrap()
-    }
+    };
+    (
+        res.0,
+        (
+            res.1.map,
+            res.1.node_order.into_pyarray(py).into(),
+            res.1.node_block_results,
+        ),
+    )
 }
 
 fn layout_trial(

crates/accelerate/src/sabre_swap/mod.rs

@@ -208,12 +208,13 @@ fn cmap_from_neighor_table(neighbor_table: &NeighborTable) -> DiGraph<(), ()> {
 /// Run sabre swap on a circuit
 ///
 /// Returns:
-///     (SwapMap, gate_order): A tuple where the first element is a mapping of
+///     (SwapMap, gate_order, node_block_results): A tuple where the first element is a mapping of
 ///     DAGCircuit node ids to a list of virtual qubit swaps that should be
 ///     added before that operation. The second element is a numpy array of
 ///     node ids that represents the traversal order used by sabre.
 #[pyfunction]
 pub fn build_swap_map(
+    py: Python,
     num_qubits: usize,
     dag: &SabreDAG,
     neighbor_table: &NeighborTable,
@@ -223,9 +224,9 @@ pub fn build_swap_map(
     num_trials: usize,
     seed: Option<u64>,
     run_in_parallel: Option<bool>,
-) -> SabreResult {
+) -> (SwapMap, PyObject, NodeBlockResults) {
     let dist = distance_matrix.as_array();
-    build_swap_map_inner(
+    let res = build_swap_map_inner(
         num_qubits,
         dag,
         neighbor_table,
@@ -235,6 +236,11 @@ pub fn build_swap_map(
         layout,
         num_trials,
         run_in_parallel,
+    );
+    (
+        res.map,
+        res.node_order.into_pyarray(py).into(),
+        res.node_block_results,
     )
 }
 

qiskit/transpiler/passes/routing/sabre_swap.py

@@ -341,11 +341,14 @@ def empty_dag(node, block):
         return out
 
     def apply_inner(out_dag, current_layout, qubit_indices_inner, result, id_to_node):
-        for node_id in result.node_order:
+        node_order = result[1]
+        swap_map = result[0]
+        node_block_results = result[2]
+        for node_id in node_order:
             node = id_to_node[node_id]
             if isinstance(node.op, ControlFlowOp):
                 # Handle control flow op and continue.
-                block_results = result.node_block_results[node_id]
+                block_results = node_block_results[node_id]
                 mapped_block_dags = []
                 idle_qubits = set(out_dag.qubits)
                 for block, block_result in zip(node.op.blocks, block_results):
@@ -360,7 +363,11 @@ def apply_inner(out_dag, current_layout, qubit_indices_inner, result, id_to_node
                         mapped_block_dag,
                         mapped_block_layout,
                         block_qubit_indices,
-                        block_result.result,
+                        (
+                            block_result.result.map,
+                            block_result.result.node_order,
+                            block_result.result.node_block_results,
+                        ),
                         block_id_to_node,
                     )
 
@@ -396,9 +403,9 @@ def apply_inner(out_dag, current_layout, qubit_indices_inner, result, id_to_node
                 continue
 
             # If we get here, the node isn't a control-flow gate.
-            if node_id in result.map:
+            if node_id in swap_map:
                 process_swaps(
-                    result.map[node_id],
+                    swap_map[node_id],
                     out_dag,
                     current_layout,
                     canonical_register,

releasenotes/notes/fix-sabre-performance-1b4503b6ecde6d13.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    Fixed a performance regression in the :class:`~.SabreLayout` and
+    :class:`~.SabreSwap` transpiler passes.
+    Fixed `#10650 <https://github.com/Qiskit/qiskit-terra/issues/10650>`__