From cb508880984d209501bb77a662450147dcadc4e5 Mon Sep 17 00:00:00 2001 From: Henry Zou Date: Thu, 11 Jul 2024 10:29:56 -0400 Subject: [PATCH 1/4] This commit ports the core logic of `star_preroute` from Python to Rust. The changes involve creating a new Rust module for the star prerouting algorithm and updating the corresponding Python code to integrate with this new Rust functionality. Details: - New Rust file: Added `star_preroute.rs` to handle the core logic of the function `star_preroute` from the python side. This file defines the type aliases for node and block representations, which matches the current block representation of the `StarBlock` (except that the center is just a bool, as we only need to know if there is a center), and the node representation matches how the nodes used in `SabreDAG`. The `star_preroute` function processes the star blocks witihin the `SabreDAG` and finds the linear routing equivalent and then returns the result as a `SabreResult`. Thus we can use the same methods we used in Sabre, such as `_build_sabre_dag` and `_apply_sabre_result`. - Node representation: A key part of this implementation is how it takes advantage of `SabreResult` and `SabreDAG`, so the node representation is a tuple of the node id, list of qubit indices, a set of classical bit indices, and a directive flag. However, once we update the regular DAG to rust, this part may change significantly. - Updates in the SABRE rust module: To use `sabre_dag.rs` and `swap_map.rs` in `star_prerouting`, I change them to be public in `crates/accelerate/src/sabre/mod.rs`. Not sure if it makes more sense to do it this way or to move `star_prerouting` to `crates/accelerate/src/sabre/` since it mimics the methods used in Sabre to change the dag. - Python side updates: Imported the necessary modules and only modified the function `star_preroute` so that now the function performs the heavy lifting of transforming the DAG within the Rust space, leveraging `_build_sabre_dag` and `_apply_sabre_result`. - Possible issues: I was not sure how correctly handle control flow from the rust side. I know that `route.rs` handles this with `route_control_flow_block` to populate the `node_block_results` for `SabreResult`, but I was not sure how to best take advantage of this function for `star_prerouting`. Currently, the `node_block_results` for `star_prerouting` essentially always empty and just there to have`SabreResult`. There also seems to be no unit tests for `star_prerouting` that includes control flow. --- crates/accelerate/src/lib.rs | 1 + crates/accelerate/src/sabre/mod.rs | 4 +- crates/accelerate/src/star_prerouting.rs | 215 ++++++++++++++++++ crates/pyext/src/lib.rs | 4 +- qiskit/__init__.py | 1 + .../passes/routing/star_prerouting.py | 183 +++++++-------- ...port_star_prerouting-13fae3ff78feb5e3.yaml | 11 + 7 files changed, 313 insertions(+), 106 deletions(-) create mode 100644 crates/accelerate/src/star_prerouting.rs create mode 100644 releasenotes/notes/port_star_prerouting-13fae3ff78feb5e3.yaml diff --git a/crates/accelerate/src/lib.rs b/crates/accelerate/src/lib.rs index dcfbdc9f1878..6e93c68e6177 100644 --- a/crates/accelerate/src/lib.rs +++ b/crates/accelerate/src/lib.rs @@ -27,6 +27,7 @@ pub mod results; pub mod sabre; pub mod sampled_exp_val; pub mod sparse_pauli_op; +pub mod star_prerouting; pub mod stochastic_swap; pub mod synthesis; pub mod two_qubit_decompose; diff --git a/crates/accelerate/src/sabre/mod.rs b/crates/accelerate/src/sabre/mod.rs index 3eb8ebb3a219..1229be16b723 100644 --- a/crates/accelerate/src/sabre/mod.rs +++ b/crates/accelerate/src/sabre/mod.rs @@ -14,8 +14,8 @@ mod layer; mod layout; mod neighbor_table; mod route; -mod sabre_dag; -mod swap_map; +pub mod sabre_dag; +pub mod swap_map; use hashbrown::HashMap; use numpy::{IntoPyArray, ToPyArray}; diff --git a/crates/accelerate/src/star_prerouting.rs b/crates/accelerate/src/star_prerouting.rs new file mode 100644 index 000000000000..ec1aacbf1f77 --- /dev/null +++ b/crates/accelerate/src/star_prerouting.rs @@ -0,0 +1,215 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2024 +// +// This code is licensed under the Apache License, Version 2.0. You may +// obtain a copy of this license in the LICENSE.txt file in the root directory +// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +// +// Any modifications or derivative works of this code must retain this +// copyright notice, and modified files need to carry a notice indicating +// that they have been altered from the originals. + +/// Type alias for a node representation. +/// Each node is represented as a tuple containing: +/// - Node id (usize) +/// - List of involved qubit indices (Vec) +/// - Set of involved classical bit indices (HashSet) +/// - Directive flag (bool) +type Nodes = (usize, Vec, HashSet, bool); + +/// Type alias for a block representation. +/// Each block is represented by a tuple containing: +/// - A boolean indicating the presence of a center (bool) +/// - A list of nodes (Vec) +type Block = (bool, Vec); + +use crate::nlayout::PhysicalQubit; +use crate::nlayout::VirtualQubit; +use crate::sabre::sabre_dag::SabreDAG; +use crate::sabre::swap_map::SwapMap; +use crate::sabre::BlockResult; +use crate::sabre::NodeBlockResults; +use crate::sabre::SabreResult; +use hashbrown::HashMap; +use hashbrown::HashSet; +use numpy::IntoPyArray; +use pyo3::prelude::*; + +/// Python function to perform star prerouting on a SabreDAG. +/// This function processes star blocks and updates the DAG and qubit mapping. +#[pyfunction] +#[pyo3(text_signature = "(dag, blocks, processing_order, /)")] +fn star_preroute( + py: Python, + dag: &mut SabreDAG, + blocks: Vec, + processing_order: Vec, +) -> (SwapMap, PyObject, NodeBlockResults, PyObject) { + let mut qubit_mapping: Vec = (0..dag.num_qubits).collect(); + let mut processed_block_ids: HashSet = HashSet::new(); + let last_2q_gate = processing_order.iter().rev().find(|node| node.1.len() > 1); + let mut is_first_star = true; + + // Structures for SabreResult + let mut out_map: HashMap> = + HashMap::with_capacity(dag.dag.node_count()); + let mut gate_order: Vec = Vec::with_capacity(dag.dag.node_count()); + let node_block_results: HashMap> = HashMap::new(); + + // Process each node in the given processing order + for node in processing_order.iter() { + if let Some(block_id) = find_block_id(&blocks, node) { + // Skip if the block has already been processed + if !processed_block_ids.insert(block_id) { + continue; + } + process_block( + &mut qubit_mapping, + &blocks[block_id], + last_2q_gate, + &mut is_first_star, + &mut gate_order, + &mut out_map, + ); + } else { + // Apply operation for nodes not part of any block + gate_order.push(node.0); + } + } + + let res = SabreResult { + map: SwapMap { map: out_map }, + node_order: gate_order, + node_block_results: NodeBlockResults { + results: node_block_results, + }, + }; + + let final_res = ( + res.map, + res.node_order.into_pyarray_bound(py).into(), + res.node_block_results, + qubit_mapping.into_pyarray_bound(py).into(), + ); + + final_res +} + +/// Finds the block ID for a given node. +/// +/// Args: +/// +/// * `blocks` - A vector of blocks to search for the node. +/// * `node` - The node for which the block ID needs to be found. +/// +/// Returns: +/// +/// An option containing the block ID if the node is part of a block, otherwise None. +fn find_block_id(blocks: &[Block], node: &Nodes) -> Option { + blocks.iter().enumerate().find_map(|(i, block)| { + if block.1.iter().any(|n| n.0 == node.0) { + Some(i) + } else { + None + } + }) +} + +/// Processes a star block, applying operations and handling swaps. +/// +/// Args: +/// +/// * `qubit_mapping` - A mutable reference to the qubit mapping vector. +/// * `block` - A tuple containing a boolean indicating the presence of a center and a vector of nodes representing the star block. +/// * `last_2q_gate` - The last two-qubit gate in the processing order. +/// * `is_first_star` - A mutable reference to a boolean indicating if this is the first star block being processed. +/// * `gate_order` - A mutable reference to the gate order vector. +/// * `out_map` - A mutable reference to the output map. +fn process_block<'a>( + qubit_mapping: &mut [usize], + block: &'a Block, + last_2q_gate: Option<&'a Nodes>, + is_first_star: &mut bool, + gate_order: &mut Vec, + out_map: &mut HashMap>, +) { + let (has_center, sequence) = block; + + // If the block contains exactly 2 nodes, apply them directly + if sequence.len() == 2 { + for inner_node in sequence { + gate_order.push(inner_node.0); + } + return; + } + + let mut prev_qargs = None; + let mut swap_source = false; + + // Process each node in the block + for (i, inner_node) in sequence.iter().enumerate() { + // Apply operation directly if it's a single-qubit operation or the same as previous qargs + if inner_node.1.len() == 1 || prev_qargs == Some(&inner_node.1) { + gate_order.push(inner_node.0); + continue; + } + + // If this is the first star and no swap source has been identified, set swap_source + if *is_first_star && !swap_source { + swap_source = *has_center; + gate_order.push(inner_node.0); + prev_qargs = Some(&inner_node.1); + continue; + } + + // Place 2q-gate and subsequent swap gate + gate_order.push(inner_node.0); + + if inner_node != last_2q_gate.unwrap() && inner_node.1.len() == 2 { + // Use the node ID of the next node in the sequence to match how SABRE applies swaps + if let Some(next_node) = sequence.get(i + 1) { + apply_swap(qubit_mapping, &inner_node.1, next_node.0, out_map); + } + } + + prev_qargs = Some(&inner_node.1); + } + *is_first_star = false; +} + +/// Applies a swap operation to the DAG and updates the qubit mapping. +/// +/// # Args: +/// +/// * `qubit_mapping` - A mutable reference to the qubit mapping vector. +/// * `qargs` - A slice containing the qubit indices for the swap operation. +/// * `next_node_id` - The ID of the next node in the sequence. +/// * `out_map` - A mutable reference to the output map. +fn apply_swap( + qubit_mapping: &mut [usize], + qargs: &[VirtualQubit], + next_node_id: usize, + out_map: &mut HashMap>, +) { + if qargs.len() == 2 { + let idx0 = qargs[0].index(); + let idx1 = qargs[1].index(); + + // Update the `qubit_mapping` and `out_map` to reflect the swap operation + qubit_mapping.swap(idx0, idx1); + out_map.insert( + next_node_id, + vec![[ + PhysicalQubit::new(qubit_mapping[idx0].try_into().unwrap()), + PhysicalQubit::new(qubit_mapping[idx1].try_into().unwrap()), + ]], + ); + } +} + +#[pymodule] +pub fn star_prerouting(m: &Bound) -> PyResult<()> { + m.add_wrapped(wrap_pyfunction!(star_preroute))?; + Ok(()) +} diff --git a/crates/pyext/src/lib.rs b/crates/pyext/src/lib.rs index 72f0d759099a..c4c6b0b0e8b6 100644 --- a/crates/pyext/src/lib.rs +++ b/crates/pyext/src/lib.rs @@ -18,7 +18,8 @@ use qiskit_accelerate::{ error_map::error_map, euler_one_qubit_decomposer::euler_one_qubit_decomposer, isometry::isometry, nlayout::nlayout, optimize_1q_gates::optimize_1q_gates, pauli_exp_val::pauli_expval, results::results, sabre::sabre, sampled_exp_val::sampled_exp_val, - sparse_pauli_op::sparse_pauli_op, stochastic_swap::stochastic_swap, synthesis::synthesis, + sparse_pauli_op::sparse_pauli_op, star_prerouting::star_prerouting, + stochastic_swap::stochastic_swap, synthesis::synthesis, two_qubit_decompose::two_qubit_decompose, uc_gate::uc_gate, utils::utils, vf2_layout::vf2_layout, }; @@ -41,6 +42,7 @@ fn _accelerate(m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pymodule!(sabre))?; m.add_wrapped(wrap_pymodule!(sampled_exp_val))?; m.add_wrapped(wrap_pymodule!(sparse_pauli_op))?; + m.add_wrapped(wrap_pymodule!(star_prerouting))?; m.add_wrapped(wrap_pymodule!(stochastic_swap))?; m.add_wrapped(wrap_pymodule!(two_qubit_decompose))?; m.add_wrapped(wrap_pymodule!(uc_gate))?; diff --git a/qiskit/__init__.py b/qiskit/__init__.py index 9aaa7a76a68e..a2863c790f92 100644 --- a/qiskit/__init__.py +++ b/qiskit/__init__.py @@ -77,6 +77,7 @@ sys.modules["qiskit._accelerate.sabre"] = _accelerate.sabre sys.modules["qiskit._accelerate.sampled_exp_val"] = _accelerate.sampled_exp_val sys.modules["qiskit._accelerate.sparse_pauli_op"] = _accelerate.sparse_pauli_op +sys.modules["qiskit._accelerate.star_prerouting"] = _accelerate.star_prerouting sys.modules["qiskit._accelerate.stochastic_swap"] = _accelerate.stochastic_swap sys.modules["qiskit._accelerate.two_qubit_decompose"] = _accelerate.two_qubit_decompose sys.modules["qiskit._accelerate.vf2_layout"] = _accelerate.vf2_layout diff --git a/qiskit/transpiler/passes/routing/star_prerouting.py b/qiskit/transpiler/passes/routing/star_prerouting.py index 3679e8bfb8e3..afc697104b40 100644 --- a/qiskit/transpiler/passes/routing/star_prerouting.py +++ b/qiskit/transpiler/passes/routing/star_prerouting.py @@ -14,11 +14,15 @@ from typing import Iterable, Union, Optional, List, Tuple from math import floor, log10 -from qiskit.circuit import Barrier +from qiskit.circuit import SwitchCaseOp, Clbit, ClassicalRegister, Barrier +from qiskit.circuit.controlflow import condition_resources, node_resources from qiskit.dagcircuit import DAGOpNode, DAGDepNode, DAGDependency, DAGCircuit from qiskit.transpiler import Layout from qiskit.transpiler.basepasses import TransformationPass -from qiskit.circuit.library import SwapGate +from qiskit.transpiler.passes.routing.sabre_swap import _build_sabre_dag, _apply_sabre_result + +from qiskit._accelerate import star_prerouting +from qiskit._accelerate.nlayout import NLayout class StarBlock: @@ -305,113 +309,86 @@ def star_preroute(self, dag, blocks, processing_order): new_dag: a dag specifying the pre-routed circuit qubit_mapping: the final qubit mapping after pre-routing """ - node_to_block_id = {} - for i, block in enumerate(blocks): - for node in block.get_nodes(): - node_to_block_id[node] = i - - new_dag = dag.copy_empty_like() - processed_block_ids = set() - qubit_mapping = list(range(len(dag.qubits))) - - def _apply_mapping(qargs, qubit_mapping, qubits): - return tuple(qubits[qubit_mapping[dag.find_bit(qubit).index]] for qubit in qargs) - - is_first_star = True - last_2q_gate = [ - op - for op in reversed(processing_order) - if ((len(op.qargs) > 1) and (op.name != "barrier")) + # Convert the DAG to a SabreDAG + num_qubits = len(dag.qubits) + canonical_register = dag.qregs["q"] + current_layout = Layout.generate_trivial_layout(canonical_register) + qubit_indices = {bit: idx for idx, bit in enumerate(canonical_register)} + layout_mapping = {qubit_indices[k]: v for k, v in current_layout.get_virtual_bits().items()} + initial_layout = NLayout(layout_mapping, num_qubits, num_qubits) + sabre_dag, circuit_to_dag_dict = _build_sabre_dag(dag, num_qubits, qubit_indices) + + # Extract the nodes from the blocks for the Rust representation + rust_blocks = [ + (block.center is not None, _extract_nodes(block.get_nodes(), dag)) for block in blocks ] - if len(last_2q_gate) > 0: - last_2q_gate = last_2q_gate[0] - else: - last_2q_gate = None + # Determine the processing order of the nodes in the DAG for the Rust representation int_digits = floor(log10(len(processing_order))) + 1 processing_order_index_map = { - node: f"a{str(index).zfill(int(int_digits))}" - for index, node in enumerate(processing_order) + node: f"a{index:0{int_digits}}" for index, node in enumerate(processing_order) } def tie_breaker_key(node): return processing_order_index_map.get(node, node.sort_key) - for node in dag.topological_op_nodes(key=tie_breaker_key): - block_id = node_to_block_id.get(node, None) - if block_id is not None: - if block_id in processed_block_ids: - continue - - processed_block_ids.add(block_id) - - # process the whole block - block = blocks[block_id] - sequence = block.nodes - center_node = block.center - - if len(sequence) == 2: - for inner_node in sequence: - new_dag.apply_operation_back( - inner_node.op, - _apply_mapping(inner_node.qargs, qubit_mapping, dag.qubits), - inner_node.cargs, - check=False, - ) - continue - swap_source = None - prev = None - for inner_node in sequence: - if (len(inner_node.qargs) == 1) or (inner_node.qargs == prev): - new_dag.apply_operation_back( - inner_node.op, - _apply_mapping(inner_node.qargs, qubit_mapping, dag.qubits), - inner_node.cargs, - check=False, - ) - continue - if is_first_star and swap_source is None: - swap_source = center_node - new_dag.apply_operation_back( - inner_node.op, - _apply_mapping(inner_node.qargs, qubit_mapping, dag.qubits), - inner_node.cargs, - check=False, - ) - - prev = inner_node.qargs - continue - # place 2q-gate and subsequent swap gate - new_dag.apply_operation_back( - inner_node.op, - _apply_mapping(inner_node.qargs, qubit_mapping, dag.qubits), - inner_node.cargs, - check=False, - ) - - if not inner_node is last_2q_gate and not isinstance(inner_node.op, Barrier): - new_dag.apply_operation_back( - SwapGate(), - _apply_mapping(inner_node.qargs, qubit_mapping, dag.qubits), - inner_node.cargs, - check=False, - ) - # Swap mapping - index_0 = dag.find_bit(inner_node.qargs[0]).index - index_1 = dag.find_bit(inner_node.qargs[1]).index - qubit_mapping[index_1], qubit_mapping[index_0] = ( - qubit_mapping[index_0], - qubit_mapping[index_1], - ) - - prev = inner_node.qargs - is_first_star = False - else: - # the node is not part of a block - new_dag.apply_operation_back( - node.op, - _apply_mapping(node.qargs, qubit_mapping, dag.qubits), - node.cargs, - check=False, - ) - return new_dag, qubit_mapping + rust_processing_order = _extract_nodes(dag.topological_op_nodes(key=tie_breaker_key), dag) + + # Run the star prerouting algorithm to obtain the new DAG and qubit mapping + *sabre_result, qubit_mapping = star_prerouting.star_preroute( + sabre_dag, rust_blocks, rust_processing_order + ) + + res_dag = _apply_sabre_result( + dag.copy_empty_like(), + dag, + sabre_result, + initial_layout, + dag.qubits, + circuit_to_dag_dict, + ) + + return res_dag, qubit_mapping + + +def _extract_nodes( + nodes: List[DAGOpNode], dag: DAGCircuit +) -> List[Tuple[int, List[int], set[int], bool]]: + """Extract and format node information for Rust representation used in SabreDAG. + + Each node is represented as a tuple containing: + - Node ID (int): The unique identifier of the node in the DAG. + - Qubit indices (list of int): Indices of qubits involved in the node's operation. + - Classical bit indices (set of int): Indices of classical bits involved in the node's operation. + - Directive flag (bool): Indicates whether the operation is a directive (True) or not (False). + + Args: + nodes (list[DAGOpNode]): List of DAGOpNode objects representing the nodes to extract information from. + dag (DAGCircuit): DAGCircuit object containing the circuit structure. + + Returns: + list of tuples: Each tuple contains (node_id, qubit_indices, classical_bit_indices, is_directive). + """ + extracted_node_info = [] + for node in nodes: + qubit_indices = [dag.find_bit(qubit).index for qubit in node.qargs] + classical_bit_indices = set() + + if node.op.condition is not None: + classical_bit_indices.update(condition_resources(node.op.condition).clbits) + + if isinstance(node.op, SwitchCaseOp): + switch_case_target = node.op.target + if isinstance(switch_case_target, Clbit): + classical_bit_indices.add(switch_case_target) + elif isinstance(switch_case_target, ClassicalRegister): + classical_bit_indices.update(switch_case_target) + else: # Assume target is an expression involving classical bits + classical_bit_indices.update(node_resources(switch_case_target).clbits) + + is_directive = getattr(node.op, "_directive", False) + extracted_node_info.append( + (node._node_id, qubit_indices, classical_bit_indices, is_directive) + ) + + return extracted_node_info diff --git a/releasenotes/notes/port_star_prerouting-13fae3ff78feb5e3.yaml b/releasenotes/notes/port_star_prerouting-13fae3ff78feb5e3.yaml new file mode 100644 index 000000000000..f8eca807bec6 --- /dev/null +++ b/releasenotes/notes/port_star_prerouting-13fae3ff78feb5e3.yaml @@ -0,0 +1,11 @@ +--- +features_transpiler: + - | + Port part of the logic from the :class:`StarPrerouting`, used to + find a star graph connectivity subcircuit and replaces it with a + linear routing equivalent. + - | + The function :func:`star_preroute` now performs the heavily lifting + to transform the dag by in the rust space by taking advantage + of the functions :func:`_build_sabre_dag` and + :func:`_apply_sabre_result`. From 2d5e276ec958843f314cfb768dd75247c99b7d8c Mon Sep 17 00:00:00 2001 From: Henry Zou Date: Fri, 12 Jul 2024 11:31:02 -0400 Subject: [PATCH 2/4] lint --- qiskit/transpiler/passes/routing/star_prerouting.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/qiskit/transpiler/passes/routing/star_prerouting.py b/qiskit/transpiler/passes/routing/star_prerouting.py index afc697104b40..f1f3f4bd42db 100644 --- a/qiskit/transpiler/passes/routing/star_prerouting.py +++ b/qiskit/transpiler/passes/routing/star_prerouting.py @@ -351,9 +351,7 @@ def tie_breaker_key(node): return res_dag, qubit_mapping -def _extract_nodes( - nodes: List[DAGOpNode], dag: DAGCircuit -) -> List[Tuple[int, List[int], set[int], bool]]: +def _extract_nodes(nodes, dag): """Extract and format node information for Rust representation used in SabreDAG. Each node is represented as a tuple containing: @@ -363,11 +361,11 @@ def _extract_nodes( - Directive flag (bool): Indicates whether the operation is a directive (True) or not (False). Args: - nodes (list[DAGOpNode]): List of DAGOpNode objects representing the nodes to extract information from. + nodes (list[DAGOpNode]): List of DAGOpNode objects to extract information from. dag (DAGCircuit): DAGCircuit object containing the circuit structure. Returns: - list of tuples: Each tuple contains (node_id, qubit_indices, classical_bit_indices, is_directive). + list of tuples: Each tuple contains information about a node in the format described above. """ extracted_node_info = [] for node in nodes: From 14640ecb70838c226e961afab08c225226e6e1e4 Mon Sep 17 00:00:00 2001 From: Henry Zou Date: Thu, 25 Jul 2024 11:30:19 -0400 Subject: [PATCH 3/4] Added additional test and adjust the swap map result - Added the additional test of qft linearization and that the resultings circuit has `n-2` swap gates where `n` is the number of cp gates. - Changed the `node_id` in `apply_swap` of `star_preroute.rs` to use the current node id as it is more efficient, but just does not match how we do it in Sabre. This makes it so that we apply the gate first then the swap, which fixes an error we had before where we never placed a swap gate at the end of the processing a block. This only affected tests where we had multiple blocks to process. To make sure we apply the results correctly from `SabreResult`, I added a flag to `_apply_sabre_result` to treat the case of `StarPrerouting` differently so that it applies the swap after applying the node. - Added a hasp map from node to block to make processing each node in the given processing order have `n + n` time complexity instead of `n^2`. As a result, we can also remove the function `find_block_id` --- crates/accelerate/src/star_prerouting.rs | 44 ++++++----------- .../transpiler/passes/routing/sabre_swap.py | 47 ++++++++++++++----- .../passes/routing/star_prerouting.py | 1 + .../python/transpiler/test_star_prerouting.py | 21 +++++++++ 4 files changed, 71 insertions(+), 42 deletions(-) diff --git a/crates/accelerate/src/star_prerouting.rs b/crates/accelerate/src/star_prerouting.rs index ec1aacbf1f77..a4ed5d42e513 100644 --- a/crates/accelerate/src/star_prerouting.rs +++ b/crates/accelerate/src/star_prerouting.rs @@ -48,7 +48,7 @@ fn star_preroute( ) -> (SwapMap, PyObject, NodeBlockResults, PyObject) { let mut qubit_mapping: Vec = (0..dag.num_qubits).collect(); let mut processed_block_ids: HashSet = HashSet::new(); - let last_2q_gate = processing_order.iter().rev().find(|node| node.1.len() > 1); + let last_2q_gate = processing_order.iter().rev().find(|node| node.1.len() == 2); let mut is_first_star = true; // Structures for SabreResult @@ -57,9 +57,17 @@ fn star_preroute( let mut gate_order: Vec = Vec::with_capacity(dag.dag.node_count()); let node_block_results: HashMap> = HashMap::new(); + // Create a HashMap to store the node-to-block mapping + let mut node_to_block: HashMap = HashMap::new(); + for (block_id, block) in blocks.iter().enumerate() { + for node in &block.1 { + node_to_block.insert(node.0, block_id); + } + } + // Process each node in the given processing order for node in processing_order.iter() { - if let Some(block_id) = find_block_id(&blocks, node) { + if let Some(&block_id) = node_to_block.get(&node.0) { // Skip if the block has already been processed if !processed_block_ids.insert(block_id) { continue; @@ -96,26 +104,6 @@ fn star_preroute( final_res } -/// Finds the block ID for a given node. -/// -/// Args: -/// -/// * `blocks` - A vector of blocks to search for the node. -/// * `node` - The node for which the block ID needs to be found. -/// -/// Returns: -/// -/// An option containing the block ID if the node is part of a block, otherwise None. -fn find_block_id(blocks: &[Block], node: &Nodes) -> Option { - blocks.iter().enumerate().find_map(|(i, block)| { - if block.1.iter().any(|n| n.0 == node.0) { - Some(i) - } else { - None - } - }) -} - /// Processes a star block, applying operations and handling swaps. /// /// Args: @@ -148,7 +136,7 @@ fn process_block<'a>( let mut swap_source = false; // Process each node in the block - for (i, inner_node) in sequence.iter().enumerate() { + for inner_node in sequence.iter() { // Apply operation directly if it's a single-qubit operation or the same as previous qargs if inner_node.1.len() == 1 || prev_qargs == Some(&inner_node.1) { gate_order.push(inner_node.0); @@ -168,9 +156,7 @@ fn process_block<'a>( if inner_node != last_2q_gate.unwrap() && inner_node.1.len() == 2 { // Use the node ID of the next node in the sequence to match how SABRE applies swaps - if let Some(next_node) = sequence.get(i + 1) { - apply_swap(qubit_mapping, &inner_node.1, next_node.0, out_map); - } + apply_swap(qubit_mapping, &inner_node.1, inner_node.0, out_map); } prev_qargs = Some(&inner_node.1); @@ -184,12 +170,12 @@ fn process_block<'a>( /// /// * `qubit_mapping` - A mutable reference to the qubit mapping vector. /// * `qargs` - A slice containing the qubit indices for the swap operation. -/// * `next_node_id` - The ID of the next node in the sequence. +/// * `node_id` - The ID of the node in the sequence. /// * `out_map` - A mutable reference to the output map. fn apply_swap( qubit_mapping: &mut [usize], qargs: &[VirtualQubit], - next_node_id: usize, + node_id: usize, out_map: &mut HashMap>, ) { if qargs.len() == 2 { @@ -199,7 +185,7 @@ fn apply_swap( // Update the `qubit_mapping` and `out_map` to reflect the swap operation qubit_mapping.swap(idx0, idx1); out_map.insert( - next_node_id, + node_id, vec![[ PhysicalQubit::new(qubit_mapping[idx0].try_into().unwrap()), PhysicalQubit::new(qubit_mapping[idx1].try_into().unwrap()), diff --git a/qiskit/transpiler/passes/routing/sabre_swap.py b/qiskit/transpiler/passes/routing/sabre_swap.py index acb23f39ab09..946b75bae540 100644 --- a/qiskit/transpiler/passes/routing/sabre_swap.py +++ b/qiskit/transpiler/passes/routing/sabre_swap.py @@ -328,6 +328,7 @@ def _apply_sabre_result( initial_layout, physical_qubits, circuit_to_dag_dict, + apply_swap_first=True, ): """Apply the ``SabreResult`` to ``out_dag``, mutating it in place. This function in effect performs the :class:`.ApplyLayout` transpiler pass with ``initial_layout`` and the Sabre routing @@ -350,6 +351,9 @@ def _apply_sabre_result( circuit_to_dag_dict (Mapping[int, DAGCircuit]): a mapping of the Python object identity (as returned by :func:`id`) of a control-flow block :class:`.QuantumCircuit` to a :class:`.DAGCircuit` that represents the same thing. + apply_swap_first (bool): a flag indicating whether to apply swaps before operations. + In SabreSwap, the swap map reflects applying the swap before applying the node, while + in StarPrerouting, the swap map reflects applying the swap after applying the node. """ # The swap gate is a singleton instance, so we don't need to waste time reconstructing it each @@ -381,19 +385,36 @@ def recurse(dest_dag, source_dag, result, root_logical_map, layout): swap_map, node_order, node_block_results = result for node_id in node_order: node = source_dag._multi_graph[node_id] - if node_id in swap_map: - apply_swaps(dest_dag, swap_map[node_id], layout) - if not isinstance(node.op, ControlFlowOp): - dest_dag.apply_operation_back( - node.op, - [ - physical_qubits[layout.virtual_to_physical(root_logical_map[q])] - for q in node.qargs - ], - node.cargs, - check=False, - ) - continue + + if apply_swap_first: + if node_id in swap_map: + apply_swaps(dest_dag, swap_map[node_id], layout) + if not isinstance(node.op, ControlFlowOp): + dest_dag.apply_operation_back( + node.op, + [ + physical_qubits[layout.virtual_to_physical(root_logical_map[q])] + for q in node.qargs + ], + node.cargs, + check=False, + ) + continue + else: + if not isinstance(node.op, ControlFlowOp): + dest_dag.apply_operation_back( + node.op, + [ + physical_qubits[layout.virtual_to_physical(root_logical_map[q])] + for q in node.qargs + ], + node.cargs, + check=False, + ) + if node_id in swap_map: + apply_swaps(dest_dag, swap_map[node_id], layout) + if not isinstance(node.op, ControlFlowOp): + continue # At this point, we have to handle a control-flow node. block_results = node_block_results[node_id] diff --git a/qiskit/transpiler/passes/routing/star_prerouting.py b/qiskit/transpiler/passes/routing/star_prerouting.py index f1f3f4bd42db..57868efce950 100644 --- a/qiskit/transpiler/passes/routing/star_prerouting.py +++ b/qiskit/transpiler/passes/routing/star_prerouting.py @@ -346,6 +346,7 @@ def tie_breaker_key(node): initial_layout, dag.qubits, circuit_to_dag_dict, + apply_swap_first=False, ) return res_dag, qubit_mapping diff --git a/test/python/transpiler/test_star_prerouting.py b/test/python/transpiler/test_star_prerouting.py index ddc8096eefd7..25162f16b0b6 100644 --- a/test/python/transpiler/test_star_prerouting.py +++ b/test/python/transpiler/test_star_prerouting.py @@ -26,6 +26,7 @@ ) from qiskit.quantum_info import Operator from qiskit.transpiler.passes import VF2Layout, ApplyLayout, SabreSwap, SabreLayout +from qiskit.transpiler.passes.layout.vf2_utils import build_interaction_graph from qiskit.transpiler.passes.routing.star_prerouting import StarPreRouting from qiskit.transpiler.coupling import CouplingMap from qiskit.transpiler.passmanager import PassManager @@ -482,3 +483,23 @@ def test_routing_after_star_prerouting(self): self.assertTrue(Operator.from_circuit(res_sabre), qc) self.assertTrue(Operator.from_circuit(res_star), qc) self.assertTrue(Operator.from_circuit(res_star), Operator.from_circuit(res_sabre)) + + @ddt.data(4, 8, 16, 32) + def test_qft_linearization(self, num_qubits): + """Test the QFT circuit to verify if it is linearized and requires n-2 swaps.""" + + qc = QFT(num_qubits, do_swaps=False, insert_barriers=True).decompose() + dag = circuit_to_dag(qc) + new_dag = StarPreRouting().run(dag) + new_qc = dag_to_circuit(new_dag) + + # Check that resulting result has n-2 swaps, where n is the number of cp gates + swap_count = new_qc.count_ops().get("swap", 0) + cp_count = new_qc.count_ops().get("cp", 0) + self.assertEqual(swap_count, cp_count - 2) + + # Confirm linearization by checking that the number of edges is equal to the number of nodes + interaction_graph = build_interaction_graph(new_dag, strict_direction=False)[0] + num_edges = interaction_graph.num_edges() + num_nodes = interaction_graph.num_nodes() + self.assertEqual(num_edges, num_nodes - 1) From 9d3904efa80a969168d1b06b3019222871296ab7 Mon Sep 17 00:00:00 2001 From: Henry Zou Date: Fri, 26 Jul 2024 16:29:48 -0400 Subject: [PATCH 4/4] Reverted changes to `_apply_sabre_result` and fixed handling on rust side - Removed `apply_swap_first` flag in `_apply_sabre_result` as it did not make sense to have it as there are no other scenario where a user may want to have control over applying the swap first. - To adjust for this and make `star_preroute` consistent with `apply_sabre_result` to apply swaps before the node id on the swap map, I adjusted `star_preroute.rs` to first process the blocks to gather the swap locations and the gate order. Once we have the full gate order, we can use the swap locations to apply the swaps while knowing the `qargs` of the node before the swap and the `node_id` of the node after the swap. - Since the above in done in the main `star_preroute` function, I removed `qubit_ampping` and `out_map` as arguments for `process_blocks`. --- crates/accelerate/src/star_prerouting.rs | 47 ++++++++++++------- .../transpiler/passes/routing/sabre_swap.py | 47 +++++-------------- .../passes/routing/star_prerouting.py | 1 - 3 files changed, 43 insertions(+), 52 deletions(-) diff --git a/crates/accelerate/src/star_prerouting.rs b/crates/accelerate/src/star_prerouting.rs index a4ed5d42e513..fd2156ad2011 100644 --- a/crates/accelerate/src/star_prerouting.rs +++ b/crates/accelerate/src/star_prerouting.rs @@ -47,7 +47,7 @@ fn star_preroute( processing_order: Vec, ) -> (SwapMap, PyObject, NodeBlockResults, PyObject) { let mut qubit_mapping: Vec = (0..dag.num_qubits).collect(); - let mut processed_block_ids: HashSet = HashSet::new(); + let mut processed_block_ids: HashSet = HashSet::with_capacity(blocks.len()); let last_2q_gate = processing_order.iter().rev().find(|node| node.1.len() == 2); let mut is_first_star = true; @@ -58,27 +58,28 @@ fn star_preroute( let node_block_results: HashMap> = HashMap::new(); // Create a HashMap to store the node-to-block mapping - let mut node_to_block: HashMap = HashMap::new(); + let mut node_to_block: HashMap = HashMap::with_capacity(processing_order.len()); for (block_id, block) in blocks.iter().enumerate() { for node in &block.1 { node_to_block.insert(node.0, block_id); } } + // Store nodes where swaps will be placed. + let mut swap_locations: Vec<&Nodes> = Vec::with_capacity(processing_order.len()); - // Process each node in the given processing order - for node in processing_order.iter() { + // Process blocks, gathering swap locations and updating the gate order + for node in &processing_order { if let Some(&block_id) = node_to_block.get(&node.0) { // Skip if the block has already been processed if !processed_block_ids.insert(block_id) { continue; } process_block( - &mut qubit_mapping, &blocks[block_id], last_2q_gate, &mut is_first_star, &mut gate_order, - &mut out_map, + &mut swap_locations, ); } else { // Apply operation for nodes not part of any block @@ -86,6 +87,22 @@ fn star_preroute( } } + // Apply the swaps based on the gathered swap locations and gate order + for (index, node_id) in gate_order.iter().enumerate() { + for swap_location in &swap_locations { + if *node_id == swap_location.0 { + if let Some(next_node_id) = gate_order.get(index + 1) { + apply_swap( + &mut qubit_mapping, + &swap_location.1, + *next_node_id, + &mut out_map, + ); + } + } + } + } + let res = SabreResult { map: SwapMap { map: out_map }, node_order: gate_order, @@ -108,19 +125,17 @@ fn star_preroute( /// /// Args: /// -/// * `qubit_mapping` - A mutable reference to the qubit mapping vector. /// * `block` - A tuple containing a boolean indicating the presence of a center and a vector of nodes representing the star block. /// * `last_2q_gate` - The last two-qubit gate in the processing order. /// * `is_first_star` - A mutable reference to a boolean indicating if this is the first star block being processed. /// * `gate_order` - A mutable reference to the gate order vector. -/// * `out_map` - A mutable reference to the output map. +/// * `swap_locations` - A mutable reference to the nodes where swaps will be placed after fn process_block<'a>( - qubit_mapping: &mut [usize], block: &'a Block, last_2q_gate: Option<&'a Nodes>, is_first_star: &mut bool, gate_order: &mut Vec, - out_map: &mut HashMap>, + swap_locations: &mut Vec<&'a Nodes>, ) { let (has_center, sequence) = block; @@ -155,10 +170,8 @@ fn process_block<'a>( gate_order.push(inner_node.0); if inner_node != last_2q_gate.unwrap() && inner_node.1.len() == 2 { - // Use the node ID of the next node in the sequence to match how SABRE applies swaps - apply_swap(qubit_mapping, &inner_node.1, inner_node.0, out_map); + swap_locations.push(inner_node); } - prev_qargs = Some(&inner_node.1); } *is_first_star = false; @@ -169,13 +182,13 @@ fn process_block<'a>( /// # Args: /// /// * `qubit_mapping` - A mutable reference to the qubit mapping vector. -/// * `qargs` - A slice containing the qubit indices for the swap operation. -/// * `node_id` - The ID of the node in the sequence. +/// * `qargs` - Qubit indices for the swap operation (node before the swap) +/// * `next_node_id` - ID of the next node in the gate order (node after the swap) /// * `out_map` - A mutable reference to the output map. fn apply_swap( qubit_mapping: &mut [usize], qargs: &[VirtualQubit], - node_id: usize, + next_node_id: usize, out_map: &mut HashMap>, ) { if qargs.len() == 2 { @@ -185,7 +198,7 @@ fn apply_swap( // Update the `qubit_mapping` and `out_map` to reflect the swap operation qubit_mapping.swap(idx0, idx1); out_map.insert( - node_id, + next_node_id, vec![[ PhysicalQubit::new(qubit_mapping[idx0].try_into().unwrap()), PhysicalQubit::new(qubit_mapping[idx1].try_into().unwrap()), diff --git a/qiskit/transpiler/passes/routing/sabre_swap.py b/qiskit/transpiler/passes/routing/sabre_swap.py index 946b75bae540..acb23f39ab09 100644 --- a/qiskit/transpiler/passes/routing/sabre_swap.py +++ b/qiskit/transpiler/passes/routing/sabre_swap.py @@ -328,7 +328,6 @@ def _apply_sabre_result( initial_layout, physical_qubits, circuit_to_dag_dict, - apply_swap_first=True, ): """Apply the ``SabreResult`` to ``out_dag``, mutating it in place. This function in effect performs the :class:`.ApplyLayout` transpiler pass with ``initial_layout`` and the Sabre routing @@ -351,9 +350,6 @@ def _apply_sabre_result( circuit_to_dag_dict (Mapping[int, DAGCircuit]): a mapping of the Python object identity (as returned by :func:`id`) of a control-flow block :class:`.QuantumCircuit` to a :class:`.DAGCircuit` that represents the same thing. - apply_swap_first (bool): a flag indicating whether to apply swaps before operations. - In SabreSwap, the swap map reflects applying the swap before applying the node, while - in StarPrerouting, the swap map reflects applying the swap after applying the node. """ # The swap gate is a singleton instance, so we don't need to waste time reconstructing it each @@ -385,36 +381,19 @@ def recurse(dest_dag, source_dag, result, root_logical_map, layout): swap_map, node_order, node_block_results = result for node_id in node_order: node = source_dag._multi_graph[node_id] - - if apply_swap_first: - if node_id in swap_map: - apply_swaps(dest_dag, swap_map[node_id], layout) - if not isinstance(node.op, ControlFlowOp): - dest_dag.apply_operation_back( - node.op, - [ - physical_qubits[layout.virtual_to_physical(root_logical_map[q])] - for q in node.qargs - ], - node.cargs, - check=False, - ) - continue - else: - if not isinstance(node.op, ControlFlowOp): - dest_dag.apply_operation_back( - node.op, - [ - physical_qubits[layout.virtual_to_physical(root_logical_map[q])] - for q in node.qargs - ], - node.cargs, - check=False, - ) - if node_id in swap_map: - apply_swaps(dest_dag, swap_map[node_id], layout) - if not isinstance(node.op, ControlFlowOp): - continue + if node_id in swap_map: + apply_swaps(dest_dag, swap_map[node_id], layout) + if not isinstance(node.op, ControlFlowOp): + dest_dag.apply_operation_back( + node.op, + [ + physical_qubits[layout.virtual_to_physical(root_logical_map[q])] + for q in node.qargs + ], + node.cargs, + check=False, + ) + continue # At this point, we have to handle a control-flow node. block_results = node_block_results[node_id] diff --git a/qiskit/transpiler/passes/routing/star_prerouting.py b/qiskit/transpiler/passes/routing/star_prerouting.py index 1d531334dc68..53bc971a268b 100644 --- a/qiskit/transpiler/passes/routing/star_prerouting.py +++ b/qiskit/transpiler/passes/routing/star_prerouting.py @@ -346,7 +346,6 @@ def tie_breaker_key(node): initial_layout, dag.qubits, circuit_to_dag_dict, - apply_swap_first=False, ) return res_dag, qubit_mapping