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

crates/accelerate/src/unitary_synthesis.rs

@@ -12,7 +12,7 @@
 #![allow(clippy::too_many_arguments)]
 
 #[cfg(feature = "cache_pygates")]
-use std::cell::OnceCell;
+use std::sync::OnceLock;
 use std::f64::consts::PI;
 
 use approx::relative_eq;
@@ -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);
     }

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::{
@@ -302,7 +302,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 {
@@ -314,7 +314,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 {
@@ -939,7 +939,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)?;
         }
@@ -1048,7 +1048,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)?;
         }
@@ -1106,7 +1106,7 @@ impl CircuitData {
             params,
             extra_attrs: ExtraInstructionAttributes::default(),
             #[cfg(feature = "cache_pygates")]
-            py_op: OnceCell::new(),
+            py_op: OnceLock::new(),
         });
         Ok(())
     }

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<Py<PyAny>>,
+    pub py_op: OnceLock<Py<PyAny>>,
 }
 
 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(),
         })
     }
 

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())

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::<PyResult<Vec<_>>>()?;
@@ -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;

crates/circuit/src/dag_node.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 std::hash::Hasher;
 
 use crate::circuit_instruction::{CircuitInstruction, OperationFromPython};
@@ -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(),
         })
     }
 

crates/circuit/src/packed_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 std::ptr::NonNull;
 
 use pyo3::intern;
@@ -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<T>`) 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<T>`) 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<Py<PyAny>>,
+    pub py_op: OnceLock<Py<PyAny>>,
 }
 
 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.