From a92753a5500858de2e7b058dc0a8b906d7849b99 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Tue, 8 Nov 2022 09:40:37 -0500 Subject: [PATCH] Oxidize SabreLayout pass This commit modifies the SabreLayout pass when run without the routing_pass argument to run primarily in Rust. This builds on top of the rust version of SabreSwap previously added in #7977, #8388, and #8572. Internally, when the routing_pass argument is not set SabreLayout will perform the full sabre algorithm both layout selection and final swap mapping in rust and return the selected initial layout, the final layout, the toplogical sorting used to traverse the circuit, and a SwapMap for any swaps inserted. This is then used to build the output circuit in place of running separate layout and routing passes. The preset pass managers are updated to handle the new combined layout and routing mode of operation for SabreLayout. The routing stage to the preset pass managers remains intact, it will just operate as if a perfect layout was selected and skip SabreSwap because the circuit is already matching the connectivity constraints. Besides just operating more quickly because the heavy lifting of the algorithm operates more efficiently in a compiled language, doing this in rust also lets change our parallelization model for running multiple seed in Sabre. Just as in #8572 we added support for SabreSwap to run multiple parallel trials with different seeds this commit adds a layout_trials argument to SabreLayout to try multiple seeds in parallel. When this is used it parallelizes at the outer layer for each layout/routing combination and the total minimal swap count seed is used. So for example if you set swap_trials=5 and layout_trails=5 that will run 5 tasks in the threadpool with 5 different seeds for the outer layout run. Inside that every time sabre swap is run (which will be multiple times as part of layout plus the final routing run) it tries 5 different seeds for each execution serially inside that parallel task. This should hopefully further improve the quality of the transpiler output and better match expectations for users who were previously calling transpile() multiple times to emulate this behavior. Implements #9090 --- qiskit/__init__.py | 1 + .../transpiler/passes/layout/sabre_layout.py | 180 ++++++++++--- .../transpiler/passes/routing/sabre_swap.py | 84 ++++-- .../transpiler/preset_passmanagers/level0.py | 9 +- .../transpiler/preset_passmanagers/level1.py | 9 +- .../transpiler/preset_passmanagers/level2.py | 9 +- .../transpiler/preset_passmanagers/level3.py | 9 +- src/lib.rs | 2 + src/sabre_layout.rs | 241 ++++++++++++++++++ src/sabre_swap/mod.rs | 37 ++- tox.ini | 2 +- 11 files changed, 510 insertions(+), 73 deletions(-) create mode 100644 src/sabre_layout.rs diff --git a/qiskit/__init__.py b/qiskit/__init__.py index e272ec304d3d..2b6d68fdd709 100644 --- a/qiskit/__init__.py +++ b/qiskit/__init__.py @@ -37,6 +37,7 @@ sys.modules["qiskit._accelerate.nlayout"] = qiskit._accelerate.nlayout 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.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 diff --git a/qiskit/transpiler/passes/layout/sabre_layout.py b/qiskit/transpiler/passes/layout/sabre_layout.py index 3c3abdaf3dce..6e5618034b8f 100644 --- a/qiskit/transpiler/passes/layout/sabre_layout.py +++ b/qiskit/transpiler/passes/layout/sabre_layout.py @@ -15,22 +15,30 @@ import logging import numpy as np +import retworkx from qiskit.converters import dag_to_circuit from qiskit.transpiler.passes.layout.set_layout import SetLayout from qiskit.transpiler.passes.layout.full_ancilla_allocation import FullAncillaAllocation from qiskit.transpiler.passes.layout.enlarge_with_ancilla import EnlargeWithAncilla from qiskit.transpiler.passes.layout.apply_layout import ApplyLayout -from qiskit.transpiler.passes.routing import SabreSwap from qiskit.transpiler.passmanager import PassManager from qiskit.transpiler.layout import Layout -from qiskit.transpiler.basepasses import AnalysisPass +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 ( + Heuristic, + NeighborTable, +) +from qiskit.transpiler.passes.routing.sabre_swap import process_swaps, apply_gate +from qiskit.tools.parallel import CPU_COUNT logger = logging.getLogger(__name__) -class SabreLayout(AnalysisPass): +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 @@ -50,7 +58,13 @@ class SabreLayout(AnalysisPass): """ def __init__( - self, coupling_map, routing_pass=None, seed=None, max_iterations=3, swap_trials=None + self, + coupling_map, + routing_pass=None, + seed=None, + max_iterations=3, + swap_trials=None, + layout_trials=None, ): """SabreLayout initializer. @@ -71,6 +85,8 @@ def __init__( on the number of trials run. This option is mutually exclusive with the ``routing_pass`` argument and an error will be raised if both are used. + layout_trials (int): The number of random seed trials to run + layout with. Raises: TranspilerError: If both ``routing_pass`` and ``swap_trials`` are @@ -78,13 +94,24 @@ def __init__( """ super().__init__() self.coupling_map = coupling_map - if routing_pass is not None and swap_trials is not None: + self._neighbor_table = None + if self.coupling_map is not None: + self._neighbor_table = NeighborTable(retworkx.adjacency_matrix(self.coupling_map.graph)) + + if routing_pass is not None and (swap_trials is not None or layout_trials is not None): raise TranspilerError("Both routing_pass and swap_trials can't be set at the same time") self.routing_pass = routing_pass self.seed = seed self.max_iterations = max_iterations self.trials = swap_trials - self.swap_trials = swap_trials + if swap_trials is None: + self.swap_trials = CPU_COUNT + else: + self.swap_trials = swap_trials + if layout_trials is None: + self.layout_trials = CPU_COUNT + else: + self.layout_trials = layout_trials def run(self, dag): """Run the SabreLayout pass on `dag`. @@ -92,6 +119,10 @@ def run(self, dag): Args: dag (DAGCircuit): DAG to find layout for. + Returns: + DAGCircuit: The output dag if swap mapping was run + (otherwise the input dag is returned unmodified). + Raises: TranspilerError: if dag wider than self.coupling_map """ @@ -101,44 +132,119 @@ def run(self, dag): # Choose a random initial_layout. if self.seed is None: self.seed = np.random.randint(0, np.iinfo(np.int32).max) - rng = np.random.default_rng(self.seed) - physical_qubits = rng.choice(self.coupling_map.size(), len(dag.qubits), replace=False) - physical_qubits = rng.permutation(physical_qubits) - initial_layout = Layout({q: dag.qubits[i] for i, q in enumerate(physical_qubits)}) + if self.routing_pass is not None: + rng = np.random.default_rng(self.seed) + + physical_qubits = rng.choice(self.coupling_map.size(), len(dag.qubits), replace=False) + physical_qubits = rng.permutation(physical_qubits) + initial_layout = Layout({q: dag.qubits[i] for i, q in enumerate(physical_qubits)}) - if self.routing_pass is None: - self.routing_pass = SabreSwap( - self.coupling_map, "decay", seed=self.seed, fake_run=True, trials=self.swap_trials - ) - else: self.routing_pass.fake_run = True - # Do forward-backward iterations. - circ = dag_to_circuit(dag) - rev_circ = circ.reverse_ops() - for _ in range(self.max_iterations): - for _ in ("forward", "backward"): - pm = self._layout_and_route_passmanager(initial_layout) - new_circ = pm.run(circ) - - # Update initial layout and reverse the unmapped circuit. - pass_final_layout = pm.property_set["final_layout"] - final_layout = self._compose_layouts( - initial_layout, pass_final_layout, new_circ.qregs - ) - initial_layout = final_layout - circ, rev_circ = rev_circ, circ + # Do forward-backward iterations. + circ = dag_to_circuit(dag) + rev_circ = circ.reverse_ops() + for _ in range(self.max_iterations): + for _ in ("forward", "backward"): + pm = self._layout_and_route_passmanager(initial_layout) + new_circ = pm.run(circ) - # Diagnostics - logger.info("new initial layout") - logger.info(initial_layout) + # Update initial layout and reverse the unmapped circuit. + pass_final_layout = pm.property_set["final_layout"] + final_layout = self._compose_layouts( + initial_layout, pass_final_layout, new_circ.qregs + ) + initial_layout = final_layout + circ, rev_circ = rev_circ, circ - for qreg in dag.qregs.values(): - initial_layout.add_register(qreg) + # Diagnostics + logger.info("new initial layout") + logger.info(initial_layout) - self.property_set["layout"] = initial_layout - self.routing_pass.fake_run = False + for qreg in dag.qregs.values(): + initial_layout.add_register(qreg) + self.property_set["layout"] = initial_layout + self.routing_pass.fake_run = False + return dag + else: + dist_matrix = self.coupling_map.distance_matrix + original_qubit_indices = {bit: index for index, bit in enumerate(dag.qubits)} + original_clbit_indices = {bit: index for index, bit in enumerate(dag.clbits)} + + dag_list = [] + for node in dag.topological_op_nodes(): + cargs = {original_clbit_indices[x] for x in node.cargs} + if node.op.condition is not None: + for clbit in dag._bits_in_condition(node.op.condition): + cargs.add(original_clbit_indices[clbit]) + + dag_list.append( + ( + node._node_id, + [original_qubit_indices[x] for x in node.qargs], + cargs, + ) + ) + ((initial_layout, final_layout), swap_map, gate_order) = sabre_layout_and_routing( + len(dag.clbits), + dag_list, + self._neighbor_table, + dist_matrix, + Heuristic.Decay, + self.seed, + self.max_iterations, + self.swap_trials, + self.layout_trials, + ) + # Apply initial layout selected. + # this is a pseudo-pass manager to avoid the repeated round trip between + # dag and circuit and just use a dag + original_dag = dag + layout_dict = {} + num_qubits = len(dag.qubits) + for k, v in initial_layout.layout_mapping(): + if k < num_qubits: + layout_dict[dag.qubits[k]] = v + initital_layout = Layout(layout_dict) + self.property_set["layout"] = initital_layout + ancilla_pass = FullAncillaAllocation(self.coupling_map) + ancilla_pass.property_set = self.property_set + dag = ancilla_pass.run(dag) + enlarge_pass = EnlargeWithAncilla() + enlarge_pass.property_set = ancilla_pass.property_set + dag = enlarge_pass.run(dag) + apply_pass = ApplyLayout() + apply_pass.property_set = enlarge_pass.property_set + dag = apply_pass.run(dag) + # Apply sabre swap ontop of circuit with sabre layout + final_layout_mapping = final_layout.layout_mapping() + self.property_set["final_layout"] = Layout( + {dag.qubits[k]: v for (k, v) in final_layout_mapping} + ) + mapped_dag = dag.copy_empty_like() + canonical_register = dag.qregs["q"] + current_layout = Layout.generate_trivial_layout(canonical_register) + qubit_indices = {bit: idx for idx, bit in enumerate(canonical_register)} + layout_mapping = { + qubit_indices[k]: v for k, v in current_layout.get_virtual_bits().items() + } + original_layout = NLayout(layout_mapping, len(dag.qubits), self.coupling_map.size()) + for node_id in gate_order: + node = original_dag._multi_graph[node_id] + process_swaps( + swap_map, + node, + mapped_dag, + original_layout, + canonical_register, + False, + qubit_indices, + ) + apply_gate( + mapped_dag, node, original_layout, canonical_register, False, layout_dict + ) + return mapped_dag def _layout_and_route_passmanager(self, initial_layout): """Return a passmanager for a full layout and routing. diff --git a/qiskit/transpiler/passes/routing/sabre_swap.py b/qiskit/transpiler/passes/routing/sabre_swap.py index 64fa06282866..85af0192b7d1 100644 --- a/qiskit/transpiler/passes/routing/sabre_swap.py +++ b/qiskit/transpiler/passes/routing/sabre_swap.py @@ -242,33 +242,63 @@ def run(self, dag): if not self.fake_run: for node_id in gate_order: node = dag._multi_graph[node_id] - self._process_swaps(swap_map, node, mapped_dag, original_layout, canonical_register) - self._apply_gate(mapped_dag, node, original_layout, canonical_register) - return mapped_dag - return dag - - def _process_swaps(self, swap_map, node, mapped_dag, current_layout, canonical_register): - if node._node_id in swap_map: - for swap in swap_map[node._node_id]: - swap_qargs = [canonical_register[swap[0]], canonical_register[swap[1]]] - self._apply_gate( + process_swaps( + swap_map, + node, mapped_dag, - DAGOpNode(op=SwapGate(), qargs=swap_qargs), - current_layout, + original_layout, canonical_register, + self.fake_run, + self._qubit_indices, ) - current_layout.swap_logical(*swap) - - def _apply_gate(self, mapped_dag, node, current_layout, canonical_register): - new_node = self._transform_gate_for_layout(node, current_layout, canonical_register) - if self.fake_run: - return new_node - return mapped_dag.apply_operation_back(new_node.op, new_node.qargs, new_node.cargs) - - def _transform_gate_for_layout(self, op_node, layout, device_qreg): - """Return node implementing a virtual op on given layout.""" - mapped_op_node = copy(op_node) - mapped_op_node.qargs = tuple( - device_qreg[layout.logical_to_physical(self._qubit_indices[x])] for x in op_node.qargs - ) - return mapped_op_node + apply_gate( + mapped_dag, + node, + original_layout, + canonical_register, + self.fake_run, + self._qubit_indices, + ) + return mapped_dag + return dag + + +def process_swaps( + swap_map, + node, + mapped_dag, + current_layout, + canonical_register, + fake_run, + qubit_indices, +): + """Process swaps from SwapMap.""" + if node._node_id in swap_map: + for swap in swap_map[node._node_id]: + swap_qargs = [canonical_register[swap[0]], canonical_register[swap[1]]] + apply_gate( + mapped_dag, + DAGOpNode(op=SwapGate(), qargs=swap_qargs), + current_layout, + canonical_register, + fake_run, + qubit_indices, + ) + current_layout.swap_logical(*swap) + + +def apply_gate(mapped_dag, node, current_layout, canonical_register, fake_run, qubit_indices): + """Apply gate given the current layout.""" + new_node = transform_gate_for_layout(node, current_layout, canonical_register, qubit_indices) + if fake_run: + return new_node + return mapped_dag.apply_operation_back(new_node.op, new_node.qargs, new_node.cargs) + + +def transform_gate_for_layout(op_node, layout, device_qreg, qubit_indices): + """Return node implementing a virtual op on given layout.""" + mapped_op_node = copy(op_node) + mapped_op_node.qargs = tuple( + device_qreg[layout.logical_to_physical(qubit_indices[x])] for x in op_node.qargs + ) + return mapped_op_node diff --git a/qiskit/transpiler/preset_passmanagers/level0.py b/qiskit/transpiler/preset_passmanagers/level0.py index 78cffd88f24a..546c5c05589e 100644 --- a/qiskit/transpiler/preset_passmanagers/level0.py +++ b/qiskit/transpiler/preset_passmanagers/level0.py @@ -109,10 +109,17 @@ def _choose_layout_condition(property_set): "layout", layout_method, pass_manager_config, optimization_level=0 ) else: + + def _swap_mapped(property_set): + return property_set["final_layout"] is None + layout = PassManager() layout.append(_given_layout) layout.append(_choose_layout, condition=_choose_layout_condition) - layout += common.generate_embed_passmanager(coupling_map) + embed = common.generate_embed_passmanager(coupling_map) + layout.append( + [pass_ for x in embed.passes() for pass_ in x["passes"]], condition=_swap_mapped + ) routing = routing_pm else: layout = None diff --git a/qiskit/transpiler/preset_passmanagers/level1.py b/qiskit/transpiler/preset_passmanagers/level1.py index b3a58fe703d8..7b4ae1774e39 100644 --- a/qiskit/transpiler/preset_passmanagers/level1.py +++ b/qiskit/transpiler/preset_passmanagers/level1.py @@ -196,12 +196,19 @@ def _opt_control(property_set): "layout", layout_method, pass_manager_config, optimization_level=1 ) else: + + def _swap_mapped(property_set): + return property_set["final_layout"] is None + layout = PassManager() layout.append(_given_layout) layout.append(_choose_layout_0, condition=_choose_layout_condition) layout.append(_choose_layout_1, condition=_layout_not_perfect) layout.append(_improve_layout, condition=_vf2_match_not_found) - layout += common.generate_embed_passmanager(coupling_map) + embed = common.generate_embed_passmanager(coupling_map) + layout.append( + [pass_ for x in embed.passes() for pass_ in x["passes"]], condition=_swap_mapped + ) routing = routing_pm diff --git a/qiskit/transpiler/preset_passmanagers/level2.py b/qiskit/transpiler/preset_passmanagers/level2.py index 1edd917c1160..6b767511b21e 100644 --- a/qiskit/transpiler/preset_passmanagers/level2.py +++ b/qiskit/transpiler/preset_passmanagers/level2.py @@ -165,11 +165,18 @@ def _opt_control(property_set): "layout", layout_method, pass_manager_config, optimization_level=2 ) else: + + def _swap_mapped(property_set): + return property_set["final_layout"] is None + layout = PassManager() layout.append(_given_layout) layout.append(_choose_layout_0, condition=_choose_layout_condition) layout.append(_choose_layout_1, condition=_vf2_match_not_found) - layout += common.generate_embed_passmanager(coupling_map) + embed = common.generate_embed_passmanager(coupling_map) + layout.append( + [pass_ for x in embed.passes() for pass_ in x["passes"]], condition=_swap_mapped + ) routing = routing_pm else: layout = None diff --git a/qiskit/transpiler/preset_passmanagers/level3.py b/qiskit/transpiler/preset_passmanagers/level3.py index 46cc03e58917..fd930303fad6 100644 --- a/qiskit/transpiler/preset_passmanagers/level3.py +++ b/qiskit/transpiler/preset_passmanagers/level3.py @@ -192,11 +192,18 @@ def _opt_control(property_set): "layout", layout_method, pass_manager_config, optimization_level=3 ) else: + + def _swap_mapped(property_set): + return property_set["final_layout"] is None + layout = PassManager() layout.append(_given_layout) layout.append(_choose_layout_0, condition=_choose_layout_condition) layout.append(_choose_layout_1, condition=_vf2_match_not_found) - layout += common.generate_embed_passmanager(coupling_map) + embed = common.generate_embed_passmanager(coupling_map) + layout.append( + [pass_ for x in embed.passes() for pass_ in x["passes"]], condition=_swap_mapped + ) routing = routing_pm else: layout = None diff --git a/src/lib.rs b/src/lib.rs index 828f444188a9..9159d487e8b4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,7 @@ mod nlayout; mod optimize_1q_gates; mod pauli_exp_val; mod results; +mod sabre_layout; mod sabre_swap; mod sampled_exp_val; mod sparse_pauli_op; @@ -51,5 +52,6 @@ fn _accelerate(_py: Python<'_>, m: &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))?; Ok(()) } diff --git a/src/sabre_layout.rs b/src/sabre_layout.rs new file mode 100644 index 000000000000..5b38399f5cdc --- /dev/null +++ b/src/sabre_layout.rs @@ -0,0 +1,241 @@ +// 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)] + +use hashbrown::HashSet; +use ndarray::prelude::*; +use numpy::IntoPyArray; +use numpy::PyReadonlyArray2; +use pyo3::prelude::*; +use pyo3::wrap_pyfunction; +use pyo3::Python; +use rand::prelude::*; +use rand_pcg::Pcg64Mcg; +use rayon::prelude::*; +use retworkx_core::petgraph::prelude::*; + +use crate::getenv_use_multiple_threads; +use crate::nlayout::NLayout; +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}; + +#[pyfunction] +pub fn sabre_layout_and_routing( + py: Python, + num_clbits: usize, + dag_nodes: Vec<(usize, Vec, HashSet)>, + neighbor_table: &NeighborTable, + distance_matrix: PyReadonlyArray2, + heuristic: &Heuristic, + seed: u64, + max_iterations: usize, + num_swap_trials: usize, + num_layout_trials: usize, +) -> ([NLayout; 2], SwapMap, PyObject) { + let run_in_parallel = getenv_use_multiple_threads(); + let outer_rng = Pcg64Mcg::seed_from_u64(seed); + let seed_vec: Vec = outer_rng + .sample_iter(&rand::distributions::Standard) + .take(num_layout_trials) + .collect(); + let dist = distance_matrix.as_array(); + let result = if run_in_parallel { + seed_vec + .into_par_iter() + .enumerate() + .map(|(index, seed_trial)| { + ( + index, + layout_trial( + num_clbits, + dag_nodes.clone(), + neighbor_table, + &dist, + heuristic, + seed_trial, + max_iterations, + num_swap_trials, + ), + ) + }) + .min_by_key(|(index, result)| { + ( + result.1.map.values().map(|x| x.len()).sum::(), + *index, + ) + }) + .unwrap() + .1 + } else { + seed_vec + .into_iter() + .map(|seed_trial| { + layout_trial( + num_clbits, + dag_nodes.clone(), + neighbor_table, + &dist, + heuristic, + seed_trial, + max_iterations, + num_swap_trials, + ) + }) + .min_by_key(|result| result.1.map.values().map(|x| x.len()).sum::()) + .unwrap() + }; + (result.0, result.1, result.2.into_pyarray(py).into()) +} + +fn layout_trial( + num_clbits: usize, + mut dag_nodes: Vec<(usize, Vec, HashSet)>, + neighbor_table: &NeighborTable, + distance_matrix: &ArrayView2, + heuristic: &Heuristic, + seed: u64, + max_iterations: usize, + num_swap_trials: usize, +) -> ([NLayout; 2], SwapMap, Vec) { + // 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 = (0..num_physical_qubits).collect(); + physical_qubits.shuffle(&mut rng); + let mut phys_to_logic = vec![0; num_physical_qubits]; + physical_qubits + .iter() + .enumerate() + .for_each(|(logic, phys)| phys_to_logic[*phys] = logic); + let mut initial_layout = NLayout { + logic_to_phys: physical_qubits, + phys_to_logic, + }; + let mut rev_dag_nodes: Vec<(usize, Vec, HashSet)> = + dag_nodes.iter().rev().cloned().collect(); + for _iter in 0..max_iterations { + // forward and reverse + for _direction in 0..2 { + let dag = apply_layout(&dag_nodes, &initial_layout, num_physical_qubits, num_clbits); + let mut pass_final_layout = NLayout { + logic_to_phys: (0..num_physical_qubits).collect(), + phys_to_logic: (0..num_physical_qubits).collect(), + }; + build_swap_map_inner( + num_physical_qubits, + &dag, + neighbor_table, + distance_matrix, + heuristic, + seed, + &mut pass_final_layout, + num_swap_trials, + Some(false), + ); + let final_layout = compose_layout(&initial_layout, &pass_final_layout); + initial_layout = final_layout; + std::mem::swap(&mut dag_nodes, &mut rev_dag_nodes); + } + } + let layout_dag = apply_layout(&dag_nodes, &initial_layout, num_physical_qubits, num_clbits); + let mut final_layout = initial_layout.clone(); + let (swap_map, gate_order) = build_swap_map_inner( + num_physical_qubits, + &layout_dag, + neighbor_table, + distance_matrix, + heuristic, + seed, + &mut final_layout, + num_swap_trials, + Some(false), + ); + ([initial_layout, final_layout], swap_map, gate_order) +} + +fn apply_layout( + dag_nodes: &[(usize, Vec, HashSet)], + layout: &NLayout, + num_qubits: usize, + num_clbits: usize, +) -> SabreDAG { + let layout_dag_nodes: Vec<(usize, Vec, HashSet)> = dag_nodes + .iter() + .map(|(node_index, qargs, cargs)| { + let new_qargs: Vec = qargs.iter().map(|n| layout.logic_to_phys[*n]).collect(); + (*node_index, new_qargs, cargs.clone()) + }) + .collect(); + build_sabre_dag(layout_dag_nodes, num_qubits, num_clbits) +} + +fn build_sabre_dag( + layout_dag_nodes: Vec<(usize, Vec, HashSet)>, + num_qubits: usize, + num_clbits: usize, +) -> SabreDAG { + let mut dag: DiGraph<(usize, Vec), ()> = + Graph::with_capacity(layout_dag_nodes.len(), 2 * layout_dag_nodes.len()); + let mut first_layer = Vec::::new(); + let mut qubit_pos: Vec> = vec![None; num_qubits]; + let mut clbit_pos: Vec> = vec![None; num_clbits]; + for node in &layout_dag_nodes { + let qargs = &node.1; + let cargs = &node.2; + let gate_index = dag.add_node((node.0, qargs.clone())); + let mut is_front = true; + for x in qargs { + if let Some(predecessor) = qubit_pos[*x] { + is_front = false; + dag.add_edge(predecessor, gate_index, ()); + } + qubit_pos[*x] = Some(gate_index); + } + for x in cargs { + if let Some(predecessor) = clbit_pos[*x] { + is_front = false; + dag.add_edge(predecessor, gate_index, ()); + } + clbit_pos[*x] = Some(gate_index); + } + if is_front { + first_layer.push(gate_index); + } + } + SabreDAG { dag, first_layer } +} + +fn compose_layout(initial_layout: &NLayout, final_layout: &NLayout) -> NLayout { + let logic_to_phys = initial_layout + .logic_to_phys + .iter() + .map(|n| final_layout.logic_to_phys[*n]) + .collect(); + let phys_to_logic = final_layout + .phys_to_logic + .iter() + .map(|n| initial_layout.phys_to_logic[*n]) + .collect(); + + NLayout { + logic_to_phys, + phys_to_logic, + } +} + +#[pymodule] +pub fn sabre_layout(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_wrapped(wrap_pyfunction!(sabre_layout_and_routing))?; + Ok(()) +} diff --git a/src/sabre_swap/mod.rs b/src/sabre_swap/mod.rs index 1798a0dd0e39..d6009ac7af95 100644 --- a/src/sabre_swap/mod.rs +++ b/src/sabre_swap/mod.rs @@ -161,9 +161,38 @@ pub fn build_swap_map( seed: u64, layout: &mut NLayout, num_trials: usize, + run_in_parallel: Option, ) -> (SwapMap, PyObject) { - let run_in_parallel = getenv_use_multiple_threads(); let dist = distance_matrix.as_array(); + let (swap_map, gate_order) = build_swap_map_inner( + num_qubits, + dag, + neighbor_table, + &dist, + heuristic, + seed, + layout, + num_trials, + run_in_parallel, + ); + (swap_map, gate_order.into_pyarray(py).into()) +} + +pub fn build_swap_map_inner( + num_qubits: usize, + dag: &SabreDAG, + neighbor_table: &NeighborTable, + dist: &ArrayView2, + heuristic: &Heuristic, + seed: u64, + layout: &mut NLayout, + num_trials: usize, + run_in_parallel: Option, +) -> (SwapMap, Vec) { + let run_in_parallel = match run_in_parallel { + Some(run_in_parallel) => run_in_parallel, + None => getenv_use_multiple_threads(), + }; let coupling_graph: DiGraph<(), ()> = cmap_from_neighor_table(neighbor_table); let outer_rng = Pcg64Mcg::seed_from_u64(seed); let seed_vec: Vec = outer_rng @@ -181,7 +210,7 @@ pub fn build_swap_map( num_qubits, dag, neighbor_table, - &dist, + dist, &coupling_graph, heuristic, seed_trial, @@ -205,7 +234,7 @@ pub fn build_swap_map( num_qubits, dag, neighbor_table, - &dist, + dist, &coupling_graph, heuristic, seed_trial, @@ -220,7 +249,7 @@ pub fn build_swap_map( SwapMap { map: result.out_map, }, - result.gate_order.into_pyarray(py).into(), + result.gate_order, ) } diff --git a/tox.ini b/tox.ini index 59cc9f6748b6..8dbd5b2d32c2 100644 --- a/tox.ini +++ b/tox.ini @@ -14,7 +14,7 @@ setenv = QISKIT_SUPRESS_PACKAGING_WARNINGS=Y QISKIT_TEST_CAPTURE_STREAMS=1 QISKIT_PARALLEL=FALSE -passenv = RAYON_NUM_THREADS OMP_NUM_THREADS QISKIT_PARALLEL RUST_BACKTRACE SETUPTOOLS_ENABLE_FEATURES +passenv = RAYON_NUM_THREADS OMP_NUM_THREADS QISKIT_PARALLEL RUST_BACKTRACE SETUPTOOLS_ENABLE_FEATURES QISKIT_IN_PARALLEL deps = setuptools_rust # This is work around for the bug of tox 3 (see #8606 for more details.) -r{toxinidir}/requirements.txt -r{toxinidir}/requirements-dev.txt