Repository: Qiskit/qiskit-terra (python)

Issue:
Refactor `EquivalenceLibrary` to internally be a graph
### What should we add?

In #7211 we refactored the search algorithm in the `BasisTranslator` transpiler pass to perform a dijkstra search using retworkx. To accomplish this the pass had to generate a retworkx graph object from the `EquivalenceLibrary` on demand and then perform the dijkstra search over that graph. However, it would be better if the equivalence library itself just internally wrapped a graph object. The data in the `EquivalenceLibrary` is inherently a graph and using a graph representation for the internal data model makes sense and will also improve the performance of the basis translator because there will be no overhead for constructing the graph object.

The output of this issue should be to change the internal `_map` attribute of the `EquivalenceLibrary` with a `retworkx.PyDiGraph` object built as the basis translator does during it's basis search: https://github.com/Qiskit/qiskit-terra/blob/8a3e760ffaf6dcd06c0a0e04f13612fdefd9ab3c/qiskit/transpiler/passes/basis/basis_translator.py#L470-L502
We should maintain the same public API that exists on the class, but just change the internal data model (and likely add helper functions to access the graph object). Then after we update the `EquivalenceLibrary` in a second pull request we can update the `BasisTranslator` to leverage the internal graph object instead of creating it on demand for each search.

Pull request:
Convert Equivalence Library to graph
<!--
⚠️ If you do not respect this template, your pull request will be closed.
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Also, please add a release note file using reno if the change needs to be
  documented in the release notes.
⚠️ If your pull request fixes an open issue, please link to the issue.

- [ ] I have added the tests to cover my changes.
- [ ] I have updated the documentation accordingly.
- [ ] I have read the CONTRIBUTING document.
-->

### Summary

Fixes #8664

### Details and comments

This PR converts the internal representation for the `EquivalenceLibrary` from a map to a graph. The existing API remains with three new attributes - `graph`, `key_to_node_index`, and `num_gates_for_rule` - added. Previously this graph was being created by the `BasisTranslator` to be used in `basis_search`. Now the graph is created as the `EquivalenceLibrary` is being filled. All accesses and mods to the `EquivalenceLibrary` are now made in the graph and the graph can now be accessed directly by the `BasisTranslator`.

The `EquivalenceLibrary` takes an optional parameter, `base`, which is another library and would be added onto the back of the new library. Since it's easier and more efficient to add to the `base` graph, the new library is now added onto the `base` graph. This is the reason for the change to `test_add_equivalence`.

A consistent slight speedup was noticed while running the `circuit` and `transpiler` CI tests. Some `BasisTranslator` benchmarks should be run to check performance changes.

- [x] Run `BasisTranslator` benchmarks
- [x] Release note?




Patch:
diff --git a/qiskit/circuit/equivalence.py b/qiskit/circuit/equivalence.py
index 80731a3dc..bd4711614 100644
--- a/qiskit/circuit/equivalence.py
+++ b/qiskit/circuit/equivalence.py
@@ -12,6 +12,7 @@
 
 """Gate equivalence library."""
 
+import copy
 from collections import namedtuple
 
 from rustworkx.visualization import graphviz_draw  # pylint: disable=no-name-in-module
@@ -22,10 +23,9 @@ from .exceptions import CircuitError
 from .parameterexpression import ParameterExpression
 
 Key = namedtuple("Key", ["name", "num_qubits"])
-
-Entry = namedtuple("Entry", ["search_base", "equivalences"])
-
 Equivalence = namedtuple("Equivalence", ["params", "circuit"])  # Ordered to match Gate.params
+NodeData = namedtuple("NodeData", ["key", "equivs"])
+EdgeData = namedtuple("EdgeData", ["index", "num_gates", "rule", "source"])
 
 
 class EquivalenceLibrary:
@@ -37,11 +37,39 @@ class EquivalenceLibrary:
 
         Args:
             base (Optional[EquivalenceLibrary]):  Base equivalence library to
