Speeds up RouteCQC by encoding qubits as integers

cirq-core/cirq/transformers/routing/mapping_manager.py

@@ -14,154 +14,196 @@
 
 """Manages the mapping from logical to physical qubits during a routing procedure."""
 
-from typing import Dict, Sequence, TYPE_CHECKING
+from typing import List, Dict, Sequence, TYPE_CHECKING
 import networkx as nx
 
-from cirq import protocols, value
+import numpy as np
 
 if TYPE_CHECKING:
     import cirq
 
 
-@value.value_equality
 class MappingManager:
     """Class that manages the mapping from logical to physical qubits.
 
-    Convenience methods over distance and mapping queries of the physical qubits are also provided.
-    All such public methods of this class expect logical qubits.
+    For efficiency, the mapping manager maps all logical and physical qubits to integers, and
+    maintains a mapping from logical qubit integers to physical qubit integers. This speedup is
+    important to avoid qubit hashing in hot-paths like querying distance of two logical qubits
+    on the device (via `dist_on_device` method).
+
+    All public methods of this class expect logical qubits (or corresponding integers that the
+    logical qubits are mapped to, via `self.logical_qid_to_int` map).
     """
 
     def __init__(
         self, device_graph: nx.Graph, initial_mapping: Dict['cirq.Qid', 'cirq.Qid']
     ) -> None:
         """Initializes MappingManager.
 
-        Sorts the nodes and edges in the device graph to guarantee graph equality. If undirected,
-        also sorts the nodes within each edge.
-
         Args:
             device_graph: connectivity graph of qubits in the hardware device.
             initial_mapping: the initial mapping of logical (keys) to physical qubits (values).
         """
-        # make sure edge insertion order is the same amongst equivalent graphs.
-        if nx.is_directed(device_graph):
-            self.device_graph = nx.DiGraph()
-            self.device_graph.add_nodes_from(sorted(list(device_graph.nodes(data=True))))
-            self.device_graph.add_edges_from(sorted(list(device_graph.edges)))
-        else:
-            self.device_graph = nx.Graph()
-            self.device_graph.add_nodes_from(sorted(list(device_graph.nodes(data=True))))
-            self.device_graph.add_edges_from(
-                sorted(list(sorted(edge) for edge in device_graph.edges))
-            )
-
-        self._map = initial_mapping.copy()
-        self._inverse_map = {v: k for k, v in self._map.items()}
-        self._induced_subgraph = nx.induced_subgraph(self.device_graph, self._map.values())
+        # Map both logical and physical qubits to integers.
+        self._logical_qid_to_int = {q: i for i, q in enumerate(sorted(initial_mapping.keys()))}
+        self._int_to_logical_qid = sorted(
+            self._logical_qid_to_int.keys(), key=lambda x: self._logical_qid_to_int[x]
+        )
+        self._physical_qid_to_int = {q: i for i, q in enumerate(sorted(initial_mapping.values()))}
+        self._int_to_physical_qid = sorted(
+            self._physical_qid_to_int.keys(), key=lambda x: self._physical_qid_to_int[x]
+        )
+        logical_qubits, physical_qubits = (
+            zip(*[(k, v) for k, v in initial_mapping.items()]) if initial_mapping else ([], [])
+        )
+        num_qubits = len(logical_qubits)
+        self._logical_to_physical = np.asarray(
+            [
+                self._physical_qid_to_int[physical_qubits[i]]
+                for i in sorted(
+                    range(num_qubits), key=lambda x: self._logical_qid_to_int[logical_qubits[x]]
+                )
+            ]
+        )
+        self._physical_to_logical = np.asarray(
+            [
+                self._logical_qid_to_int[logical_qubits[i]]
+                for i in sorted(
+                    range(num_qubits), key=lambda x: self._physical_qid_to_int[physical_qubits[x]]
+                )
+            ]
+        )
+        # Construct the induced subgraph (on integers) and corresponding distance matrix.
+        self._induced_subgraph_int = nx.relabel_nodes(
+            nx.induced_subgraph(device_graph, initial_mapping.values()),
+            {q: self._physical_qid_to_int[q] for q in initial_mapping.values()},
+        )
+        # Compute floyd warshall dictionary.
         self._predecessors, self._distances = nx.floyd_warshall_predecessor_and_distance(
-            self._induced_subgraph
+            self._induced_subgraph_int
         )
 
     @property
