Skip to content

Commit

Permalink
Refactor Sabre to an explicitly stateful router (#11977)
Browse files Browse the repository at this point in the history
This commit is an overdue tidy up of the Sabre code, which had been
through a few growth spurts with the addition of the combined
layout-and-routing pass in Rust, and the support for control-flow
operations and directives. I have a few things I'd like to try with
Sabre, and the code was getting quite unwieldy to modify and extend.

This refactors the Sabre routing internals, encapsulating a "routing
target" into a single view object that is used to define the hardware
target, and the stateful components of the routing algorithm into a
formal `RoutingState` object.  The functions that build up the routing
algorithm then become stateful instance methods, avoiding needing to
pass many things through several internal function calls.

In addition to the non-trivial lines-of-code savings, this also made it
clearer to me (while doing the refactor) that routing-state methods were
not all really at similar levels of abstraction, meaning that things
like the escape-valve mechanism took up oversized space in the
description of the main algorithm, and the control-flow block handling
was not as neatly split from the rest of the logic as it could have
been.  This reorganises some of the methods to make the important
components of the algorithms clearer; the top level of the algorithm now
fits on one screen.

Lastly, this moves both layout and routing into a unified `sabre`
module, mostly just to simplify all the `use` statements and to put
logically grouped code in the same place.
  • Loading branch information
jakelishman committed Apr 11, 2024
1 parent 4ff0fa2 commit 3af3cf5
Show file tree
Hide file tree
Showing 12 changed files with 763 additions and 852 deletions.
6 changes: 2 additions & 4 deletions crates/accelerate/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,7 @@ mod optimize_1q_gates;
mod pauli_exp_val;
mod quantum_circuit;
mod results;
mod sabre_layout;
mod sabre_swap;
mod sabre;
mod sampled_exp_val;
mod sparse_pauli_op;
mod stochastic_swap;
Expand All @@ -51,7 +50,7 @@ pub fn getenv_use_multiple_threads() -> bool {
fn _accelerate(m: &Bound<PyModule>) -> PyResult<()> {
m.add_wrapped(wrap_pymodule!(nlayout::nlayout))?;
m.add_wrapped(wrap_pymodule!(stochastic_swap::stochastic_swap))?;
m.add_wrapped(wrap_pymodule!(sabre_swap::sabre_swap))?;
m.add_wrapped(wrap_pymodule!(sabre::sabre))?;
m.add_wrapped(wrap_pymodule!(pauli_exp_val::pauli_expval))?;
m.add_wrapped(wrap_pymodule!(dense_layout::dense_layout))?;
m.add_wrapped(wrap_pymodule!(quantum_circuit::quantum_circuit))?;
Expand All @@ -60,7 +59,6 @@ fn _accelerate(m: &Bound<PyModule>) -> PyResult<()> {
m.add_wrapped(wrap_pymodule!(results::results))?;
m.add_wrapped(wrap_pymodule!(optimize_1q_gates::optimize_1q_gates))?;
m.add_wrapped(wrap_pymodule!(sampled_exp_val::sampled_exp_val))?;
m.add_wrapped(wrap_pymodule!(sabre_layout::sabre_layout))?;
m.add_wrapped(wrap_pymodule!(vf2_layout::vf2_layout))?;
m.add_wrapped(wrap_pymodule!(two_qubit_decompose::two_qubit_decompose))?;
m.add_wrapped(wrap_pymodule!(utils::utils))?;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
// copyright notice, and modified files need to carry a notice indicating
// that they have been altered from the originals.

use ahash;
use indexmap::IndexMap;
use ndarray::prelude::*;
use rustworkx_core::petgraph::prelude::*;
Expand All @@ -29,7 +28,7 @@ use crate::nlayout::PhysicalQubit;
/// extension, not for every swap trialled.
pub struct FrontLayer {
/// Map of the (index to the) node to the qubits it acts on.
nodes: IndexMap<NodeIndex, [PhysicalQubit; 2], ahash::RandomState>,
nodes: IndexMap<NodeIndex, [PhysicalQubit; 2], ::ahash::RandomState>,
/// Map of each qubit to the node that acts on it and the other qubit that node acts on, if this
/// qubit is active (otherwise `None`).
qubits: Vec<Option<(NodeIndex, PhysicalQubit)>>,
Expand All @@ -42,12 +41,18 @@ impl FrontLayer {
// pair, and can only have one gate in the layer.
nodes: IndexMap::with_capacity_and_hasher(
num_qubits as usize / 2,
ahash::RandomState::default(),
::ahash::RandomState::default(),
),
qubits: vec![None; num_qubits as usize],
}
}

/// View onto the mapping between qubits and their `(node, other_qubit)` pair. Index `i`
/// corresponds to physical qubit `i`.
pub fn qubits(&self) -> &[Option<(NodeIndex, PhysicalQubit)>] {
&self.qubits
}

/// Add a node into the front layer, with the two qubits it operates on.
pub fn insert(&mut self, index: NodeIndex, qubits: [PhysicalQubit; 2]) {
let [a, b] = qubits;
Expand Down Expand Up @@ -105,27 +110,6 @@ impl FrontLayer {
/ self.nodes.len() as f64
}

/// Populate a of nodes that would be routable if the given swap was applied to a layout. This
/// mutates `routable` to avoid heap allocations in the main logic loop.
pub fn routable_after(
&self,
routable: &mut Vec<NodeIndex>,
swap: &[PhysicalQubit; 2],
coupling: &DiGraph<(), ()>,
) {
let [a, b] = *swap;
if let Some((node, c)) = self.qubits[a.index()] {
if coupling.contains_edge(NodeIndex::new(b.index()), NodeIndex::new(c.index())) {
routable.push(node);
}
}
if let Some((node, c)) = self.qubits[b.index()] {
if coupling.contains_edge(NodeIndex::new(a.index()), NodeIndex::new(c.index())) {
routable.push(node);
}
}
}

/// Apply a physical swap to the current layout data structure.
pub fn apply_swap(&mut self, swap: [PhysicalQubit; 2]) {
let [a, b] = swap;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,22 +11,23 @@
// that they have been altered from the originals.
#![allow(clippy::too_many_arguments)]

use hashbrown::HashSet;
use ndarray::prelude::*;
use numpy::{IntoPyArray, PyArray, PyReadonlyArray2};
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use pyo3::Python;

use hashbrown::HashSet;
use numpy::{IntoPyArray, PyArray, PyReadonlyArray2};
use rand::prelude::*;
use rand_pcg::Pcg64Mcg;
use rayon::prelude::*;

use crate::getenv_use_multiple_threads;
use crate::nlayout::{NLayout, PhysicalQubit};
use crate::sabre_swap::neighbor_table::NeighborTable;
use crate::sabre_swap::sabre_dag::SabreDAG;
use crate::sabre_swap::swap_map::SwapMap;
use crate::sabre_swap::{build_swap_map_inner, Heuristic, NodeBlockResults, SabreResult};

use super::neighbor_table::NeighborTable;
use super::route::{swap_map, swap_map_trial, RoutingTargetView};
use super::sabre_dag::SabreDAG;
use super::swap_map::SwapMap;
use super::{Heuristic, NodeBlockResults, SabreResult};

#[pyfunction]
#[pyo3(signature = (dag, neighbor_table, distance_matrix, heuristic, max_iterations, num_swap_trials, num_random_trials, seed=None, partial_layouts=vec![]))]
Expand All @@ -35,14 +36,19 @@ pub fn sabre_layout_and_routing(
dag: &SabreDAG,
neighbor_table: &NeighborTable,
distance_matrix: PyReadonlyArray2<f64>,
heuristic: &Heuristic,
heuristic: Heuristic,
max_iterations: usize,
num_swap_trials: usize,
num_random_trials: usize,
seed: Option<u64>,
mut partial_layouts: Vec<Vec<Option<u32>>>,
) -> (NLayout, PyObject, (SwapMap, PyObject, NodeBlockResults)) {
let run_in_parallel = getenv_use_multiple_threads();
let target = RoutingTargetView {
neighbors: neighbor_table,
coupling: &neighbor_table.coupling_graph(),
distance: distance_matrix.as_array(),
};
let mut starting_layouts: Vec<Vec<Option<u32>>> =
(0..num_random_trials).map(|_| vec![]).collect();
starting_layouts.append(&mut partial_layouts);
Expand All @@ -54,7 +60,6 @@ pub fn sabre_layout_and_routing(
.sample_iter(&rand::distributions::Standard)
.take(starting_layouts.len())
.collect();
let dist = distance_matrix.as_array();
let res = if run_in_parallel && starting_layouts.len() > 1 {
seed_vec
.into_par_iter()
Expand All @@ -63,9 +68,8 @@ pub fn sabre_layout_and_routing(
(
index,
layout_trial(
&target,
dag,
neighbor_table,
&dist,
heuristic,
seed_trial,
max_iterations,
Expand All @@ -89,9 +93,8 @@ pub fn sabre_layout_and_routing(
.enumerate()
.map(|(index, seed_trial)| {
layout_trial(
&target,
dag,
neighbor_table,
&dist,
heuristic,
seed_trial,
max_iterations,
Expand All @@ -115,19 +118,21 @@ pub fn sabre_layout_and_routing(
}

fn layout_trial(
target: &RoutingTargetView,
dag: &SabreDAG,
neighbor_table: &NeighborTable,
distance_matrix: &ArrayView2<f64>,
heuristic: &Heuristic,
heuristic: Heuristic,
seed: u64,
max_iterations: usize,
num_swap_trials: usize,
run_swap_in_parallel: bool,
starting_layout: &[Option<u32>],
) -> (NLayout, Vec<PhysicalQubit>, SabreResult) {
let num_physical_qubits: u32 = distance_matrix.shape()[0].try_into().unwrap();
let num_physical_qubits: u32 = target.neighbors.num_qubits().try_into().unwrap();
let mut rng = Pcg64Mcg::seed_from_u64(seed);

// This is purely for RNG compatibility during a refactor.
let routing_seed = Pcg64Mcg::seed_from_u64(seed).next_u64();

// Pick a random initial layout including a full ancilla allocation.
let mut initial_layout = {
let physical_qubits: Vec<PhysicalQubit> = if !starting_layout.is_empty() {
Expand Down Expand Up @@ -182,29 +187,18 @@ fn layout_trial(

for _iter in 0..max_iterations {
for dag in [&dag_no_control_forward, &dag_no_control_reverse] {
let (_result, final_layout) = build_swap_map_inner(
num_physical_qubits,
dag,
neighbor_table,
distance_matrix,
heuristic,
Some(seed),
&initial_layout,
1,
Some(false),
);
let (_result, final_layout) =
swap_map_trial(target, dag, heuristic, &initial_layout, routing_seed);
initial_layout = final_layout;
}
}

let (sabre_result, final_layout) = build_swap_map_inner(
num_physical_qubits,
let (sabre_result, final_layout) = swap_map(
target,
dag,
neighbor_table,
distance_matrix,
heuristic,
Some(seed),
&initial_layout,
Some(seed),
num_swap_trials,
Some(run_swap_in_parallel),
);
Expand All @@ -214,9 +208,3 @@ fn layout_trial(
.collect();
(initial_layout, final_permutation, sabre_result)
}

#[pymodule]
pub fn sabre_layout(m: &Bound<PyModule>) -> PyResult<()> {
m.add_wrapped(wrap_pyfunction!(sabre_layout_and_routing))?;
Ok(())
}
128 changes: 128 additions & 0 deletions crates/accelerate/src/sabre/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
// This code is part of Qiskit.
//
// (C) Copyright IBM 2022
//
// This code is licensed under the Apache License, Version 2.0. You may
// obtain a copy of this license in the LICENSE.txt file in the root directory
// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
//
// Any modifications or derivative works of this code must retain this
// copyright notice, and modified files need to carry a notice indicating
// that they have been altered from the originals.

mod layer;
mod layout;
mod neighbor_table;
mod route;
mod sabre_dag;
mod swap_map;

use hashbrown::HashMap;
use numpy::{IntoPyArray, ToPyArray};
use pyo3::exceptions::PyIndexError;
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use pyo3::Python;

use crate::nlayout::PhysicalQubit;
use neighbor_table::NeighborTable;
use sabre_dag::SabreDAG;
use swap_map::SwapMap;

#[pyclass]
#[derive(Clone, Copy)]
pub enum Heuristic {
Basic,
Lookahead,
Decay,
}

/// A container for Sabre mapping results.
#[pyclass(module = "qiskit._accelerate.sabre")]
#[derive(Clone, Debug)]
pub struct SabreResult {
#[pyo3(get)]
pub map: SwapMap,
pub node_order: Vec<usize>,
#[pyo3(get)]
pub node_block_results: NodeBlockResults,
}

#[pymethods]
impl SabreResult {
#[getter]
fn node_order(&self, py: Python) -> PyObject {
self.node_order.to_pyarray_bound(py).into()
}
}

#[pyclass(mapping, module = "qiskit._accelerate.sabre")]
#[derive(Clone, Debug)]
pub struct NodeBlockResults {
pub results: HashMap<usize, Vec<BlockResult>>,
}

#[pymethods]
impl NodeBlockResults {
// Mapping Protocol
pub fn __len__(&self) -> usize {
self.results.len()
}

pub fn __contains__(&self, object: usize) -> bool {
self.results.contains_key(&object)
}

pub fn __getitem__(&self, py: Python, object: usize) -> PyResult<PyObject> {
match self.results.get(&object) {
Some(val) => Ok(val
.iter()
.map(|x| x.clone().into_py(py))
.collect::<Vec<_>>()
.into_pyarray_bound(py)
.into()),
None => Err(PyIndexError::new_err(format!(
"Node index {object} has no block results",
))),
}
}

pub fn __str__(&self) -> PyResult<String> {
Ok(format!("{:?}", self.results))
}
}

#[pyclass(module = "qiskit._accelerate.sabre")]
#[derive(Clone, Debug)]
pub struct BlockResult {
#[pyo3(get)]
pub result: SabreResult,
pub swap_epilogue: Vec<[PhysicalQubit; 2]>,
}

#[pymethods]
impl BlockResult {
#[getter]
fn swap_epilogue(&self, py: Python) -> PyObject {
self.swap_epilogue
.iter()
.map(|x| x.into_py(py))
.collect::<Vec<_>>()
.into_pyarray_bound(py)
.into()
}
}

#[pymodule]
pub fn sabre(m: &Bound<PyModule>) -> PyResult<()> {
m.add_wrapped(wrap_pyfunction!(route::sabre_routing))?;
m.add_wrapped(wrap_pyfunction!(layout::sabre_layout_and_routing))?;
m.add_class::<Heuristic>()?;
m.add_class::<NeighborTable>()?;
m.add_class::<SabreDAG>()?;
m.add_class::<SwapMap>()?;
m.add_class::<BlockResult>()?;
m.add_class::<NodeBlockResults>()?;
m.add_class::<SabreResult>()?;
Ok(())
}
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ use crate::nlayout::PhysicalQubit;
///
/// and used solely to represent neighbors of each node in qiskit-terra's rust
/// module.
#[pyclass(module = "qiskit._accelerate.sabre_swap")]
#[pyclass(module = "qiskit._accelerate.sabre")]
#[derive(Clone, Debug)]
pub struct NeighborTable {
// The choice of 4 `PhysicalQubit`s in the stack-allocated region is because a) this causes the
Expand All @@ -50,6 +50,10 @@ impl NeighborTable {
.map(move |v| (NodeIndex::new(u), NodeIndex::new(v.index())))
}))
}

pub fn num_qubits(&self) -> usize {
self.neighbors.len()
}
}

impl std::ops::Index<PhysicalQubit> for NeighborTable {
Expand Down
Loading

0 comments on commit 3af3cf5

Please sign in to comment.