scnc
/
qiskit
mirror of https://github.com/Qiskit/qiskit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor Sabre to an explicitly stateful router (#11977) 

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.
This commit is contained in:
Jake Lishman 2024-04-11 19:22:31 +01:00 committed by GitHub
parent 4ff0fa2f62
commit 3af3cf5880
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 763 additions and 852 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
crates/accelerate/src/lib.rs
View File

@ -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;
@ -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))?;
@ -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))?;

32
crates/accelerate/src/sabre_swap/layer.rs → crates/accelerate/src/sabre/layer.rs
View File

@ -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::*;
@ -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)>>,
@ -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;
@ -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;

68
crates/accelerate/src/sabre_layout.rs → crates/accelerate/src/sabre/layout.rs
View File

@ -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![]))]
@ -35,7 +36,7 @@ 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,
@ -43,6 +44,11 @@ pub fn sabre_layout_and_routing(
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);
@ -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()
@ -63,9 +68,8 @@ pub fn sabre_layout_and_routing(
(
index,
layout_trial(
&target,
dag,
neighbor_table,
&dist,
heuristic,
seed_trial,
max_iterations,
@ -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,
@ -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() {
@ -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),
);
@ -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
crates/accelerate/src/sabre/mod.rs Normal file
View File

@ -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(())
}

6
crates/accelerate/src/sabre_swap/neighbor_table.rs → crates/accelerate/src/sabre/neighbor_table.rs
View File

@ -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
@ -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 {

584
crates/accelerate/src/sabre/route.rs Normal file
View File

@ -0,0 +1,584 @@
// This code is part of Qiskit.
//
// (C) Copyright IBM 2024
//
// 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.
use std::cmp::Ordering;
use pyo3::prelude::*;
use pyo3::Python;
use hashbrown::HashMap;
use indexmap::IndexMap;
use ndarray::prelude::*;
use numpy::{IntoPyArray, PyArray, PyReadonlyArray2};
use rand::prelude::*;
use rand_pcg::Pcg64Mcg;
use rayon::prelude::*;
use rustworkx_core::dictmap::*;
use rustworkx_core::petgraph::prelude::*;
use rustworkx_core::petgraph::visit::EdgeRef;
use rustworkx_core::shortest_path::dijkstra;
use rustworkx_core::token_swapper::token_swapper;
use crate::getenv_use_multiple_threads;
use crate::nlayout::{NLayout, PhysicalQubit};
use super::layer::{ExtendedSet, FrontLayer};
use super::neighbor_table::NeighborTable;
use super::sabre_dag::SabreDAG;
use super::swap_map::SwapMap;
use super::{BlockResult, Heuristic, NodeBlockResults, SabreResult};
/// Epsilon used in minimum-score calculations.
const BEST_EPSILON: f64 = 1e-10;
/// Size of lookahead window.
const EXTENDED_SET_SIZE: usize = 20;
/// Decay coefficient for penalizing serial swaps.
const DECAY_RATE: f64 = 0.001;
/// How often to reset all decay rates to 1.
const DECAY_RESET_INTERVAL: u8 = 5;
/// Weight of lookahead window compared to front_layer.
const EXTENDED_SET_WEIGHT: f64 = 0.5;
/// Number of trials for control flow block swap epilogues.
const SWAP_EPILOGUE_TRIALS: usize = 4;
/// A view object onto a full routing target. This is cheap to clone and to replace components
/// within it; cloning only duplicates the inner references and not the data objects beneath. This
/// struct doesn't own its data because it's typically a view onto data generated from Python, and
/// we want to avoid copies.
pub struct RoutingTargetView<'a> {
pub neighbors: &'a NeighborTable,
pub coupling: &'a DiGraph<(), ()>,
pub distance: ArrayView2<'a, f64>,
}
/// Long-term internal state of the Sabre routing algorithm. This includes all the scratch space
/// and tracking that we use over the course of many swap insertions, but doesn't include ephemeral
/// state that never needs to leave the main loop. This is mostly just a convenience, so we don't
/// have to pass everything from function to function.
struct RoutingState<'a, 'b> {
target: &'a RoutingTargetView<'b>,
dag: &'a SabreDAG,
heuristic: Heuristic,
/// Mapping of instructions (node indices) to swaps that precede them.
out_map: HashMap<usize, Vec<[PhysicalQubit; 2]>>,
/// Order of the instructions (node indices) in the problem DAG in the output.
gate_order: Vec<usize>,
/// Map from node index of a control-flow op to its inner result.
node_block_results: HashMap<usize, Vec<BlockResult>>,
front_layer: FrontLayer,
extended_set: ExtendedSet,
/// How many predecessors still need to be satisfied for each node index before it is at the
/// front of the topological iteration through the nodes as they're routed.
required_predecessors: &'a mut [u32],
layout: NLayout,
/// Tracking for the 'decay' heuristic on each qubit.
qubits_decay: &'a mut [f64],
/// Reusable allocated storage space for choosing the best swap. This is owned outside of the
/// `choose_best_swap` function so that we don't need to reallocate and then re-grow the
/// collection on every entry.
swap_scratch: Vec<[PhysicalQubit; 2]>,
rng: Pcg64Mcg,
seed: u64,
}
impl<'a, 'b> RoutingState<'a, 'b> {
/// Apply a swap to the program-state structures (front layer, extended set and current
/// layout).
#[inline]
fn apply_swap(&mut self, swap: [PhysicalQubit; 2]) {
self.front_layer.apply_swap(swap);
self.extended_set.apply_swap(swap);
self.layout.swap_physical(swap[0], swap[1]);
}
/// Return the node, if any, that is on this qubit and is routable with the current layout.
#[inline]
fn routable_node_on_qubit(&self, qubit: PhysicalQubit) -> Option<NodeIndex> {
self.front_layer.qubits()[qubit.index()].and_then(|(node, other)| {
self.target
.coupling
.contains_edge(NodeIndex::new(qubit.index()), NodeIndex::new(other.index()))
.then_some(node)
})
}
/// Update the system state as the given `nodes` are added to the routing order, preceded by
/// the given `swaps`. This involves updating the output values `gate_order` and `out_map`,
/// but also the tracking objects `front_layer`, `extended_set` and `required_predecessors` by
/// removing the routed nodes and adding any now-reachable ones.
fn update_route(&mut self, nodes: &[NodeIndex], swaps: Vec<[PhysicalQubit; 2]>) {
// First node gets the swaps attached. We don't add to the `gate_order` here because
// `route_reachable_nodes` is responsible for that part.
self.out_map
.insert(self.dag.dag[nodes[0]].py_node_id, swaps);
for node in nodes {
self.front_layer.remove(node);
}
self.route_reachable_nodes(nodes);
// Ideally we'd know how to mutate the extended set directly, but since its limited size
// ties its construction strongly to the iteration order through the front layer, it's not
// easy to do better than just emptying it and rebuilding.
self.extended_set.clear();
self.populate_extended_set();
}
/// Search forwards in the DAG from all the nodes in `nodes`, adding them to the `gate_order`
/// or the current `front_layer` as appropriate, and continue inspecting gates until there is
/// nothing further with no required predecessors.
///
/// The nodes in `nodes` should all already have no further required predecessors.
fn route_reachable_nodes(&mut self, nodes: &[NodeIndex]) {
let mut to_visit = nodes.to_vec();
let mut i = 0;
let dag = &self.dag;
// Iterate through `to_visit`, except we often push new nodes onto the end of it.
while i < to_visit.len() {
let node_id = to_visit[i];
let node = &dag.dag[node_id];
i += 1;
// If the node is a directive that means it can be placed anywhere.
if !node.directive {
if let Some(blocks) = dag.node_blocks.get(&node.py_node_id) {
let block_results = blocks
.iter()
.map(|block| self.route_control_flow_block(block))
.collect::<Vec<_>>();
self.node_block_results
.insert(node.py_node_id, block_results);
} else {
match node.qubits[..] {
// A gate op whose connectivity must match the device to be placed in the
// gate order.
[a, b]
if !self.target.coupling.contains_edge(
NodeIndex::new(a.to_phys(&self.layout).index()),
NodeIndex::new(b.to_phys(&self.layout).index()),
) =>
{
// 2Q op that cannot be placed. Add it to the front layer and move on.
self.front_layer.insert(
node_id,
[a.to_phys(&self.layout), b.to_phys(&self.layout)],
);
continue;
}
_ => {}
}
}
}
// If we reach here, the node is routable.
self.gate_order.push(node.py_node_id);
for edge in dag.dag.edges_directed(node_id, Direction::Outgoing) {
let successor_node = edge.target();
let successor_index = successor_node.index();
self.required_predecessors[successor_index] -= 1;
if self.required_predecessors[successor_index] == 0 {
to_visit.push(successor_node);
}
}
}
}
/// Inner worker to route a control-flow block. Since control-flow blocks are routed to
/// restore the layout at the end of themselves, and the recursive calls spawn their own
/// tracking states, this does not affect our own state.
fn route_control_flow_block(&self, block: &SabreDAG) -> BlockResult {
let (result, mut block_final_layout) =
swap_map_trial(self.target, block, self.heuristic, &self.layout, self.seed);
// For now, we always append a swap circuit that gets the inner block back to the
// parent's layout.
let swap_epilogue = {
// Map physical location in the final layout from the inner routing to the current
// location in the outer routing.
let mapping: HashMap<NodeIndex, NodeIndex> = block_final_layout
.iter_physical()
.map(|(p, v)| {
(
NodeIndex::new(p.index()),
NodeIndex::new(v.to_phys(&self.layout).index()),
)
})
.collect();
let swaps = token_swapper(
&self.target.coupling,
mapping,
Some(SWAP_EPILOGUE_TRIALS),
Some(self.seed),
None,
)
.unwrap();
// Convert physical swaps to virtual swaps
swaps
.into_iter()
.map(|(l, r)| {
let p_l = PhysicalQubit::new(l.index().try_into().unwrap());
let p_r = PhysicalQubit::new(r.index().try_into().unwrap());
block_final_layout.swap_physical(p_l, p_r);
[p_l, p_r]
})
.collect()
};
BlockResult {
result,
swap_epilogue,
}
}
/// Fill the given `extended_set` with the next nodes that would be reachable after the front
/// layer (and themselves). This uses `required_predecessors` as scratch space for efficiency,
/// but returns it to the same state as the input on return.
fn populate_extended_set(&mut self) {
let mut to_visit = self.front_layer.iter_nodes().copied().collect::<Vec<_>>();
let mut decremented: IndexMap<usize, u32, ahash::RandomState> =
IndexMap::with_hasher(ahash::RandomState::default());
let mut i = 0;
let mut visit_now: Vec<NodeIndex> = Vec::new();
let dag = &self.dag;
while i < to_visit.len() && self.extended_set.len() < EXTENDED_SET_SIZE {
// Visit runs of non-2Q gates fully before moving on to children of 2Q gates. This way,
// traversal order is a BFS of 2Q gates rather than of all gates.
visit_now.push(to_visit[i]);
let mut j = 0;
while let Some(node) = visit_now.get(j) {
for edge in dag.dag.edges_directed(*node, Direction::Outgoing) {
let successor_node = edge.target();
let successor_index = successor_node.index();
*decremented.entry(successor_index).or_insert(0) += 1;
self.required_predecessors[successor_index] -= 1;
if self.required_predecessors[successor_index] == 0 {
if !dag.dag[successor_node].directive
&& !dag.node_blocks.contains_key(&successor_index)
{
if let [a, b] = dag.dag[successor_node].qubits[..] {
self.extended_set
.push([a.to_phys(&self.layout), b.to_phys(&self.layout)]);
to_visit.push(successor_node);
continue;
}
}
visit_now.push(successor_node);
}
}
j += 1;
}
visit_now.clear();
i += 1;
}
for (node, amount) in decremented.iter() {
self.required_predecessors[*node] += *amount;
}
}
/// Add swaps to the current set that greedily bring the nearest node together. This is a
/// "release valve" mechanism; it ignores all the Sabre heuristics and forces progress, so we
/// can't get permanently stuck.
fn force_enable_closest_node(
&mut self,
current_swaps: &mut Vec<[PhysicalQubit; 2]>,
) -> NodeIndex {
let (&closest_node, &qubits) = {
let dist = &self.target.distance;
self.front_layer
.iter()
.min_by(|(_, qubits_a), (_, qubits_b)| {
dist[[qubits_a[0].index(), qubits_a[1].index()]]
.partial_cmp(&dist[[qubits_b[0].index(), qubits_b[1].index()]])
.unwrap_or(Ordering::Equal)
})
.unwrap()
};
let shortest_path = {
let mut shortest_paths: DictMap<NodeIndex, Vec<NodeIndex>> = DictMap::new();
(dijkstra(
&self.target.coupling,
NodeIndex::new(qubits[0].index()),
Some(NodeIndex::new(qubits[1].index())),
|_| Ok(1.),
Some(&mut shortest_paths),
) as PyResult<Vec<Option<f64>>>)
.unwrap();
shortest_paths
.get(&NodeIndex::new(qubits[1].index()))
.unwrap()
.iter()
.map(|n| PhysicalQubit::new(n.index() as u32))
.collect::<Vec<_>>()
};
// Insert greedy swaps along that shortest path, splitting them between moving the left side
// and moving the right side to minimise the depth. One side needs to move up to the split
// point and the other can stop one short because the gate will be routable then.
let split: usize = shortest_path.len() / 2;
current_swaps.reserve(shortest_path.len() - 2);
for i in 0..split {
current_swaps.push([shortest_path[i], shortest_path[i + 1]]);
}
for i in 0..split - 1 {
let end = shortest_path.len() - 1 - i;
current_swaps.push([shortest_path[end], shortest_path[end - 1]]);
}
current_swaps.iter().for_each(|&swap| self.apply_swap(swap));
closest_node
}
/// Return the swap of two virtual qubits that produces the best score of all possible swaps.
fn choose_best_swap(&mut self) -> [PhysicalQubit; 2] {
self.swap_scratch.clear();
let mut min_score = f64::MAX;
// The decay heuristic is the only one that actually needs the absolute score.
let dist = &self.target.distance;
let absolute_score = match self.heuristic {
Heuristic::Decay => {
self.front_layer.total_score(dist)
+ EXTENDED_SET_WEIGHT * self.extended_set.total_score(dist)
}
_ => 0.0,
};
for swap in obtain_swaps(&self.front_layer, self.target.neighbors) {
let score = match self.heuristic {
Heuristic::Basic => self.front_layer.score(swap, dist),
Heuristic::Lookahead => {
self.front_layer.score(swap, dist)
+ EXTENDED_SET_WEIGHT * self.extended_set.score(swap, dist)
}
Heuristic::Decay => {
self.qubits_decay[swap[0].index()].max(self.qubits_decay[swap[1].index()])
* (absolute_score
+ self.front_layer.score(swap, dist)
+ EXTENDED_SET_WEIGHT * self.extended_set.score(swap, dist))
}
};
if score < min_score - BEST_EPSILON {
min_score = score;
self.swap_scratch.clear();
self.swap_scratch.push(swap);
} else if (score - min_score).abs() < BEST_EPSILON {
self.swap_scratch.push(swap);
}
}
*self.swap_scratch.choose(&mut self.rng).unwrap()
}
}
/// Return a set of candidate swaps that affect qubits in front_layer.
///
/// For each virtual qubit in `front_layer`, find its current location on hardware and the physical
/// qubits in that neighborhood. Every swap on virtual qubits that corresponds to one of those
/// physical couplings is a candidate swap.
fn obtain_swaps<'a>(
front_layer: &'a FrontLayer,
neighbors: &'a NeighborTable,
) -> impl Iterator<Item = [PhysicalQubit; 2]> + 'a {
front_layer.iter_active().flat_map(move |&p| {
neighbors[p].iter().filter_map(move |&neighbor| {
if neighbor > p || !front_layer.is_active(neighbor) {
Some([p, neighbor])
} else {
None
}
})
})
}
/// Run sabre swap on a circuit
///
/// Returns:
/// (SwapMap, gate_order, node_block_results, final_permutation): 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. The third is inner results for the blocks of
/// control flow, and the fourth is a permutation, where `final_permution[i]` is the final
/// logical position of the qubit that began in position `i`.
#[pyfunction]
#[allow(clippy::too_many_arguments)]
pub fn sabre_routing(
py: Python,
dag: &SabreDAG,
neighbor_table: &NeighborTable,
distance_matrix: PyReadonlyArray2<f64>,
heuristic: Heuristic,
initial_layout: &NLayout,
num_trials: usize,
seed: Option<u64>,
run_in_parallel: Option<bool>,
) -> (SwapMap, PyObject, NodeBlockResults, PyObject) {
let target = RoutingTargetView {
neighbors: neighbor_table,
coupling: &neighbor_table.coupling_graph(),
distance: distance_matrix.as_array(),
};
let (res, final_layout) = swap_map(
&target,
dag,
heuristic,
initial_layout,
seed,
num_trials,
run_in_parallel,
);
(
res.map,
res.node_order.into_pyarray_bound(py).into(),
res.node_block_results,
PyArray::from_iter_bound(
py,
(0u32..neighbor_table.num_qubits().try_into().unwrap()).map(|phys| {
PhysicalQubit::new(phys)
.to_virt(initial_layout)
.to_phys(&final_layout)
}),
)
.into(),
)
}
/// Run (potentially in parallel) several trials of the Sabre routing algorithm on the given
/// problem and return the one with fewest swaps.
pub fn swap_map(
target: &RoutingTargetView,
dag: &SabreDAG,
heuristic: Heuristic,
initial_layout: &NLayout,
seed: Option<u64>,
num_trials: usize,
run_in_parallel: Option<bool>,
) -> (SabreResult, NLayout) {
let run_in_parallel = match run_in_parallel {
Some(run_in_parallel) => run_in_parallel,
None => getenv_use_multiple_threads() && num_trials > 1,
};
let outer_rng = match seed {
Some(seed) => Pcg64Mcg::seed_from_u64(seed),
None => Pcg64Mcg::from_entropy(),
};
let seed_vec: Vec<u64> = outer_rng
.sample_iter(&rand::distributions::Standard)
.take(num_trials)
.collect();
if run_in_parallel {
seed_vec
.into_par_iter()
.enumerate()
.map(|(index, seed_trial)| {
(
index,
swap_map_trial(target, dag, heuristic, initial_layout, seed_trial),
)
})
.min_by_key(|(index, (result, _))| {
[
result.map.map.values().map(|x| x.len()).sum::<usize>(),
*index,
]
})
.unwrap()
.1
} else {
seed_vec
.into_iter()
.map(|seed_trial| swap_map_trial(target, dag, heuristic, initial_layout, seed_trial))
.min_by_key(|(result, _)| result.map.map.values().map(|x| x.len()).sum::<usize>())
.unwrap()
}
}
/// Run a single trial of the Sabre routing algorithm.
pub fn swap_map_trial(
target: &RoutingTargetView,
dag: &SabreDAG,
heuristic: Heuristic,
initial_layout: &NLayout,
seed: u64,
) -> (SabreResult, NLayout) {
let num_qubits: u32 = target.neighbors.num_qubits().try_into().unwrap();
let mut state = RoutingState {
target,
dag,
heuristic,
out_map: HashMap::new(),
gate_order: Vec::with_capacity(dag.dag.node_count()),
node_block_results: HashMap::with_capacity(dag.node_blocks.len()),
front_layer: FrontLayer::new(num_qubits),
extended_set: ExtendedSet::new(num_qubits),
required_predecessors: &mut vec![0; dag.dag.node_count()],
layout: initial_layout.clone(),
qubits_decay: &mut vec![1.; num_qubits as usize],
swap_scratch: Vec::new(),
rng: Pcg64Mcg::seed_from_u64(seed),
seed,
};
for node in dag.dag.node_indices() {
for edge in dag.dag.edges(node) {
state.required_predecessors[edge.target().index()] += 1;
}
}
state.route_reachable_nodes(&dag.first_layer);
state.populate_extended_set();
// Main logic loop; the front layer only becomes empty when all nodes have been routed. At
// each iteration of this loop, we route either one or two gates.
let max_iterations_without_progress = 10 * num_qubits as usize;
let mut num_search_steps: u8 = 0;
let mut routable_nodes = Vec::<NodeIndex>::with_capacity(2);
while !state.front_layer.is_empty() {
let mut current_swaps: Vec<[PhysicalQubit; 2]> = Vec::new();
// Swap-mapping loop. This is the main part of the algorithm, which we repeat until we
// either successfully route a node, or exceed the maximum number of attempts.
while routable_nodes.is_empty() && current_swaps.len() <= max_iterations_without_progress {
let best_swap = state.choose_best_swap();
state.apply_swap(best_swap);
current_swaps.push(best_swap);
if let Some(node) = state.routable_node_on_qubit(best_swap[0]) {
routable_nodes.push(node);
}
if let Some(node) = state.routable_node_on_qubit(best_swap[1]) {
routable_nodes.push(node);
}
num_search_steps += 1;
if num_search_steps >= DECAY_RESET_INTERVAL {
state.qubits_decay.fill(1.);
num_search_steps = 0;
} else {
state.qubits_decay[best_swap[0].index()] += DECAY_RATE;
state.qubits_decay[best_swap[1].index()] += DECAY_RATE;
}
}
if routable_nodes.is_empty() {
// If we exceeded the max number of heuristic-chosen swaps without making progress,
// unwind to the last progress point and greedily swap to bring a ndoe together.
// Efficiency doesn't matter much; this path never gets taken unless we're unlucky.
current_swaps
.drain(..)
.rev()
.for_each(|swap| state.apply_swap(swap));
let force_routed = state.force_enable_closest_node(&mut current_swaps);
routable_nodes.push(force_routed);
}
state.update_route(&routable_nodes, current_swaps);
state.qubits_decay.fill(1.);
routable_nodes.clear();
}
(
SabreResult {
map: SwapMap { map: state.out_map },
node_order: state.gate_order,
node_block_results: NodeBlockResults {
results: state.node_block_results,
},
},
state.layout,
)
}

