diff --git a/crates/transpiler/src/passes/sabre/layer.rs b/crates/transpiler/src/passes/sabre/layer.rs index 13c5225684f5..2a3a1ab46d03 100644 --- a/crates/transpiler/src/passes/sabre/layer.rs +++ b/crates/transpiler/src/passes/sabre/layer.rs @@ -16,6 +16,8 @@ use rustworkx_core::petgraph::prelude::*; use qiskit_circuit::PhysicalQubit; +use super::vec_map::VecMap; + /// A container for the current non-routable parts of the front layer. This only ever holds /// two-qubit gates; the only reason a 0q- or 1q operation can be unroutable is because it has an /// unsatisfied 2q predecessor, which disqualifies it from being in the front layer. @@ -31,7 +33,7 @@ pub struct FrontLayer { nodes: IndexMap, /// 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>, + qubits: VecMap>, } impl FrontLayer { @@ -43,7 +45,7 @@ impl FrontLayer { num_qubits as usize / 2, ::ahash::RandomState::default(), ), - qubits: vec![None; num_qubits as usize], + qubits: vec![None; num_qubits as usize].into(), } } @@ -54,15 +56,15 @@ impl FrontLayer { /// 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)>] { + pub fn qubits(&self) -> &VecMap> { &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; - self.qubits[a.index()] = Some((index, b)); - self.qubits[b.index()] = Some((index, a)); + self.qubits[a] = Some((index, b)); + self.qubits[b] = Some((index, a)); self.nodes.insert(index, qubits); } @@ -74,14 +76,14 @@ impl FrontLayer { .nodes .swap_remove(index) .expect("Tried removing index that does not exist."); - self.qubits[a.index()] = None; - self.qubits[b.index()] = None; + self.qubits[a] = None; + self.qubits[b] = None; } /// Query whether a qubit has an active node. #[inline] pub fn is_active(&self, qubit: PhysicalQubit) -> bool { - self.qubits[qubit.index()].is_some() + self.qubits[qubit].is_some() } /// Calculate the score _difference_ caused by this swap, compared to not making the swap. @@ -95,10 +97,10 @@ impl FrontLayer { // equal anyway, so not affect the score. let [a, b] = swap; let mut total = 0.0; - if let Some((_, c)) = self.qubits[a.index()] { + if let Some((_, c)) = self.qubits[a] { total += dist[[b.index(), c.index()]] - dist[[a.index(), c.index()]] } - if let Some((_, c)) = self.qubits[b.index()] { + if let Some((_, c)) = self.qubits[b] { total += dist[[a.index(), c.index()]] - dist[[b.index(), c.index()]] } total @@ -114,7 +116,7 @@ impl FrontLayer { /// Apply a physical swap to the current layout data structure. pub fn apply_swap(&mut self, swap: [PhysicalQubit; 2]) { let [a, b] = swap; - match (self.qubits[a.index()], self.qubits[b.index()]) { + match (self.qubits[a], self.qubits[b]) { (Some((index1, _)), Some((index2, _))) if index1 == index2 => { let entry = self.nodes.get_mut(&index1).unwrap(); *entry = [entry[1], entry[0]]; @@ -122,17 +124,17 @@ impl FrontLayer { } _ => {} } - if let Some((index, c)) = self.qubits[a.index()] { - self.qubits[c.index()] = Some((index, b)); + if let Some((index, c)) = self.qubits[a] { + self.qubits[c] = Some((index, b)); let entry = self.nodes.get_mut(&index).unwrap(); *entry = if *entry == [a, c] { [b, c] } else { [c, b] }; } - if let Some((index, c)) = self.qubits[b.index()] { - self.qubits[c.index()] = Some((index, a)); + if let Some((index, c)) = self.qubits[b] { + self.qubits[c] = Some((index, a)); let entry = self.nodes.get_mut(&index).unwrap(); *entry = if *entry == [b, c] { [a, c] } else { [c, a] }; } - self.qubits.swap(a.index(), b.index()); + self.qubits.swap(a, b); } /// True if there are no nodes in the current layer. diff --git a/crates/transpiler/src/passes/sabre/mod.rs b/crates/transpiler/src/passes/sabre/mod.rs index d6d8b15e9c07..ebe0c559a42e 100644 --- a/crates/transpiler/src/passes/sabre/mod.rs +++ b/crates/transpiler/src/passes/sabre/mod.rs @@ -15,6 +15,7 @@ pub mod heuristic; mod layer; mod layout; pub(crate) mod route; +mod vec_map; use pyo3::prelude::*; use pyo3::wrap_pyfunction; diff --git a/crates/transpiler/src/passes/sabre/route.rs b/crates/transpiler/src/passes/sabre/route.rs index 70eb876aa5f3..269c5c9ad557 100644 --- a/crates/transpiler/src/passes/sabre/route.rs +++ b/crates/transpiler/src/passes/sabre/route.rs @@ -37,6 +37,7 @@ use smallvec::{SmallVec, smallvec}; use super::dag::{InteractionKind, SabreDAG}; use super::heuristic::{BasicHeuristic, DecayHeuristic, Heuristic, LookaheadHeuristic, SetScaling}; use super::layer::{ExtendedSet, FrontLayer}; +use super::vec_map::VecMap; use crate::TranspilerError; use crate::neighbors::Neighbors; use crate::target::{Target, TargetCouplingError}; @@ -450,8 +451,8 @@ struct State { 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: Vec, - decay: Vec, + required_predecessors: VecMap, + decay: VecMap, /// Reusable allocated storage space for accumulating and scoring swaps. This is owned as part /// of the general state to avoid reallocation costs. swap_scores: Vec<([PhysicalQubit; 2], f64)>, @@ -479,7 +480,7 @@ impl State { problem: RoutingProblem, qubit: PhysicalQubit, ) -> Option { - self.front_layer.qubits()[qubit.index()].and_then(|(node, other)| { + self.front_layer.qubits()[qubit].and_then(|(node, other)| { problem .target .neighbors @@ -563,11 +564,10 @@ impl State { .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_back(successor_node); + let successor = edge.target(); + self.required_predecessors[successor] -= 1; + if self.required_predecessors[successor] == 0 { + to_visit.push_back(successor); } } } @@ -625,25 +625,24 @@ impl State { return; }; let mut to_visit = self.front_layer.iter_nodes().copied().collect::>(); - let mut decremented: IndexMap = + let mut decremented: IndexMap = IndexMap::with_hasher(ahash::RandomState::default()); let mut i = 0; while i < to_visit.len() && self.extended_set.len() < extended_set_size { let node = to_visit[i]; for edge in problem.sabre.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 { + let successor = edge.target(); + *decremented.entry(successor).or_insert(0) += 1; + self.required_predecessors[successor] -= 1; + if self.required_predecessors[successor] == 0 { // TODO: this looks "through" control-flow ops without seeing them, but we // actually eagerly route control-flow blocks as soon as they're eligible, so // they should be reflected in the extended set. - if let InteractionKind::TwoQ([a, b]) = &problem.sabre.dag[successor_node].kind { + if let InteractionKind::TwoQ([a, b]) = &problem.sabre.dag[successor].kind { self.extended_set .push([a.to_phys(&self.layout), b.to_phys(&self.layout)]); } - to_visit.push(successor_node); + to_visit.push(successor); } } i += 1; @@ -703,27 +702,12 @@ impl State { // If we apply a single swap it could be that we route 2 nodes; that is a setup like // A - B - A - B // and we swap the middle two qubits. This cannot happen if we apply 2 or more swaps. - if current_swaps.len() > 1 { - smallvec![closest_node] - } else { - // check if the closest node has neighbors that are now routable -- for that we get - // the other physical qubit that was swapped and check whether the node on it - // is now routable - let mut possible_other_qubit = current_swaps[0] + match current_swaps.as_slice() { + [swap] => swap .iter() - // check if other nodes are in the front layer that are connected by this swap - .filter_map(|&swap_qubit| self.front_layer.qubits()[swap_qubit.index()]) - // remove the closest_node, which we know we already routed - .filter(|(node_index, _other_qubit)| *node_index != closest_node) - .map(|(_node_index, other_qubit)| other_qubit); - - // if there is indeed another candidate, check if that gate is routable - if let Some(other_qubit) = possible_other_qubit.next() { - if let Some(also_routed) = self.routable_node_on_qubit(problem, other_qubit) { - return smallvec![closest_node, also_routed]; - } - } - smallvec![closest_node] + .filter_map(|q| self.routable_node_on_qubit(problem, *q)) + .collect(), + _ => smallvec![closest_node], } } @@ -781,8 +765,7 @@ impl State { if let Some(DecayHeuristic { .. }) = problem.heuristic.decay { for (swap, score) in self.swap_scores.iter_mut() { - *score = (absolute_score + *score) - * self.decay[swap[0].index()].max(self.decay[swap[1].index()]); + *score = (absolute_score + *score) * self.decay[swap[0]].max(self.decay[swap[1]]); } } @@ -875,14 +858,15 @@ pub fn swap_map_trial<'a>( let mut order = Order::for_problem(problem); let num_qubits: u32 = problem.target.num_qubits().try_into().unwrap(); - let mut required_predecessors = vec![0; problem.sabre.dag.node_count()]; + let mut required_predecessors = + VecMap::::from(vec![0; problem.sabre.dag.node_count()]); for edge in problem.sabre.dag.edge_references() { - required_predecessors[edge.target().index()] += 1; + required_predecessors[edge.target()] += 1; } let mut state = State { front_layer: FrontLayer::new(num_qubits), extended_set: ExtendedSet::new(num_qubits), - decay: vec![1.; num_qubits as usize], + decay: vec![1.; num_qubits as usize].into(), required_predecessors, layout: initial_layout.clone(), swap_scores: Vec::with_capacity(problem.target.neighbors.edge_count() / 2), @@ -918,8 +902,8 @@ pub fn swap_map_trial<'a>( state.decay.fill(1.); num_search_steps = 0; } else { - state.decay[best_swap[0].index()] += increment; - state.decay[best_swap[1].index()] += increment; + state.decay[best_swap[0]] += increment; + state.decay[best_swap[1]] += increment; } } } diff --git a/crates/transpiler/src/passes/sabre/vec_map.rs b/crates/transpiler/src/passes/sabre/vec_map.rs new file mode 100644 index 000000000000..552f58719554 --- /dev/null +++ b/crates/transpiler/src/passes/sabre/vec_map.rs @@ -0,0 +1,64 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2026 +// +// 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 https://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::marker::PhantomData; +use std::ops; + +use rustworkx_core::petgraph::graph::IndexType; + +/// Internal helper struct that represents a `Box<[T]>`, but is indexed by a petgraph-like indexer. +/// +/// The layer structures involve several flat arrays, each representing a map from an index to a +/// value, but the index types are often different, so it's less error-prone if we _enforce_ that +/// you index using the object, rather than calling its `.index` method and erasing the type. +#[derive(Clone, Debug)] +pub struct VecMap { + phantom: PhantomData, + data: Box<[T]>, +} +impl VecMap { + /// Swap the values of two indices. + #[inline] + pub fn swap(&mut self, a: Idx, b: Idx) { + self.data.swap(a.index(), b.index()) + } + + /// Fill all entries in the slice with a given value. + #[inline] + pub fn fill(&mut self, val: T) + where + T: Clone, + { + self.data.fill(val) + } +} + +impl ops::Index for VecMap { + type Output = <[T] as ops::Index>::Output; + fn index(&self, index: Idx) -> &Self::Output { + &self.data[index.index()] + } +} +impl ops::IndexMut for VecMap { + fn index_mut(&mut self, index: Idx) -> &mut Self::Output { + &mut self.data[index.index()] + } +} + +impl From> for VecMap { + fn from(value: Vec) -> Self { + Self { + phantom: PhantomData, + data: value.into_boxed_slice(), + } + } +}