Merge remote-tracking branch 'origin/main' into partial-layout-interface

Author: mtreinish

crates/accelerate/src/sabre_layout.rs

@@ -13,8 +13,7 @@
 
 use hashbrown::HashSet;
 use ndarray::prelude::*;
-use numpy::IntoPyArray;
-use numpy::PyReadonlyArray2;
+use numpy::{IntoPyArray, PyArray, PyReadonlyArray2};
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
 use pyo3::Python;
@@ -42,7 +41,7 @@ pub fn sabre_layout_and_routing(
     num_random_trials: usize,
     seed: Option<u64>,
     mut partial_layouts: Vec<Vec<Option<usize>>>,
-) -> ([NLayout; 2], (SwapMap, PyObject, NodeBlockResults)) {
+) -> (NLayout, PyObject, (SwapMap, PyObject, NodeBlockResults)) {
     let run_in_parallel = getenv_use_multiple_threads();
     let mut starting_layouts: Vec<Vec<Option<usize>>> =
         (0..num_random_trials).map(|_| vec![]).collect();
@@ -76,7 +75,7 @@ pub fn sabre_layout_and_routing(
                     ),
                 )
             })
-            .min_by_key(|(index, (_, result))| {
+            .min_by_key(|(index, (_, _, result))| {
                 (
                     result.map.map.values().map(|x| x.len()).sum::<usize>(),
                     *index,
@@ -101,15 +100,16 @@ pub fn sabre_layout_and_routing(
                     &starting_layouts[index],
                 )
             })
-            .min_by_key(|(_, result)| result.map.map.values().map(|x| x.len()).sum::<usize>())
+            .min_by_key(|(_, _, result)| result.map.map.values().map(|x| x.len()).sum::<usize>())
             .unwrap()
     };
     (
         res.0,
+        PyArray::from_vec(py, res.1).into(),
         (
-            res.1.map,
-            res.1.node_order.into_pyarray(py).into(),
-            res.1.node_block_results,
+            res.2.map,
+            res.2.node_order.into_pyarray(py).into(),
+            res.2.node_block_results,
         ),
     )
 }
@@ -124,133 +124,91 @@ fn layout_trial(
     num_swap_trials: usize,
     run_swap_in_parallel: bool,
     starting_layout: &[Option<usize>],
-) -> ([NLayout; 2], SabreResult) {
-    // Pick a random initial layout and fully populate ancillas in that layout too
+) -> (NLayout, Vec<usize>, SabreResult) {
     let num_physical_qubits = distance_matrix.shape()[0];
     let mut rng = Pcg64Mcg::seed_from_u64(seed);
-    let physical_qubits = if !starting_layout.is_empty() {
-        let used_bits: HashSet<usize> = starting_layout
-            .iter()
-            .filter_map(|x| x.as_ref())
-            .copied()
-            .collect();
-        let mut free_bits: Vec<usize> = (0..num_physical_qubits)
-            .filter(|x| !used_bits.contains(x))
-            .collect();
-        free_bits.shuffle(&mut rng);
-        (0..num_physical_qubits)
-            .map(|x| match starting_layout.get(x) {
-                Some(phys) => phys.unwrap_or_else(|| free_bits.pop().unwrap()),
-                None => free_bits.pop().unwrap(),
-            })
-            .collect()
-    } else {
-        let mut physical_qubits: Vec<usize> = (0..num_physical_qubits).collect();
-        physical_qubits.shuffle(&mut rng);
-        physical_qubits
+
+    // Pick a random initial layout including a full ancilla allocation.
+    let mut initial_layout = {
+        let physical_qubits = if !starting_layout.is_empty() {
+            let used_bits: HashSet<usize> = starting_layout
+                .iter()
+                .filter_map(|x| x.as_ref())
+                .copied()
+                .collect();
+            let mut free_bits: Vec<usize> = (0..num_physical_qubits)
+                .filter(|x| !used_bits.contains(x))
+                .collect();
+            free_bits.shuffle(&mut rng);
+            (0..num_physical_qubits)
+                .map(|x| match starting_layout.get(x) {
+                    Some(phys) => phys.unwrap_or_else(|| free_bits.pop().unwrap()),
+                    None => free_bits.pop().unwrap(),
+                })
+                .collect()
+        } else {
+            let mut physical_qubits: Vec<usize> = (0..num_physical_qubits).collect();
+            physical_qubits.shuffle(&mut rng);
+            physical_qubits
+        };
+        NLayout::from_logical_to_physical(physical_qubits)
     };
 
-    let mut initial_layout = NLayout::from_logical_to_physical(physical_qubits);
-    let new_dag_fn = |nodes| {
-        // Because the current implementation of Sabre swap doesn't permute
-        // the layout when placing control flow ops, there's no need to
-        // recurse into blocks. We remove them here, but still map control
-        // flow node IDs to an empty block list so Sabre treats these ops
-        // as control flow nodes, but doesn't route their blocks.
-        let node_blocks_empty = dag
+    // Sabre routing currently enforces that control-flow blocks return to their starting layout,
+    // which means they don't actually affect any heuristics that affect our layout choice.
+    let dag_no_control_forward = SabreDAG {
+        num_qubits: dag.num_qubits,
+        num_clbits: dag.num_clbits,
+        dag: dag.dag.clone(),
+        nodes: dag.nodes.clone(),
+        first_layer: dag.first_layer.clone(),
+        node_blocks: dag
             .node_blocks
-            .iter()
-            .map(|(node_index, _)| (*node_index, Vec::with_capacity(0)));
-        SabreDAG::new(
-            dag.num_qubits,
-            dag.num_clbits,
-            nodes,
-            node_blocks_empty.collect(),
-        )
-        .unwrap()
+            .keys()
+            .map(|index| (*index, Vec::new()))
+            .collect(),
     };
+    let dag_no_control_reverse = SabreDAG::new(
+        dag_no_control_forward.num_qubits,
+        dag_no_control_forward.num_clbits,
+        dag_no_control_forward.nodes.iter().rev().cloned().collect(),
+        dag_no_control_forward.node_blocks.clone(),
+    );
 
-    // Create forward and reverse dags (without node blocks).
-    // Once we've settled on a layout, we recursively apply it to the original
-    // DAG and its node blocks.
-    let mut dag_forward: SabreDAG = new_dag_fn(dag.nodes.clone());
-    let mut dag_reverse: SabreDAG = new_dag_fn(dag.nodes.iter().rev().cloned().collect());
     for _iter in 0..max_iterations {
-        // forward and reverse
-        for _direction in 0..2 {
-            let layout_dag = apply_layout(&dag_forward, &initial_layout);
-            let mut pass_final_layout = NLayout::generate_trivial_layout(num_physical_qubits);
-            build_swap_map_inner(
+        for dag in [&dag_no_control_forward, &dag_no_control_reverse] {
+            let (_result, final_layout) = build_swap_map_inner(
                 num_physical_qubits,
-                &layout_dag,
+                dag,
                 neighbor_table,
                 distance_matrix,
                 heuristic,
                 Some(seed),
-                &mut pass_final_layout,
+                &initial_layout,
                 num_swap_trials,
                 Some(run_swap_in_parallel),
             );
-            let final_layout = compose_layout(&initial_layout, &pass_final_layout);
             initial_layout = final_layout;
-            std::mem::swap(&mut dag_forward, &mut dag_reverse);
         }
     }
 
-    // Apply the layout to the original DAG.
-    let layout_dag = apply_layout(dag, &initial_layout);
-
-    let mut final_layout = NLayout::generate_trivial_layout(num_physical_qubits);
-    let sabre_result = build_swap_map_inner(
+    let (sabre_result, final_layout) = build_swap_map_inner(
         num_physical_qubits,
-        &layout_dag,
+        dag,
         neighbor_table,
         distance_matrix,
         heuristic,
         Some(seed),
-        &mut final_layout,
+        &initial_layout,
         num_swap_trials,
         Some(run_swap_in_parallel),
     );
-    ([initial_layout, final_layout], sabre_result)
-}
-
-fn apply_layout(dag: &SabreDAG, layout: &NLayout) -> SabreDAG {
-    let layout_nodes = dag.nodes.iter().map(|(node_index, qargs, cargs)| {
-        let new_qargs: Vec<usize> = qargs.iter().map(|n| layout.logic_to_phys[*n]).collect();
-        (*node_index, new_qargs, cargs.clone())
-    });
-    let node_blocks = dag.node_blocks.iter().map(|(node_index, blocks)| {
-        (
-            *node_index,
-            blocks.iter().map(|d| apply_layout(d, layout)).collect(),
-        )
-    });
-    SabreDAG::new(
-        dag.num_qubits,
-        dag.num_clbits,
-        layout_nodes.collect(),
-        node_blocks.collect(),
-    )
-    .unwrap()
-}
-
-fn compose_layout(initial_layout: &NLayout, final_layout: &NLayout) -> NLayout {
-    let logic_to_phys = initial_layout
-        .logic_to_phys
-        .iter()
-        .map(|n| final_layout.logic_to_phys[*n])
-        .collect();
-    let phys_to_logic = final_layout
+    let final_permutation = initial_layout
         .phys_to_logic
         .iter()
-        .map(|n| initial_layout.phys_to_logic[*n])
+        .map(|initial| final_layout.logic_to_phys[*initial])
         .collect();
-
-    NLayout {
-        logic_to_phys,
-        phys_to_logic,
-    }
+    (initial_layout, final_permutation, sabre_result)
 }
 
 #[pymodule]