From 364ffa1eb824a8a5a0e400c1a5f2d725f0d316cb Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Thu, 25 Sep 2025 15:17:20 +0530 Subject: [PATCH] Fix pickling of SabreSwap (#15074) This commit fixes the support for pickling SabreSwap. In #14317 a new RoutingTarget rust struct was added to encapsulate the target details, and this was exposed to Python so that the Python class for the transpiler pass was able to reuse the object between multiple runs. However, this new type didn't implement pickle support and it would cause a failure when trying to pickle a SabreSwap instance that had a routing target populated. This commit fixes this oversight and implements pickle support for the RoutingTarget so that SabreSwap can always be pickled. Fixes #15071 (cherry picked from commit a816e64bb972ad31a7e43cb42ef905e452610dde) --- .../transpiler/src/passes/sabre/neighbors.rs | 4 +- crates/transpiler/src/passes/sabre/route.rs | 45 +++++++++++++ ...ix-sabre-swap-pickle-a1215f58853f81d1.yaml | 7 +++ test/python/transpiler/test_sabre_swap.py | 63 +++++++++++++++++++ 4 files changed, 117 insertions(+), 2 deletions(-) create mode 100644 releasenotes/notes/fix-sabre-swap-pickle-a1215f58853f81d1.yaml 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. ┌───┐┌───┐