2
crates/accelerate/src/sabre_swap/sabre_dag.rs → crates/accelerate/src/sabre/sabre_dag.rs
View File

@ -35,7 +35,7 @@ pub struct DAGNode {
/// Control-flow operations are represented by the presence of the Python [DAGCircuit]'s node id
/// (the [DAGNode.py_node_id] field) as a key in [node_blocks], where the value is an array of the
/// inner dataflow graphs.
#[pyclass(module = "qiskit._accelerate.sabre_swap")]
#[pyclass(module = "qiskit._accelerate.sabre")]
#[derive(Clone, Debug)]
pub struct SabreDAG {
pub num_qubits: usize,

2
crates/accelerate/src/sabre_swap/swap_map.rs → crates/accelerate/src/sabre/swap_map.rs
View File

@ -17,7 +17,7 @@ use pyo3::prelude::*;
use crate::nlayout::PhysicalQubit;
/// A container for required swaps before a gate qubit
#[pyclass(module = "qiskit._accelerate.sabre_swap")]
#[pyclass(module = "qiskit._accelerate.sabre")]
#[derive(Clone, Debug)]
pub struct SwapMap {
pub map: HashMap<usize, Vec<[PhysicalQubit; 2]>>,

773
crates/accelerate/src/sabre_swap/mod.rs
View File

@ -1,773 +0,0 @@
// 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.
#![allow(clippy::too_many_arguments)]
pub mod layer;
pub mod neighbor_table;
pub mod sabre_dag;
pub mod swap_map;
use hashbrown::HashMap;
use indexmap::IndexMap;
use ndarray::prelude::*;
use numpy::PyReadonlyArray2;
use numpy::{IntoPyArray, PyArray, ToPyArray};
use pyo3::exceptions::PyIndexError;
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use pyo3::Python;
use rand::prelude::SliceRandom;
use rand::prelude::*;
use rand_pcg::Pcg64Mcg;
use rayon::prelude::*;
use rustworkx_core::dictmap::*;
use rustworkx_core::petgraph::prelude::*;
use rustworkx_core::petgraph::visit::EdgeRef;
use rustworkx_core::shortest_path::dijkstra;
use rustworkx_core::token_swapper::token_swapper;
use std::cmp::Ordering;
use crate::getenv_use_multiple_threads;
use crate::nlayout::{NLayout, PhysicalQubit};
use layer::{ExtendedSet, FrontLayer};
use neighbor_table::NeighborTable;
use sabre_dag::SabreDAG;
use swap_map::SwapMap;
const BEST_EPSILON: f64 = 1e-10; // Epsilon used in minimum-score calculations.
const EXTENDED_SET_SIZE: usize = 20; // Size of lookahead window.
const DECAY_RATE: f64 = 0.001; // Decay coefficient for penalizing serial swaps.
const DECAY_RESET_INTERVAL: u8 = 5; // How often to reset all decay rates to 1.
const EXTENDED_SET_WEIGHT: f64 = 0.5; // Weight of lookahead window compared to front_layer.
const SWAP_EPILOGUE_TRIALS: usize = 4; // Number of trials for control flow block swap epilogues.
#[pyclass]
pub enum Heuristic {
Basic,
Lookahead,
Decay,
}
/// A container for Sabre mapping results.
#[pyclass(module = "qiskit._accelerate.sabre_swap")]
#[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_swap")]
#[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_swap")]
#[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()
}
}
/// Return a set of candidate swaps that affect qubits in front_layer.
///
/// For each virtual qubit in front_layer, find its current location
/// on hardware and the physical qubits in that neighborhood. Every SWAP
/// on virtual qubits that corresponds to one of those physical couplings
/// is a candidate SWAP.
fn obtain_swaps<'a>(
front_layer: &'a FrontLayer,
neighbors: &'a NeighborTable,
) -> impl Iterator<Item = [PhysicalQubit; 2]> + 'a {
front_layer.iter_active().flat_map(move |&p| {
neighbors[p].iter().filter_map(move |&neighbor| {
if neighbor > p || !front_layer.is_active(neighbor) {
Some([p, neighbor])
} else {
None
}
})
})
}
/// Fill the given `extended_set` with the next nodes that would be reachable after the front layer
/// (and themselves). This uses `required_predecessors` as scratch space for efficiency, but
/// returns it to the same state as the input on return.
fn populate_extended_set(
extended_set: &mut ExtendedSet,
dag: &SabreDAG,
front_layer: &FrontLayer,
layout: &NLayout,
required_predecessors: &mut [u32],
) {
let mut to_visit = front_layer.iter_nodes().copied().collect::<Vec<_>>();
let mut decremented: IndexMap<usize, u32, ahash::RandomState> =
IndexMap::with_hasher(ahash::RandomState::default());
let mut i = 0;
let mut visit_now: Vec<NodeIndex> = Vec::new();
while i < to_visit.len() && extended_set.len() < EXTENDED_SET_SIZE {
// Visit runs of non-2Q gates fully before moving on to children
// of 2Q gates. This way, traversal order is a BFS of 2Q gates rather
// than of all gates.
visit_now.push(to_visit[i]);
let mut j = 0;
while let Some(node) = visit_now.get(j) {
for edge in dag.dag.edges_directed(*node, Direction::Outgoing) {
let successor_node = edge.target();
let successor_index = successor_node.index();
*decremented.entry(successor_index).or_insert(0) += 1;
required_predecessors[successor_index] -= 1;
if required_predecessors[successor_index] == 0 {
if !dag.dag[successor_node].directive
&& !dag.node_blocks.contains_key(&successor_index)
{
if let [a, b] = dag.dag[successor_node].qubits[..] {
extended_set.push([a.to_phys(layout), b.to_phys(layout)]);
to_visit.push(successor_node);
continue;
}
}
visit_now.push(successor_node);
}
}
j += 1;
}
visit_now.clear();
i += 1;
}
for (node, amount) in decremented.iter() {
required_predecessors[*node] += *amount;
}
}
/// Run sabre swap on a circuit
///
/// Returns:
/// (SwapMap, gate_order, node_block_results, final_permutation): 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. The third is inner results for the blocks of
/// control flow, and the fourth is a permutation, where `final_permution[i]` is the final
/// logical position of the qubit that began in position `i`.
#[pyfunction]
pub fn build_swap_map(
py: Python,
num_qubits: u32,
dag: &SabreDAG,
neighbor_table: &NeighborTable,
distance_matrix: PyReadonlyArray2<f64>,
heuristic: &Heuristic,
initial_layout: &NLayout,
num_trials: usize,
seed: Option<u64>,
run_in_parallel: Option<bool>,
) -> (SwapMap, PyObject, NodeBlockResults, PyObject) {
let dist = distance_matrix.as_array();
let (res, final_layout) = build_swap_map_inner(
num_qubits,
dag,
neighbor_table,
&dist,
heuristic,
seed,
initial_layout,
num_trials,
run_in_parallel,
);
(
res.map,
res.node_order.into_pyarray_bound(py).into(),
res.node_block_results,
PyArray::from_iter_bound(
py,
(0..num_qubits).map(|phys| {
PhysicalQubit::new(phys)
.to_virt(initial_layout)
.to_phys(&final_layout)
}),
)
.into(),
)
}
pub fn build_swap_map_inner(
num_qubits: u32,
dag: &SabreDAG,
neighbor_table: &NeighborTable,
dist: &ArrayView2<f64>,
heuristic: &Heuristic,
seed: Option<u64>,
initial_layout: &NLayout,
num_trials: usize,
run_in_parallel: Option<bool>,
) -> (SabreResult, NLayout) {
let run_in_parallel = match run_in_parallel {
Some(run_in_parallel) => run_in_parallel,
None => getenv_use_multiple_threads() && num_trials > 1,
};
let coupling_graph = neighbor_table.coupling_graph();
let outer_rng = match seed {
Some(seed) => Pcg64Mcg::seed_from_u64(seed),
None => Pcg64Mcg::from_entropy(),
};
let seed_vec: Vec<u64> = outer_rng
.sample_iter(&rand::distributions::Standard)
.take(num_trials)
.collect();
if run_in_parallel {
seed_vec
.into_par_iter()
.enumerate()
.map(|(index, seed_trial)| {
(
index,
swap_map_trial(
num_qubits,
dag,
neighbor_table,
dist,
&coupling_graph,
heuristic,
seed_trial,
initial_layout,
),
)
})
.min_by_key(|(index, (result, _))| {
[
result.map.map.values().map(|x| x.len()).sum::<usize>(),
*index,
]
})
.unwrap()
.1
} else {
seed_vec
.into_iter()
.map(|seed_trial| {
swap_map_trial(
num_qubits,
dag,
neighbor_table,
dist,
&coupling_graph,
heuristic,
seed_trial,
initial_layout,
)
})
.min_by_key(|(result, _)| result.map.map.values().map(|x| x.len()).sum::<usize>())
.unwrap()
}
}
fn swap_map_trial(
num_qubits: u32,
dag: &SabreDAG,
neighbor_table: &NeighborTable,
dist: &ArrayView2<f64>,
coupling_graph: &DiGraph<(), ()>,
heuristic: &Heuristic,
seed: u64,
initial_layout: &NLayout,
) -> (SabreResult, NLayout) {
let max_iterations_without_progress = 10 * num_qubits as usize;
let mut out_map: HashMap<usize, Vec<[PhysicalQubit; 2]>> = HashMap::new();
let mut gate_order = Vec::with_capacity(dag.dag.node_count());
let mut front_layer = FrontLayer::new(num_qubits);
let mut extended_set = ExtendedSet::new(num_qubits);
let mut required_predecessors: Vec<u32> = vec![0; dag.dag.node_count()];
let mut layout = initial_layout.clone();
let mut num_search_steps: u8 = 0;
let mut qubits_decay: Vec<f64> = vec![1.; num_qubits as usize];
let mut rng = Pcg64Mcg::seed_from_u64(seed);
let mut node_block_results: HashMap<usize, Vec<BlockResult>> =
HashMap::with_capacity(dag.node_blocks.len());
for node in dag.dag.node_indices() {
for edge in dag.dag.edges(node) {
required_predecessors[edge.target().index()] += 1;
}
}
// This closure is used to curry parameters so we can avoid passing
// everything and the kitchen sink to update_routes and
// route_reachable_nodes.
let route_block_dag = |block_dag: &SabreDAG, current_layout: &NLayout| {
swap_map_trial(
num_qubits,
block_dag,
neighbor_table,
dist,
coupling_graph,
heuristic,
seed,
current_layout,
)
};
route_reachable_nodes(
seed,
&dag.first_layer,
dag,
&layout,
coupling_graph,
&mut gate_order,
&mut front_layer,
&mut required_predecessors,
&mut node_block_results,
&route_block_dag,
);
populate_extended_set(
&mut extended_set,
dag,
&front_layer,
&layout,
&mut required_predecessors,
);
// Main logic loop; the front layer only becomes empty when all nodes have been routed. At
// each iteration of this loop, we route either one or two gates.
let mut routable_nodes = Vec::<NodeIndex>::with_capacity(2);
// Reusable allocated storage space for choosing the best swap. This is owned outside of the
// `choose_best_swap` function so that we don't need to reallocate and then re-grow the
// collection on every entry.
let mut swap_scratch = Vec::<[PhysicalQubit; 2]>::new();
while !front_layer.is_empty() {
let mut current_swaps: Vec<[PhysicalQubit; 2]> = Vec::new();
// Swap-mapping loop. This is the main part of the algorithm, which we repeat until we
// either successfully route a node, or exceed the maximum number of attempts.
while routable_nodes.is_empty() && current_swaps.len() <= max_iterations_without_progress {
let best_swap = choose_best_swap(
&front_layer,
&extended_set,
neighbor_table,
dist,
&qubits_decay,
heuristic,
&mut rng,
&mut swap_scratch,
);
front_layer.routable_after(&mut routable_nodes, &best_swap, coupling_graph);
front_layer.apply_swap(best_swap);
extended_set.apply_swap(best_swap);
layout.swap_physical(best_swap[0], best_swap[1]);
current_swaps.push(best_swap);
num_search_steps += 1;
if num_search_steps >= DECAY_RESET_INTERVAL {
qubits_decay.fill(1.);
num_search_steps = 0;
} else {
qubits_decay[best_swap[0].index()] += DECAY_RATE;
qubits_decay[best_swap[1].index()] += DECAY_RATE;
}
}
// If we exceeded the number of allowed attempts without successfully routing a node, we
// reset back to the state we were in last time we routed a node, then find the node in the
// front layer whose qubits are the closest in the coupling map, and greedily insert swaps
// to make the node routable. We could clone the layout each time we route a gate, but
// this path is only an escape mechansim for the algorithm getting stuck, so it should
// ideally never be taken, and it doesn't matter if it's not the speediest---it's better to
// keep the other path faster.
if routable_nodes.is_empty() {
undo_swaps(
&mut current_swaps,
&mut front_layer,
&mut extended_set,
&mut layout,
);
let (&node, &qubits) = closest_operation(&front_layer, dist);
swaps_to_route(&mut current_swaps, &qubits, coupling_graph);
for &[a, b] in current_swaps.iter() {
front_layer.apply_swap([a, b]);
extended_set.apply_swap([a, b]);
layout.swap_physical(a, b);
}
routable_nodes.push(node);
}
update_route(
seed,
&routable_nodes,
current_swaps,
dag,
&layout,
coupling_graph,
&mut gate_order,
&mut out_map,
&mut front_layer,
&mut extended_set,
&mut required_predecessors,
&mut node_block_results,
&route_block_dag,
);
qubits_decay.fill(1.);
routable_nodes.clear();
}
(
SabreResult {
map: SwapMap { map: out_map },
node_order: gate_order,
node_block_results: NodeBlockResults {
results: node_block_results,
},
},
layout,
)
}
/// Update the system state as the given `nodes` are added to the routing order, preceded by the
/// given `swaps`. This involves updating the output values `gate_order` and `out_map`, but also
/// the tracking objects `front_layer`, `extended_set` and `required_predecessors` by removing the
/// routed nodes and adding any now-reachable ones.
fn update_route<F>(
seed: u64,
nodes: &[NodeIndex],
swaps: Vec<[PhysicalQubit; 2]>,
dag: &SabreDAG,
layout: &NLayout,
coupling: &DiGraph<(), ()>,
gate_order: &mut Vec<usize>,
out_map: &mut HashMap<usize, Vec<[PhysicalQubit; 2]>>,
front_layer: &mut FrontLayer,
extended_set: &mut ExtendedSet,
required_predecessors: &mut [u32],
node_block_results: &mut HashMap<usize, Vec<BlockResult>>,
route_block_dag: &F,
) where
F: Fn(&SabreDAG, &NLayout) -> (SabreResult, NLayout),
{
// First node gets the swaps attached. We don't add to the `gate_order` here because
// `route_reachable_nodes` is responsible for that part.
out_map.insert(dag.dag[nodes[0]].py_node_id, swaps);
for node in nodes {
front_layer.remove(node);
}
route_reachable_nodes(
seed,
nodes,
dag,
layout,
coupling,
gate_order,
front_layer,
required_predecessors,
node_block_results,
route_block_dag,
);
// Ideally we'd know how to mutate the extended set directly, but since its limited size ties
// its construction strongly to the iteration order through the front layer, it's not easy to
// do better than just emptying it and rebuilding.
extended_set.clear();
populate_extended_set(
extended_set,
dag,
front_layer,
layout,
required_predecessors,
);
}
fn gen_swap_epilogue(
coupling: &DiGraph<(), ()>,
mut from_layout: NLayout,
to_layout: &NLayout,
seed: u64,
) -> Vec<[PhysicalQubit; 2]> {
// Map physical location in from_layout to physical location in to_layout
let mapping: HashMap<NodeIndex, NodeIndex> = from_layout
.iter_physical()
.map(|(p, v)| {
(
NodeIndex::new(p.index()),
NodeIndex::new(v.to_phys(to_layout).index()),
)
})
.collect();
let swaps = token_swapper(
coupling,
mapping,
Some(SWAP_EPILOGUE_TRIALS),
Some(seed),
None,
)
.unwrap();
// Convert physical swaps to virtual swaps
swaps
.into_iter()
.map(|(l, r)| {
let p_l = PhysicalQubit::new(l.index().try_into().unwrap());
let p_r = PhysicalQubit::new(r.index().try_into().unwrap());
from_layout.swap_physical(p_l, p_r);
[p_l, p_r]
})
.collect()
}
/// Search forwards in the `dag` from all the nodes in `to_visit`, adding them to the `gate_order`
/// or the current `front_layer` as appropriate, and continue inspecting gates until there is
/// nothing further with no required predecessors.
///
/// The nodes in `to_visit` should all already have no further required predecessors.
fn route_reachable_nodes<F>(
seed: u64,
to_visit: &[NodeIndex],
dag: &SabreDAG,
layout: &NLayout,
coupling: &DiGraph<(), ()>,
gate_order: &mut Vec<usize>,
front_layer: &mut FrontLayer,
required_predecessors: &mut [u32],
node_block_results: &mut HashMap<usize, Vec<BlockResult>>,
route_block_dag: &F,
) where
F: Fn(&SabreDAG, &NLayout) -> (SabreResult, NLayout),
{
let mut to_visit = to_visit.to_vec();
let mut i = 0;
// Iterate through `to_visit`, except we often push new nodes onto the end of it.
while i < to_visit.len() {
let node_id = to_visit[i];
let node = &dag.dag[node_id];
i += 1;
// If the node is a directive that means it can be placed anywhere
if !node.directive {
match dag.node_blocks.get(&node.py_node_id) {
Some(blocks) => {
// Control flow op. Route all blocks for current layout.
let mut block_results: Vec<BlockResult> = Vec::with_capacity(blocks.len());
for inner_dag in blocks {
let (inner_dag_routed, inner_final_layout) =
route_block_dag(inner_dag, layout);
// For now, we always append a swap circuit that gets the inner block
// back to the parent's layout.
let swap_epilogue =
gen_swap_epilogue(coupling, inner_final_layout, layout, seed);
let block_result = BlockResult {
result: inner_dag_routed,
swap_epilogue,
};
block_results.push(block_result);
}
node_block_results.insert_unique_unchecked(node.py_node_id, block_results);
}
None => match node.qubits[..] {
// A gate op whose connectivity must match the device to be
// placed in the gate order.
[a, b]
if !coupling.contains_edge(
NodeIndex::new(a.to_phys(layout).index()),
NodeIndex::new(b.to_phys(layout).index()),
) =>
{
// 2Q op that cannot be placed. Add it to the front layer
// and move on.
front_layer.insert(node_id, [a.to_phys(layout), b.to_phys(layout)]);
continue;
}
_ => {}
},
}
}
gate_order.push(node.py_node_id);
for edge in dag.dag.edges_directed(node_id, Direction::Outgoing) {
let successor_node = edge.target();
let successor_index = successor_node.index();
required_predecessors[successor_index] -= 1;
if required_predecessors[successor_index] == 0 {
to_visit.push(successor_node);
}
}
}
}
/// Walk through the swaps in the given vector, undoing them on the layout and removing them.
fn undo_swaps(
swaps: &mut Vec<[PhysicalQubit; 2]>,
front_layer: &mut FrontLayer,
extended_set: &mut ExtendedSet,
layout: &mut NLayout,
) {
swaps.drain(..).rev().for_each(|swap| {
front_layer.apply_swap(swap);
extended_set.apply_swap(swap);
layout.swap_physical(swap[0], swap[1]);
});
}
/// Find the node index and its associated virtual qubits that is currently the closest to being
/// routable in terms of number of swaps.
fn closest_operation<'a>(
front_layer: &'a FrontLayer,
dist: &'_ ArrayView2<f64>,
) -> (&'a NodeIndex, &'a [PhysicalQubit; 2]) {
front_layer
.iter()
.min_by(|(_, qubits_a), (_, qubits_b)| {
dist[[qubits_a[0].index(), qubits_a[1].index()]]
.partial_cmp(&dist[[qubits_b[0].index(), qubits_b[1].index()]])
.unwrap_or(Ordering::Equal)
})
.unwrap()
}
/// Add the minimal set of swaps to the `swaps` vector that bring the two `qubits` together so that
/// a 2q gate on them could be routed.
fn swaps_to_route(
swaps: &mut Vec<[PhysicalQubit; 2]>,
qubits: &[PhysicalQubit; 2],
coupling_graph: &DiGraph<(), ()>,
) {
let mut shortest_paths: DictMap<NodeIndex, Vec<NodeIndex>> = DictMap::new();
(dijkstra(
coupling_graph,
NodeIndex::new(qubits[0].index()),
Some(NodeIndex::new(qubits[1].index())),
|_| Ok(1.),
Some(&mut shortest_paths),
) as PyResult<Vec<Option<f64>>>)
.unwrap();
let shortest_path = shortest_paths
.get(&NodeIndex::new(qubits[1].index()))
.unwrap()
.iter()
.map(|n| PhysicalQubit::new(n.index() as u32))
.collect::<Vec<_>>();
// Insert greedy swaps along that shortest path, splitting them between moving the left side
// and moving the right side to minimise the depth. One side needs to move up to the split
// point and the other can stop one short because the gate will be routable then.
let split: usize = shortest_path.len() / 2;
swaps.reserve(shortest_path.len() - 2);
for i in 0..split {
swaps.push([shortest_path[i], shortest_path[i + 1]]);
}
for i in 0..split - 1 {
let end = shortest_path.len() - 1 - i;
swaps.push([shortest_path[end], shortest_path[end - 1]]);
}
}
/// Return the swap of two virtual qubits that produces the best score of all possible swaps.
fn choose_best_swap(
layer: &FrontLayer,
extended_set: &ExtendedSet,
neighbor_table: &NeighborTable,
dist: &ArrayView2<f64>,
qubits_decay: &[f64],
heuristic: &Heuristic,
rng: &mut Pcg64Mcg,
best_swaps: &mut Vec<[PhysicalQubit; 2]>,
) -> [PhysicalQubit; 2] {
best_swaps.clear();
let mut min_score = f64::MAX;
// The decay heuristic is the only one that actually needs the absolute score.
let absolute_score = match heuristic {
Heuristic::Decay => {
layer.total_score(dist) + EXTENDED_SET_WEIGHT * extended_set.total_score(dist)
}
_ => 0.0,
};
for swap in obtain_swaps(layer, neighbor_table) {
let score = match heuristic {
Heuristic::Basic => layer.score(swap, dist),
Heuristic::Lookahead => {
layer.score(swap, dist) + EXTENDED_SET_WEIGHT * extended_set.score(swap, dist)
}
Heuristic::Decay => {
qubits_decay[swap[0].index()].max(qubits_decay[swap[1].index()])
* (absolute_score
+ layer.score(swap, dist)
+ EXTENDED_SET_WEIGHT * extended_set.score(swap, dist))
}
};
if score < min_score - BEST_EPSILON {
min_score = score;
best_swaps.clear();
best_swaps.push(swap);
} else if (score - min_score).abs() < BEST_EPSILON {
best_swaps.push(swap);
}
}
*best_swaps.choose(rng).unwrap()
}
#[pymodule]
pub fn sabre_swap(m: &Bound<PyModule>) -> PyResult<()> {
m.add_wrapped(wrap_pyfunction!(build_swap_map))?;
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(())
}

