[WIP] Add collect_multi_blocks function #461
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.
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.
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.
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 <[email protected]>
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: