Fix pickling of SabreSwap (Qiskit#15074) (Qiskit#15076)

This commit fixes the support for pickling SabreSwap. In Qiskit#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 Qiskit#15071

(cherry picked from commit a816e64)

Co-authored-by: Matthew Treinish <[email protected]>

Author: mergify[bot]

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<PhysicalQubit>,
-    partition: Vec<usize>,
+    pub(crate) neighbors: Vec<PhysicalQubit>,
+    pub(crate) partition: Vec<usize>,
 }
 impl Neighbors {
     /// Construct the neighbor adjacency table from a coupling graph.

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<RoutingTarget>);
 #[pymethods]
 impl PyRoutingTarget {
+    #[new]
+    fn py_new() -> Self {
+        PyRoutingTarget(None)
+    }
+
+    fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
+        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<PyDict>) -> PyResult<()> {
+        let neighbors_array: Option<Vec<PhysicalQubit>> = value
+            .get_item("neighbors")?
+            .map(|x| x.extract())
+            .transpose()?;
+        if let Some(neighbors_array) = neighbors_array {
+            let partition: Vec<usize> = 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<Self> {
         let coupling = match target.coupling_graph() {

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 <https://github.com/Qiskit/qiskit/issues/15071>`__.

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.
                   ┌───┐┌───┐