Skip to content

Commit

Permalink
Merge branch 'main' into bit_u32_sequel
Browse files Browse the repository at this point in the history
  • Loading branch information
@raynelfss
raynelfss authored Jan 30, 2025
2 parents ae3385e + 315d6ad commit b18c5ab
Show file tree
Hide file tree
Showing 3 changed files with 596 additions and 73 deletions.
2 changes: 1 addition & 1 deletion crates/circuit/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ features = ["rayon"]

[dependencies.smallvec]
workspace = true
features = ["union"]
features = ["union", "const_generics"]

[features]
cache_pygates = []
Expand Down
160 changes: 110 additions & 50 deletions crates/circuit/src/dag_circuit.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use crate::dag_node::{DAGInNode, DAGNode, DAGOpNode, DAGOutNode};
use crate::dot_utils::build_dot;
use crate::error::DAGCircuitError;
use crate::imports;
use crate::interner::{Interned, Interner};
use crate::interner::{Interned, InternedMap, Interner};
use crate::operations::{Operation, OperationRef, Param, PyInstruction, StandardGate};
use crate::packed_instruction::{PackedInstruction, PackedOperation};
use crate::rustworkx_core_vnext::isomorphism;
Expand Down Expand Up @@ -4810,6 +4810,76 @@ impl DAGCircuit {
&self.vars
}

/// Merge the `qargs` in a different [Interner] into this DAG, remapping the qubits.
///
/// This is useful for simplifying the direct mapping of [PackedInstruction]s from one DAG to
/// another, like in substitution methods, or rebuilding a new DAG out of a lot of smaller ones.
/// See [Interner::merge_map_slice] for more information on the mapping function.
///
/// The input [InternedMap] is cleared of its previous entries by this method, and then we
/// re-use the allocation.
pub fn merge_qargs_using(
&mut self,
other: &Interner<[Qubit]>,
map_fn: impl FnMut(&Qubit) -> Option<Qubit>,
map: &mut InternedMap<[Qubit]>,
) {
// 4 is an arbitrary guess for the amount of stack space to allocate for mapping the
// `qargs`, but it doesn't matter if it's too short because it'll safely spill to the heap.
self.qargs_interner
.merge_map_slice_using::<4>(other, map_fn, map);
}

/// Merge the `qargs` in a different [Interner] into this DAG, remapping the qubits.
///
/// This is useful for simplifying the direct mapping of [PackedInstruction]s from one DAG to
/// another, like in substitution methods, or rebuilding a new DAG out of a lot of smaller ones.
/// See [Interner::merge_map_slice] for more information on the mapping function.
pub fn merge_qargs(
&mut self,
other: &Interner<[Qubit]>,
map_fn: impl FnMut(&Qubit) -> Option<Qubit>,
) -> InternedMap<[Qubit]> {
let mut out = InternedMap::new();
self.merge_qargs_using(other, map_fn, &mut out);
out
}

/// Merge the `cargs` in a different [Interner] into this DAG, remapping the clbits.
///
/// This is useful for simplifying the direct mapping of [PackedInstruction]s from one DAG to
/// another, like in substitution methods, or rebuilding a new DAG out of a lot of smaller ones.
/// See [Interner::merge_map_slice] for more information on the mapping function.
///
/// The input [InternedMap] is cleared of its previous entries by this method, and then we
/// re-use the allocation.
pub fn merge_cargs_using(
&mut self,
other: &Interner<[Clbit]>,
map_fn: impl FnMut(&Clbit) -> Option<Clbit>,
map: &mut InternedMap<[Clbit]>,
) {
// 4 is an arbitrary guess for the amount of stack space to allocate for mapping the
// `cargs`, but it doesn't matter if it's too short because it'll safely spill to the heap.
self.cargs_interner
.merge_map_slice_using::<4>(other, map_fn, map);
}

/// Merge the `cargs` in a different [Interner] into this DAG, remapping the clbits.
///
/// This is useful for simplifying the direct mapping of [PackedInstruction]s from one DAG to
/// another, like in substitution methods, or rebuilding a new DAG out of a lot of smaller ones.
/// See [Interner::merge_map_slice] for more information on the mapping function.
pub fn merge_cargs(
&mut self,
other: &Interner<[Clbit]>,
map_fn: impl FnMut(&Clbit) -> Option<Clbit>,
) -> InternedMap<[Clbit]> {
let mut out = InternedMap::new();
self.merge_cargs_using(other, map_fn, &mut out);
out
}

