Optimize shortestpath in SABRE

src/qibo/transpiler/blocks.py

@@ -169,6 +169,15 @@ def remove_block(self, block: "Block"):
                 "The block you are trying to remove is not present in the circuit blocks.",
             )
 
+    def return_last_block(self):
+        """Return the last block in the circuit blocks."""
+        if len(self.block_list) == 0:
+            raise_error(
+                BlockingError,
+                "No blocks found in the circuit blocks.",
+            )
+        return self.block_list[-1]
+
 
 def block_decomposition(circuit: Circuit, fuse: bool = True):
     """Decompose a circuit into blocks of gates acting on two qubits.

src/qibo/transpiler/router.py

@@ -257,6 +257,32 @@ def update(self, swap: tuple):
         ), self._circuit_logical.index(swap[1])
         self._circuit_logical[idx_0], self._circuit_logical[idx_1] = swap[1], swap[0]
 
+    def undo(self, swap: tuple):
+        """Undo the last swap.
+
+        If the last block is not a swap or does not match the swap to undo, an error is raised.
+        Method works in-place.
+
+        Args:
+            swap (tuple): tuple containing the logical qubits to be swapped.
+        """
+        physical_swap = self.logical_to_physical(swap, index=True)
+        last_swap_block = self._routed_blocks.return_last_block()
+        if last_swap_block.gates[0].__class__ != gates.SWAP or sorted(
+            last_swap_block.qubits
+        ) != sorted(physical_swap):
+            raise_error(
+                ConnectivityError,
+                "The last block does not match the swap to undo.",
+            )
+        self._routed_blocks.remove_block(last_swap_block)
+        self._swaps -= 1
+
+        idx_0, idx_1 = self._circuit_logical.index(
+            swap[0]
+        ), self._circuit_logical.index(swap[1])
+        self._circuit_logical[idx_0], self._circuit_logical[idx_1] = swap[1], swap[0]
+
     def get_logical_qubits(self, block: Block):
         """Returns the current logical qubits where a block is acting on.
 
