Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Use VF2 to find a partial layout for seeding a SabreLayout trial #10169

Closed
wants to merge 9 commits into from
80 changes: 76 additions & 4 deletions crates/accelerate/src/sabre_layout.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
// that they have been altered from the originals.
#![allow(clippy::too_many_arguments)]

use hashbrown::HashSet;
use ndarray::prelude::*;
use numpy::PyReadonlyArray2;
use pyo3::prelude::*;
Expand All @@ -36,6 +37,7 @@ pub fn sabre_layout_and_routing(
num_swap_trials: usize,
num_layout_trials: usize,
seed: Option<u64>,
partial_layout: Option<Vec<Option<usize>>>,
) -> ([NLayout; 2], SabreResult) {
let run_in_parallel = getenv_use_multiple_threads();
let outer_rng = match seed {
Expand All @@ -44,14 +46,15 @@ pub fn sabre_layout_and_routing(
};
let seed_vec: Vec<u64> = outer_rng
.sample_iter(&rand::distributions::Standard)
.take(num_layout_trials)
.take(num_layout_trials + 1)
.collect();
let dist = distance_matrix.as_array();
if run_in_parallel && num_layout_trials > 1 {
seed_vec
.into_par_iter()
.enumerate()
.map(|(index, seed_trial)| {
let partial = if index == num_layout_trials { &partial_layout } else { &None };
(
index,
layout_trial(
Expand All @@ -63,6 +66,7 @@ pub fn sabre_layout_and_routing(
max_iterations,
num_swap_trials,
run_in_parallel,
partial,
),
)
})
Expand All @@ -77,7 +81,9 @@ pub fn sabre_layout_and_routing(
} else {
seed_vec
.into_iter()
.map(|seed_trial| {
.enumerate()
.map(|(index, seed_trial)| {
let partial = if index == num_layout_trials { &partial_layout } else { &None };
layout_trial(
dag,
neighbor_table,
Expand All @@ -87,6 +93,7 @@ pub fn sabre_layout_and_routing(
max_iterations,
num_swap_trials,
run_in_parallel,
partial,
)
})
.min_by_key(|(_, result)| result.map.map.values().map(|x| x.len()).sum::<usize>())
Expand All @@ -103,12 +110,77 @@ fn layout_trial(
max_iterations: usize,
num_swap_trials: usize,
run_swap_in_parallel: bool,
partial_layout: &Option<Vec<Option<usize>>>,
) -> ([NLayout; 2], SabreResult) {
// Pick a random initial layout and fully populate ancillas in that layout too
let num_physical_qubits = distance_matrix.shape()[0];
let mut rng = Pcg64Mcg::seed_from_u64(seed);
let mut physical_qubits: Vec<usize> = (0..num_physical_qubits).collect();
physical_qubits.shuffle(&mut rng);
let physical_qubits: Vec<usize> = match partial_layout {
Some(partial_layout_bits) => {
let mut used_bits: HashSet<usize> = partial_layout_bits
.iter()
.filter_map(|x| x.as_ref())
.copied()
.collect();
// Compute nearest neighbors to use for free bits
let mut free_bits_distance: Vec<(usize, f64)> = distance_matrix
.axis_iter(Axis(0))
.enumerate()
.filter_map(|(x, row)| {
// If starting from free bit don't check distance
if !used_bits.contains(&x) {
None
} else {
let index_distance =
row.into_iter().enumerate().filter_map(|(y, value)| {
if used_bits.contains(&y) {
None
} else {
Some((y, *value))
}
});
Some(index_distance)
}
})
.flatten()
.collect();
free_bits_distance.par_sort_by(|a, b| {
// Reverse arg order so lowest distance is at the end of Vec
// and when we pop below we get the closest nodes first
b.1.partial_cmp(&a.1).unwrap()
});

let mut get_free_bit = || -> usize {
// As the free_bits_distance Vec will have multiple
// entries for each bit we need to loop until we find
// an unused bit
let mut new_bit = free_bits_distance.pop().unwrap().0;
while used_bits.contains(&new_bit) {
new_bit = free_bits_distance.pop().unwrap().0;
}
used_bits.insert(new_bit);
new_bit
};

(0..num_physical_qubits)
.map(|x| {
if x >= partial_layout_bits.len() {
get_free_bit()
} else {
match partial_layout_bits[x] {
Some(phys) => phys,
None => get_free_bit(),
}
}
})
.collect()
}
None => {
let mut physical_qubits: Vec<usize> = (0..num_physical_qubits).collect();
physical_qubits.shuffle(&mut rng);
physical_qubits
}
};
let mut initial_layout = NLayout::from_logical_to_physical(physical_qubits);
let new_dag_fn = |nodes| {
// Because the current implementation of Sabre swap doesn't permute
Expand Down
2 changes: 1 addition & 1 deletion crates/accelerate/src/sabre_swap/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ 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 EXTENDED_SET_SIZE: usize = 72; // 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.
Expand Down
168 changes: 163 additions & 5 deletions qiskit/transpiler/passes/layout/sabre_layout.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,11 @@
"""Layout selection using the SABRE bidirectional search approach from Li et al.
"""

from collections import defaultdict
import copy
import logging
import time

import numpy as np
import rustworkx as rx

Expand All @@ -24,6 +27,7 @@
from qiskit.transpiler.passes.layout.enlarge_with_ancilla import EnlargeWithAncilla
from qiskit.transpiler.passes.layout.apply_layout import ApplyLayout
from qiskit.transpiler.passes.layout import disjoint_utils
from qiskit.transpiler.passes.layout import vf2_utils
from qiskit.transpiler.passmanager import PassManager
from qiskit.transpiler.layout import Layout
from qiskit.transpiler.basepasses import TransformationPass
Expand All @@ -38,18 +42,25 @@
from qiskit.transpiler.target import Target
from qiskit.transpiler.coupling import CouplingMap
from qiskit.tools.parallel import CPU_COUNT
from qiskit.circuit.controlflow import ControlFlowOp, ForLoopOp
from qiskit.converters import circuit_to_dag

logger = logging.getLogger(__name__)


class SabreLayout(TransformationPass):
"""Choose a Layout via iterative bidirectional routing of the input circuit.

Starting with a random initial `Layout`, the algorithm does a full routing
of the circuit (via the `routing_pass` method) to end up with a
`final_layout`. This final_layout is then used as the initial_layout for
routing the reverse circuit. The algorithm iterates a number of times until
it finds an initial_layout that reduces full routing cost.
The algorithm does a full routing of the circuit (via the `routing_pass`
method) to end up with a `final_layout`. This final_layout is then used as
the initial_layout for routing the reverse circuit. The algorithm iterates a
number of times until it finds an initial_layout that reduces full routing cost.

Prior to running the SABRE algorithm this transpiler pass will try to find the layout
for deepest layer that is has an isomorphic subgraph in the coupling graph. This is
done by progressively using the algorithm from :class:`~.VF2Layout` on the circuit
until a mapping is not found. This partial layout is then used to seed the SABRE algorithm
and then random physical bits are selected for the remaining elements in the mapping.

This method exploits the reversibility of quantum circuits, and tries to
include global circuit information in the choice of initial_layout.
Expand Down Expand Up @@ -85,6 +96,10 @@ def __init__(
swap_trials=None,
layout_trials=None,
skip_routing=False,
vf2_partial_layout=True,
vf2_call_limit=None,
vf2_time_limit=None,
vf2_max_trials=None,
):
"""SabreLayout initializer.

Expand Down Expand Up @@ -121,6 +136,16 @@ def __init__(
will be returned in the property set. This is a tradeoff to run custom
routing with multiple layout trials, as using this option will cause
SabreLayout to run the routing stage internally but not use that result.
vf2_partial_layout (bool): Run vf2 partial layout
vf2_call_limit (int): The number of state visits to attempt in each execution of
VF2 to attempt to find a partial layout.
vf2_time_limit (float): The total time limit in seconds to run VF2 to find a partial
layout
vf2_max_trials (int): The maximum number of trials to run VF2 to find
a partial layout. If this is not specified the number of trials will be limited
based on the number of edges in the interaction graph or the coupling graph
(whichever is larger) if no other limits are set. If set to a value <= 0 no
limit on the number of trials will be set.

Raises:
TranspilerError: If both ``routing_pass`` and ``swap_trials`` or
Expand Down Expand Up @@ -158,6 +183,11 @@ def __init__(
self.coupling_map = copy.deepcopy(self.coupling_map)
self.coupling_map.make_symmetric()
self._neighbor_table = NeighborTable(rx.adjacency_matrix(self.coupling_map.graph))
self.avg_error_map = None
self.vf2_partial_layout = vf2_partial_layout
self.call_limit = vf2_call_limit
self.time_limit = vf2_time_limit
self.max_trials = vf2_max_trials

def run(self, dag):
"""Run the SabreLayout pass on `dag`.
Expand Down Expand Up @@ -299,6 +329,13 @@ def _inner_run(self, dag, coupling_map):
coupling_map.size(),
original_qubit_indices,
)
partial_layout = None
if self.vf2_partial_layout:
partial_layout_virtual_bits = self._vf2_partial_layout(
dag, coupling_map
).get_virtual_bits()
partial_layout = [partial_layout_virtual_bits.get(i, None) for i in dag.qubits]

((initial_layout, final_layout), sabre_result) = sabre_layout_and_routing(
sabre_dag,
neighbor_table,
Expand All @@ -308,6 +345,7 @@ def _inner_run(self, dag, coupling_map):
self.swap_trials,
self.layout_trials,
self.seed,
partial_layout,
)

# Apply initial layout selected.
Expand Down Expand Up @@ -370,3 +408,123 @@ def _compose_layouts(self, initial_layout, pass_final_layout, qregs):
qubit_map = Layout.combine_into_edge_map(initial_layout, trivial_layout)
final_layout = {v: pass_final_layout._v2p[qubit_map[v]] for v in initial_layout._v2p}
return Layout(final_layout)

# TODO: Migrate this to rust as part of sabre_layout.rs after
# https://github.com/Qiskit/rustworkx/issues/741 is implemented and released
def _vf2_partial_layout(self, dag, coupling_map):
"""Find a partial layout using vf2 on the deepest subgraph that is isomorphic to
the coupling graph."""
im_graph_node_map = {}
reverse_im_graph_node_map = {}
im_graph = rx.PyGraph(multigraph=False)
logger.debug("Buidling interaction graphs")
largest_im_graph = None
best_mapping = None
first_mapping = None
if self.avg_error_map is None:
self.avg_error_map = vf2_utils.build_average_error_map(self.target, None, coupling_map)

cm_graph, cm_nodes = vf2_utils.shuffle_coupling_graph(coupling_map, self.seed, False)
# To avoid trying to over optimize the result by default limit the number
# of trials based on the size of the graphs. For circuits with simple layouts
# like an all 1q circuit we don't want to sit forever trying every possible
# mapping in the search space if no other limits are set
if self.max_trials is None and self.call_limit is None and self.time_limit is None:
im_graph_edge_count = len(im_graph.edge_list())
cm_graph_edge_count = len(coupling_map.graph.edge_list())
self.max_trials = max(im_graph_edge_count, cm_graph_edge_count) + 15

start_time = time.time()

# A more efficient search pattern would be to do a binary search
# and find, but to conserve memory and avoid a large number of
# unecessary graphs this searchs from the beginning and continues
# until there is no vf2 match
def _visit(dag, weight, wire_map):
for node in dag.topological_op_nodes():
nonlocal largest_im_graph
largest_im_graph = im_graph.copy()
if getattr(node.op, "_directive", False):
continue
if isinstance(node.op, ControlFlowOp):
if isinstance(node.op, ForLoopOp):
inner_weight = len(node.op.params[0]) * weight
else:
inner_weight = weight
for block in node.op.blocks:
inner_wire_map = {
inner: wire_map[outer] for outer, inner in zip(node.qargs, block.qubits)
}
_visit(circuit_to_dag(block), inner_weight, inner_wire_map)
continue
len_args = len(node.qargs)
qargs = [wire_map[q] for q in node.qargs]
if len_args == 1:
if qargs[0] not in im_graph_node_map:
weights = defaultdict(int)
weights[node.name] += weight
im_graph_node_map[qargs[0]] = im_graph.add_node(weights)
reverse_im_graph_node_map[im_graph_node_map[qargs[0]]] = qargs[0]
else:
im_graph[ # pylint: disable=unsubscriptable-object
im_graph_node_map[qargs[0]]
][node.op.name] += weight
if len_args == 2:
if qargs[0] not in im_graph_node_map:
im_graph_node_map[qargs[0]] = im_graph.add_node(defaultdict(int))
reverse_im_graph_node_map[im_graph_node_map[qargs[0]]] = qargs[0]
if qargs[1] not in im_graph_node_map:
im_graph_node_map[qargs[1]] = im_graph.add_node(defaultdict(int))
reverse_im_graph_node_map[im_graph_node_map[qargs[1]]] = qargs[1]
edge = (im_graph_node_map[qargs[0]], im_graph_node_map[qargs[1]])
if im_graph.has_edge(*edge):
im_graph.get_edge_data(*edge)[node.name] += weight
else:
weights = defaultdict(int)
weights[node.name] += weight
im_graph.add_edge(*edge, weights)
if len_args > 2:
raise TranspilerError(
"Encountered an instruction operating on more than 2 qubits, this pass "
"only functions with 1 or 2 qubit operations."
)
vf2_mapping = rx.vf2_mapping(
cm_graph,
im_graph,
subgraph=True,
id_order=False,
induced=False,
call_limit=self.call_limit,
)
try:
nonlocal first_mapping
first_mapping = next(vf2_mapping)
except StopIteration:
break
nonlocal best_mapping
best_mapping = vf2_mapping
elapsed_time = time.time() - start_time
if (
self.time_limit is not None
and best_mapping is not None
and elapsed_time >= self.time_limit
):
logger.debug(
"SabreLayout VF2 heuristic has taken %s which exceeds configured max time: %s",
elapsed_time,
self.time_limit,
)
break

_visit(dag, 1, {bit: bit for bit in dag.qubits})
logger.debug("Finding best mappings of largest partial subgraph")
im_graph = largest_im_graph

def mapping_to_layout(layout_mapping):
return Layout({reverse_im_graph_node_map[k]: v for k, v in layout_mapping.items()})

layout_mapping = {im_i: cm_nodes[cm_i] for cm_i, im_i in first_mapping.items()}
chosen_layout = mapping_to_layout(layout_mapping)
for reg in dag.qregs.values():
chosen_layout.add_register(reg)
return chosen_layout
Loading