diff --git a/crates/transpiler/src/passes/sabre/neighbors.rs b/crates/transpiler/src/passes/sabre/neighbors.rs index 67eb597a4152..8d664f198199 100644 --- a/crates/transpiler/src/passes/sabre/neighbors.rs +++ b/crates/transpiler/src/passes/sabre/neighbors.rs @@ -26,8 +26,8 @@ use rustworkx_core::petgraph::visit::*; /// small that a linear search is faster). #[derive(Clone, Debug)] pub struct Neighbors { - neighbors: Vec, - partition: Vec, + pub(crate) neighbors: Vec, + pub(crate) partition: Vec, } impl Neighbors { /// Construct the neighbor adjacency table from a coupling graph. diff --git a/crates/transpiler/src/passes/sabre/route.rs b/crates/transpiler/src/passes/sabre/route.rs index b417625ae3f3..724ec3a1db98 100644 --- a/crates/transpiler/src/passes/sabre/route.rs +++ b/crates/transpiler/src/passes/sabre/route.rs @@ -17,6 +17,7 @@ use std::num::NonZero; use numpy::{PyArray2, ToPyArray}; use pyo3::prelude::*; +use pyo3::types::PyDict; use pyo3::Python; use hashbrown::HashSet; @@ -315,6 +316,50 @@ impl RoutingTarget { pub struct PyRoutingTarget(pub Option); #[pymethods] impl PyRoutingTarget { + #[new] + fn py_new() -> Self { + PyRoutingTarget(None) + } + + fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult> { + let out_dict = PyDict::new(py); + out_dict.set_item( + "neighbors", + self.0.as_ref().map(|x| x.neighbors.neighbors.clone()), + )?; + out_dict.set_item( + "partition", + self.0.as_ref().map(|x| x.neighbors.partition.clone()), + )?; + Ok(out_dict) + } + + fn __setstate__(&mut self, value: Bound) -> PyResult<()> { + let neighbors_array: Option> = value + .get_item("neighbors")? + .map(|x| x.extract()) + .transpose()?; + if let Some(neighbors_array) = neighbors_array { + let partition: Vec = value + .get_item("partition")? + .map(|x| x.extract()) + .transpose()? + .unwrap(); + let neighbors = Neighbors { + neighbors: neighbors_array, + partition, + }; + if self.0.is_none() { + self.0 = Some(RoutingTarget::from_neighbors(neighbors)); + } else { + self.0.as_mut().unwrap().distance = + distance_matrix(&neighbors, usize::MAX, f64::NAN); + self.0.as_mut().unwrap().neighbors = neighbors; + } + } + Ok(()) + } + #[staticmethod] pub(crate) fn from_target(target: &Target) -> PyResult { let coupling = match target.coupling_graph() { diff --git a/releasenotes/notes/fix-sabre-swap-pickle-a1215f58853f81d1.yaml b/releasenotes/notes/fix-sabre-swap-pickle-a1215f58853f81d1.yaml new file mode 100644 index 000000000000..62a0e9d2776f --- /dev/null +++ b/releasenotes/notes/fix-sabre-swap-pickle-a1215f58853f81d1.yaml @@ -0,0 +1,7 @@ +--- +fixes: + - | + Fixed an issue with :mod:`pickle` support for the :class:`.SabreSwap` where a + :class:`.SabreSwap` instance would error when being pickled after the + :meth:`.SabreSwap.run` method was run. + Fixed `#15071 `__. diff --git a/test/python/transpiler/test_sabre_swap.py b/test/python/transpiler/test_sabre_swap.py index f8fe9d5a6937..fa9aa8cec740 100644 --- a/test/python/transpiler/test_sabre_swap.py +++ b/test/python/transpiler/test_sabre_swap.py @@ -14,6 +14,9 @@ import unittest import itertools +import pickle +from copy import deepcopy +import io import ddt import numpy.random @@ -101,6 +104,66 @@ def looping_circuit(uphill_swaps=1, additional_local_minimum_gates=0): class TestSabreSwap(QiskitTestCase): """Tests the SabreSwap pass.""" + def test_sabre_swap_pickle(self): + """Test the pass can be pickled.""" + coupling = CouplingMap.from_ring(5) + target = Target.from_configuration(["u", "cx"], coupling_map=coupling) + sabre_swap = SabreSwap(target, "lookahead", seed=42, trials=1024) + with io.BytesIO() as buf: + pickle.dump(sabre_swap, buf) + buf.seek(0) + output = pickle.load(buf) + self.assertIsInstance(output, SabreSwap) + self.assertIsNone(output._routing_target) + self.assertEqual(sabre_swap.heuristic, output.heuristic) + self.assertEqual(sabre_swap.trials, output.trials) + self.assertEqual(sabre_swap.fake_run, output.fake_run) + + test_circuit = QuantumCircuit(5) + test_circuit.cx(0, 1) + test_circuit.cx(0, 2) + test_circuit.cx(0, 3) + test_circuit.cx(0, 4) + before_result = sabre_swap(test_circuit) + with io.BytesIO() as buf: + pickle.dump(sabre_swap, buf) + buf.seek(0) + output = pickle.load(buf) + self.assertIsInstance(output, SabreSwap) + self.assertIsNotNone(output._routing_target) + self.assertEqual(sabre_swap.heuristic, output.heuristic) + self.assertEqual(sabre_swap.trials, output.trials) + self.assertEqual(sabre_swap.fake_run, output.fake_run) + after_result = output(test_circuit) + self.assertEqual(before_result, after_result) + + def test_sabre_swap_deepcopy(self): + """Test the pass can be deepcopied.""" + coupling = CouplingMap.from_ring(5) + target = Target.from_configuration(["u", "cx"], coupling_map=coupling) + sabre_swap = SabreSwap(target, "lookahead", seed=42, trials=1024) + output = deepcopy(sabre_swap) + self.assertIsInstance(output, SabreSwap) + self.assertIsNone(output._routing_target) + self.assertEqual(sabre_swap.heuristic, output.heuristic) + self.assertEqual(sabre_swap.trials, output.trials) + self.assertEqual(sabre_swap.fake_run, output.fake_run) + + test_circuit = QuantumCircuit(5) + test_circuit.cx(0, 1) + test_circuit.cx(0, 2) + test_circuit.cx(0, 3) + test_circuit.cx(0, 4) + before_result = sabre_swap(test_circuit) + output = deepcopy(sabre_swap) + self.assertIsInstance(output, SabreSwap) + self.assertIsNotNone(output._routing_target) + self.assertEqual(sabre_swap.heuristic, output.heuristic) + self.assertEqual(sabre_swap.trials, output.trials) + self.assertEqual(sabre_swap.fake_run, output.fake_run) + after_result = output(test_circuit) + self.assertEqual(before_result, after_result) + def test_trivial_case(self): """Test that an already mapped circuit is unchanged. ┌───┐┌───┐