-
Notifications
You must be signed in to change notification settings - Fork 2.4k
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
Oxidize Commutation Analysis #12995
Oxidize Commutation Analysis #12995
Conversation
Co-authored-by: Raynel Sanchez <[email protected]>
This commit migrates the entirety of the `DAGCircuit` class to Rust. It fully replaces the Python version of the class. The primary advantage of this migration is moving from a Python space rustworkx directed graph representation to a Rust space petgraph (the upstream library for rustworkx) directed graph. Moving the graph data structure to rust enables us to directly interact with the DAG directly from transpiler passes in Rust in the future. This will enable a significant speed-up in those transpiler passes. Additionally, this should also improve the memory footprint as the DAGCircuit no longer stores `DAGNode` instances, and instead stores a lighter enum NodeType, which simply contains a `PackedInstruction` or the wire objects directly. Internally, the new Rust-based `DAGCircuit` uses a `petgraph::StableGraph` with node weights of type `NodeType` and edge weights of type `Wire`. The NodeType enum contains variants for `QubitIn`, `QubitOut`, `ClbitIn`, `ClbitOut`, and `Operation`, which should save us from all of the `isinstance` checking previously needed when working with `DAGNode` Python instances. The `Wire` enum contains variants `Qubit`, `Clbit`, and `Var`. As the full Qiskit data model is not rust-native at this point while all the class code in the `DAGCircuit` exists in Rust now, there are still sections that rely on Python or actively run Python code via Rust to function. These typically involve anything that uses `condition`, control flow, classical vars, calibrations, bit/register manipulation, etc. In the future as we either migrate this functionality to Rust or deprecate and remove it this can be updated in place to avoid the use of Python. API access from Python-space remains in terms of `DAGNode` instances to maintain API compatibility with the Python implementation. However, internally, we convert to and deal in terms of NodeType. When the user requests a particular node via lookup or iteration, we inflate an ephemeral `DAGNode` based on the internal `NodeType` and give them that. This is very similar to what was done in Qiskit#10827 when porting CircuitData to Rust. As part of this porting there are a few small differences to keep in mind with the new Rust implementation of DAGCircuit. The first is that the topological ordering is slightly different with the new DAGCircuit. Previously, the Python version of `DAGCircuit` using a lexicographical topological sort key which was basically `"0,1,0,2"` where the first `0,1` are qargs on qubit indices `0,1` for nodes and `0,2` are cargs on clbit indices `0,2`. However, the sort key has now changed to be `(&[Qubit(0), Qubit(1)], &[Clbit(0), Clbit(2)])` in rust in this case which for the most part should behave identically, but there are some edge cases that will appear where the sort order is different. It will always be a valid topological ordering as the lexicographical key is used as a tie breaker when generating a topological sort. But if you're relaying on the exact same sort order there will be differences after this PR. The second is that a lot of undocumented functionality in the DAGCircuit which previously worked because of Python's implicit support for interacting with data structures is no longer functional. For example, previously the `DAGCircuit.qubits` list could be set directly (as the circuit visualizers previously did), but this was never documented as supported (and would corrupt the DAGCircuit). Any functionality like this we'd have to explicit include in the Rust implementation and as they were not included in the documented public API this PR opted to remove the vast majority of this type of functionality. The last related thing might require future work to mitigate is that this PR breaks the linkage between `DAGNode` and the underlying `DAGCirucit` object. In the Python implementation the `DAGNode` objects were stored directly in the `DAGCircuit` and when an API method returned a `DAGNode` from the DAG it was a shared reference to the underlying object in the `DAGCircuit`. This meant if you mutated the `DAGNode` it would be reflected in the `DAGCircuit`. This was not always a sound usage of the API as the `DAGCircuit` was implicitly caching many attributes of the DAG and you should always be using the `DAGCircuit` API to mutate any nodes to prevent any corruption of the `DAGCircuit`. However, now as the underlying data store for nodes in the DAG are no longer the python space objects returned by `DAGCircuit` methods mutating a `DAGNode` will not make any change in the underlying `DAGCircuit`. This can come as quite the surprise at first, especially if you were relying on this side effect, even if it was unsound. It's also worth noting that 2 large pieces of functionality from rustworkx are included in this PR. These are the new files `rustworkx_core_vnext` and `dot_utils` which are rustworkx's VF2 implementation and its dot file generation. As there was not a rust interface exposed for this functionality from rustworkx-core there was no way to use these functions in rustworkx. Until these interfaces added to rustworkx-core in future releases we'll have to keep these local copies. The vf2 implementation is in progress in Qiskit/rustworkx#1235, but `dot_utils` might make sense to keep around longer term as it is slightly modified from the upstream rustworkx implementation to directly interface with `DAGCircuit` instead of a generic graph. Co-authored-by: Matthew Treinish <[email protected]> Co-authored-by: Raynel Sanchez <[email protected]> Co-authored-by: Elena Peña Tapia <[email protected]> Co-authored-by: Alexander Ivrii <[email protected]> Co-authored-by: Eli Arbel <[email protected]> Co-authored-by: John Lapeyre <[email protected]> Co-authored-by: Jake Lishman <[email protected]>
Right now there is a bug in the matplotlib circuit visualizer likely caused by the new `__eq__` implementation for `DAGOpNode` that didn't exist before were some gates are missing from the visualization. In the interest of unblocking this PR this commit updates the references for these cases temporarily until this issue is fixed.
Co-authored-by: Raynel Sanchez <[email protected]>
Pull Request Test Coverage Report for Build 10705679280Warning: This coverage report may be inaccurate.This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.
Details
💛 - Coveralls |
such as the relative placement and the parameter hash
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is looking good thanks for doing this. I left some small mostly mechanical comments inline. I still want to review the pass logic to double check it all matches up exactly. Especially as there seem to be a fair amount of panic!()
calls which I want to double check and make sure we're not masking potential user actionable error conditions (or indicate a logic issue).
use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType, Wire}; | ||
use rustworkx_core::petgraph::stable_graph::NodeIndex; | ||
|
||
type AIndexSet<T> = IndexSet<T, BuildHasherDefault<AHasher>>; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This can be expressed a bit more succinctly with:
type AIndexSet<T> = IndexSet<T, BuildHasherDefault<AHasher>>; | |
type AIndexSet<T> = IndexSet<T, ::ahash::RandomState>; |
That being said I'm not sure I'm a fan of using a type alias for this instead of just using IndexSet<NodeIndex, RandomState>
(assuming a use ahash::RandomState
is added) where we need to define a new indexset.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Overall this LGTM, after the most recent update the code looks good to me. I had a few questions inline nothing really blocking. But at least for the data structure I'd like to know the rationale behind using a HashMap<Wire, Vec<IndexSet<NodeIndex>>
as the type when in python it was all lists. I can believe using a Vec<IndexSet>
as the weight here is the better choice but I'd like to know the rationale before approving.
let mut commutation_set: HashMap<Wire, CommutingNodes> = HashMap::new(); | ||
let mut node_indices: HashMap<(NodeIndex, Wire), usize> = HashMap::new(); | ||
|
||
let max_num_qubits = 3; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not that it matters, but I'd have done this with:
const MAX_NUM_QUBITS: u32 = 3;
outside the function definition because it is never changed. But realistically the compiler will likely end up with the same generated code in both cases so it doesn't matter.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Still consistency is nice -- I'll move it outside 👍🏻
for qubit in 0..dag.num_qubits() { | ||
let wire = Wire::Qubit(Qubit(qubit as u32)); | ||
|
||
for current_gate_idx in dag.nodes_on_wire(py, &wire, false) { | ||
// get the commutation set associated with the current wire, or create a new | ||
// index set containing the current gate | ||
let commutation_entry = commutation_set | ||
.entry(wire.clone()) | ||
.or_insert_with(|| vec![AIndexSet::from_iter([current_gate_idx])]); | ||
|
||
// we can unwrap as we know the commutation entry has at least one element | ||
let last = commutation_entry.last_mut().unwrap(); | ||
|
||
// if the current gate index is not in the set, check whether it commutes with | ||
// the previous nodes -- if yes, add it to the commutation set | ||
if !last.contains(¤t_gate_idx) { | ||
let mut all_commute = true; | ||
|
||
for prev_gate_idx in last.iter() { | ||
// if the node is an input/output node, they do not commute, so we only | ||
// continue if the nodes are operation nodes | ||
if let (NodeType::Operation(packed_inst0), NodeType::Operation(packed_inst1)) = | ||
(&dag.dag[current_gate_idx], &dag.dag[*prev_gate_idx]) | ||
{ | ||
let op1 = packed_inst0.op.view(); | ||
let op2 = packed_inst1.op.view(); | ||
let params1 = packed_inst0.params_view(); | ||
let params2 = packed_inst1.params_view(); | ||
let qargs1 = dag.get_qargs(packed_inst0.qubits); | ||
let qargs2 = dag.get_qargs(packed_inst1.qubits); | ||
let cargs1 = dag.get_cargs(packed_inst0.clbits); | ||
let cargs2 = dag.get_cargs(packed_inst1.clbits); | ||
|
||
all_commute = commutation_checker.commute_inner( | ||
py, | ||
&op1, | ||
params1, | ||
packed_inst0.extra_attrs.as_deref(), | ||
qargs1, | ||
cargs1, | ||
&op2, | ||
params2, | ||
packed_inst1.extra_attrs.as_deref(), | ||
qargs2, | ||
cargs2, | ||
max_num_qubits, | ||
)?; | ||
if !all_commute { | ||
break; | ||
} | ||
} else { | ||
all_commute = false; | ||
break; | ||
} | ||
} | ||
|
||
if all_commute { | ||
// all commute, add to current list | ||
last.insert(current_gate_idx); | ||
} else { | ||
// does not commute, create new list | ||
commutation_entry.push(AIndexSet::from_iter([current_gate_idx])) | ||
} | ||
} | ||
|
||
node_indices.insert( | ||
(current_gate_idx, wire.clone()), | ||
commutation_entry.len() - 1, | ||
); | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is a probably a good candidate for parallelization. We can naively make this a parallel iterator for the outer loop like:
0..dag.num_qubits().into_par_iter().map(|qubit| Wire::Qubit(Qubit(qubit as u32))).for_each(|wire| {
We'll probably have to play some games around synchronization, although we can probably change it into an iterator collector. But we should do this in a follow up because it might require some deeper changes (as the compiler points them out).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah good point -- I'll have a look in parallel 🙂 I was trying to propagate the Result coming from commute_inner
but if we collect then we can still do that 👍🏻
// we could cache the py_wires to avoid this match and the python object creation, | ||
// but this didn't make a noticable difference in runtime |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Realistically the BitData
you're accessing here is the cache. It's basically just doing a vec lookup and returning the python object.
That being said I'm wondering can wire here be anything besides a Qubit
? Not that it really matters because this code is correct and will handle the other cases as expected.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
No you're right, it cannot, since we explicitly only iterate over 0..dag.num_qubits().map(|i| Wire::Qubit(i))
. We can change the code to expect a Qubit, that might be a bit cleaner to the readers
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I was wondering what happens if a dag has non-continuous qubits, i.e. gaps in the set of qubits such as qubits 0...3 and 6..10 with no qubits at index 4, 5 - can this happen and was this supported in the Python implementation?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The indices can't ever be non-continuous in Rust or Python. They're literally the index in the list
/Vec
for dag.qubits
. So if you remove a qubit everything has a shift penalty. The rust code just makes this a bit more explicit than it ever was in Python because there was a layer of indirection around it in Python with the bit objects.
- vec<vec> is slightly faster than vec<indexset> - add custom types to satisfies clippy's complex type complaint - don't handle Clbit/Var
I ran ASV comparing this against main to see what the performance looks like and it's really good:
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This LGTM, one super small thing in the Python code that is out of date with what got merged in for the commutation checker. It's not a bug but we just have an extra condition and comment that aren't needed anymore.
Since this PR was first written the split between the python side and rust side of the CommutationChecker class has changed so that there are no longer separate classes anymore. The implementations are unified and the python space class just wraps an inner rust object. However, the construction of the CommutationAnalysis pass was still written assuming there was the possibility to get either a rust or Python object. This commit fixes this and the type change on the `comm_checker` attribute by removing the unnecessary logic.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I fixed that last little nit and now this should be ready to merge. Thanks @sbrandhsn and @Cryoris for doing this!
Summary
Details and comments
Part of #12208, oxidises
CommutationAnalysis
. This is based on PR #12959 (and thus #12550) and PR #12870 and shows promising runtime improvements in preliminary tests. When #12550 is merged, the interface toCommutationChecker
needs a little thought, then this should be ready to go, in principle.