@@ -659,8 +685,7 @@ def __init__(
         self.circuit = None
         self._memory_map = None
         self._final_measurements = None
-        self._temporary_added_swaps = 0
-        self._saved_circuit = None
+        self._temp_added_swaps = []
         random.seed(seed)
 
     def __call__(self, circuit: Circuit, initial_layout: dict):
@@ -674,7 +699,6 @@ def __call__(self, circuit: Circuit, initial_layout: dict):
             (:class:`qibo.models.circuit.Circuit`, dict): routed circuit and final layout.
         """
         self._preprocessing(circuit=circuit, initial_layout=initial_layout)
-        self._saved_circuit = deepcopy(self.circuit)
         longest_path = np.max(self._dist_matrix)
 
         while self._dag.number_of_nodes() != 0:
@@ -687,9 +711,12 @@ def __call__(self, circuit: Circuit, initial_layout: dict):
             # If the number of added swaps is too high, the algorithm is stuck.
             # Reset the circuit to the last saved state and make the nearest gate executable by manually adding SWAPs.
             if (
-                self._temporary_added_swaps > self.swap_threshold * longest_path
+                len(self._temp_added_swaps) > self.swap_threshold * longest_path
             ):  # threshold is arbitrary
-                self.circuit = deepcopy(self._saved_circuit)
+                while self._temp_added_swaps:
+                    swap = self._temp_added_swaps.pop()
+                    self.circuit.undo(swap)
+                self._temp_added_swaps = []
                 self._shortest_path_routing()
 
         circuit_kwargs = circuit.init_kwargs
@@ -800,7 +827,7 @@ def _find_new_mapping(self):
         for qubit in self.circuit.logical_to_physical(best_candidate, index=True):
             self._delta_register[qubit] += self.delta
         self.circuit.update(best_candidate)
-        self._temporary_added_swaps += 1
+        self._temp_added_swaps.append(best_candidate)
 
     def _compute_cost(self, candidate: int):
         """Compute the cost associated to a possible SWAP candidate."""
@@ -897,8 +924,7 @@ def _execute_blocks(self, blocklist: list):
         self._update_front_layer()
         self._memory_map = []
         self._delta_register = [1.0 for _ in self._delta_register]
-        self._temporary_added_swaps = 0
-        self._saved_circuit = deepcopy(self.circuit)
+        self._temp_added_swaps = []
 
     def _shortest_path_routing(self):
         """Route a gate in the front layer using the shortest path. This method is executed when the standard SABRE fails to find an optimized solution.

tests/test_transpiler_blocks.py

@@ -367,3 +367,27 @@ def test_block_on_qubits():
     assert new_block.gates[2].qubits == (3,)
     assert new_block.gates[3].qubits == (3, 2)
     assert new_block.gates[4].qubits == (3,)
+
+
+def test_return_last_block():
+    circ = Circuit(4)
+    circ.add(gates.CZ(0, 1))
+    circ.add(gates.CZ(1, 3))
+    circ.add(gates.CZ(1, 2))
+    circ.add(gates.CZ(2, 3))
+    circuit_blocks = CircuitBlocks(circ)
+    last_block = circuit_blocks.return_last_block()
+    assert_gates_equality(last_block.gates, [gates.CZ(2, 3)])
+
+    circuit_blocks.remove_block(last_block)
+    last_block_2 = circuit_blocks.return_last_block()
+    assert_gates_equality(last_block_2.gates, [gates.CZ(1, 2)])
+
+
+def test_return_last_block_error():
+    circ = Circuit(4)
+    circuit_blocks = CircuitBlocks(circ)
+
+    # No blocks in the circuit
+    with pytest.raises(BlockingError):
+        last_block = circuit_blocks.return_last_block()

tests/test_transpiler_router.py

@@ -9,6 +9,7 @@
 from qibo.models import Circuit
 from qibo.quantum_info.random_ensembles import random_unitary
 from qibo.transpiler._exceptions import ConnectivityError
+from qibo.transpiler.blocks import Block
 from qibo.transpiler.optimizer import Preprocessing
 from qibo.transpiler.pipeline import (
     assert_circuit_equivalence,
@@ -472,3 +473,49 @@ def test_star_router(nqubits, depth, middle_qubit, measurements, unitaries):
         final_map=final_qubit_map,
         initial_map=initial_layout,
     )
+
+
+def test_undo():
+    circ = Circuit(4)
+    initial_layout = {"q0": 0, "q1": 1, "q2": 2, "q3": 3}
+    circuit_map = CircuitMap(initial_layout=initial_layout, circuit=circ)
+
+    # Two SWAP gates are added
+    circuit_map.update((1, 2))
+    circuit_map.update((2, 3))
+    assert circuit_map._circuit_logical == [0, 3, 1, 2]
+    assert len(circuit_map._routed_blocks.block_list) == 2
+
+    # Undo the last SWAP gate
+    circuit_map.undo((2, 3))
+    assert circuit_map._circuit_logical == [0, 2, 1, 3]
+    assert circuit_map._swaps == 1
+    assert len(circuit_map._routed_blocks.block_list) == 1
+
+    # Undo the first SWAP gate
+    circuit_map.undo((1, 2))
+    assert circuit_map._circuit_logical == [0, 1, 2, 3]
+    assert circuit_map._swaps == 0
+    assert len(circuit_map._routed_blocks.block_list) == 0
+
+
+def test_undo_error():
+    circ = Circuit(4)
+    initial_layout = {"q0": 0, "q1": 1, "q2": 2, "q3": 3}
+    circuit_map = CircuitMap(initial_layout=initial_layout, circuit=circ)
+
+    circuit_map.update((1, 2))
+    circuit_map.update((2, 3))
+
+    # The last block is a SWAP gate on qubits (2, 3)
+    with pytest.raises(ConnectivityError):
+        circuit_map.undo((1, 2))
+
+    circ_cz = Circuit(4)
+    circuit_map_cz = CircuitMap(initial_layout=initial_layout, circuit=circ_cz)
+    block = Block(qubits=(0, 1), gates=[gates.CZ(0, 1)])
+    circuit_map_cz._routed_blocks.add_block(block)
+
+    # The last block is a CZ gate
+    with pytest.raises(ConnectivityError):
+        circuit_map_cz.undo((0, 1))