-
Notifications
You must be signed in to change notification settings - Fork 2.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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.
- Loading branch information
1 parent
4ff0fa2
commit 3af3cf5
Showing
12 changed files
with
763 additions
and
852 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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(()) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.