Port star_preroute to rust (backport #12761)

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;

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};

crates/accelerate/src/star_prerouting.rs

@@ -0,0 +1,214 @@
+// 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<VirtualQubit>)
+/// - Set of involved classical bit indices (HashSet<usize>)
+/// - Directive flag (bool)
+type Nodes = (usize, Vec<VirtualQubit>, HashSet<usize>, 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<Nodes>)
+type Block = (bool, Vec<Nodes>);
+
+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<Block>,
+    processing_order: Vec<Nodes>,
+) -> (SwapMap, PyObject, NodeBlockResults, PyObject) {
+    let mut qubit_mapping: Vec<usize> = (0..dag.num_qubits).collect();
+    let mut processed_block_ids: HashSet<usize> = 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;
+
+    // Structures for SabreResult
+    let mut out_map: HashMap<usize, Vec<[PhysicalQubit; 2]>> =
+        HashMap::with_capacity(dag.dag.node_count());
+    let mut gate_order: Vec<usize> = Vec::with_capacity(dag.dag.node_count());
+    let node_block_results: HashMap<usize, Vec<BlockResult>> = HashMap::new();
+
+    // Create a HashMap to store the node-to-block mapping
+    let mut node_to_block: HashMap<usize, usize> = 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 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(
+                &blocks[block_id],
+                last_2q_gate,
+                &mut is_first_star,
+                &mut gate_order,
+                &mut swap_locations,
+            );
+        } else {
+            // Apply operation for nodes not part of any block
+            gate_order.push(node.0);
+        }
+    }
+
+    // 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,
+        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
+}
+
+/// Processes a star block, applying operations and handling swaps.
+///
+/// Args:
+///
+/// * `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.
+/// * `swap_locations` - A mutable reference to the nodes where swaps will be placed after
+fn process_block<'a>(
+    block: &'a Block,
+    last_2q_gate: Option<&'a Nodes>,
+    is_first_star: &mut bool,
+    gate_order: &mut Vec<usize>,
+    swap_locations: &mut Vec<&'a Nodes>,
+) {
+    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 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);
+            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 {
+            swap_locations.push(inner_node);
+        }
+        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` - 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],
+    next_node_id: usize,
+    out_map: &mut HashMap<usize, Vec<[PhysicalQubit; 2]>>,
+) {
+    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<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(star_preroute))?;
+    Ok(())
+}

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<PyModule>) -> 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))?;

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