-                will be referenced if an entry is not found in this library.
+                be referenced if an entry is not found in this library.
         """
         self._base = base
 
-        self._map = {}
+        if base is None:
+            self._graph = rx.PyDiGraph()
+            self._key_to_node_index = dict()
+            self._rule_count = 0
+        else:
+            self._graph = base._graph.copy()
+            self._key_to_node_index = copy.deepcopy(base._key_to_node_index)
+            self._rule_count = base._rule_count
+
+    @property
+    def graph(self) -> rx.PyDiGraph:
+        """Return graph representing the equivalence library data.
+
+        This property should be treated as read-only as it provides
+        a reference to the internal state of the :class:`~.EquivalenceLibrary` object.
+        If the graph returned by this property is mutated it could corrupt the
+        the contents of the object. If you need to modify the output ``PyDiGraph``
+        be sure to make a copy prior to any modification.
+
+        Returns:
+            PyDiGraph: A graph object with equivalence data in each node.
+        """
+        return self._graph
+
+    def _set_default_node(self, key):
+        """Create a new node if key not found"""
+        if key not in self._key_to_node_index:
+            self._key_to_node_index[key] = self._graph.add_node(NodeData(key=key, equivs=[]))
+        return self._key_to_node_index[key]
 
     def add_equivalence(self, gate, equivalent_circuit):
         """Add a new equivalence to the library. Future queries for the Gate
@@ -62,13 +90,25 @@ class EquivalenceLibrary:
         _raise_if_param_mismatch(gate.params, equivalent_circuit.parameters)
 
         key = Key(name=gate.name, num_qubits=gate.num_qubits)
-
         equiv = Equivalence(params=gate.params.copy(), circuit=equivalent_circuit.copy())
 
-        if key not in self._map:
-            self._map[key] = Entry(search_base=True, equivalences=[])
+        target = self._set_default_node(key)
+        self._graph[target].equivs.append(equiv)
 
-        self._map[key].equivalences.append(equiv)
+        sources = {
+            Key(name=instruction.operation.name, num_qubits=len(instruction.qubits))
+            for instruction in equivalent_circuit
+        }
+        edges = [
+            (
+                self._set_default_node(source),
+                target,
+                EdgeData(index=self._rule_count, num_gates=len(sources), rule=equiv, source=source),
+            )
+            for source in sources
+        ]
+        self._graph.add_edges_from(edges)
+        self._rule_count += 1
 
     def has_entry(self, gate):
         """Check if a library contains any decompositions for gate.
@@ -80,10 +120,9 @@ class EquivalenceLibrary:
             Bool: True if gate has a known decomposition in the library.
                 False otherwise.
         """
-
         key = Key(name=gate.name, num_qubits=gate.num_qubits)
 
-        return key in self._map or (self._base.has_entry(gate) if self._base is not None else False)
+        return key in self._key_to_node_index
 
     def set_entry(self, gate, entry):
         """Set the equivalence record for a Gate. Future queries for the Gate
@@ -98,16 +137,14 @@ class EquivalenceLibrary:
             entry (List['QuantumCircuit']) : A list of QuantumCircuits, each
                 equivalently implementing the given Gate.
         """
-
         for equiv in entry:
             _raise_if_shape_mismatch(gate, equiv)
             _raise_if_param_mismatch(gate.params, equiv.parameters)
 
         key = Key(name=gate.name, num_qubits=gate.num_qubits)
-
         equivs = [Equivalence(params=gate.params.copy(), circuit=equiv.copy()) for equiv in entry]
 
-        self._map[key] = Entry(search_base=False, equivalences=equivs)
+        self._graph[self._set_default_node(key)] = NodeData(key=key, equivs=equivs)
 
     def get_entry(self, gate):
         """Gets the set of QuantumCircuits circuits from the library which
@@ -128,13 +165,30 @@ class EquivalenceLibrary:
                 ordering of the StandardEquivalenceLibrary will not generally be
                 consistent across Qiskit versions.
         """