/// Return an iterator of gate runs with non-conditional op nodes of given names
pub fn collect_runs(
&self,
Expand Down Expand Up @@ -6383,10 +6453,24 @@ impl DAGCircuit {
&self.op_names
}

/// Extends the DAG with valid instances of [PackedInstruction]
/// Extends the DAG with valid instances of [PackedInstruction].
pub fn extend<I>(&mut self, py: Python, iter: I) -> PyResult<Vec<NodeIndex>>
where
I: IntoIterator<Item = PackedInstruction>,
{
self.try_extend(
py,
iter.into_iter()
.map(|inst| -> Result<PackedInstruction, Infallible> { Ok(inst) }),
)
}

/// Extends the DAG with valid instances of [PackedInstruction], where the iterator produces the
/// results in a fallible manner.
pub fn try_extend<I, E>(&mut self, py: Python, iter: I) -> PyResult<Vec<NodeIndex>>
where
I: IntoIterator<Item = Result<PackedInstruction, E>>,
PyErr: From<E>,
{
// Create HashSets to keep track of each bit/var's last node
let mut qubit_last_nodes: HashMap<Qubit, NodeIndex> = HashMap::default();
Expand All @@ -6400,6 +6484,7 @@ impl DAGCircuit {
// Store new nodes to return
let mut new_nodes = Vec::with_capacity(iter.size_hint().1.unwrap_or_default());
for instr in iter {
let instr = instr?;
let op_name = instr.op.name();
let (all_cbits, vars): (Vec<Clbit>, Option<Vec<PyObject>>) = {
if self.may_have_additional_wires(py, &instr) {
Expand Down Expand Up @@ -6571,8 +6656,8 @@ impl DAGCircuit {

new_dag.metadata = qc.metadata.map(|meta| meta.unbind());

// Add the qubits depending on order.
let qubit_map: Option<Vec<Qubit>> = if let Some(qubit_ordering) = qubit_order {
// Add the qubits depending on order, and produce the qargs map.
let qarg_map = if let Some(qubit_ordering) = qubit_order {
let mut ordered_vec = Vec::from_iter((0..num_qubits as u32).map(Qubit));
qubit_ordering
.into_iter()
Expand All @@ -6587,7 +6672,11 @@ impl DAGCircuit {
ordered_vec[qubit_index.index()] = new_dag.add_qubit_unchecked(py, &qubit)?;
Ok(())
})?;
Some(ordered_vec)
// The `Vec::get` use is because an arbitrary interner might contain old references to
// bit instances beyond `num_qubits`, such as if it's from a DAG that had wires removed.
new_dag.merge_qargs(qc_data.qargs_interner(), |bit| {
ordered_vec.get(bit.index()).copied()
})
} else {
qc_data
.qubits()
Expand All @@ -6597,11 +6686,11 @@ impl DAGCircuit {
new_dag.add_qubit_unchecked(py, qubit.bind(py))?;
Ok(())
})?;
None
new_dag.merge_qargs(qc_data.qargs_interner(), |bit| Some(*bit))
};

// Add the clbits depending on order.
let clbit_map: Option<Vec<Clbit>> = if let Some(clbit_ordering) = clbit_order {
// Add the clbits depending on order, and produce the cargs map.
let carg_map = if let Some(clbit_ordering) = clbit_order {
let mut ordered_vec = Vec::from_iter((0..num_clbits as u32).map(Clbit));
clbit_ordering
.into_iter()
Expand All @@ -6616,7 +6705,11 @@ impl DAGCircuit {
ordered_vec[clbit_index.index()] = new_dag.add_clbit_unchecked(py, &clbit)?;
Ok(())
})?;
Some(ordered_vec)
// The `Vec::get` use is because an arbitrary interner might contain old references to
// bit instances beyond `num_clbits`, such as if it's from a DAG that had wires removed.
new_dag.merge_cargs(qc_data.cargs_interner(), |bit| {
ordered_vec.get(bit.index()).copied()
})
} else {
qc_data
.clbits()
Expand All @@ -6626,7 +6719,7 @@ impl DAGCircuit {
new_dag.add_clbit_unchecked(py, clbit.bind(py))?;
Ok(())
})?;
None
new_dag.merge_cargs(qc_data.cargs_interner(), |bit| Some(*bit))
};

// Add all of the new vars.
Expand All @@ -6651,57 +6744,24 @@ impl DAGCircuit {
new_dag.add_creg(py, &creg)?;
}

// Pre-process and re-intern all indices again.
let instructions: Vec<PackedInstruction> = qc_data
.iter()
.map(|instr| -> PyResult<PackedInstruction> {
// Re-map the qubits
let new_qargs = if let Some(qubit_mapping) = &qubit_map {
let qargs = qc_data
.get_qargs(instr.qubits)
.iter()
.map(|bit| qubit_mapping[bit.index()])
.collect();
new_dag.qargs_interner.insert_owned(qargs)
} else {
new_dag
.qargs_interner
.insert(qc_data.get_qargs(instr.qubits))
};
// Remap the clbits
let new_cargs = if let Some(clbit_mapping) = &clbit_map {
let qargs = qc_data
.get_cargs(instr.clbits)
.iter()
.map(|bit| clbit_mapping[bit.index()])
.collect();
new_dag.cargs_interner.insert_owned(qargs)
} else {
new_dag
.cargs_interner
.insert(qc_data.get_cargs(instr.clbits))
};
// Copy the operations

new_dag.try_extend(
py,
qc_data.iter().map(|instr| -> PyResult<PackedInstruction> {
Ok(PackedInstruction {
op: if copy_op {
instr.op.py_deepcopy(py, None)?
} else {
instr.op.clone()
},
qubits: new_qargs,
clbits: new_cargs,
qubits: qarg_map[instr.qubits],
clbits: carg_map[instr.clbits],
params: instr.params.clone(),
extra_attrs: instr.extra_attrs.clone(),
#[cfg(feature = "cache_pygates")]
py_op: OnceLock::new(),
})
})
.collect::<PyResult<Vec<_>>>()?;

// Finally add all the instructions back
new_dag.extend(py, instructions)?;

}),
)?;
Ok(new_dag)
}

Expand Down
Loading

0 comments on commit b18c5ab

Please sign in to comment.