Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use OnceLock in imports. #13776

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 1 addition & 5 deletions crates/accelerate/src/gate_direction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -219,11 +219,7 @@ fn py_fix_direction_target(
py,
None,
Some(qargs),
Some(
get_std_gate_class(py, std_gate)
.expect("These gates should have Python classes")
.bind(py),
),
Some(get_std_gate_class(py, std_gate).bind(py)),
Some(inst.params_view().to_vec()),
)
.unwrap_or(false)
Expand Down
129 changes: 67 additions & 62 deletions crates/circuit/src/imports.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,32 +15,33 @@
// python side casting

use pyo3::prelude::*;
use pyo3::sync::GILOnceCell;
use pyo3::sync::OnceLockExt;
use std::sync::OnceLock;

use crate::operations::{StandardGate, STANDARD_GATE_SIZE};

/// Helper wrapper around `GILOnceCell` instances that are just intended to store a Python object
/// Helper wrapper around `OnceLock` instances that are just intended to store a Python object
/// that is lazily imported.
pub struct ImportOnceCell {
module: &'static str,
object: &'static str,
cell: GILOnceCell<Py<PyAny>>,
cell: OnceLock<Py<PyAny>>,
}

impl ImportOnceCell {
pub const fn new(module: &'static str, object: &'static str) -> Self {
Self {
module,
object,
cell: GILOnceCell::new(),
cell: OnceLock::new(),
}
}

/// Get the underlying GIL-independent reference to the contained object, importing if
/// required.
#[inline]
pub fn get(&self, py: Python) -> &Py<PyAny> {
self.cell.get_or_init(py, || {
self.cell.get_or_init_py_attached(py, || {
py.import(self.module)
.unwrap()
.getattr(self.object)
Expand Down Expand Up @@ -254,67 +255,71 @@ static STDGATE_IMPORT_PATHS: [[&str; 2]; STANDARD_GATE_SIZE] = [
/// class that matches it. This is typically used when we need to convert from the internal rust
/// representation to a Python object for a python user to interact with.
///
/// NOTE: the order here is significant it must match the StandardGate variant's number must match
/// NOTE: the order here is significant. The StandardGate variant's number must match the
/// index of it's entry in this table. This is all done statically for performance
static STDGATE_PYTHON_GATES: [GILOnceCell<PyObject>; STANDARD_GATE_SIZE] = [
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
GILOnceCell::new(),
static STDGATE_PYTHON_GATES: [OnceLock<PyObject>; STANDARD_GATE_SIZE] = [
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
];

#[inline]
pub fn get_std_gate_class(py: Python, rs_gate: StandardGate) -> PyResult<&'static Py<PyAny>> {
STDGATE_PYTHON_GATES[rs_gate as usize].get_or_try_init(py, || {
pub fn get_std_gate_class(py: Python, rs_gate: StandardGate) -> &'static Py<PyAny> {
STDGATE_PYTHON_GATES[rs_gate as usize].get_or_init_py_attached(py, || {
let [py_mod, py_class] = STDGATE_IMPORT_PATHS[rs_gate as usize];
Ok(py.import(py_mod)?.getattr(py_class)?.unbind())
py.import(py_mod)
.unwrap()
.getattr(py_class)
.unwrap()
.unbind()
})
}
4 changes: 2 additions & 2 deletions crates/circuit/src/operations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -419,7 +419,7 @@ impl StandardGate {
params: Option<&[Param]>,
extra_attrs: &ExtraInstructionAttributes,
) -> PyResult<Py<PyAny>> {
let gate_class = get_std_gate_class(py, *self)?;
let gate_class = get_std_gate_class(py, *self);
let args = match params.unwrap_or(&[]) {
&[] => PyTuple::empty(py),
params => PyTuple::new(py, params.iter().map(|x| x.into_pyobject(py).unwrap()))?,
Expand Down Expand Up @@ -715,7 +715,7 @@ impl StandardGate {
}

#[getter]
pub fn get_gate_class(&self, py: Python) -> PyResult<&'static Py<PyAny>> {
pub fn get_gate_class(&self, py: Python) -> &'static Py<PyAny> {
get_std_gate_class(py, *self)
}

Expand Down
2 changes: 1 addition & 1 deletion crates/circuit/src/packed_instruction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,7 @@ impl PackedOperation {
let py = py_type.py();
let py_op = match self.view() {
OperationRef::Standard(standard) => {
return get_std_gate_class(py, standard)?
return get_std_gate_class(py, standard)
.bind(py)
.downcast::<PyType>()?
.is_subclass(py_type)
Expand Down