-
         key = Key(name=gate.name, num_qubits=gate.num_qubits)
-
         query_params = gate.params
 
         return [_rebind_equiv(equiv, query_params) for equiv in self._get_equivalences(key)]
 
+    def keys(self):
+        """Return list of keys to key to node index map.
+
+        Returns:
+            List: Keys to the key to node index map.
+        """
+        return self._key_to_node_index.keys()
+
+    def node_index(self, key):
+        """Return node index for a given key.
+
+        Args:
+            key (Key): Key to an equivalence.
+
+        Returns:
+            Int: Index to the node in the graph for the given key.
+        """
+        return self._key_to_node_index[key]
+
     def draw(self, filename=None):
         """Draws the equivalence relations available in the library.
 
@@ -166,7 +220,7 @@ class EquivalenceLibrary:
         graph = rx.PyDiGraph()
 
         node_map = {}
-        for key in self._get_all_keys():
+        for key in self._key_to_node_index:
             name, num_qubits = key
             equivalences = self._get_equivalences(key)
 
@@ -197,23 +251,13 @@ class EquivalenceLibrary:
 
         return graph
 
-    def _get_all_keys(self):
-        base_keys = self._base._get_all_keys() if self._base is not None else set()
-
-        self_keys = set(self._map.keys())
-
-        return self_keys | {
-            base_key
-            for base_key in base_keys
-            if base_key not in self._map or self._map[base_key].search_base
-        }
-
     def _get_equivalences(self, key):
-        search_base, equivalences = self._map.get(key, (True, []))
-
-        if search_base and self._base is not None:
-            return equivalences + self._base._get_equivalences(key)
-        return equivalences
+        """Get all the equivalences for the given key"""
+        return (
+            self._graph[self._key_to_node_index[key]].equivs
+            if key in self._key_to_node_index
+            else []
+        )
 
 
 def _raise_if_param_mismatch(gate_params, circuit_parameters):
diff --git a/qiskit/transpiler/passes/basis/basis_translator.py b/qiskit/transpiler/passes/basis/basis_translator.py
index 62121bce9..215fba774 100644
--- a/qiskit/transpiler/passes/basis/basis_translator.py
+++ b/qiskit/transpiler/passes/basis/basis_translator.py
@@ -26,7 +26,7 @@ import rustworkx
 from qiskit.circuit import Gate, ParameterVector, QuantumRegister, ControlFlowOp, QuantumCircuit
 from qiskit.dagcircuit import DAGCircuit
 from qiskit.converters import circuit_to_dag, dag_to_circuit
-from qiskit.circuit.equivalence import Key
+from qiskit.circuit.equivalence import Key, NodeData
 from qiskit.transpiler.basepasses import TransformationPass
 from qiskit.transpiler.exceptions import TranspilerError
 
@@ -370,17 +370,24 @@ class StopIfBasisRewritable(Exception):
 class BasisSearchVisitor(rustworkx.visit.DijkstraVisitor):  # pylint: disable=no-member
     """Handles events emitted during `rustworkx.dijkstra_search`."""
 
-    def __init__(self, graph, source_basis, target_basis, num_gates_for_rule):
+    def __init__(self, graph, source_basis, target_basis):
         self.graph = graph
         self.target_basis = set(target_basis)
         self._source_gates_remain = set(source_basis)
-        self._num_gates_remain_for_rule = dict(num_gates_for_rule)
+        self._num_gates_remain_for_rule = dict()
+        save_index = -1
+        for edata in self.graph.edges():
+            if save_index == edata.index:
+                continue
+            self._num_gates_remain_for_rule[edata.index] = edata.num_gates
+            save_index = edata.index
+
         self._basis_transforms = []
         self._predecessors = dict()
         self._opt_cost_map = dict()
 
     def discover_vertex(self, v, score):
-        gate = self.graph[v]
+        gate = self.graph[v].key
         self._source_gates_remain.discard(gate)
         self._opt_cost_map[gate] = score
         rule = self._predecessors.get(gate, None)
@@ -406,22 +413,21 @@ class BasisSearchVisitor(rustworkx.visit.DijkstraVisitor):  # pylint: disable=no
         if edata is None:
             return
 
-        index = edata["index"]
-        self._num_gates_remain_for_rule[index] -= 1
+        self._num_gates_remain_for_rule[edata.index] -= 1
 
-        target = self.graph[target]
+        target = self.graph[target].key
         # if there are gates in this `rule` that we have not yet generated, we can't apply
         # this `rule`. if `target` is already in basis, it's not beneficial to use this rule.
-        if self._num_gates_remain_for_rule[index] > 0 or target in self.target_basis:
+        if self._num_gates_remain_for_rule[edata.index] > 0 or target in self.target_basis:
             raise rustworkx.visit.PruneSearch  # pylint: disable=no-member
 
     def edge_relaxed(self, edge):
         _, target, edata = edge
         if edata is not None:
-            gate = self.graph[target]
-            self._predecessors[gate] = edata["rule"]
+            gate = self.graph[target].key
+            self._predecessors[gate] = edata.rule
 
-    def edge_cost(self, edge):
+    def edge_cost(self, edge_data):
         """Returns the cost of an edge.
 
         This function computes the cost of this edge rule by summing
