From 223cea7f614c8d8d61319271b68be1b86d241400 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Wed, 29 Jan 2025 17:52:04 -0500 Subject: [PATCH] Add unitary gate to rust --- Cargo.lock | 74 ++++++++++- Cargo.toml | 1 + .../accelerate/src/target_transpiler/mod.rs | 3 + crates/circuit/Cargo.toml | 6 +- crates/circuit/src/circuit_instruction.rs | 49 +++++++- crates/circuit/src/dag_circuit.rs | 9 +- crates/circuit/src/operations.rs | 117 +++++++++++++++++- crates/circuit/src/packed_instruction.rs | 28 ++++- 8 files changed, 271 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a20ef58bcd4d..40f083bc5469 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -85,7 +85,7 @@ checksum = "6d7a33e7b9505a52e33ed0ad66db6434f18cda0b1c72665fabf14e85cdd39e43" dependencies = [ "proc-macro2", "quote", - "syn 2.0.79", + "syn 2.0.96", ] [[package]] @@ -709,6 +709,33 @@ dependencies = [ "autocfg", ] +[[package]] +name = "nalgebra" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26aecdf64b707efd1310e3544d709c5c0ac61c13756046aaaba41be5c4f66a3b" +dependencies = [ + "approx 0.5.1", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.96", +] + [[package]] name = "nano-gemm" version = "0.1.2" @@ -848,6 +875,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -865,6 +903,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b94caae805f998a07d33af06e6a3891e38556051b8045c615470a71590e13e78" dependencies = [ "libc", + "nalgebra", "ndarray", "num-complex", "num-integer", @@ -1226,6 +1265,7 @@ dependencies = [ "hashbrown 0.14.5", "indexmap", "itertools 0.13.0", + "nalgebra", "ndarray", "num-complex", "numpy", @@ -1475,6 +1515,15 @@ dependencies = [ "rayon-cond", ] +[[package]] +name = "safe_arch" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -1521,6 +1570,19 @@ dependencies = [ "digest", ] +[[package]] +name = "simba" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3a386a501cd104797982c15ae17aafe8b9261315b5d07e3ec803f2ea26be0fa" +dependencies = [ + "approx 0.5.1", + "num-complex", + "num-traits", + "paste", + "wide", +] + [[package]] name = "smallvec" version = "1.13.2" @@ -1694,6 +1756,16 @@ version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +[[package]] +name = "wide" +version = "0.7.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41b5576b9a81633f3e8df296ce0063042a73507636cbe956c61133dd7034ab22" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi-util" version = "0.1.9" diff --git a/Cargo.toml b/Cargo.toml index 50276f437ec5..4c898a2710ba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ indexmap.version = "2.7.1" hashbrown.version = "0.14.5" num-bigint = "0.4" num-complex = "0.4" +nalgebra = "0.33" ndarray = "0.15" numpy = "0.23" smallvec = "1.13" diff --git a/crates/accelerate/src/target_transpiler/mod.rs b/crates/accelerate/src/target_transpiler/mod.rs index 0a14a3dec72f..9df6047a494b 100644 --- a/crates/accelerate/src/target_transpiler/mod.rs +++ b/crates/accelerate/src/target_transpiler/mod.rs @@ -783,6 +783,9 @@ impl Target { OperationRef::Gate(gate) => gate.gate.clone_ref(py), OperationRef::Instruction(instruction) => instruction.instruction.clone_ref(py), OperationRef::Operation(operation) => operation.operation.clone_ref(py), + OperationRef::Unitary(unitary) => unitary + .create_py_op(py, &ExtraInstructionAttributes::default())? + .into_any(), }, TargetOperation::Variadic(op_cls) => op_cls.clone_ref(py), }; diff --git a/crates/circuit/Cargo.toml b/crates/circuit/Cargo.toml index 9d5691be41e2..cf3339334a2c 100644 --- a/crates/circuit/Cargo.toml +++ b/crates/circuit/Cargo.toml @@ -20,10 +20,10 @@ bytemuck.workspace = true bitfield-struct.workspace = true num-complex.workspace = true ndarray.workspace = true -numpy.workspace = true thiserror.workspace = true approx.workspace = true itertools.workspace = true +nalgebra.workspace = true [dependencies.pyo3] workspace = true @@ -41,6 +41,10 @@ features = ["rayon"] workspace = true features = ["union"] +[dependencies.numpy] +workspace = true +features = ["nalgebra"] + [features] cache_pygates = [] diff --git a/crates/circuit/src/circuit_instruction.rs b/crates/circuit/src/circuit_instruction.rs index 3fccccdefafb..d4b3496aefdc 100644 --- a/crates/circuit/src/circuit_instruction.rs +++ b/crates/circuit/src/circuit_instruction.rs @@ -13,7 +13,7 @@ #[cfg(feature = "cache_pygates")] use std::sync::OnceLock; -use numpy::{IntoPyArray, PyArray2}; +use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2}; use pyo3::basic::CompareOp; use pyo3::exceptions::{PyDeprecationWarning, PyTypeError}; use pyo3::prelude::*; @@ -21,6 +21,7 @@ use pyo3::types::{PyBool, PyList, PyString, PyTuple, PyType}; use pyo3::IntoPyObjectExt; use pyo3::{intern, PyObject, PyResult}; +use nalgebra::{MatrixView2, MatrixView4}; use num_complex::Complex64; use smallvec::SmallVec; @@ -28,8 +29,8 @@ use crate::imports::{ CONTROLLED_GATE, CONTROL_FLOW_OP, GATE, INSTRUCTION, OPERATION, WARNINGS_WARN, }; use crate::operations::{ - Operation, OperationRef, Param, PyGate, PyInstruction, PyOperation, StandardGate, - StandardInstruction, StandardInstructionType, + ArrayType, Operation, OperationRef, Param, PyGate, PyInstruction, PyOperation, StandardGate, + StandardInstruction, StandardInstructionType, UnitaryGate, }; use crate::packed_instruction::PackedOperation; @@ -341,6 +342,9 @@ impl CircuitInstruction { OperationRef::Gate(gate) => gate.gate.clone_ref(py), OperationRef::Instruction(instruction) => instruction.instruction.clone_ref(py), OperationRef::Operation(operation) => operation.operation.clone_ref(py), + OperationRef::Unitary(unitary) => { + unitary.create_py_op(py, &self.extra_attrs)?.into_any() + } }; #[cfg(feature = "cache_pygates")] @@ -742,6 +746,45 @@ impl<'py> FromPyObject<'py> for OperationFromPython { }); } + // We need to check by name here to avoid a circular import during initial loading + if ob.getattr(intern!(py, "name"))?.extract::()? == "unitary" { + let params = extract_params()?; + if let Param::Obj(data) = ¶ms[0] { + let py_matrix: PyReadonlyArray2 = data.extract(py)?; + let matrix: Option> = py_matrix.try_as_matrix(); + if let Some(x) = matrix { + let unitary_gate = UnitaryGate { + array: ArrayType::OneQ(x.into_owned()), + }; + return Ok(OperationFromPython { + operation: PackedOperation::from_unitary(Box::new(unitary_gate)), + params: SmallVec::new(), + extra_attrs: extract_extra()?, + }); + } + let matrix: Option> = py_matrix.try_as_matrix(); + if let Some(x) = matrix { + let unitary_gate = UnitaryGate { + array: ArrayType::TwoQ(x.into_owned()), + }; + return Ok(OperationFromPython { + operation: PackedOperation::from_unitary(Box::new(unitary_gate)), + params: SmallVec::new(), + extra_attrs: extract_extra()?, + }); + } else { + let unitary_gate = UnitaryGate { + array: ArrayType::NDArray(py_matrix.as_array().to_owned()), + }; + return Ok(OperationFromPython { + operation: PackedOperation::from_unitary(Box::new(unitary_gate)), + params: SmallVec::new(), + extra_attrs: extract_extra()?, + }); + }; + } + } + if ob_type.is_subclass(GATE.get_bound(py))? { let params = extract_params()?; let gate = Box::new(PyGate { diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs index c165d33fc033..b80572b83e09 100644 --- a/crates/circuit/src/dag_circuit.rs +++ b/crates/circuit/src/dag_circuit.rs @@ -3295,7 +3295,8 @@ def _format(operand): py_op.operation.setattr(py, "condition", new_condition)?; } OperationRef::StandardGate(_) - | OperationRef::StandardInstruction(_) => {} + | OperationRef::StandardInstruction(_) + | OperationRef::Unitary(_) => {} } } } @@ -6245,9 +6246,9 @@ impl DAGCircuit { }; #[cfg(feature = "cache_pygates")] let py_op = match new_op.operation.view() { - OperationRef::StandardGate(_) | OperationRef::StandardInstruction(_) => { - OnceLock::new() - } + OperationRef::StandardGate(_) + | OperationRef::StandardInstruction(_) + | OperationRef::Unitary(_) => OnceLock::new(), OperationRef::Gate(gate) => OnceLock::from(gate.gate.clone_ref(py)), OperationRef::Instruction(instruction) => { OnceLock::from(instruction.instruction.clone_ref(py)) diff --git a/crates/circuit/src/operations.rs b/crates/circuit/src/operations.rs index a3a917815150..b30ea57e3e1a 100644 --- a/crates/circuit/src/operations.rs +++ b/crates/circuit/src/operations.rs @@ -17,19 +17,21 @@ use std::{fmt, vec}; use crate::circuit_data::CircuitData; use crate::circuit_instruction::ExtraInstructionAttributes; use crate::imports::{get_std_gate_class, BARRIER, DELAY, MEASURE, RESET}; -use crate::imports::{PARAMETER_EXPRESSION, QUANTUM_CIRCUIT}; +use crate::imports::{PARAMETER_EXPRESSION, QUANTUM_CIRCUIT, UNITARY_GATE}; use crate::{gate_matrix, impl_intopyobject_for_copy_pyclass, Qubit}; -use ndarray::{aview2, Array2}; +use nalgebra::{Matrix2, Matrix4}; +use ndarray::{array, aview2, Array2}; use num_complex::Complex64; use smallvec::{smallvec, SmallVec}; use numpy::IntoPyArray; use numpy::PyArray2; use numpy::PyReadonlyArray2; +use numpy::ToPyArray; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::{IntoPyDict, PyFloat, PyIterator, PyList, PyTuple}; +use pyo3::types::{IntoPyDict, PyDict, PyFloat, PyIterator, PyList, PyTuple}; use pyo3::{intern, IntoPyObjectExt, Python}; #[derive(Clone, Debug, IntoPyObject, IntoPyObjectRef)] @@ -164,6 +166,7 @@ pub enum OperationRef<'a> { Gate(&'a PyGate), Instruction(&'a PyInstruction), Operation(&'a PyOperation), + Unitary(&'a UnitaryGate), } impl Operation for OperationRef<'_> { @@ -175,6 +178,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.name(), Self::Instruction(instruction) => instruction.name(), Self::Operation(operation) => operation.name(), + Self::Unitary(unitary) => unitary.name(), } } #[inline] @@ -185,6 +189,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.num_qubits(), Self::Instruction(instruction) => instruction.num_qubits(), Self::Operation(operation) => operation.num_qubits(), + Self::Unitary(unitary) => unitary.num_qubits(), } } #[inline] @@ -195,6 +200,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.num_clbits(), Self::Instruction(instruction) => instruction.num_clbits(), Self::Operation(operation) => operation.num_clbits(), + Self::Unitary(unitary) => unitary.num_clbits(), } } #[inline] @@ -205,6 +211,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.num_params(), Self::Instruction(instruction) => instruction.num_params(), Self::Operation(operation) => operation.num_params(), + Self::Unitary(unitary) => unitary.num_params(), } } #[inline] @@ -215,6 +222,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.control_flow(), Self::Instruction(instruction) => instruction.control_flow(), Self::Operation(operation) => operation.control_flow(), + Self::Unitary(unitary) => unitary.control_flow(), } } #[inline] @@ -225,6 +233,7 @@ impl Operation for OperationRef<'_> { OperationRef::Gate(gate) => gate.blocks(), OperationRef::Instruction(instruction) => instruction.blocks(), OperationRef::Operation(operation) => operation.blocks(), + Self::Unitary(unitary) => unitary.blocks(), } } #[inline] @@ -235,6 +244,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.matrix(params), Self::Instruction(instruction) => instruction.matrix(params), Self::Operation(operation) => operation.matrix(params), + Self::Unitary(unitary) => unitary.matrix(params), } } #[inline] @@ -245,6 +255,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.definition(params), Self::Instruction(instruction) => instruction.definition(params), Self::Operation(operation) => operation.definition(params), + Self::Unitary(unitary) => unitary.definition(params), } } #[inline] @@ -255,6 +266,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.standard_gate(), Self::Instruction(instruction) => instruction.standard_gate(), Self::Operation(operation) => operation.standard_gate(), + Self::Unitary(unitary) => unitary.standard_gate(), } } #[inline] @@ -265,6 +277,7 @@ impl Operation for OperationRef<'_> { Self::Gate(gate) => gate.directive(), Self::Instruction(instruction) => instruction.directive(), Self::Operation(operation) => operation.directive(), + Self::Unitary(unitary) => unitary.directive(), } } } @@ -2792,3 +2805,101 @@ impl Operation for PyOperation { }) } } + +#[derive(Clone, Debug)] +pub enum ArrayType { + NDArray(Array2), + OneQ(Matrix2), + TwoQ(Matrix4), +} + +/// This class is used to wrap a Python side Operation that is not in the standard library +#[derive(Clone, Debug)] +// We bit-pack pointers to this, so having a known alignment even on 32-bit systems is good. +#[repr(align(8))] +pub struct UnitaryGate { + pub array: ArrayType, +} + +impl Operation for UnitaryGate { + fn name(&self) -> &str { + "unitary" + } + fn num_qubits(&self) -> u32 { + match &self.array { + ArrayType::NDArray(arr) => arr.shape()[0].ilog2() as u32, + ArrayType::OneQ(_) => 1, + ArrayType::TwoQ(_) => 2, + } + } + fn num_clbits(&self) -> u32 { + 0 + } + fn num_params(&self) -> u32 { + 0 + } + fn control_flow(&self) -> bool { + false + } + fn blocks(&self) -> Vec { + vec![] + } + fn matrix(&self, _params: &[Param]) -> Option> { + match &self.array { + ArrayType::NDArray(arr) => Some(arr.clone()), + ArrayType::OneQ(mat) => Some(array!( + [mat[(0, 0)], mat[(0, 1)]], + [mat[(1, 0)], mat[(1, 1)]], + )), + ArrayType::TwoQ(mat) => Some(array!( + [mat[(0, 0)], mat[(0, 1)], mat[(0, 2)], mat[(0, 3)]], + [mat[(1, 0)], mat[(1, 1)], mat[(1, 2)], mat[(1, 3)]], + [mat[(2, 0)], mat[(2, 1)], mat[(2, 2)], mat[(2, 3)]], + [mat[(3, 0)], mat[(3, 1)], mat[(3, 2)], mat[(3, 3)]], + )), + } + } + fn definition(&self, _params: &[Param]) -> Option { + None + } + fn standard_gate(&self) -> Option { + None + } + + fn directive(&self) -> bool { + false + } +} + +impl UnitaryGate { + pub fn create_py_op( + &self, + py: Python, + extra_attrs: &ExtraInstructionAttributes, + ) -> PyResult> { + let (label, _unit, _duration, condition) = ( + extra_attrs.label(), + extra_attrs.unit(), + extra_attrs.duration(), + extra_attrs.condition(), + ); + let kwargs = PyDict::new(py); + if let Some(label) = label { + kwargs.set_item(intern!(py, "label"), label.into_py_any(py)?)?; + } + let out_array = match &self.array { + ArrayType::NDArray(arr) => arr.to_pyarray(py), + ArrayType::OneQ(arr) => arr.to_pyarray(py), + ArrayType::TwoQ(arr) => arr.to_pyarray(py), + }; + kwargs.set_item(intern!(py, "check_input"), false)?; + kwargs.set_item(intern!(py, "num_qubits"), self.num_qubits())?; + let mut gate = UNITARY_GATE + .get_bound(py) + .call((out_array,), Some(&kwargs))?; + if let Some(cond) = condition { + gate = gate.call_method1(intern!(py, "c_if"), (cond,))?; + } + Ok(gate.unbind()) + } +} diff --git a/crates/circuit/src/packed_instruction.rs b/crates/circuit/src/packed_instruction.rs index c7ec40aeef3c..fe4a5523a7dc 100644 --- a/crates/circuit/src/packed_instruction.rs +++ b/crates/circuit/src/packed_instruction.rs @@ -23,11 +23,11 @@ use smallvec::SmallVec; use crate::circuit_data::CircuitData; use crate::circuit_instruction::ExtraInstructionAttributes; -use crate::imports::{get_std_gate_class, BARRIER, DEEPCOPY, DELAY, MEASURE, RESET}; +use crate::imports::{get_std_gate_class, BARRIER, DEEPCOPY, DELAY, MEASURE, RESET, UNITARY_GATE}; use crate::interner::Interned; use crate::operations::{ Operation, OperationRef, Param, PyGate, PyInstruction, PyOperation, StandardGate, - StandardInstruction, + StandardInstruction, UnitaryGate, }; use crate::{Clbit, Qubit}; @@ -43,13 +43,14 @@ enum PackedOperationType { PyGate = 2, PyInstruction = 3, PyOperation = 4, + UnitaryGate = 5, } unsafe impl ::bytemuck::CheckedBitPattern for PackedOperationType { type Bits = u8; fn is_valid_bit_pattern(bits: &Self::Bits) -> bool { - *bits < 5 + *bits < 6 } } unsafe impl ::bytemuck::NoUninit for PackedOperationType {} @@ -65,6 +66,7 @@ unsafe impl ::bytemuck::NoUninit for PackedOperationType {} /// Gate(Box), /// Instruction(Box), /// Operation(Box), +/// UnitaryGate(Box), /// } /// ``` /// @@ -253,7 +255,7 @@ mod standard_instruction { /// A private module to encapsulate the encoding of pointer types. mod pointer { - use crate::operations::{PyGate, PyInstruction, PyOperation}; + use crate::operations::{PyGate, PyInstruction, PyOperation, UnitaryGate}; use crate::packed_instruction::{PackedOperation, PackedOperationType}; use std::ptr::NonNull; @@ -294,6 +296,7 @@ mod pointer { impl_packable_pointer!(PyGate, PackedOperationType::PyGate); impl_packable_pointer!(PyInstruction, PackedOperationType::PyInstruction); impl_packable_pointer!(PyOperation, PackedOperationType::PyOperation); + impl_packable_pointer!(UnitaryGate, PackedOperationType::UnitaryGate); impl From> for PackedOperation { fn from(value: Box) -> Self { @@ -338,6 +341,7 @@ mod pointer { PackedOperationType::PyGate => drop_pointer_as::(self), PackedOperationType::PyInstruction => drop_pointer_as::(self), PackedOperationType::PyOperation => drop_pointer_as::(self), + PackedOperationType::UnitaryGate => drop_pointer_as::(self), } } } @@ -396,6 +400,7 @@ impl PackedOperation { OperationRef::Instruction(self.try_into().unwrap()) } PackedOperationType::PyOperation => OperationRef::Operation(self.try_into().unwrap()), + PackedOperationType::UnitaryGate => OperationRef::Unitary(self.try_into().unwrap()), } } @@ -425,6 +430,10 @@ impl PackedOperation { operation.into() } + pub fn from_unitary(unitary: Box) -> Self { + unitary.into() + } + /// Check equality of the operation, including Python-space checks, if appropriate. pub fn py_eq(&self, py: Python, other: &PackedOperation) -> PyResult { match (self.view(), other.view()) { @@ -484,6 +493,7 @@ impl PackedOperation { op_name: operation.op_name.clone(), } .into()), + OperationRef::Unitary(unitary) => Ok(unitary.clone().into()), } } @@ -521,6 +531,7 @@ impl PackedOperation { op_name: operation.op_name.clone(), }) .into()), + OperationRef::Unitary(unitary) => Ok(unitary.clone().into()), } } @@ -559,6 +570,12 @@ impl PackedOperation { OperationRef::Gate(gate) => gate.gate.bind(py), OperationRef::Instruction(instruction) => instruction.instruction.bind(py), OperationRef::Operation(operation) => operation.operation.bind(py), + OperationRef::Unitary(_) => { + return UNITARY_GATE + .get_bound(py) + .downcast::()? + .is_subclass(py_type); + } }; py_op.is_instance(py_type) } @@ -573,6 +590,7 @@ impl Operation for PackedOperation { OperationRef::Gate(gate) => gate.name(), OperationRef::Instruction(instruction) => instruction.name(), OperationRef::Operation(operation) => operation.name(), + OperationRef::Unitary(unitary) => unitary.name(), }; // SAFETY: all of the inner parts of the view are owned by `self`, so it's valid for us to // forcibly reborrowing up to our own lifetime. We avoid using `` @@ -636,6 +654,7 @@ impl Clone for PackedOperation { OperationRef::Operation(operation) => { Self::from_operation(Box::new(operation.to_owned())) } + OperationRef::Unitary(unitary) => Self::from_unitary(Box::new(unitary.clone())), } } } @@ -740,6 +759,7 @@ impl PackedInstruction { OperationRef::Gate(gate) => Ok(gate.gate.clone_ref(py)), OperationRef::Instruction(instruction) => Ok(instruction.instruction.clone_ref(py)), OperationRef::Operation(operation) => Ok(operation.operation.clone_ref(py)), + OperationRef::Unitary(unitary) => unitary.create_py_op(py, &self.extra_attrs), } };