diff --git a/crates/accelerate/src/target_transpiler/mod.rs b/crates/accelerate/src/target_transpiler/mod.rs index 4f6b14f09c77..a1c4823dc116 100644 --- a/crates/accelerate/src/target_transpiler/mod.rs +++ b/crates/accelerate/src/target_transpiler/mod.rs @@ -14,6 +14,7 @@ mod errors; mod instruction_properties; +mod nullable_index_map; use std::ops::Index; @@ -22,6 +23,7 @@ use ahash::RandomState; use ahash::HashSet; use indexmap::{IndexMap, IndexSet}; use itertools::Itertools; +use nullable_index_map::NullableIndexMap; use pyo3::{ exceptions::{PyAttributeError, PyIndexError, PyKeyError, PyValueError}, prelude::*, @@ -31,7 +33,6 @@ use pyo3::{ use qiskit_circuit::circuit_instruction::convert_py_to_operation_type; use qiskit_circuit::operations::{Operation, OperationType, Param}; -use rustworkx_core::dictmap::InitWithHasher; use smallvec::SmallVec; use crate::nlayout::PhysicalQubit; @@ -50,11 +51,11 @@ mod exceptions { // Custom types type Qargs = SmallVec<[PhysicalQubit; 2]>; type GateMap = IndexMap; -type PropsMap = IndexMap, Option, RandomState>; +type PropsMap = NullableIndexMap>; type GateMapState = Vec<(String, Vec<(Option, Option)>)>; #[derive(Debug, Clone, FromPyObject)] -enum TargetOperation { +pub enum TargetOperation { Normal(NormalOperation), Variadic(PyObject), } @@ -94,7 +95,7 @@ impl TargetOperation { } #[derive(Debug, Clone)] -struct NormalOperation { +pub struct NormalOperation { operation: OperationType, params: SmallVec<[Param; 3]>, op_object: PyObject, @@ -167,7 +168,7 @@ pub struct Target { _gate_name_map: IndexMap, global_operations: IndexMap, RandomState>, variable_class_operations: IndexSet, - qarg_gate_map: IndexMap, Option>, RandomState>, + qarg_gate_map: NullableIndexMap>>, non_global_strict_basis: Option>, non_global_basis: Option>, } @@ -264,7 +265,7 @@ impl Target { _gate_name_map: IndexMap::default(), variable_class_operations: IndexSet::default(), global_operations: IndexMap::default(), - qarg_gate_map: IndexMap::default(), + qarg_gate_map: NullableIndexMap::default(), non_global_basis: None, non_global_strict_basis: None, }) @@ -340,7 +341,7 @@ impl Target { &mut self, instruction: TargetOperation, name: &str, - mut properties: Option, + properties: Option, ) -> PyResult<()> { if self.gate_map.contains_key(name) { return Err(PyAttributeError::new_err(format!( @@ -356,10 +357,10 @@ impl Target { self.variable_class_operations.insert(name.to_string()); } TargetOperation::Normal(_) => { - if let Some(properties) = properties.as_mut() { + if let Some(mut properties) = properties { qargs_val = PropsMap::with_capacity(properties.len()); let inst_num_qubits = instruction.num_qubits(); - if properties.contains_key(&None) { + if properties.contains_key(None) { self.global_operations .entry(inst_num_qubits) .and_modify(|e| { @@ -367,7 +368,8 @@ impl Target { }) .or_insert(HashSet::from_iter([name.to_string()])); } - let property_keys: Vec> = properties.keys().cloned().collect(); + let property_keys: Vec> = + properties.keys().map(|qargs| qargs.cloned()).collect(); for qarg in property_keys { if let Some(qarg) = qarg.as_ref() { if qarg.len() != inst_num_qubits as usize { @@ -388,16 +390,14 @@ impl Target { }) + 1, )); } - let inst_properties = properties.swap_remove(&qarg).unwrap(); + let inst_properties = properties.swap_remove(qarg.as_ref()).unwrap(); qargs_val.insert(qarg.clone(), inst_properties); - self.qarg_gate_map - .entry(qarg) - .and_modify(|e| { - if let Some(e) = e { - e.insert(name.to_string()); - } - }) - .or_insert(Some(HashSet::from_iter([name.to_string()]))); + if let Some(Some(value)) = self.qarg_gate_map.get_mut(qarg.as_ref()) { + value.insert(name.to_string()); + } else { + self.qarg_gate_map + .insert(qarg, Some(HashSet::from_iter([name.to_string()]))); + } } } else { qargs_val = PropsMap::with_capacity(0); @@ -433,14 +433,16 @@ impl Target { ))); }; let mut prop_map = self[&instruction].clone(); - if !(prop_map.contains_key(&qargs)) { + if !(prop_map.contains_key(qargs.as_ref())) { return Err(PyKeyError::new_err(format!( "Provided qarg {:?} not in this Target for {:?}.", &qargs.unwrap_or_default(), &instruction ))); } - prop_map.entry(qargs).and_modify(|e| *e = properties); + if let Some(e) = prop_map.get_mut(qargs.as_ref()) { + *e = properties; + } self.gate_map .entry(instruction) .and_modify(|e| *e = prop_map); @@ -522,7 +524,7 @@ impl Target { /// KeyError: If ``qargs`` is not in target #[pyo3(name = "operation_names_for_qargs", signature=(qargs=None, /))] pub fn py_operation_names_for_qargs(&self, qargs: Option) -> PyResult> { - match self.operation_names_for_qargs(&qargs) { + match self.operation_names_for_qargs(qargs.as_ref()) { Ok(set) => Ok(set), Err(e) => Err(PyKeyError::new_err(e.message)), } @@ -632,10 +634,10 @@ impl Target { if let Some(_qargs) = &qargs { if self.gate_map.contains_key(op_name) { let gate_map_name = &self.gate_map[op_name]; - if gate_map_name.contains_key(&qargs) { + if gate_map_name.contains_key(qargs.as_ref()) { return Ok(true); } - if gate_map_name.contains_key(&None) { + if gate_map_name.contains_key(None) { let qubit_comparison = self._gate_name_map[op_name].num_qubits(); return Ok(qubit_comparison == _qargs.len() as u32 @@ -692,7 +694,7 @@ impl Target { return Ok(true); } } - Ok(self.instruction_supported(&operation_name, &qargs)) + Ok(self.instruction_supported(&operation_name, qargs.as_ref())) } else { Ok(false) } @@ -914,7 +916,7 @@ impl Target { .unwrap() .extract::()? .into_iter() - .map(|(name, prop_map)| (name, IndexMap::from_iter(prop_map.into_iter()))), + .map(|(name, prop_map)| (name, PropsMap::from_iter(prop_map.into_iter()))), ); self._gate_name_map = state .get_item("gate_name_map")? @@ -924,7 +926,7 @@ impl Target { .get_item("global_operations")? .unwrap() .extract::, RandomState>>()?; - self.qarg_gate_map = IndexMap::from_iter( + self.qarg_gate_map = NullableIndexMap::from_iter( state .get_item("qarg_gate_map")? .unwrap() @@ -962,7 +964,7 @@ impl Target { self.gate_map.iter().flat_map(move |(op, props_map)| { props_map .keys() - .map(move |qargs| (&self._gate_name_map[op], qargs.as_ref())) + .map(move |qargs| (&self._gate_name_map[op], qargs)) }) } /// Returns an iterator over the operation names in the target. @@ -971,15 +973,7 @@ impl Target { } /// Get the operation objects in the target. - pub fn operations(&self) -> impl Iterator)> { - return self._operations().filter_map(|operation| match operation { - TargetOperation::Normal(normal) => Some((&normal.operation, &normal.params)), - _ => None, - }); - } - - /// Get the operation objects in the target. - fn _operations(&self) -> impl Iterator { + pub fn operations(&self) -> impl Iterator { return self._gate_name_map.values(); } @@ -1063,13 +1057,13 @@ impl Target { /// Gets all the operation names that use these qargs. Rust native equivalent of ``BaseTarget.operation_names_for_qargs()`` pub fn operation_names_for_qargs( &self, - qargs: &Option, + qargs: Option<&Qargs>, ) -> Result, TargetKeyError> { // When num_qubits == 0 we return globally defined operators let mut res: HashSet<&str> = HashSet::default(); let mut qargs = qargs; if self.num_qubits.unwrap_or_default() == 0 || self.num_qubits.is_none() { - qargs = &None; + qargs = None; } if let Some(qargs) = qargs.as_ref() { if qargs @@ -1107,7 +1101,7 @@ impl Target { /// Returns an iterator rust-native operation instances and parameters present in the Target that affect the provided qargs. pub fn ops_from_qargs( &self, - qargs: &Option, + qargs: Option<&Qargs>, ) -> Result)>, TargetKeyError> { match self.operation_names_for_qargs(qargs) { Ok(operations) => { @@ -1132,7 +1126,7 @@ impl Target { operation: &str, ) -> Result>, TargetKeyError> { if let Some(gate_map_oper) = self.gate_map.get(operation) { - if gate_map_oper.contains_key(&None) { + if gate_map_oper.contains_key(None) { return Ok(None); } let qargs = gate_map_oper.keys().flatten(); @@ -1183,18 +1177,18 @@ impl Target { if qargs.len() == 1 && is_none { return None; } - Some(qargs.map(|qarg| qarg.as_ref())) + Some(qargs) } - pub fn instruction_supported(&self, operation_name: &str, qargs: &Option) -> bool { + pub fn instruction_supported(&self, operation_name: &str, qargs: Option<&Qargs>) -> bool { if self.gate_map.contains_key(operation_name) { - if let Some(_qargs) = qargs.as_ref() { - let qarg_set: HashSet = _qargs.iter().cloned().collect(); + if let Some(_qargs) = qargs { + let qarg_set: HashSet<&PhysicalQubit> = _qargs.iter().collect(); if let Some(gate_prop_name) = self.gate_map.get(operation_name) { if gate_prop_name.contains_key(qargs) { return true; } - if gate_prop_name.contains_key(&None) { + if gate_prop_name.contains_key(None) { let obj = &self._gate_name_map[operation_name]; if self.variable_class_operations.contains(operation_name) { return qargs.is_none() diff --git a/crates/accelerate/src/target_transpiler/nullable_index_map.rs b/crates/accelerate/src/target_transpiler/nullable_index_map.rs new file mode 100644 index 000000000000..35d747660857 --- /dev/null +++ b/crates/accelerate/src/target_transpiler/nullable_index_map.rs @@ -0,0 +1,455 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2024 +// +// This code is licensed under the Apache License, Version 2.0. You may +// obtain a copy of this license in the LICENSE.txt file in the root directory +// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +// +// Any modifications or derivative works of this code must retain this +// copyright notice, and modified files need to carry a notice indicating +// that they have been altered from the originals. + +use ahash::RandomState; +use indexmap::{ + map::{IntoIter as BaseIntoIter, Iter as BaseIter, Keys as BaseKeys, Values as BaseValues}, + IndexMap, +}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use pyo3::IntoPy; +use rustworkx_core::dictmap::InitWithHasher; +use std::ops::Index; +use std::{hash::Hash, mem::swap}; + +type BaseMap = IndexMap; + +/// +/// An `IndexMap`-like structure thet can be used when one of the keys can have a `None` value. +/// +/// This structure is essentially a wrapper around the `IndexMap` struct that allows the +/// storage of `Option` key values as `K`` and keep an extra slot reserved only for the +/// `None` instance. There are some upsides to this including: +/// +/// The ability to index using Option<&K> to index a specific key. +/// Store keys as non option wrapped to obtain references to K instead of reference to Option. +/// +/// **Warning:** This is an experimental feature and should be used with care as it does not +/// fully implement all the methods present in `IndexMap` due to API limitations. +#[derive(Debug, Clone)] +pub struct NullableIndexMap +where + K: Eq + Hash + Clone, + V: Clone, +{ + map: BaseMap, + null_val: Option, +} + +impl NullableIndexMap +where + K: Eq + Hash + Clone, + V: Clone, +{ + /// Returns a reference to the value stored at `key`, if it does not exist + /// `None` is returned instead. + pub fn get(&self, key: Option<&K>) -> Option<&V> { + match key { + Some(key) => self.map.get(key), + None => self.null_val.as_ref(), + } + } + + /// Returns a mutable reference to the value stored at `key`, if it does not + /// exist `None` is returned instead. + pub fn get_mut(&mut self, key: Option<&K>) -> Option<&mut V> { + match key { + Some(key) => self.map.get_mut(key), + None => self.null_val.as_mut(), + } + } + + /// Inserts a `value` in the slot alotted to `key`. + /// + /// If a previous value existed there previously it will be returned, otherwise + /// `None` will be returned. + pub fn insert(&mut self, key: Option, value: V) -> Option { + match key { + Some(key) => self.map.insert(key, value), + None => { + let mut old_val = Some(value); + swap(&mut old_val, &mut self.null_val); + old_val + } + } + } + + /// Creates an instance of `NullableIndexMap` with capacity to hold `n`+1 key-value + /// pairs. + /// + /// Notice that an extra space needs to be alotted to store the instance of `None` a + /// key. + pub fn with_capacity(n: usize) -> Self { + Self { + map: BaseMap::with_capacity(n), + null_val: None, + } + } + + /// Creates an instance of `NullableIndexMap` from an iterator over instances of + /// `(Option, V)`. + pub fn from_iter<'a, I>(iter: I) -> Self + where + I: IntoIterator, V)> + 'a, + { + let mut null_val = None; + let filtered = iter.into_iter().filter_map(|item| match item { + (Some(key), value) => Some((key, value)), + (None, value) => { + null_val = Some(value); + None + } + }); + Self { + map: IndexMap::from_iter(filtered), + null_val, + } + } + + /// Returns `true` if the map contains a slot indexed by `key`, otherwise `false`. + pub fn contains_key(&self, key: Option<&K>) -> bool { + match key { + Some(key) => self.map.contains_key(key), + None => self.null_val.is_some(), + } + } + + /// Extends the key-value pairs in the map with the contents of an iterator over + /// `(Option, V)`. + /// + /// If an already existent key is provided, it will be replaced by the entry provided + /// in the iterator. + pub fn extend<'a, I>(&mut self, iter: I) + where + I: IntoIterator, V)> + 'a, + { + let filtered = iter.into_iter().filter_map(|item| match item { + (Some(key), value) => Some((key, value)), + (None, value) => { + self.null_val = Some(value); + None + } + }); + self.map.extend(filtered) + } + + /// Removes the entry allotted to `key` from the map and returns it. The index of + /// this entry is then replaced by the entry located at the last index. + /// + /// `None` will be returned if the `key` is not present in the map. + pub fn swap_remove(&mut self, key: Option<&K>) -> Option { + match key { + Some(key) => self.map.swap_remove(key), + None => { + let mut ret_val = None; + swap(&mut ret_val, &mut self.null_val); + ret_val + } + } + } + + /// Returns an iterator over references of the key-value pairs of the map. + pub fn iter(&self) -> Iter { + Iter { + map: self.map.iter(), + null_value: &self.null_val, + } + } + + /// Returns an iterator over references of the keys present in the map. + pub fn keys(&self) -> Keys { + Keys { + map_keys: self.map.keys(), + null_value: self.null_val.is_some(), + } + } + + /// Returns an iterator over references of all the values present in the map. + pub fn values(&self) -> Values { + Values { + map_values: self.map.values(), + null_value: &self.null_val, + } + } + + /// Returns the number of key-value pairs present in the map. + pub fn len(&self) -> usize { + self.map.len() + self.null_val.is_some() as usize + } +} + +impl IntoIterator for NullableIndexMap +where + K: Eq + Hash + Clone, + V: Clone, +{ + type Item = (Option, V); + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter { + map: self.map.into_iter(), + null_value: self.null_val, + } + } +} + +/// Iterator for the key-value pairs in `NullableIndexMap`. +pub struct Iter<'a, K, V> { + map: BaseIter<'a, K, V>, + null_value: &'a Option, +} + +impl<'a, K, V> Iterator for Iter<'a, K, V> { + type Item = (Option<&'a K>, &'a V); + + fn next(&mut self) -> Option { + if let Some((key, val)) = self.map.next() { + Some((Some(key), val)) + } else if let Some(value) = self.null_value { + let value = value; + self.null_value = &None; + Some((None, value)) + } else { + None + } + } + + fn size_hint(&self) -> (usize, Option) { + ( + self.map.size_hint().0 + self.null_value.is_some() as usize, + self.map + .size_hint() + .1 + .map(|hint| hint + self.null_value.is_some() as usize), + ) + } +} + +impl<'a, K, V> ExactSizeIterator for Iter<'a, K, V> { + fn len(&self) -> usize { + self.map.len() + self.null_value.is_some() as usize + } +} + +/// Owned iterator over the key-value pairs in `NullableIndexMap`. +pub struct IntoIter +where + V: Clone, +{ + map: BaseIntoIter, + null_value: Option, +} + +impl Iterator for IntoIter +where + V: Clone, +{ + type Item = (Option, V); + + fn next(&mut self) -> Option { + if let Some((key, val)) = self.map.next() { + Some((Some(key), val)) + } else if self.null_value.is_some() { + let mut value = None; + swap(&mut value, &mut self.null_value); + Some((None, value.unwrap())) + } else { + None + } + } + + fn size_hint(&self) -> (usize, Option) { + ( + self.map.size_hint().0 + self.null_value.is_some() as usize, + self.map + .size_hint() + .1 + .map(|hint| hint + self.null_value.is_some() as usize), + ) + } +} + +impl ExactSizeIterator for IntoIter +where + V: Clone, +{ + fn len(&self) -> usize { + self.map.len() + self.null_value.is_some() as usize + } +} + +/// Iterator over the keys of a `NullableIndexMap`. +pub struct Keys<'a, K, V> { + map_keys: BaseKeys<'a, K, V>, + null_value: bool, +} + +impl<'a, K, V> Iterator for Keys<'a, K, V> { + type Item = Option<&'a K>; + + fn next(&mut self) -> Option { + if let Some(key) = self.map_keys.next() { + Some(Some(key)) + } else if self.null_value { + self.null_value = false; + Some(None) + } else { + None + } + } + + fn size_hint(&self) -> (usize, Option) { + ( + self.map_keys.size_hint().0 + self.null_value as usize, + self.map_keys + .size_hint() + .1 + .map(|hint| hint + self.null_value as usize), + ) + } +} + +impl<'a, K, V> ExactSizeIterator for Keys<'a, K, V> { + fn len(&self) -> usize { + self.map_keys.len() + self.null_value as usize + } +} + +/// Iterator over the values of a `NullableIndexMap`. +pub struct Values<'a, K, V> { + map_values: BaseValues<'a, K, V>, + null_value: &'a Option, +} + +impl<'a, K, V> Iterator for Values<'a, K, V> { + type Item = &'a V; + + fn next(&mut self) -> Option { + if let Some(value) = self.map_values.next() { + Some(value) + } else if self.null_value.is_some() { + let return_value = self.null_value; + self.null_value = &None; + return_value.as_ref() + } else { + None + } + } + + fn size_hint(&self) -> (usize, Option) { + ( + self.map_values.size_hint().0 + self.null_value.is_some() as usize, + self.map_values + .size_hint() + .1 + .map(|hint| hint + self.null_value.is_some() as usize), + ) + } +} + +impl<'a, K, V> ExactSizeIterator for Values<'a, K, V> { + fn len(&self) -> usize { + self.map_values.len() + self.null_value.is_some() as usize + } +} + +impl Index> for NullableIndexMap +where + K: Eq + Hash + Clone, + V: Clone, +{ + type Output = V; + fn index(&self, index: Option<&K>) -> &Self::Output { + match index { + Some(k) => self.map.index(k), + None => match &self.null_val { + Some(val) => val, + None => panic!("The provided key is not present in map: None"), + }, + } + } +} + +impl Default for NullableIndexMap +where + K: Eq + Hash + Clone, + V: Clone, +{ + fn default() -> Self { + Self { + map: IndexMap::default(), + null_val: None, + } + } +} + +impl<'py, K, V> FromPyObject<'py> for NullableIndexMap +where + K: IntoPy + FromPyObject<'py> + Eq + Hash + Clone, + V: IntoPy + FromPyObject<'py> + Clone, +{ + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let map: IndexMap, V, RandomState> = ob.extract()?; + let mut null_val: Option = None; + let filtered = map + .into_iter() + .filter_map(|(key, value)| match (key, value) { + (Some(key), value) => Some((key, value)), + (None, value) => { + null_val = Some(value); + None + } + }); + Ok(Self { + map: filtered.collect(), + null_val, + }) + } +} + +impl IntoPy for NullableIndexMap +where + K: IntoPy + Eq + Hash + Clone, + V: IntoPy + Clone, +{ + fn into_py(self, py: Python<'_>) -> PyObject { + let map_object = self.map.into_py(py); + let bound_map_obj = map_object.bind(py); + let downcast_dict: &Bound = bound_map_obj.downcast().unwrap(); + if let Some(null_val) = self.null_val { + downcast_dict + .set_item(py.None(), null_val.into_py(py)) + .unwrap(); + } + map_object + } +} + +impl ToPyObject for NullableIndexMap +where + K: ToPyObject + Eq + Hash + Clone, + V: ToPyObject + Clone, +{ + fn to_object(&self, py: Python<'_>) -> PyObject { + let map_object = self.map.to_object(py); + let bound_map_obj = map_object.bind(py); + let downcast_dict: &Bound = bound_map_obj.downcast().unwrap(); + if let Some(null_val) = &self.null_val { + downcast_dict + .set_item(py.None(), null_val.to_object(py)) + .unwrap(); + } + map_object + } +}