@@ -430,19 +436,17 @@ class BasisSearchVisitor(rustworkx.visit.DijkstraVisitor):  # pylint: disable=no
         will later add it.
         """
 
-        if edge is None:
+        if edge_data is None:
             # the target of the edge is a gate in the target basis,
             # so we return a default value of 1.
             return 1
 
         cost_tot = 0
-        rule = edge["rule"]
-        for instruction in rule.circuit:
+        for instruction in edge_data.rule.circuit:
             key = Key(name=instruction.operation.name, num_qubits=len(instruction.qubits))
             cost_tot += self._opt_cost_map[key]
 
-        source = edge["source"]
-        return cost_tot - self._opt_cost_map[source]
+        return cost_tot - self._opt_cost_map[edge_data.source]
 
     @property
     def basis_transforms(self):
@@ -477,62 +481,33 @@ def _basis_search(equiv_lib, source_basis, target_basis):
     if not source_basis:
         return []
 
-    all_gates_in_lib = set()
-
-    graph = rustworkx.PyDiGraph()
-    nodes_to_indices = dict()
-    num_gates_for_rule = dict()
-
-    def lazy_setdefault(key):
-        if key not in nodes_to_indices:
-            nodes_to_indices[key] = graph.add_node(key)
-        return nodes_to_indices[key]
-
-    rcounter = 0  # running sum of the number of equivalence rules in the library.
-    for key in equiv_lib._get_all_keys():
-        target = lazy_setdefault(key)
-        all_gates_in_lib.add(key)
-        for equiv in equiv_lib._get_equivalences(key):
-            sources = {
-                Key(name=instruction.operation.name, num_qubits=len(instruction.qubits))
-                for instruction in equiv.circuit
-            }
-            all_gates_in_lib |= sources
-            edges = [
-                (
-                    lazy_setdefault(source),
-                    target,
-                    {"index": rcounter, "rule": equiv, "source": source},
-                )
-                for source in sources
-            ]
-
-            num_gates_for_rule[rcounter] = len(sources)
-            graph.add_edges_from(edges)
-            rcounter += 1
-
     # This is only neccessary since gates in target basis are currently reported by
     # their names and we need to have in addition the number of qubits they act on.
-    target_basis_keys = [
-        key
-        for gate in target_basis
-        for key in filter(lambda key, name=gate: key.name == name, all_gates_in_lib)
-    ]
+    target_basis_keys = [key for key in equiv_lib.keys() if key.name in target_basis]
+
+    graph = equiv_lib.graph
+    vis = BasisSearchVisitor(graph, source_basis, target_basis_keys)
 
-    vis = BasisSearchVisitor(graph, source_basis, target_basis_keys, num_gates_for_rule)
     # we add a dummy node and connect it with gates in the target basis.
     # we'll start the search from this dummy node.
-    dummy = graph.add_node("dummy starting node")
-    graph.add_edges_from_no_data([(dummy, nodes_to_indices[key]) for key in target_basis_keys])
-    rtn = None
-    try:
-        rustworkx.digraph_dijkstra_search(graph, [dummy], vis.edge_cost, vis)
-    except StopIfBasisRewritable:
-        rtn = vis.basis_transforms
+    dummy = graph.add_node(NodeData(key="key", equivs=[("dummy starting node", 0)]))
 
-        logger.debug("Transformation path:")
-        for gate_name, gate_num_qubits, params, equiv in rtn:
-            logger.debug("%s/%s => %s\n%s", gate_name, gate_num_qubits, params, equiv)
+    try:
+        graph.add_edges_from_no_data(
+            [(dummy, equiv_lib.node_index(key)) for key in target_basis_keys]
+        )
+        rtn = None
+        try:
+            rustworkx.digraph_dijkstra_search(graph, [dummy], vis.edge_cost, vis)
+        except StopIfBasisRewritable:
+            rtn = vis.basis_transforms
+
+            logger.debug("Transformation path:")
+            for gate_name, gate_num_qubits, params, equiv in rtn:
+                logger.debug("%s/%s => %s\n%s", gate_name, gate_num_qubits, params, equiv)
+    finally:
+        # Remove dummy node in order to return graph to original state
+        graph.remove_node(dummy)
 
     return rtn
 
diff --git a/releasenotes/notes/equivalence-to-graph-3b52912ecb542db8.yaml b/releasenotes/notes/equivalence-to-graph-3b52912ecb542db8.yaml
new file mode 100644
index 000000000..360a93b5b
--- /dev/null
+++ b/releasenotes/notes/equivalence-to-graph-3b52912ecb542db8.yaml
@@ -0,0 +1,47 @@
+---
+features:
+  - |
+    The :class:`~.EquivalenceLibrary` is now
+    represented internally as a ``PyDiGraph``. As part of this change,
+    there is a new public-facing attribute - ``graph``: the graph
+    representation of the EquivalenceLibrary. This attribute is
+    intended for use internally in qiskit and therefore should always
+    be copied before being modified by the user to prevent possible
+    corruption of the internal equivalence graph.
+upgrade:
+  - |
+    :class:`~.EquivalenceLibrary` objects that are initialized with the ``base``
+    attribute will no long have a shared reference with the
+    :class:`~.EquivalenceLibrary` passed in. In earlier releases if you mutated
+    ``base`` after it was used to create a new :class:`~.EquivalenceLibrary`
+    instance both instances would reflect that change. This no longer is the case
+    and updates to ``base`` will no longer be reflected in the new
+    :class:`~.EquivalenceLibrary`. For example, if you created an equivalence library
+    with::
+
+        import math
+
+        from qiskit.circuit import QuantumCircuit
+        from qiskit.circuit.library import XGate
+        from qiskit.circuit.equivalence import EquivalenceLibrary
+
+        original_lib = EquivalenceLibrary()
+        qc = QuantumCircuit(1)
+        qc.rx(math.pi, 0)
+        original_lib.add_equivalence(XGate(), qc)
+        new_lib = EquivalenceLibrary(base=original_lib)
+
+    if you modified ``original_lib`` with::
+
+        import from qiskit.circuit.library import SXGate
+
+        qc = QuantumCircuit(1)
+        qc.rx(math.pi / 2, 0)
+        original_lib.add_equivalence(SXGate(), qc)
+
+    in previous releases ``new_lib`` would also include the definition of ``SXGate``
+    after it was added to ``original_lib``, but in this release this no longer will
+    be the case. This change was made because     of the change in internal data
+    structure to be a graph, which improved performance of the
+    :class:`~.EquivalenceLibrary` class, especially when there are multiple runs of
+    the :class:`~.BasisTranslator` transpiler pass.


Test patch:
diff --git a/test/python/circuit/test_equivalence.py b/test/python/circuit/test_equivalence.py
index f83a9f150..ec3022ce5 100644
--- a/test/python/circuit/test_equivalence.py
+++ b/test/python/circuit/test_equivalence.py
@@ -16,6 +16,8 @@
 import unittest
 import numpy as np
 
+import rustworkx as rx
+
 from qiskit.test import QiskitTestCase
 
 from qiskit.circuit import QuantumCircuit, Parameter, Gate
@@ -24,6 +26,7 @@ from qiskit.circuit.exceptions import CircuitError
 from qiskit.converters import circuit_to_instruction, circuit_to_gate
 from qiskit.circuit import EquivalenceLibrary
 from qiskit.utils import optionals
+from qiskit.circuit.equivalence import Key, Equivalence, NodeData, EdgeData
 
 from ..visualization.visualization import QiskitVisualizationTestCase, path_to_diagram_reference
 
@@ -154,6 +157,47 @@ class TestEquivalenceLibraryWithoutBase(QiskitTestCase):
 
         self.assertFalse(eq_lib.has_entry(OneQubitZeroParamGate()))
 
+    def test_equivalence_graph(self):
+        """Verify valid graph created by add_equivalence"""
+
+        eq_lib = EquivalenceLibrary()
+
+        gate = OneQubitZeroParamGate()
+        first_equiv = QuantumCircuit(1)
+        first_equiv.h(0)
+        eq_lib.add_equivalence(gate, first_equiv)
+
+        equiv_copy = eq_lib._get_equivalences(Key(name="1q0p", num_qubits=1))[0].circuit
+
+        egraph = rx.PyDiGraph()
+        node_wt = NodeData(
+            key=Key(name="1q0p", num_qubits=1), equivs=[Equivalence(params=[], circuit=equiv_copy)]
+        )
+
+        egraph.add_node(node_wt)
+
+        node_wt = NodeData(key=Key(name="h", num_qubits=1), equivs=[])
+        egraph.add_node(node_wt)
+
+        edge_wt = EdgeData(
+            index=0,
+            num_gates=1,
+            rule=Equivalence(params=[], circuit=equiv_copy),
+            source=Key(name="h", num_qubits=1),
+        )
+        egraph.add_edge(0, 1, edge_wt)
+
+        for node in eq_lib.graph.nodes():
+            self.assertTrue(node in egraph.nodes())
+            for edge in eq_lib.graph.edges():
+                self.assertTrue(edge in egraph.edges())
+
+        self.assertEqual(len(eq_lib.graph.nodes()), len(egraph.nodes()))
+        self.assertEqual(len(eq_lib.graph.edges()), len(egraph.edges()))
+
+        keys = {Key(name="1q0p", num_qubits=1): 0, Key(name="h", num_qubits=1): 1}.keys()
+        self.assertEqual(keys, eq_lib.keys())
+
 
 class TestEquivalenceLibraryWithBase(QiskitTestCase):
     """Test cases for EquivalenceLibrary with base library."""
@@ -207,8 +251,9 @@ class TestEquivalenceLibraryWithBase(QiskitTestCase):
         entry = eq_lib.get_entry(gate)
 
         self.assertEqual(len(entry), 2)
-        self.assertEqual(entry[0], second_equiv)
-        self.assertEqual(entry[1], first_equiv)
+        self.assertNotEqual(entry[0], entry[1])
+        self.assertTrue(entry[0] in [first_equiv, second_equiv])
+        self.assertTrue(entry[1] in [first_equiv, second_equiv])
 
     def test_set_entry(self):
         """Verify we find only equivalences from top when explicitly set."""