Initial: Try to speed up Target rust-side mapping access.

crates/accelerate/src/target_transpiler/mod.rs

@@ -24,6 +24,8 @@ use ahash::HashSet;
 use indexmap::{IndexMap, IndexSet};
 use itertools::Itertools;
 use nullable_index_map::NullableIndexMap;
+use pyo3::types::IntoPyDict;
+use pyo3::types::PyIterator;
 use pyo3::{
     exceptions::{PyAttributeError, PyIndexError, PyKeyError, PyValueError},
     prelude::*,
@@ -177,6 +179,7 @@ pub(crate) struct Target {
     qarg_gate_map: NullableIndexMap<Qargs, Option<HashSet<String>>>,
     non_global_strict_basis: Option<Vec<String>>,
     non_global_basis: Option<Vec<String>>,
+    _py_gate_map_cache: IndexMap<String, Py<PyDict>, RandomState>,
 }
 
 #[pymethods]
@@ -274,6 +277,7 @@ impl Target {
             qarg_gate_map: NullableIndexMap::default(),
             non_global_basis: None,
             non_global_strict_basis: None,
+            _py_gate_map_cache: IndexMap::default(),
         })
     }
 
@@ -360,6 +364,7 @@ impl Target {
         self.gate_map.insert(name.to_string(), qargs_val);
         self.non_global_basis = None;
         self.non_global_strict_basis = None;
+        self._py_gate_map_cache.swap_remove(name);
         Ok(())
     }
 