3
qiskit/__init__.py
View File

@ -62,8 +62,7 @@ import qiskit._accelerate
sys.modules["qiskit._accelerate.nlayout"] = qiskit._accelerate.nlayout
sys.modules["qiskit._accelerate.quantum_circuit"] = qiskit._accelerate.quantum_circuit
sys.modules["qiskit._accelerate.stochastic_swap"] = qiskit._accelerate.stochastic_swap
sys.modules["qiskit._accelerate.sabre_swap"] = qiskit._accelerate.sabre_swap
sys.modules["qiskit._accelerate.sabre_layout"] = qiskit._accelerate.sabre_layout
sys.modules["qiskit._accelerate.sabre"] = qiskit._accelerate.sabre
sys.modules["qiskit._accelerate.pauli_expval"] = qiskit._accelerate.pauli_expval
sys.modules["qiskit._accelerate.dense_layout"] = qiskit._accelerate.dense_layout
sys.modules["qiskit._accelerate.sparse_pauli_op"] = qiskit._accelerate.sparse_pauli_op

4
qiskit/transpiler/passes/layout/sabre_layout.py
View File

@ -35,8 +35,8 @@ from qiskit.transpiler.layout import Layout
from qiskit.transpiler.basepasses import TransformationPass
from qiskit.transpiler.exceptions import TranspilerError
from qiskit._accelerate.nlayout import NLayout
from qiskit._accelerate.sabre_layout import sabre_layout_and_routing
from qiskit._accelerate.sabre_swap import (
from qiskit._accelerate.sabre import (
sabre_layout_and_routing,
Heuristic,
NeighborTable,
)

7
qiskit/transpiler/passes/routing/sabre_swap.py
View File

@ -31,8 +31,8 @@ from qiskit.transpiler.passes.layout import disjoint_utils
from qiskit.dagcircuit import DAGCircuit
from qiskit.utils.parallel import CPU_COUNT
from qiskit._accelerate.sabre_swap import (
build_swap_map,
from qiskit._accelerate.sabre import (
sabre_routing,
Heuristic,
NeighborTable,
SabreDAG,
@ -239,8 +239,7 @@ class SabreSwap(TransformationPass):
self._qubit_indices,
)
sabre_start = time.perf_counter()
*sabre_result, final_permutation = build_swap_map(
len(dag.qubits),
*sabre_result, final_permutation = sabre_routing(
sabre_dag,
self._neighbor_table,
self.dist_matrix,