Skip to content

Commit

Permalink
Oxidize SabreLayout pass (#9116)
Browse files Browse the repository at this point in the history 
* 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

* Use deepcopy for coupling map copy

Previously this PR was using copy() to copy the coupling map before we
mutated it to be symmetric (a requirement for the sabre algorithm).
However, this modification of the object was leaking out causing test
failures. This commit switches it to a deepcopy to ensure there are no
shared references (and a comment added to explain it's needed).

* Fix failing unitary synthesis tests

This PR branch modifies the default behavior of the SabreLayout pass so
it is now a transformation pass that computes a layout, applies it, and
then performs routing. This means when using sabre layout in a custom
pass manager we no longer need to embed a layout after computing the
layout. The failing unitary synthesis tests were using a custom pass
manager and trying to apply the layout again after SabreLayout already
did. This commit just removes this now unecessary steps from the test
code.

* Add release note

* Run BarrierBeforeMeasurement before new SabreLayout

Now that the routing stage is integrated into the SabreLayout pass we
should be running the BarrierBeforeMeasurement pass prior to layout in
the preset pass managers instead of before routing. The goal of the pass
is to prevent the routing algorithms for accidentally reusing a qubit
after a final measurement which would be invalid by inserting a barrier
before the measurements to ensure all qubits are swap mapped prior to
adding the measurements during routing. While this might not strictly be
necessary (it didn't affect any test output) it feels like best practice
to ensure we're doing this prior to potentially routing to prevent
issues.

* Improve docstrings

* Set a fixed number of layout trials in preset pass managers

For reproducible results with a fixed seed this commit sets a fixed
number of layout_trials for the SabreLayout pass in the preset pass
managers. If we did not set a fixed value than the output of the
transpiler with a fixed seed will vary based on the number of
physical cores that is running the compilation. To start
optimization levels 0 and 1 use 5, level 2 uses 10, and level
3 uses 20 which matches the swap_trials argument we used. This is just a
starting point, we can adjust these values later if needed.

* Update tests for layout changes

This commit updates the tests which are checking exact layouts with a
fixed seed when running SabreLayout. The changes to SabreLayout breaks
exact seed reproducibility from the earlier version of the pass. So we
need to update these tests for their new layout assignment from the
improved pass. One exception is a test which was trying to assert that
transpile() preserves a swap if it's in the basis set. However, the new
layout and routing output from SabreLayout for that test was resulting
in all the swaps getting optimized away at optimization level 3
(resulting in 13 cx gates instead of ~4 cx gates and 5 swaps before,
which would be more efficient on real hardware). So the test was removed
and only run at lower optimziation levels.

* Set a fixed number of layout trials in SabreLayout tests

The dedicated tests for SabreLayout were not running a fixed number of
trials. This was causing a different layout to be returned in tests when
run across multiple systems as the number of trials defaults to the
number of physical CPUs. This commit fixes the trial count to the number
of cores on the local system where the layout was updated. This should
fix the non-determinism in the tests causing failures in CI and on
different local systems.

* Run SabreSwap in parallel if only a single layout trial

If there is only a single layout trial being run we don't have to worry
about trying to do too much work in parallel at once by parallelizing
the inner sabre swap execution. This commit updates the threading logic
to enable running the inner sabre swap trials in parallel if there is
only a single layout trial.

* Remove duplicated SabreDAG creation

* Correctly apply selected layout on dag nodes

This commit corrects a bug in the PR branch that was caused by applying
the selected initial layout in a trial to the swapped order node list.
This was causing unexpected results when applying the circuit because
the intent was to apply it only to the original input not the reversed
input.

* Remove unnecessary clone from serial layout trials

In the case we're evaluating the layout trials serially instead of in a
parallel iterator we don't need to clone the dag nodes list. This is
because nothing will be modifying it in parallel, so we don't need a
thread local copy. Each call to layout_trial() will keep the dag nodes
vector intact (see previous commit for fixing this) so it can just be
passed by reference if there are no parallel threads involved.

* Fix seed setup when no user seed specified

This commit fixes an issue prevent seed randomization when no seed is
specified. On subsequent uses of a pass SabreLayout would not randomize
the seed between runs because it was setting the seed to instance state.
This commit fixes this issue by relying on initializing the RNG from
entropy each time run() is called if no user specified seed is provided.

* Start from trivial layout for routing stage

This commit fixes the routing run to run from a trivial layout instead
of the initial layout. By the time we do final routing for a trial we've
already applied the selected initial layout to the SabreDAG. So the
correct layout to use for running final swap mapping is a trivial layout
where logical bit 0 is at physical bit 0. Using initial layout twice
means we end up mapping more than is needed resulting in incorrect
results.

* Revert "Correctly apply selected layout on dag nodes"

This change was incorrect, the output was already in the correct order
and this was causing the behavior it strived to fix. This commit reverts
the addition of the extra mem::swap() call to fix things.

This reverts commit d98ef6c.

* Deduplicate NLayout trivial layout creation

This commit deduplicates the trivial layout generation for the NLayout
class. Previously there were a few places both in rust and python that
sabre layout was manually generating a trivial NLayout object. THis
commit adds a static method to the NLayout class that allows both Python
and Rust to easily create a new trivial NLayout object instead of
manually creating the object.

* Fix fixed layout tests after updates

Since more recent commits fixed a few bugs in the behavior of the
SabreLayout pass, the previously updated fixed layout tests were no
longer correct. This commit updates the tests which were now failing
because the layout changed again after fixing bugs in the new pass code.

* Try nesting parallelism in the sabres

Looking at profiles for running the new SabreLayout pass, as expected
the runtime of the rust SabreSwap routines is dominating. This is
because we've basically serialized the sabre swap routines and are
running multiple seed trials. As an experiment this commit sets the
inner SabreSwap routines to run in parallel too. Since the rayon
algorithm uses a work stealing algorithm this hopefully shouldn't cause
too much extra overhead, especially because the layout trials are quite
fast. This ideally means we're just scheduling each sabre swap trial in
a big parallel work queue and rayon does the rest of the magic to figure
out how to execute things. Initial testing is showing an improvement for
large circuits and a more modest improvement for more modest circuits.

* Add skip_routing argument to preserve custom user provided routing

This commit adds a new argument, skip_routing, to the SabreLayout
constructor. The intent of this new option is to enable mixing custom
routing_method user arguments with SabreLayout in it's new accelerated
mode of operation. In the earlier commits no matter what users specified
the preset pass manager construction would use sabreswap for routing as
it was run internally as part of layout. This meant doing something
like:

transpile(qc, backend, routing_method='stochastic')

would really run SabreSwap which is clearly not the user intent. To
provide the layout benefits with multiple seed trials this new argument
allows disabling the application of the routing found. This comes with a
runtime penalty because effectively we end up running routing twice and
only using one of the results. But for custom user provided methods or
plugins this seems like a reasonable tradeoff.

* Fix typo in docstring

* Update random seed usage in rust code

In #9132 we updated the random seed parameters in the rust code for
sabre swap to make the seed optional and default to initializing from
entropy if it's not specified. This commit updates the usage to account
for this change on main.

* s/retworkx/rustworkx/g

* Add alternate constructor for NLayout from a logic_to_phys vec

This commit adds a new constructor method to the NLayout class that
builds an NLayout object from just a logic_to_phys Vec. This constructor
can be accessed from either rust or python (although it's not as
efficient from Python). This is used to simplify some of the SabreLayout
rust code that was doing this inline manually.

* Move layout embedding into a method

This commit moves the code the optimized SabreLayout pass was using to
embed the found layout from the Rust code into a method. This will make
it easier to refactor later if a more efficient pass manager path is
added.

* Simplify pass logic and update comments

This commit removes an unnecessary else branch in the SabreLayout.run()
code to make it slightly easier to read. At the same time some comments
are updated to better explain the logic of the code.

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
  • Loading branch information
@mtreinish @mergify
mtreinish and mergify[bot] authored Dec 8, 2022
1 parent 45db5ec commit 7243c71
Show file tree
Hide file tree
Showing 16 changed files with 791 additions and 142 deletions.
1 change: 1 addition & 0 deletions qiskit/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
236 changes: 193 additions & 43 deletions qiskit/transpiler/passes/layout/sabre_layout.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,24 +13,33 @@
"""Layout selection using the SABRE bidirectional search approach from Li et al.
"""

import copy
import logging
import numpy as np
import rustworkx as rx

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
Expand All @@ -42,6 +51,21 @@ class SabreLayout(AnalysisPass):
This method exploits the reversibility of quantum circuits, and tries to
include global circuit information in the choice of initial_layout.
By default this pass will run both layout and routing and will transform the
circuit so that the layout is applied to the input dag (meaning that the output
circuit will have ancilla qubits allocated for unused qubits on the coupling map
and the qubits will be reordered to match the mapped physical qubits) and then
routing will be applied (inserting :class:`~.SwapGate`s to account for limited
connectivity). This is unlike most other layout passes which are :class:`~.AnalysisPass`
objects and just find an initial layout and set that on the property set. This is
done because by default the pass will run parallel seed trials with different random
seeds for selecting the random initial layout and then selecting the routed output
which results in the least number of swap gates needed.
You can use the ``routing_pass`` argument to have this pass operate as a typical
layout pass. When specified this will use the specified routing pass to select an
initial layout only and will not run multiple seed trials.
**References:**
[1] Li, Gushu, Yufei Ding, and Yuan Xie. "Tackling the qubit mapping problem
Expand All @@ -50,15 +74,26 @@ 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,
skip_routing=False,
):
"""SabreLayout initializer.
Args:
coupling_map (Coupling): directed graph representing a coupling map.
routing_pass (BasePass): the routing pass to use while iterating.
This is mutually exclusive with the ``swap_trials`` argument and
if both are set an error will be raised.
If specified this pass operates as an :class:`~.AnalysisPass` and
will only populate the ``layout`` field in the property set and
the input dag is returned unmodified. This argument is mutually
exclusive with the ``swap_trials`` and the ``layout_trials``
arguments and if this is specified at the same time as either
argument an error will be raised.
seed (int): seed for setting a random first trial layout.
max_iterations (int): number of forward-backward iterations.
swap_trials (int): The number of trials to run of
Expand All @@ -71,74 +106,189 @@ 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. When > 1 the trial that resuls in the output with
the fewest swap gates will be selected. If this is not specified
(and ``routing_pass`` is not set) then the number of local
physical CPUs will be used as the default value. This option is
mutually exclusive with the ``routing_pass`` argument and an error
will be raised if both are used.
skip_routing (bool): If this is set ``True`` and ``routing_pass`` is not used
then routing will not be applied to the output circuit. Only the layout
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.
Raises:
TranspilerError: If both ``routing_pass`` and ``swap_trials`` are
specified
TranspilerError: If both ``routing_pass`` and ``swap_trials`` or
both ``routing_pass`` and ``layout_trials`` are specified
"""
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:
if not self.coupling_map.is_symmetric:
# deepcopy is needed here to avoid modifications updating
# shared references in passes which require directional
# constraints
self.coupling_map = copy.deepcopy(self.coupling_map)
self.coupling_map.make_symmetric()
self._neighbor_table = NeighborTable(rx.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
self.skip_routing = skip_routing

def run(self, dag):
"""Run the SabreLayout pass on `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
"""
if len(dag.qubits) > self.coupling_map.size():
raise TranspilerError("More virtual qubits exist than physical.")

# 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)
if self.routing_pass is not None:
if self.seed is None:
seed = np.random.randint(0, np.iinfo(np.int32).max)
else:
seed = self.seed
rng = np.random.default_rng(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)})
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)

# 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

# Diagnostics
logger.info("new initial layout")
logger.info(initial_layout)

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
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)}

# Diagnostics
logger.info("new initial layout")
logger.info(initial_layout)
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])

for qreg in dag.qregs.values():
initial_layout.add_register(qreg)
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.
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
# If skip_routing is set then return the layout in the property set
# and throwaway the extra work we did to compute the swap map
if self.skip_routing:
return dag
# After this point the pass is no longer an analysis pass and the
# output circuit returned is transformed with the layout applied
# and swaps inserted
dag = self._apply_layout_no_pass_manager(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"]
qubit_indices = {bit: idx for idx, bit in enumerate(canonical_register)}
original_layout = NLayout.generate_trivial_layout(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

self.property_set["layout"] = initial_layout
self.routing_pass.fake_run = False
def _apply_layout_no_pass_manager(self, dag):
"""Apply and embed a layout into a dagcircuit without using a ``PassManager`` to
avoid circuit<->dag conversion.
"""
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)
return dag

def _layout_and_route_passmanager(self, initial_layout):
"""Return a passmanager for a full layout and routing.
Expand Down
Loading

0 comments on commit 7243c71

Please sign in to comment.