@@ -374,30 +379,29 @@ impl Target {
     #[pyo3(text_signature = "(instruction, qargs, properties, /,)")]
     fn update_instruction_properties(
         &mut self,
-        instruction: String,
+        instruction: &str,
         qargs: Option<Qargs>,
         properties: Option<InstructionProperties>,
     ) -> PyResult<()> {
-        if !self.contains_key(&instruction) {
+        if !self.contains_key(instruction) {
             return Err(PyKeyError::new_err(format!(
                 "Provided instruction: '{:?}' not in this Target.",
                 &instruction
             )));
         };
-        let mut prop_map = self[&instruction].clone();
-        if !(prop_map.contains_key(qargs.as_ref())) {
+        if !(self[instruction].contains_key(qargs.as_ref())) {
             return Err(PyKeyError::new_err(format!(
                 "Provided qarg {:?} not in this Target for {:?}.",
                 &qargs.unwrap_or_default(),
                 &instruction
             )));
         }
-        if let Some(e) = prop_map.get_mut(qargs.as_ref()) {
-            *e = properties;
+        if let Some(obj) = self.gate_map.get_mut(instruction) {
+            if let Some(e) = obj.get_mut(qargs.as_ref()) {
+                *e = properties;
+            }
         }
-        self.gate_map
-            .entry(instruction)
-            .and_modify(|e| *e = prop_map);
+        self._py_gate_map_cache.swap_remove(instruction);
         Ok(())
     }
 
@@ -790,8 +794,8 @@ impl Target {
 
     // Magic methods:
 
-    fn __len__(&self) -> PyResult<usize> {
-        Ok(self.gate_map.len())
+    fn __len__(&self) -> usize {
+        self.gate_map.len()
     }
 
     fn __getstate__(&self, py: Python<'_>) -> PyResult<Py<PyDict>> {
@@ -898,6 +902,86 @@ impl Target {
             .extract::<Option<Vec<String>>>()?;
         Ok(())
     }
+
+    /// Performs caching of python side dict objects for better access performance.
+    fn __getitem__(&mut self, key: &Bound<PyAny>) -> PyResult<PyObject> {
+        let py: Python = key.py();
+        let Ok(key) = key.extract::<String>() else {return Err(PyKeyError::new_err(
+            "Invalid Key Type for Target."
+        ));};
+        if let Some(val) = self._py_gate_map_cache.get(key.as_str()) {
+            Ok(val.as_any().clone_ref(py))
+        } else if let Some(val) = self.gate_map.get(key.as_str()) {
+            let new_py: Vec<(Option<Bound<'_, PyTuple>>, PyObject)> = val
+                .clone()
+                .into_iter()
+                .map(|(key, val)| (key.map(|key| PyTuple::new_bound(py, key)), val.into_py(py)))
+                .collect();
+            let dict_py = new_py.into_py_dict_bound(py).unbind();
+            self._py_gate_map_cache.insert(key, dict_py.clone_ref(py));
+            Ok(dict_py.into())
+        } else {
+            Err(PyKeyError::new_err(format!(
+                "Key: '{key}' not present in target"
+            )))
+        }
+    }
+
+    /// Retrieves all the keys in the gate_map
+    #[pyo3(name = "keys")]
+    fn py_keys(slf: PyRef<Self>) -> PyResult<PyObject> {
+        let dict = PyDict::new_bound(slf.py());
+        let mapped_keys = slf.keys().map(|key| key.to_string());
+        for key in mapped_keys {
+            dict.set_item(key, slf.py().None())?;
+        }
+        Ok(dict.as_any().call_method0("keys")?.into())
+    }
+
+    #[pyo3(name = "values")]
+    fn py_values(slf: PyRef<Self>) -> PyResult<Py<PyList>> {
+        let py = slf.py();
+        let values = PyList::empty_bound(py);
+        let mapped_values = slf.values().map(|value| {
+            let vec_dict: Vec<_> = value
+                .clone()
+                .into_iter()
+                .map(|(key, val)| (key.map(|key| PyTuple::new_bound(py, key)), val.into_py(py)))
+                .collect();
+            vec_dict.into_py_dict_bound(py)
+        });
+        for key in mapped_values {
+            values.append(key)?;
+        }
+        Ok(values.into())
+    }
+
+    fn items(slf: PyRef<Self>) -> PyResult<Py<PyList>> {
+        let py = slf.py();
+        let items_list = PyList::empty_bound(py);
+        for (key, value) in slf.gate_map.iter().map(|(keys, values)| {
+            (keys.to_string(), {
+                let values_vec: Vec<_> = values
+                    .clone()
+                    .into_iter()
+                    .map(|(key, val)| (key.map(|key| PyTuple::new_bound(py, key)), val.into_py(py)))
+                    .collect();
+                values_vec.into_py_dict_bound(py)
+            })
+        }) {
+            items_list.append((key, value))?;
+        }
+        Ok(items_list.into())
+    }
+
+    fn __contains__(&self, item: &Bound<PyAny>) -> PyResult<bool> {
+        let Ok(item) = item.extract::<String>() else {return Err(PyKeyError::new_err("Invalid key type, not supported in Target."));};
+        Ok(self.gate_map.contains_key(&item))
+    }
+
+    fn __iter__(slf: &Bound<Self>) -> PyResult<Py<PyIterator>> {
+        Ok(slf.call_method0("keys")?.iter()?.into())
+    }
 }
 
 // Rust native methods

qiskit/transpiler/target.py

@@ -334,7 +334,6 @@ def __init__(
         concurrent_measurements=None,  # pylint: disable=unused-argument
     ):
         # A nested mapping of gate name -> qargs -> properties
-        self._gate_map = {}
         self._coupling_graph = None
         self._instruction_durations = None
         self._instruction_schedule_map = None
@@ -422,10 +421,7 @@ def add_instruction(self, instruction, properties=None, name=None):
             instruction_name = name
         if properties is None or is_class:
             properties = {None: None}
-        if instruction_name in self._gate_map:
-            raise AttributeError(f"Instruction {instruction_name} is already in the target")
         super().add_instruction(instruction, instruction_name, properties)
-        self._gate_map[instruction_name] = properties
         self._coupling_graph = None
         self._instruction_durations = None
         self._instruction_schedule_map = None
@@ -441,7 +437,6 @@ def update_instruction_properties(self, instruction, qargs, properties):
             KeyError: If ``instruction`` or ``qarg`` are not in the target
         """
         super().update_instruction_properties(instruction, qargs, properties)
-        self._gate_map[instruction][qargs] = properties
         self._instruction_durations = None
         self._instruction_schedule_map = None
 
@@ -487,7 +482,7 @@ def update_from_instruction_schedule_map(self, inst_map, inst_name_map=None, err
                 except TypeError:
                     qargs = (qargs,)
                 try:
-                    props = self._gate_map[inst_name][qargs]
+                    props = self[inst_name][qargs]
                 except (KeyError, TypeError):
                     props = None
 
@@ -521,7 +516,7 @@ def update_from_instruction_schedule_map(self, inst_map, inst_name_map=None, err
             if not out_props:
                 continue
             # Prepare Qiskit Gate object assigned to the entries
-            if inst_name not in self._gate_map:
+            if inst_name not in self:
                 # Entry not found: Add new instruction
                 if inst_name in qiskit_inst_name_map:
                     # Remove qargs with length that doesn't match with instruction qubit number
@@ -559,7 +554,7 @@ def update_from_instruction_schedule_map(self, inst_map, inst_name_map=None, err
             else:
                 # Entry found: Update "existing" instructions.
                 for qargs, prop in out_props.items():
-                    if qargs not in self._gate_map[inst_name]:
+                    if qargs not in self[inst_name]:
                         continue
                     self.update_instruction_properties(inst_name, qargs, prop)
 
@@ -571,9 +566,9 @@ def qargs_for_operation_name(self, operation):
         Returns:
             set: The set of qargs the gate instance applies to.
         """
-        if None in self._gate_map[operation]:
+        if None in self[operation]:
             return None
-        return self._gate_map[operation].keys()
+        return self[operation].keys()
 
     def durations(self):
         """Get an InstructionDurations object from the target
@@ -585,7 +580,7 @@ def durations(self):
         if self._instruction_durations is not None:
             return self._instruction_durations
         out_durations = []
-        for instruction, props_map in self._gate_map.items():
+        for instruction, props_map in self.items():
             for qarg, properties in props_map.items():
                 if properties is not None and properties.duration is not None:
                     out_durations.append((instruction, list(qarg), properties.duration, "s"))
@@ -613,7 +608,7 @@ def instruction_schedule_map(self):
         if self._instruction_schedule_map is not None:
             return self._instruction_schedule_map
         out_inst_schedule_map = InstructionScheduleMap()
-        for instruction, qargs in self._gate_map.items():
+        for instruction, qargs in self.items():
             for qarg, properties in qargs.items():
                 # Directly getting CalibrationEntry not to invoke .get_schedule().
                 # This keeps PulseQobjDef un-parsed.
@@ -639,11 +634,11 @@ def has_calibration(
             Returns ``True`` if the calibration is supported and ``False`` if it isn't.
         """
         qargs = tuple(qargs)
-        if operation_name not in self._gate_map:
+        if operation_name not in self:
             return False
-        if qargs not in self._gate_map[operation_name]:
+        if qargs not in self[operation_name]:
             return False
-        return getattr(self._gate_map[operation_name][qargs], "_calibration", None) is not None
+        return getattr(self[operation_name][qargs], "_calibration", None) is not None
 
     def get_calibration(
         self,
@@ -670,29 +665,9 @@ def get_calibration(
             raise KeyError(
                 f"Calibration of instruction {operation_name} for qubit {qargs} is not defined."
             )
-        cal_entry = getattr(self._gate_map[operation_name][qargs], "_calibration")
+        cal_entry = getattr(self[operation_name][qargs], "_calibration")
         return cal_entry.get_schedule(*args, **kwargs)
 
-    @property
-    def operation_names(self):
-        """Get the operation names in the target."""
-        return self._gate_map.keys()
-
-    @property
-    def instructions(self):
-        """Get the list of tuples ``(:class:`~qiskit.circuit.Instruction`, (qargs))``
-        for the target
-
-        For globally defined variable width operations the tuple will be of the form
-        ``(class, None)`` where class is the actual operation class that
-        is globally defined.
-        """
-        return [
-            (self._gate_name_map[op], qarg)
-            for op, qargs in self._gate_map.items()
-            for qarg in qargs
-        ]
-
     def instruction_properties(self, index):
         """Get the instruction properties for a specific instruction tuple
 
@@ -729,15 +704,15 @@ def instruction_properties(self, index):
             InstructionProperties: The instruction properties for the specified instruction tuple
         """
         instruction_properties = [
-            inst_props for qargs in self._gate_map.values() for inst_props in qargs.values()
+            inst_props for qargs in self.values() for inst_props in qargs.values()
         ]
         return instruction_properties[index]
 
     def _build_coupling_graph(self):
         self._coupling_graph = rx.PyDiGraph(multigraph=False)
         if self.num_qubits is not None:
             self._coupling_graph.add_nodes_from([{} for _ in range(self.num_qubits)])
-        for gate, qarg_map in self._gate_map.items():
+        for gate, qarg_map in self.items():
             if qarg_map is None:
                 if self._gate_name_map[gate].num_qubits == 2:
                     self._coupling_graph = None  # pylint: disable=attribute-defined-outside-init
@@ -841,37 +816,13 @@ def _filter_coupling_graph(self):
             graph.remove_nodes_from(list(to_remove))
         return graph
 
-    def __iter__(self):
-        return iter(self._gate_map)
-
-    def __getitem__(self, key):
-        return self._gate_map[key]
-
     def get(self, key, default=None):
         """Gets an item from the Target. If not found return a provided default or `None`."""
         try:
             return self[key]
         except KeyError:
             return default
 
-    def __len__(self):
-        return len(self._gate_map)
-
-    def __contains__(self, item):
-        return item in self._gate_map
-
-    def keys(self):
-        """Return the keys (operation_names) of the Target"""
-        return self._gate_map.keys()
-
-    def values(self):
-        """Return the Property Map (qargs -> InstructionProperties) of every instruction in the Target"""
-        return self._gate_map.values()
-
-    def items(self):
-        """Returns pairs of Gate names and its property map (str, dict[tuple, InstructionProperties])"""
-        return self._gate_map.items()
-
     def __str__(self):
         output = io.StringIO()
         if self.description is not None:
@@ -880,7 +831,7 @@ def __str__(self):
             output.write("Target\n")
         output.write(f"Number of qubits: {self.num_qubits}\n")
         output.write("Instructions:\n")
-        for inst, qarg_props in self._gate_map.items():
+        for inst, qarg_props in self.items():
             output.write(f"\t{inst}\n")
             for qarg, props in qarg_props.items():
                 if qarg is None:
@@ -910,15 +861,13 @@ def __str__(self):
 
     def __getstate__(self) -> dict:
         return {
-            "_gate_map": self._gate_map,
             "coupling_graph": self._coupling_graph,
             "instruction_durations": self._instruction_durations,
             "instruction_schedule_map": self._instruction_schedule_map,
             "base": super().__getstate__(),
         }
 
     def __setstate__(self, state: tuple):
-        self._gate_map = state["_gate_map"]
         self._coupling_graph = state["coupling_graph"]
         self._instruction_durations = state["instruction_durations"]
         self._instruction_schedule_map = state["instruction_schedule_map"]