From 3a9993a461df63b80564280ac34e51a43d784ad3 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Mon, 11 Nov 2024 13:12:27 -0500 Subject: [PATCH] Use OnceLock instead of OnceCell (#13410) * Use OnceLock instead of OnceCell OnceLock is a thread-safe version of OnceCell that enables us to use PackedInstruction from a threaded environment. There is some overhead associated with this, primarily in memory as the OnceLock is a larger type than a OnceCell. But the tradeoff is worth it to start leverage multithreading for circuits. Fixes #13219 * Update twirling too * Fix rustfmt --- crates/accelerate/src/twirling.rs | 10 +++++----- crates/accelerate/src/unitary_synthesis.rs | 6 +++--- crates/circuit/src/circuit_data.rs | 12 ++++++------ crates/circuit/src/circuit_instruction.rs | 6 +++--- crates/circuit/src/converters.rs | 4 ++-- crates/circuit/src/dag_circuit.rs | 18 +++++++++--------- crates/circuit/src/dag_node.rs | 8 ++++---- crates/circuit/src/packed_instruction.rs | 14 +++++++------- 8 files changed, 39 insertions(+), 39 deletions(-) diff --git a/crates/accelerate/src/twirling.rs b/crates/accelerate/src/twirling.rs index 0867bc161555..29c9da3671bc 100644 --- a/crates/accelerate/src/twirling.rs +++ b/crates/accelerate/src/twirling.rs @@ -209,7 +209,7 @@ fn twirl_gate( params: None, extra_attrs: ExtraInstructionAttributes::new(None, None, None, None), #[cfg(feature = "cache_pygates")] - py_op: std::cell::OnceCell::new(), + py_op: std::sync::OnceLock::new(), }, )?; out_circ.push( @@ -221,7 +221,7 @@ fn twirl_gate( params: None, extra_attrs: ExtraInstructionAttributes::new(None, None, None, None), #[cfg(feature = "cache_pygates")] - py_op: std::cell::OnceCell::new(), + py_op: std::sync::OnceLock::new(), }, )?; @@ -235,7 +235,7 @@ fn twirl_gate( params: None, extra_attrs: ExtraInstructionAttributes::new(None, None, None, None), #[cfg(feature = "cache_pygates")] - py_op: std::cell::OnceCell::new(), + py_op: std::sync::OnceLock::new(), }, )?; out_circ.push( @@ -247,7 +247,7 @@ fn twirl_gate( params: None, extra_attrs: ExtraInstructionAttributes::new(None, None, None, None), #[cfg(feature = "cache_pygates")] - py_op: std::cell::OnceCell::new(), + py_op: std::sync::OnceLock::new(), }, )?; @@ -361,7 +361,7 @@ fn generate_twirled_circuit( )), extra_attrs: inst.extra_attrs.clone(), #[cfg(feature = "cache_pygates")] - py_op: std::cell::OnceCell::new(), + py_op: std::sync::OnceLock::new(), }; #[cfg(feature = "cache_pygates")] new_inst.py_op.set(new_inst_obj).unwrap(); diff --git a/crates/accelerate/src/unitary_synthesis.rs b/crates/accelerate/src/unitary_synthesis.rs index b5ef66db4f1e..62f41c78084c 100644 --- a/crates/accelerate/src/unitary_synthesis.rs +++ b/crates/accelerate/src/unitary_synthesis.rs @@ -11,9 +11,9 @@ // that they have been altered from the originals. #![allow(clippy::too_many_arguments)] -#[cfg(feature = "cache_pygates")] -use std::cell::OnceCell; use std::f64::consts::PI; +#[cfg(feature = "cache_pygates")] +use std::sync::OnceLock; use approx::relative_eq; use hashbrown::{HashMap, HashSet}; @@ -149,7 +149,7 @@ fn apply_synth_sequence( params: new_params, extra_attrs: ExtraInstructionAttributes::default(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }; instructions.push(instruction); } diff --git a/crates/circuit/src/circuit_data.rs b/crates/circuit/src/circuit_data.rs index f24277e18c7e..0025664e4d94 100644 --- a/crates/circuit/src/circuit_data.rs +++ b/crates/circuit/src/circuit_data.rs @@ -11,7 +11,7 @@ // that they have been altered from the originals. #[cfg(feature = "cache_pygates")] -use std::cell::OnceCell; +use std::sync::OnceLock; use crate::bit_data::BitData; use crate::circuit_instruction::{ @@ -303,7 +303,7 @@ impl CircuitData { params: inst.params.clone(), extra_attrs: inst.extra_attrs.clone(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }); } } else if copy_instructions { @@ -315,7 +315,7 @@ impl CircuitData { params: inst.params.clone(), extra_attrs: inst.extra_attrs.clone(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }); } } else { @@ -940,7 +940,7 @@ impl CircuitData { params, extra_attrs: ExtraInstructionAttributes::default(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }); res.track_instruction_parameters(py, res.data.len() - 1)?; } @@ -1049,7 +1049,7 @@ impl CircuitData { params, extra_attrs: ExtraInstructionAttributes::default(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }); res.track_instruction_parameters(py, res.data.len() - 1)?; } @@ -1107,7 +1107,7 @@ impl CircuitData { params, extra_attrs: ExtraInstructionAttributes::default(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }); Ok(()) } diff --git a/crates/circuit/src/circuit_instruction.rs b/crates/circuit/src/circuit_instruction.rs index e1bf3cfc2e55..e449ad660e38 100644 --- a/crates/circuit/src/circuit_instruction.rs +++ b/crates/circuit/src/circuit_instruction.rs @@ -11,7 +11,7 @@ // that they have been altered from the originals. #[cfg(feature = "cache_pygates")] -use std::cell::OnceCell; +use std::sync::OnceLock; use numpy::IntoPyArray; use pyo3::basic::CompareOp; @@ -236,7 +236,7 @@ pub struct CircuitInstruction { pub params: SmallVec<[Param; 3]>, pub extra_attrs: ExtraInstructionAttributes, #[cfg(feature = "cache_pygates")] - pub py_op: OnceCell>, + pub py_op: OnceLock>, } impl CircuitInstruction { @@ -301,7 +301,7 @@ impl CircuitInstruction { params, extra_attrs: ExtraInstructionAttributes::new(label, None, None, None), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }) } diff --git a/crates/circuit/src/converters.rs b/crates/circuit/src/converters.rs index dea366d02ff6..030a582bdad7 100644 --- a/crates/circuit/src/converters.rs +++ b/crates/circuit/src/converters.rs @@ -11,7 +11,7 @@ // that they have been altered from the originals. #[cfg(feature = "cache_pygates")] -use std::cell::OnceCell; +use std::sync::OnceLock; use ::pyo3::prelude::*; use hashbrown::HashMap; @@ -126,7 +126,7 @@ pub fn dag_to_circuit( )), extra_attrs: instr.extra_attrs.clone(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }) } else { Ok(instr.clone()) diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs index eb49130fcd37..25e52831683c 100644 --- a/crates/circuit/src/dag_circuit.rs +++ b/crates/circuit/src/dag_circuit.rs @@ -68,7 +68,7 @@ use std::convert::Infallible; use std::f64::consts::PI; #[cfg(feature = "cache_pygates")] -use std::cell::OnceCell; +use std::sync::OnceLock; static CONTROL_FLOW_OP_NAMES: [&str; 4] = ["for_loop", "while_loop", "if_else", "switch_case"]; static SEMANTIC_EQ_SYMMETRIC: [&str; 4] = ["barrier", "swap", "break_loop", "continue_loop"]; @@ -5147,7 +5147,7 @@ impl DAGCircuit { let py_op = if let Some(py_op) = py_op { py_op.into() } else { - OnceCell::new() + OnceLock::new() }; let packed_instruction = PackedInstruction { op, @@ -6193,7 +6193,7 @@ impl DAGCircuit { .then(|| Box::new(new_gate.1.iter().map(|x| Param::Float(*x)).collect())), extra_attrs: ExtraInstructionAttributes::default(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), } } else { panic!("This method only works if provided index is an op node"); @@ -6276,12 +6276,12 @@ impl DAGCircuit { }; #[cfg(feature = "cache_pygates")] let py_op = match new_op.operation.view() { - OperationRef::Standard(_) => OnceCell::new(), - OperationRef::Gate(gate) => OnceCell::from(gate.gate.clone_ref(py)), + OperationRef::Standard(_) => OnceLock::new(), + OperationRef::Gate(gate) => OnceLock::from(gate.gate.clone_ref(py)), OperationRef::Instruction(instruction) => { - OnceCell::from(instruction.instruction.clone_ref(py)) + OnceLock::from(instruction.instruction.clone_ref(py)) } - OperationRef::Operation(op) => OnceCell::from(op.operation.clone_ref(py)), + OperationRef::Operation(op) => OnceLock::from(op.operation.clone_ref(py)), }; let inst = PackedInstruction { op: new_op.operation, @@ -6732,7 +6732,7 @@ impl DAGCircuit { params: instr.params.clone(), extra_attrs: instr.extra_attrs.clone(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }) }) .collect::>>()?; @@ -6994,7 +6994,7 @@ impl DAGCircuit { params: (!new_op.params.is_empty()).then(|| new_op.params.into()), extra_attrs, #[cfg(feature = "cache_pygates")] - py_op: py_op_cache.map(OnceCell::from).unwrap_or_default(), + py_op: py_op_cache.map(OnceLock::from).unwrap_or_default(), }); if let Some(weight) = self.dag.node_weight_mut(node_index) { *weight = new_weight; diff --git a/crates/circuit/src/dag_node.rs b/crates/circuit/src/dag_node.rs index 6c3a2d15fbad..2fdfcdcbaef2 100644 --- a/crates/circuit/src/dag_node.rs +++ b/crates/circuit/src/dag_node.rs @@ -10,9 +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::OnceCell; use std::hash::Hasher; +#[cfg(feature = "cache_pygates")] +use std::sync::OnceLock; use crate::circuit_instruction::{CircuitInstruction, OperationFromPython}; use crate::imports::QUANTUM_CIRCUIT; @@ -241,7 +241,7 @@ impl DAGOpNode { instruction.operation = instruction.operation.py_deepcopy(py, None)?; #[cfg(feature = "cache_pygates")] { - instruction.py_op = OnceCell::new(); + instruction.py_op = OnceLock::new(); } } let base = PyClassInitializer::from(DAGNode { node: None }); @@ -293,7 +293,7 @@ impl DAGOpNode { params: self.instruction.params.clone(), extra_attrs: self.instruction.extra_attrs.clone(), #[cfg(feature = "cache_pygates")] - py_op: OnceCell::new(), + py_op: OnceLock::new(), }) } diff --git a/crates/circuit/src/packed_instruction.rs b/crates/circuit/src/packed_instruction.rs index af72b3226a7c..4da706280336 100644 --- a/crates/circuit/src/packed_instruction.rs +++ b/crates/circuit/src/packed_instruction.rs @@ -10,9 +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::OnceCell; use std::ptr::NonNull; +#[cfg(feature = "cache_pygates")] +use std::sync::OnceLock; use pyo3::intern; use pyo3::prelude::*; @@ -504,17 +504,17 @@ pub struct PackedInstruction { pub extra_attrs: ExtraInstructionAttributes, #[cfg(feature = "cache_pygates")] - /// This is hidden in a `OnceCell` because it's just an on-demand cache; we don't create this - /// unless asked for it. A `OnceCell` of a non-null pointer type (like `Py`) is the same + /// This is hidden in a `OnceLock` because it's just an on-demand cache; we don't create this + /// unless asked for it. A `OnceLock` of a non-null pointer type (like `Py`) is the same /// size as a pointer and there are no runtime checks on access beyond the initialisation check, /// which is a simple null-pointer check. /// - /// WARNING: remember that `OnceCell`'s `get_or_init` method is no-reentrant, so the initialiser + /// WARNING: remember that `OnceLock`'s `get_or_init` method is no-reentrant, so the initialiser /// must not yield the GIL to Python space. We avoid using `GILOnceCell` here because it /// requires the GIL to even `get` (of course!), which makes implementing `Clone` hard for us. /// We can revisit once we're on PyO3 0.22+ and have been able to disable its `py-clone` /// feature. - pub py_op: OnceCell>, + pub py_op: OnceLock>, } impl PackedInstruction { @@ -581,7 +581,7 @@ impl PackedInstruction { } }; - // `OnceCell::get_or_init` and the non-stabilised `get_or_try_init`, which would otherwise + // `OnceLock::get_or_init` and the non-stabilised `get_or_try_init`, which would otherwise // be nice here are both non-reentrant. This is a problem if the init yields control to the // Python interpreter as this one does, since that can allow CPython to freeze the thread // and for another to attempt the initialisation.