From fdbf962af6751666e6725089c5de8c3ba293b98e Mon Sep 17 00:00:00 2001 From: Jake Lishman Date: Fri, 21 Jun 2024 11:15:07 +0100 Subject: [PATCH] Localise `py_op` caching data in `RefCell` (#12594) This localises the caching concerns of the `PackedInstruction::py_op` field into a method `unpack_py_op`, which can now update the cache through an immutable reference (if no other immutable references are taken out). Having the new method to encapsulate the `cache_pyops` feature simplifies the large `#[cfg(feature = "cache_pyop")]` gates. --- crates/circuit/src/circuit_data.rs | 363 ++-------------------- crates/circuit/src/circuit_instruction.rs | 82 ++++- 2 files changed, 102 insertions(+), 343 deletions(-) diff --git a/crates/circuit/src/circuit_data.rs b/crates/circuit/src/circuit_data.rs index da35787e3207..07f4579a4cd3 100644 --- a/crates/circuit/src/circuit_data.rs +++ b/crates/circuit/src/circuit_data.rs @@ -10,10 +10,13 @@ // copyright notice, and modified files need to carry a notice indicating // that they have been altered from the originals. +#[cfg(feature = "cache_pygates")] +use std::cell::RefCell; + use crate::bit_data::BitData; use crate::circuit_instruction::{ - convert_py_to_operation_type, operation_type_and_data_to_py, CircuitInstruction, - ExtraInstructionAttributes, OperationInput, PackedInstruction, + convert_py_to_operation_type, CircuitInstruction, ExtraInstructionAttributes, OperationInput, + PackedInstruction, }; use crate::imports::{BUILTIN_LIST, QUBIT}; use crate::interner::{IndexedInterner, Interner, InternerKey}; @@ -489,66 +492,40 @@ impl CircuitData { .getattr(intern!(py, "deepcopy"))?; for inst in &mut res.data { match &mut inst.op { - OperationType::Standard(_) => { - #[cfg(feature = "cache_pygates")] - { - inst.py_op = None; - } - } + OperationType::Standard(_) => {} OperationType::Gate(ref mut op) => { op.gate = deepcopy.call1((&op.gate,))?.unbind(); - #[cfg(feature = "cache_pygates")] - { - inst.py_op = None; - } } OperationType::Instruction(ref mut op) => { op.instruction = deepcopy.call1((&op.instruction,))?.unbind(); - #[cfg(feature = "cache_pygates")] - { - inst.py_op = None; - } } OperationType::Operation(ref mut op) => { op.operation = deepcopy.call1((&op.operation,))?.unbind(); - #[cfg(feature = "cache_pygates")] - { - inst.py_op = None; - } } }; + #[cfg(feature = "cache_pygates")] + { + *inst.py_op.borrow_mut() = None; + } } } else if copy_instructions { for inst in &mut res.data { match &mut inst.op { - OperationType::Standard(_) => { - #[cfg(feature = "cache_pygates")] - { - inst.py_op = None; - } - } + OperationType::Standard(_) => {} OperationType::Gate(ref mut op) => { op.gate = op.gate.call_method0(py, intern!(py, "copy"))?; - #[cfg(feature = "cache_pygates")] - { - inst.py_op = None; - } } OperationType::Instruction(ref mut op) => { op.instruction = op.instruction.call_method0(py, intern!(py, "copy"))?; - #[cfg(feature = "cache_pygates")] - { - inst.py_op = None; - } } OperationType::Operation(ref mut op) => { op.operation = op.operation.call_method0(py, intern!(py, "copy"))?; - #[cfg(feature = "cache_pygates")] - { - inst.py_op = None; - } } }; + #[cfg(feature = "cache_pygates")] + { + *inst.py_op.borrow_mut() = None; + } } } Ok(res) @@ -589,87 +566,10 @@ impl CircuitData { /// Args: /// func (Callable[[:class:`~.Operation`], None]): /// The callable to invoke. - #[cfg(not(feature = "cache_pygates"))] #[pyo3(signature = (func))] pub fn foreach_op(&self, py: Python<'_>, func: &Bound) -> PyResult<()> { for inst in self.data.iter() { - let label; - let duration; - let unit; - let condition; - match &inst.extra_attrs { - Some(extra_attrs) => { - label = &extra_attrs.label; - duration = &extra_attrs.duration; - unit = &extra_attrs.unit; - condition = &extra_attrs.condition; - } - None => { - label = &None; - duration = &None; - unit = &None; - condition = &None; - } - } - - let op = operation_type_and_data_to_py( - py, - &inst.op, - &inst.params, - label, - duration, - unit, - condition, - )?; - func.call1((op,))?; - } - Ok(()) - } - - /// Invokes callable ``func`` with each instruction's operation. - /// - /// Args: - /// func (Callable[[:class:`~.Operation`], None]): - /// The callable to invoke. - #[cfg(feature = "cache_pygates")] - #[pyo3(signature = (func))] - pub fn foreach_op(&mut self, py: Python<'_>, func: &Bound) -> PyResult<()> { - for inst in self.data.iter_mut() { - let op = match &inst.py_op { - Some(op) => op.clone_ref(py), - None => { - let label; - let duration; - let unit; - let condition; - match &inst.extra_attrs { - Some(extra_attrs) => { - label = &extra_attrs.label; - duration = &extra_attrs.duration; - unit = &extra_attrs.unit; - condition = &extra_attrs.condition; - } - None => { - label = &None; - duration = &None; - unit = &None; - condition = &None; - } - } - let new_op = operation_type_and_data_to_py( - py, - &inst.op, - &inst.params, - label, - duration, - unit, - condition, - )?; - inst.py_op = Some(new_op.clone_ref(py)); - new_op - } - }; - func.call1((op,))?; + func.call1((inst.unpack_py_op(py)?,))?; } Ok(()) } @@ -680,88 +580,10 @@ impl CircuitData { /// Args: /// func (Callable[[int, :class:`~.Operation`], None]): /// The callable to invoke. - #[cfg(not(feature = "cache_pygates"))] #[pyo3(signature = (func))] pub fn foreach_op_indexed(&self, py: Python<'_>, func: &Bound) -> PyResult<()> { for (index, inst) in self.data.iter().enumerate() { - let label; - let duration; - let unit; - let condition; - match &inst.extra_attrs { - Some(extra_attrs) => { - label = &extra_attrs.label; - duration = &extra_attrs.duration; - unit = &extra_attrs.unit; - condition = &extra_attrs.condition; - } - None => { - label = &None; - duration = &None; - unit = &None; - condition = &None; - } - } - - let op = operation_type_and_data_to_py( - py, - &inst.op, - &inst.params, - label, - duration, - unit, - condition, - )?; - func.call1((index, op))?; - } - Ok(()) - } - - /// Invokes callable ``func`` with the positional index and operation - /// of each instruction. - /// - /// Args: - /// func (Callable[[int, :class:`~.Operation`], None]): - /// The callable to invoke. - #[cfg(feature = "cache_pygates")] - #[pyo3(signature = (func))] - pub fn foreach_op_indexed(&mut self, py: Python<'_>, func: &Bound) -> PyResult<()> { - for (index, inst) in self.data.iter_mut().enumerate() { - let op = match &inst.py_op { - Some(op) => op.clone_ref(py), - None => { - let label; - let duration; - let unit; - let condition; - match &inst.extra_attrs { - Some(extra_attrs) => { - label = &extra_attrs.label; - duration = &extra_attrs.duration; - unit = &extra_attrs.unit; - condition = &extra_attrs.condition; - } - None => { - label = &None; - duration = &None; - unit = &None; - condition = &None; - } - } - let new_op = operation_type_and_data_to_py( - py, - &inst.op, - &inst.params, - label, - duration, - unit, - condition, - )?; - inst.py_op = Some(new_op.clone_ref(py)); - new_op - } - }; - func.call1((index, op))?; + func.call1((index, inst.unpack_py_op(py)?))?; } Ok(()) } @@ -779,49 +601,23 @@ impl CircuitData { /// func (Callable[[:class:`~.Operation`], :class:`~.Operation`]): /// A callable used to map original operation to their /// replacements. - #[cfg(not(feature = "cache_pygates"))] #[pyo3(signature = (func))] pub fn map_ops(&mut self, py: Python<'_>, func: &Bound) -> PyResult<()> { for inst in self.data.iter_mut() { - let old_op = match &inst.op { - OperationType::Standard(op) => { - let label; - let duration; - let unit; - let condition; - match &inst.extra_attrs { - Some(extra_attrs) => { - label = &extra_attrs.label; - duration = &extra_attrs.duration; - unit = &extra_attrs.unit; - condition = &extra_attrs.condition; - } - None => { - label = &None; - duration = &None; - unit = &None; - condition = &None; - } - } - if condition.is_some() { - operation_type_and_data_to_py( - py, - &inst.op, - &inst.params, - label, - duration, - unit, - condition, - )? - } else { - op.into_py(py) + let py_op = { + if let OperationType::Standard(op) = inst.op { + match inst.extra_attrs.as_deref() { + None + | Some(ExtraInstructionAttributes { + condition: None, .. + }) => op.into_py(py), + _ => inst.unpack_py_op(py)?, } + } else { + inst.unpack_py_op(py)? } - OperationType::Gate(op) => op.gate.clone_ref(py), - OperationType::Instruction(op) => op.instruction.clone_ref(py), - OperationType::Operation(op) => op.operation.clone_ref(py), }; - let result: OperationInput = func.call1((old_op,))?.extract()?; + let result: OperationInput = func.call1((py_op,))?.extract()?; match result { OperationInput::Standard(op) => { inst.op = OperationType::Standard(op); @@ -836,7 +632,7 @@ impl CircuitData { inst.op = OperationType::Operation(op); } OperationInput::Object(new_op) => { - let new_inst_details = convert_py_to_operation_type(py, new_op)?; + let new_inst_details = convert_py_to_operation_type(py, new_op.clone_ref(py))?; inst.op = new_inst_details.operation; inst.params = new_inst_details.params; if new_inst_details.label.is_some() @@ -851,103 +647,10 @@ impl CircuitData { condition: new_inst_details.condition, })) } - } - } - } - Ok(()) - } - - /// Invokes callable ``func`` with each instruction's operation, - /// replacing the operation with the result. - /// - /// .. note:: - /// - /// This is only to be used by map_vars() in quantumcircuit.py it - /// assumes that a full Python instruction will only be returned from - /// standard gates iff a condition is set. - /// - /// Args: - /// func (Callable[[:class:`~.Operation`], :class:`~.Operation`]): - /// A callable used to map original operation to their - /// replacements. - #[cfg(feature = "cache_pygates")] - #[pyo3(signature = (func))] - pub fn map_ops(&mut self, py: Python<'_>, func: &Bound) -> PyResult<()> { - for inst in self.data.iter_mut() { - let old_op = match &inst.py_op { - Some(op) => op.clone_ref(py), - None => match &inst.op { - OperationType::Standard(op) => { - let label; - let duration; - let unit; - let condition; - match &inst.extra_attrs { - Some(extra_attrs) => { - label = &extra_attrs.label; - duration = &extra_attrs.duration; - unit = &extra_attrs.unit; - condition = &extra_attrs.condition; - } - None => { - label = &None; - duration = &None; - unit = &None; - condition = &None; - } - } - if condition.is_some() { - let new_op = operation_type_and_data_to_py( - py, - &inst.op, - &inst.params, - label, - duration, - unit, - condition, - )?; - inst.py_op = Some(new_op.clone_ref(py)); - new_op - } else { - op.into_py(py) - } - } - OperationType::Gate(op) => op.gate.clone_ref(py), - OperationType::Instruction(op) => op.instruction.clone_ref(py), - OperationType::Operation(op) => op.operation.clone_ref(py), - }, - }; - let result: OperationInput = func.call1((old_op,))?.extract()?; - match result { - OperationInput::Standard(op) => { - inst.op = OperationType::Standard(op); - } - OperationInput::Gate(op) => { - inst.op = OperationType::Gate(op); - } - OperationInput::Instruction(op) => { - inst.op = OperationType::Instruction(op); - } - OperationInput::Operation(op) => { - inst.op = OperationType::Operation(op); - } - OperationInput::Object(new_op) => { - let new_inst_details = convert_py_to_operation_type(py, new_op.clone_ref(py))?; - inst.op = new_inst_details.operation; - inst.params = new_inst_details.params; - if new_inst_details.label.is_some() - || new_inst_details.duration.is_some() - || new_inst_details.unit.is_some() - || new_inst_details.condition.is_some() + #[cfg(feature = "cache_pygates")] { - inst.extra_attrs = Some(Box::new(ExtraInstructionAttributes { - label: new_inst_details.label, - duration: new_inst_details.duration, - unit: new_inst_details.unit, - condition: new_inst_details.condition, - })) + *inst.py_op.borrow_mut() = Some(new_op); } - inst.py_op = Some(new_op); } } } @@ -1537,7 +1240,7 @@ impl CircuitData { params: inst.params.clone(), extra_attrs: inst.extra_attrs.clone(), #[cfg(feature = "cache_pygates")] - py_op: inst.py_op.clone(), + py_op: RefCell::new(inst.py_op.clone()), }) } @@ -1557,7 +1260,7 @@ impl CircuitData { params: inst.params.clone(), extra_attrs: inst.extra_attrs.clone(), #[cfg(feature = "cache_pygates")] - py_op: inst.py_op.clone(), + py_op: RefCell::new(inst.py_op.clone()), }) } } diff --git a/crates/circuit/src/circuit_instruction.rs b/crates/circuit/src/circuit_instruction.rs index 2bb90367082d..5179190d8aa2 100644 --- a/crates/circuit/src/circuit_instruction.rs +++ b/crates/circuit/src/circuit_instruction.rs @@ -10,6 +10,9 @@ // copyright notice, and modified files need to carry a notice indicating // that they have been altered from the originals. +#[cfg(feature = "cache_pygates")] +use std::cell::RefCell; + use pyo3::basic::CompareOp; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; @@ -47,8 +50,61 @@ pub(crate) struct PackedInstruction { pub clbits_id: Index, pub params: SmallVec<[Param; 3]>, pub extra_attrs: Option>, + #[cfg(feature = "cache_pygates")] - pub py_op: Option, + /// This is hidden in a `RefCell` because, while that has additional memory-usage implications + /// while we're still building with the feature enabled, we intend to remove the feature in the + /// future, and hiding the cache within a `RefCell` lets us keep the cache transparently in our + /// interfaces, without needing various functions to unnecessarily take `&mut` references. + pub py_op: RefCell>, +} + +impl PackedInstruction { + /// Build a reference to the Python-space operation object (the `Gate`, etc) packed into this + /// instruction. This may construct the reference if the `PackedInstruction` is a standard + /// gate with no already stored operation. + /// + /// A standard-gate operation object returned by this function is disconnected from the + /// containing circuit; updates to its label, duration, unit and condition will not be + /// propagated back. + pub fn unpack_py_op(&self, py: Python) -> PyResult> { + #[cfg(feature = "cache_pygates")] + { + if let Some(cached_op) = self.py_op.borrow().as_ref() { + return Ok(cached_op.clone_ref(py)); + } + } + let (label, duration, unit, condition) = match self.extra_attrs.as_deref() { + Some(ExtraInstructionAttributes { + label, + duration, + unit, + condition, + }) => ( + label.as_deref(), + duration.as_ref(), + unit.as_deref(), + condition.as_ref(), + ), + None => (None, None, None, None), + }; + let out = operation_type_and_data_to_py( + py, + &self.op, + &self.params, + label, + duration, + unit, + condition, + )?; + #[cfg(feature = "cache_pygates")] + { + if let Ok(mut cell) = self.py_op.try_borrow_mut() { + cell.get_or_insert_with(|| out.clone_ref(py)); + } + } + Ok(out) + } } /// A single instruction in a :class:`.QuantumCircuit`, comprised of the :attr:`operation` and @@ -666,20 +722,20 @@ pub(crate) fn operation_type_to_py( let (label, duration, unit, condition) = match &circuit_inst.extra_attrs { None => (None, None, None, None), Some(extra_attrs) => ( - extra_attrs.label.clone(), - extra_attrs.duration.clone(), - extra_attrs.unit.clone(), - extra_attrs.condition.clone(), + extra_attrs.label.as_deref(), + extra_attrs.duration.as_ref(), + extra_attrs.unit.as_deref(), + extra_attrs.condition.as_ref(), ), }; operation_type_and_data_to_py( py, &circuit_inst.operation, &circuit_inst.params, - &label, - &duration, - &unit, - &condition, + label, + duration, + unit, + condition, ) } @@ -692,10 +748,10 @@ pub(crate) fn operation_type_and_data_to_py( py: Python, operation: &OperationType, params: &[Param], - label: &Option, - duration: &Option, - unit: &Option, - condition: &Option, + label: Option<&str>, + duration: Option<&PyObject>, + unit: Option<&str>, + condition: Option<&PyObject>, ) -> PyResult { match &operation { OperationType::Standard(op) => {