-    def map(self) -> Dict['cirq.Qid', 'cirq.Qid']:
-        """The mapping of logical qubits (keys) to physical qubits (values)."""
-        return self._map
+    def physical_qid_to_int(self) -> Dict['cirq.Qid', int]:
+        """Mapping of physical qubits, that were part of the initial mapping, to unique integers."""
+        return self._physical_qid_to_int
+
+    @property
+    def int_to_physical_qid(self) -> List['cirq.Qid']:
+        """Inverse mapping of unique integers to corresponding physical qubits.
+
+        `self.physical_qid_to_int[self.int_to_physical_qid[i]] == i` for each i.
+        """
+        return self._int_to_physical_qid
+
+    @property
+    def logical_qid_to_int(self) -> Dict['cirq.Qid', int]:
+        """Mapping of logical qubits, that were part of the initial mapping, to unique integers."""
+        return self._logical_qid_to_int
 
     @property
-    def inverse_map(self) -> Dict['cirq.Qid', 'cirq.Qid']:
-        """The mapping of physical qubits (keys) to logical qubits (values)."""
-        return self._inverse_map
+    def int_to_logical_qid(self) -> List['cirq.Qid']:
+        """Inverse mapping of unique integers to corresponding physical qubits.
+
+        `self.logical_qid_to_int[self.int_to_logical_qid[i]] == i` for each i.
+        """
+        return self._int_to_logical_qid
+
+    @property
+    def logical_to_physical(self) -> np.ndarray:
+        """The mapping of logical qubit integers to physical qubit integers.
+
+        Let `lq: cirq.Qid` be a logical qubit. Then the corresponding physical qubit that it
+        maps to can be obtained by:
+        `self.int_to_physical_qid[self.logical_to_physical[self.logical_qid_to_int[lq]]]`
+        """
+        return self._logical_to_physical
+
+    @property
+    def physical_to_logical(self) -> np.ndarray:
+        """The mapping of physical qubits integers to logical qubits integers.
+
+        Let `pq: cirq.Qid` be a physical qubit. Then the corresponding logical qubit that it
+        maps to can be obtained by:
+        `self.int_to_logical_qid[self.physical_to_logical[self.physical_qid_to_int[pq]]]`
+        """
+        return self._physical_to_logical
 
     @property
-    def induced_subgraph(self) -> nx.Graph:
-        """The induced subgraph on the set of physical qubits which are part of `self.map`."""
-        return self._induced_subgraph
+    def induced_subgraph_int(self) -> nx.Graph:
+        """Induced subgraph on physical qubit integers present in `self.logical_to_physical`."""
+        return self._induced_subgraph_int
 
-    def dist_on_device(self, lq1: 'cirq.Qid', lq2: 'cirq.Qid') -> int:
+    def dist_on_device(self, lq1: int, lq2: int) -> int:
         """Finds distance between logical qubits 'lq1' and 'lq2' on the device.
 
         Args:
-            lq1: the first logical qubit.
-            lq2: the second logical qubit.
+            lq1: integer corresponding to the first logical qubit.
+            lq2: integer corresponding to the second logical qubit.
 
         Returns:
             The shortest path distance.
         """
-        return self._distances[self._map[lq1]][self._map[lq2]]
+        return self._distances[self.logical_to_physical[lq1]][self.logical_to_physical[lq2]]
 
-    def can_execute(self, op: 'cirq.Operation') -> bool:
-        """Finds whether the given operation acts on qubits that are adjacent on the device.
+    def is_adjacent(self, lq1: int, lq2: int) -> bool:
+        """Finds whether logical qubits `lq1` and `lq2` are adjacent on the device.
 
         Args:
-            op: an operation on logical qubits.
+            lq1: integer corresponding to the first logical qubit.
+            lq2: integer corresponding to the second logical qubit.
 
         Returns:
-            True, if physical qubits corresponding to logical qubits `op.qubits` are adjacent on
+            True, if physical qubits corresponding to `lq1` and `lq2` are adjacent on
             the device.
         """
-        return protocols.num_qubits(op) < 2 or self.dist_on_device(*op.qubits) == 1
+        return self.dist_on_device(lq1, lq2) == 1
 
-    def apply_swap(self, lq1: 'cirq.Qid', lq2: 'cirq.Qid') -> None:
+    def apply_swap(self, lq1: int, lq2: int) -> None:
         """Updates the mapping to simulate inserting a swap operation between `lq1` and `lq2`.
 
         Args:
-            lq1: the first logical qubit.
-            lq2: the second logical qubit.
+            lq1: integer corresponding to the first logical qubit.
+            lq2: integer corresponding to the second logical qubit.
 
         Raises:
-            ValueError: whenever lq1 and lq2 are no adjacent on the device.
+            ValueError: whenever lq1 and lq2 are not adjacent on the device.
         """
         if self.dist_on_device(lq1, lq2) > 1:
             raise ValueError(
                 f"q1: {lq1} and q2: {lq2} are not adjacent on the device. Cannot swap them."
             )
 
-        pq1, pq2 = self._map[lq1], self._map[lq2]
-        self._map[lq1], self._map[lq2] = self._map[lq2], self._map[lq1]
-
-        self._inverse_map[pq1], self._inverse_map[pq2] = (
-            self._inverse_map[pq2],
-            self._inverse_map[pq1],
-        )
+        pq1, pq2 = self.logical_to_physical[lq1], self.logical_to_physical[lq2]
+        self._logical_to_physical[[lq1, lq2]] = self._logical_to_physical[[lq2, lq1]]
+        self._physical_to_logical[[pq1, pq2]] = self._physical_to_logical[[pq2, pq1]]
 
     def mapped_op(self, op: 'cirq.Operation') -> 'cirq.Operation':
-        """Transforms the given operation with the qubits in self._map.
+        """Transforms the given logical operation to act on corresponding physical qubits.
 
         Args:
-            op: an operation on logical qubits.
+            op: logical operation acting on logical qubits.
 
         Returns:
-            The same operation on corresponding physical qubits."""
-        return op.transform_qubits(self._map)
+            The same operation acting on corresponding physical qubits.
+        """
+        logical_ints = [self._logical_qid_to_int[q] for q in op.qubits]
+        physical_ints = self.logical_to_physical[logical_ints]
+        qubit_map: Dict['cirq.Qid', 'cirq.Qid'] = {
+            q: self._int_to_physical_qid[physical_ints[i]] for i, q in enumerate(op.qubits)
+        }
+        return op.transform_qubits(qubit_map)
 
-    def shortest_path(self, lq1: 'cirq.Qid', lq2: 'cirq.Qid') -> Sequence['cirq.Qid']:
-        """Find the shortest path between two logical qubits on the device given their mapping.
+    def shortest_path(self, lq1: int, lq2: int) -> Sequence[int]:
+        """Find the shortest path between two logical qubits on the device, given their mapping.
 
         Args:
-            lq1: the first logical qubit.
-            lq2: the second logical qubit.
+            lq1: integer corresponding to the first logical qubit.
+            lq2: integer corresponding to the second logical qubit.
 
         Returns:
-            A sequence of logical qubits on the shortest path from lq1 to lq2.
+            A sequence of logical qubit integers on the shortest path from `lq1` to `lq2`.
         """
-        return [
-            self._inverse_map[pq]
-            for pq in nx.reconstruct_path(self._map[lq1], self._map[lq2], self._predecessors)
+        return self.physical_to_logical[
+            nx.reconstruct_path(*self.logical_to_physical[[lq1, lq2]], self._predecessors)
         ]
-
-    def _value_equality_values_(self):
-        graph_equality = (
-            tuple(self.device_graph.nodes),
-            tuple(self.device_graph.edges),
-            nx.is_directed(self.device_graph),
-        )
-        map_equality = tuple(sorted(self._map.items()))
-        return (graph_equality, map_equality)
-
-    def __repr__(self) -> str:
-        graph_type = type(self.device_graph).__name__
-        return (
-            f'cirq.MappingManager('
-            f'nx.{graph_type}({dict(self.device_graph.adjacency())}),'
-            f' {self._map})'
-        )