Skip to content

Commit c33a092

Browse files
authored
Add type-safe VecMap for slices indexed by newtypes (#14910)
* Simply Sabre release-valve calculation (#15897) There is an edge case in release-valve handling where, if the shortest path is a single swap (but the heuristics have been chosen such that it is unselectable normally), it is possible for the release valve to cause _two_ gates to be routed. The previous handling of this path was written in quite a complicated manner, likely to be able to use `closest_node` by name from its discovery earlier in the release-valve process. Instead, it is easier to recognise that, given a single swap: * the two qubits cannot be part of the same gate, or the gate would have been routable _without_ the swap; * the `closest_node` must touch one of these qubits because of the the Dijkstra search; * therefore we can simply add any routable gate that touches _either_ qubit without risk of duplication or forgetting `closest_swap`. * Add explicit turbofish typing
1 parent eab94f2 commit c33a092

4 files changed

Lines changed: 105 additions & 39 deletions

File tree

crates/transpiler/src/passes/sabre/layer.rs

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ use rustworkx_core::petgraph::prelude::*;
1616

1717
use qiskit_circuit::PhysicalQubit;
1818

19+
use super::vec_map::VecMap;
20+
1921
/// A container for the current non-routable parts of the front layer. This only ever holds
2022
/// two-qubit gates; the only reason a 0q- or 1q operation can be unroutable is because it has an
2123
/// unsatisfied 2q predecessor, which disqualifies it from being in the front layer.
@@ -31,7 +33,7 @@ pub struct FrontLayer {
3133
nodes: IndexMap<NodeIndex, [PhysicalQubit; 2], ::ahash::RandomState>,
3234
/// Map of each qubit to the node that acts on it and the other qubit that node acts on, if this
3335
/// qubit is active (otherwise `None`).
34-
qubits: Vec<Option<(NodeIndex, PhysicalQubit)>>,
36+
qubits: VecMap<PhysicalQubit, Option<(NodeIndex, PhysicalQubit)>>,
3537
}
3638

3739
impl FrontLayer {
@@ -43,7 +45,7 @@ impl FrontLayer {
4345
num_qubits as usize / 2,
4446
::ahash::RandomState::default(),
4547
),
46-
qubits: vec![None; num_qubits as usize],
48+
qubits: vec![None; num_qubits as usize].into(),
4749
}
4850
}
4951

@@ -54,15 +56,15 @@ impl FrontLayer {
5456

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

6163
/// Add a node into the front layer, with the two qubits it operates on.
6264
pub fn insert(&mut self, index: NodeIndex, qubits: [PhysicalQubit; 2]) {
6365
let [a, b] = qubits;
64-
self.qubits[a.index()] = Some((index, b));
65-
self.qubits[b.index()] = Some((index, a));
66+
self.qubits[a] = Some((index, b));
67+
self.qubits[b] = Some((index, a));
6668
self.nodes.insert(index, qubits);
6769
}
6870

@@ -74,14 +76,14 @@ impl FrontLayer {
7476
.nodes
7577
.swap_remove(index)
7678
.expect("Tried removing index that does not exist.");
77-
self.qubits[a.index()] = None;
78-
self.qubits[b.index()] = None;
79+
self.qubits[a] = None;
80+
self.qubits[b] = None;
7981
}
8082

8183
/// Query whether a qubit has an active node.
8284
#[inline]
8385
pub fn is_active(&self, qubit: PhysicalQubit) -> bool {
84-
self.qubits[qubit.index()].is_some()
86+
self.qubits[qubit].is_some()
8587
}
8688

8789
/// Calculate the score _difference_ caused by this swap, compared to not making the swap.
@@ -95,10 +97,10 @@ impl FrontLayer {
9597
// equal anyway, so not affect the score.
9698
let [a, b] = swap;
9799
let mut total = 0.0;
98-
if let Some((_, c)) = self.qubits[a.index()] {
100+
if let Some((_, c)) = self.qubits[a] {
99101
total += dist[[b.index(), c.index()]] - dist[[a.index(), c.index()]]
100102
}
101-
if let Some((_, c)) = self.qubits[b.index()] {
103+
if let Some((_, c)) = self.qubits[b] {
102104
total += dist[[a.index(), c.index()]] - dist[[b.index(), c.index()]]
103105
}
104106
total
@@ -114,25 +116,25 @@ impl FrontLayer {
114116
/// Apply a physical swap to the current layout data structure.
115117
pub fn apply_swap(&mut self, swap: [PhysicalQubit; 2]) {
116118
let [a, b] = swap;
117-
match (self.qubits[a.index()], self.qubits[b.index()]) {
119+
match (self.qubits[a], self.qubits[b]) {
118120
(Some((index1, _)), Some((index2, _))) if index1 == index2 => {
119121
let entry = self.nodes.get_mut(&index1).unwrap();
120122
*entry = [entry[1], entry[0]];
121123
return;
122124
}
123125
_ => {}
124126
}
125-
if let Some((index, c)) = self.qubits[a.index()] {
126-
self.qubits[c.index()] = Some((index, b));
127+
if let Some((index, c)) = self.qubits[a] {
128+
self.qubits[c] = Some((index, b));
127129
let entry = self.nodes.get_mut(&index).unwrap();
128130
*entry = if *entry == [a, c] { [b, c] } else { [c, b] };
129131
}
130-
if let Some((index, c)) = self.qubits[b.index()] {
131-
self.qubits[c.index()] = Some((index, a));
132+
if let Some((index, c)) = self.qubits[b] {
133+
self.qubits[c] = Some((index, a));
132134
let entry = self.nodes.get_mut(&index).unwrap();
133135
*entry = if *entry == [b, c] { [a, c] } else { [c, a] };
134136
}
135-
self.qubits.swap(a.index(), b.index());
137+
self.qubits.swap(a, b);
136138
}
137139

138140
/// True if there are no nodes in the current layer.

crates/transpiler/src/passes/sabre/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ pub mod heuristic;
1515
mod layer;
1616
mod layout;
1717
pub(crate) mod route;
18+
mod vec_map;
1819

1920
use pyo3::prelude::*;
2021
use pyo3::wrap_pyfunction;

crates/transpiler/src/passes/sabre/route.rs

Lines changed: 22 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ use smallvec::{SmallVec, smallvec};
3737
use super::dag::{InteractionKind, SabreDAG};
3838
use super::heuristic::{BasicHeuristic, DecayHeuristic, Heuristic, LookaheadHeuristic, SetScaling};
3939
use super::layer::{ExtendedSet, FrontLayer};
40+
use super::vec_map::VecMap;
4041
use crate::TranspilerError;
4142
use crate::neighbors::Neighbors;
4243
use crate::target::{Target, TargetCouplingError};
@@ -451,8 +452,8 @@ struct State {
451452
extended_set: ExtendedSet,
452453
/// How many predecessors still need to be satisfied for each node index before it is at the
453454
/// front of the topological iteration through the nodes as they're routed.
454-
required_predecessors: Vec<u32>,
455-
decay: Vec<f64>,
455+
required_predecessors: VecMap<NodeIndex, u32>,
456+
decay: VecMap<PhysicalQubit, f64>,
456457
/// Reusable allocated storage space for accumulating and scoring swaps. This is owned as part
457458
/// of the general state to avoid reallocation costs.
458459
swap_scores: Vec<([PhysicalQubit; 2], f64)>,
@@ -480,7 +481,7 @@ impl State {
480481
problem: RoutingProblem,
481482
qubit: PhysicalQubit,
482483
) -> Option<NodeIndex> {
483-
self.front_layer.qubits()[qubit.index()].and_then(|(node, other)| {
484+
self.front_layer.qubits()[qubit].and_then(|(node, other)| {
484485
problem
485486
.target
486487
.neighbors
@@ -564,11 +565,10 @@ impl State {
564565
.dag
565566
.edges_directed(node_id, Direction::Outgoing)
566567
{
567-
let successor_node = edge.target();
568-
let successor_index = successor_node.index();
569-
self.required_predecessors[successor_index] -= 1;
570-
if self.required_predecessors[successor_index] == 0 {
571-
to_visit.push_back(successor_node);
568+
let successor = edge.target();
569+
self.required_predecessors[successor] -= 1;
570+
if self.required_predecessors[successor] == 0 {
571+
to_visit.push_back(successor);
572572
}
573573
}
574574
}
@@ -626,25 +626,24 @@ impl State {
626626
return;
627627
};
628628
let mut to_visit = self.front_layer.iter_nodes().copied().collect::<Vec<_>>();
629-
let mut decremented: IndexMap<usize, u32, ahash::RandomState> =
629+
let mut decremented: IndexMap<NodeIndex, u32, ahash::RandomState> =
630630
IndexMap::with_hasher(ahash::RandomState::default());
631631
let mut i = 0;
632632
while i < to_visit.len() && self.extended_set.len() < extended_set_size {
633633
let node = to_visit[i];
634634
for edge in problem.sabre.dag.edges_directed(node, Direction::Outgoing) {
635-
let successor_node = edge.target();
636-
let successor_index = successor_node.index();
637-
*decremented.entry(successor_index).or_insert(0) += 1;
638-
self.required_predecessors[successor_index] -= 1;
639-
if self.required_predecessors[successor_index] == 0 {
635+
let successor = edge.target();
636+
*decremented.entry(successor).or_insert(0) += 1;
637+
self.required_predecessors[successor] -= 1;
638+
if self.required_predecessors[successor] == 0 {
640639
// TODO: this looks "through" control-flow ops without seeing them, but we
641640
// actually eagerly route control-flow blocks as soon as they're eligible, so
642641
// they should be reflected in the extended set.
643-
if let InteractionKind::TwoQ([a, b]) = &problem.sabre.dag[successor_node].kind {
642+
if let InteractionKind::TwoQ([a, b]) = &problem.sabre.dag[successor].kind {
644643
self.extended_set
645644
.push([a.to_phys(&self.layout), b.to_phys(&self.layout)]);
646645
}
647-
to_visit.push(successor_node);
646+
to_visit.push(successor);
648647
}
649648
}
650649
i += 1;
@@ -767,8 +766,7 @@ impl State {
767766

768767
if let Some(DecayHeuristic { .. }) = problem.heuristic.decay {
769768
for (swap, score) in self.swap_scores.iter_mut() {
770-
*score = (absolute_score + *score)
771-
* self.decay[swap[0].index()].max(self.decay[swap[1].index()]);
769+
*score = (absolute_score + *score) * self.decay[swap[0]].max(self.decay[swap[1]]);
772770
}
773771
}
774772

@@ -861,14 +859,15 @@ pub fn swap_map_trial<'a>(
861859
let mut order = Order::for_problem(problem);
862860

863861
let num_qubits: u32 = problem.target.num_qubits().try_into().unwrap();
864-
let mut required_predecessors = vec![0; problem.sabre.dag.node_count()];
862+
let mut required_predecessors =
863+
VecMap::<NodeIndex, u32>::from(vec![0; problem.sabre.dag.node_count()]);
865864
for edge in problem.sabre.dag.edge_references() {
866-
required_predecessors[edge.target().index()] += 1;
865+
required_predecessors[edge.target()] += 1;
867866
}
868867
let mut state = State {
869868
front_layer: FrontLayer::new(num_qubits),
870869
extended_set: ExtendedSet::new(num_qubits),
871-
decay: vec![1.; num_qubits as usize],
870+
decay: vec![1.; num_qubits as usize].into(),
872871
required_predecessors,
873872
layout: initial_layout.clone(),
874873
swap_scores: Vec::with_capacity(problem.target.neighbors.edge_count() / 2),
@@ -904,8 +903,8 @@ pub fn swap_map_trial<'a>(
904903
state.decay.fill(1.);
905904
num_search_steps = 0;
906905
} else {
907-
state.decay[best_swap[0].index()] += increment;
908-
state.decay[best_swap[1].index()] += increment;
906+
state.decay[best_swap[0]] += increment;
907+
state.decay[best_swap[1]] += increment;
909908
}
910909
}
911910
}
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// This code is part of Qiskit.
2+
//
3+
// (C) Copyright IBM 2026
4+
//
5+
// This code is licensed under the Apache License, Version 2.0. You may
6+
// obtain a copy of this license in the LICENSE.txt file in the root directory
7+
// of this source tree or at https://www.apache.org/licenses/LICENSE-2.0.
8+
//
9+
// Any modifications or derivative works of this code must retain this
10+
// copyright notice, and modified files need to carry a notice indicating
11+
// that they have been altered from the originals.
12+
13+
use std::marker::PhantomData;
14+
use std::ops;
15+
16+
use rustworkx_core::petgraph::graph::IndexType;
17+
18+
/// Internal helper struct that represents a `Box<[T]>`, but is indexed by a petgraph-like indexer.
19+
///
20+
/// The layer structures involve several flat arrays, each representing a map from an index to a
21+
/// value, but the index types are often different, so it's less error-prone if we _enforce_ that
22+
/// you index using the object, rather than calling its `.index` method and erasing the type.
23+
#[derive(Clone, Debug)]
24+
pub struct VecMap<Idx: IndexType, T> {
25+
phantom: PhantomData<Idx>,
26+
data: Box<[T]>,
27+
}
28+
impl<Idx: IndexType, T> VecMap<Idx, T> {
29+
/// Swap the values of two indices.
30+
#[inline]
31+
pub fn swap(&mut self, a: Idx, b: Idx) {
32+
self.data.swap(a.index(), b.index())
33+
}
34+
35+
/// Fill all entries in the slice with a given value.
36+
#[inline]
37+
pub fn fill(&mut self, val: T)
38+
where
39+
T: Clone,
40+
{
41+
self.data.fill(val)
42+
}
43+
}
44+
45+
impl<Idx: IndexType, T> ops::Index<Idx> for VecMap<Idx, T> {
46+
type Output = <[T] as ops::Index<usize>>::Output;
47+
fn index(&self, index: Idx) -> &Self::Output {
48+
&self.data[index.index()]
49+
}
50+
}
51+
impl<Idx: IndexType, T> ops::IndexMut<Idx> for VecMap<Idx, T> {
52+
fn index_mut(&mut self, index: Idx) -> &mut Self::Output {
53+
&mut self.data[index.index()]
54+
}
55+
}
56+
57+
impl<Idx: IndexType, T> From<Vec<T>> for VecMap<Idx, T> {
58+
fn from(value: Vec<T>) -> Self {
59+
Self {
60+
phantom: PhantomData,
61+
data: value.into_boxed_slice(),
62+
}
63+
}
64+
}

0 commit comments

Comments
 (0)