-
Notifications
You must be signed in to change notification settings - Fork 145
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
[WIP] Add collect_multi_blocks function #461
base: main
Are you sure you want to change the base?
Conversation
This commit adds a new function collect_multi_blocks to retworkx. It's designed primarily for use by Qiskit's compiler to collect n qubit blocks from a DAGCircuit. This is similar to how collect_runs() is used to collect 1q blocks and collect_bicolor_runs is used to collect 2q blocks. This function is basically a mechanical porting of: Qiskit/qiskit#4747 which added this functionality as a Qiskit transpiler pass. With this API in retworkx the pass can still be added to Qiskit but just leverage retworkx to compute the blocks.
src/dag_algo/multiblock.rs
Outdated
/// DSU function for finding root of set of items | ||
/// If my parent is myself, I am the root. Otherwise we recursively | ||
/// find the root for my parent. After that, we assign my parent to be | ||
/// my root, saving recursion in the future. | ||
fn find_set( | ||
index: usize, | ||
parent: &mut Vec<Option<usize>>, | ||
groups: &mut HashMap<usize, Vec<usize>>, | ||
op_groups: &mut HashMap<usize, Vec<usize>>, | ||
) -> usize { | ||
if parent[index].is_none() { | ||
parent[index] = Some(index); | ||
groups.insert(index, vec![index]); | ||
op_groups.insert(index, Vec::new()); | ||
} | ||
if parent[index] == Some(index) { | ||
return index; | ||
} | ||
parent[index] = | ||
Some(find_set(parent[index].unwrap(), parent, groups, op_groups)); | ||
parent[index].unwrap() | ||
} | ||
|
||
/// DSU function for unioning two sets together | ||
/// Find the roots of each set. Then assign one to have the other | ||
/// as its parent, thus liking the sets. | ||
/// Merges smaller set into larger set in order to have better runtime | ||
fn union_set( | ||
set_a_ind: usize, | ||
set_b_ind: usize, | ||
parent: &mut Vec<Option<usize>>, | ||
groups: &mut HashMap<usize, Vec<usize>>, | ||
op_groups: &mut HashMap<usize, Vec<usize>>, | ||
) { | ||
let mut set_a = find_set(set_a_ind, parent, groups, op_groups); | ||
let mut set_b = find_set(set_b_ind, parent, groups, op_groups); | ||
|
||
if set_a == set_b { | ||
return; | ||
} | ||
if op_groups[&set_a].len() < op_groups[&set_b].len() { | ||
mem::swap(&mut set_a, &mut set_b); | ||
} | ||
parent[set_b] = Some(set_a); | ||
let mut other_op_groups = op_groups.get_mut(&set_b).unwrap().clone(); | ||
op_groups | ||
.get_mut(&set_a) | ||
.unwrap() | ||
.append(&mut other_op_groups); | ||
op_groups.get_mut(&set_b).unwrap().clear(); | ||
let mut other_groups = groups.get_mut(&set_b).unwrap().clone(); | ||
groups.get_mut(&set_a).unwrap().append(&mut other_groups); | ||
groups.get_mut(&set_b).unwrap().clear(); | ||
} |
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.
You can avoid reimplementing the whole DSU structure by using petgraph::unionfind
like we did for the Minimum-Spanning-Tree code. You probably just need to port the code from union_set
inside an if with .union(a, b)
. And keeping a hashmap with op_groups
of course.
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.
UnionFind
in petgraph merges two sets based on the rank
of the items. This is important to keep the tree structure well balanced. But in our case, since the height of the tree is at most the block_size
which i imagine to be just a small constant, the heuristic that @mtreinish used, i.e based on the size of op_groups
, will probably perform better in practice.
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.
Nice algorithm overall and especially the trick with lexicographical sorting. It's also nice that we managed to port this directly in Rust. That being said, it might be worth putting the effort to improve the readability of the code a bit (it took me a while to figure out how it's working).
Co-authored-by: georgios-ts <45130028+georgios-ts@users.noreply.github.com>
This commit adds a new function collect_multi_blocks to retworkx. It's
designed primarily for use by Qiskit's compiler to collect n qubit blocks
from a DAGCircuit. This is similar to how collect_runs() is used to
collect 1q blocks and collect_bicolor_runs is used to collect 2q blocks.
This PR is basically a mechanical porting of:
Qiskit/qiskit#4747
which added this functionality as a Qiskit transpiler pass. With this
API in retworkx the pass can still be added to Qiskit but just leverage
retworkx to compute the blocks.
TODO: