Oxidize Commutation Analysis

crates/accelerate/src/commutation_analysis.rs

@@ -0,0 +1,179 @@
+// 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.
+
+use pyo3::prelude::PyModule;
+use pyo3::{pyfunction, pymodule, wrap_pyfunction, Bound, PyResult, Python};
+use qiskit_circuit::operations::Param;
+use qiskit_circuit::Qubit;
+use smallvec::{smallvec, SmallVec};
+use std::hash::BuildHasherDefault;
+
+use crate::commutation_checker::CommutationChecker;
+use ahash::AHasher;
+use hashbrown::HashMap;
+use indexmap::IndexSet;
+use pyo3::prelude::*;
+
+use pyo3::types::{PyDict, PyList};
+use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType, Wire};
+use rustworkx_core::petgraph::stable_graph::NodeIndex;
+
+type AIndexSet<T> = IndexSet<T, BuildHasherDefault<AHasher>>;
+#[derive(Clone, Debug)]
+pub enum CommutationSetEntry {
+    Index(usize),
+    SetExists(Vec<AIndexSet<NodeIndex>>),
+}
+
+fn analyze_commutations_inner(
+    py: Python,
+    dag: &mut DAGCircuit,
+    commutation_checker: &mut CommutationChecker,
+) -> HashMap<(Option<NodeIndex>, Wire), CommutationSetEntry> {
+    let mut commutation_set: HashMap<(Option<NodeIndex>, Wire), CommutationSetEntry> =
+        HashMap::new();
+    let max_num_qubits = 3;
+
+    (0..dag.num_qubits()).for_each(|qubit| {
+        let wire = Wire::Qubit(Qubit(qubit as u32));
+        dag.nodes_on_wire(py, &wire, false)
+            .iter()
+            .for_each(|current_gate_idx| {
+                if let CommutationSetEntry::SetExists(ref mut commutation_entry) = commutation_set
+                    .entry((None, wire.clone()))
+                    .or_insert_with(|| {
+                        CommutationSetEntry::SetExists(vec![AIndexSet::from_iter([
+                            *current_gate_idx,
+                        ])])
+                    })
+                {
+                    let last = commutation_entry.last_mut().unwrap();
+
+                    if !last.contains(current_gate_idx) {
+                        if last.iter().all(|prev_gate_idx| {
+                            if let (
+                                NodeType::Operation(packed_inst0),
+                                NodeType::Operation(packed_inst1),
+                            ) = (&dag.dag[*current_gate_idx], &dag.dag[*prev_gate_idx])
+                            {
+                                let empty_params: Box<SmallVec<[Param; 3]>> = Box::new(smallvec![]);
+                                let op1 = packed_inst0.op.view();
+                                let op2 = packed_inst1.op.view();
+                                let params1 = match packed_inst0.params.as_ref() {
+                                    Some(params) => params,
+                                    None => &empty_params,
+                                };
+                                let params2 = match packed_inst1.params.as_ref() {
+                                    Some(params) => params,
+                                    None => &empty_params,
+                                };
+                                let qargs1 = dag.qargs_interner.get(packed_inst0.qubits);
+                                let qargs2 = dag.qargs_interner.get(packed_inst1.qubits);
+                                let cargs1 = dag.cargs_interner.get(packed_inst0.clbits);
+                                let cargs2 = dag.cargs_interner.get(packed_inst1.clbits);
+
+                                // TODO preliminary interface, change this when dagcircuit merges
+                                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,
+                                    )
+                                    .unwrap()
+                            } else {
+                                false
+                            }
+                        }) {
+                            // 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]))
+                        }
+                    }
+                } else {
+                    panic!("Wrong type in dictionary!");
+                }
+                if let CommutationSetEntry::SetExists(last_entry) =
+                    commutation_set.get(&(None, wire.clone())).unwrap()
+                {
+                    commutation_set.insert(
+                        (Some(*current_gate_idx), wire.clone()),
+                        CommutationSetEntry::Index(last_entry.len() - 1),
+                    );
+                }
+            })
+    });
+    commutation_set
+}
+
+#[pyfunction]
+#[pyo3(signature = (dag, commutation_checker))]
+pub(crate) fn analyze_commutations(
+    py: Python,
+    dag: &mut DAGCircuit,
+    commutation_checker: &mut CommutationChecker,
+) -> PyResult<Py<PyDict>> {
+    let commutations = analyze_commutations_inner(py, dag, commutation_checker);
+    let out_dict = PyDict::new_bound(py);
+    for (k, comms) in commutations {
+        let nidx = k.0;
+        let wire = match k.1 {
+            Wire::Qubit(q) => dag.qubits.get(q).unwrap().to_object(py),
+            Wire::Clbit(c) => dag.clbits.get(c).unwrap().to_object(py),
+            Wire::Var(v) => v,
+        };
+
+        if nidx.is_some() {
+            match comms {
+                CommutationSetEntry::Index(idx) => {
+                    out_dict.set_item((dag.get_node(py, nidx.unwrap())?, wire), idx)?
+                }
+                _ => panic!("Wrong format in commutation analysis"),
+            };
+        } else {
+            match comms {
+                CommutationSetEntry::SetExists(comm_set) => out_dict.set_item(
+                    wire,
+                    PyList::new_bound(
+                        py,
+                        comm_set.iter().map(|inner| {
+                            PyList::new_bound(
+                                py,
+                                inner
+                                    .into_iter()
+                                    .map(|ndidx| dag.get_node(py, *ndidx).unwrap()),
+                            )
+                        }),
+                    ),
+                )?,
+                _ => panic!("Wrong format in commutation analysis"),
+            }
+        }
+    }
+    Ok(out_dict.unbind())
+}
+
+#[pymodule]
+pub fn commutation_analysis(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(analyze_commutations))?;
+    Ok(())
+}

crates/accelerate/src/commutation_checker.rs

@@ -69,7 +69,7 @@ where
 /// lookups. It's not meant to be a public facing Python object though and only used
 /// internally by the Python class.
 #[pyclass(module = "qiskit._accelerate.commutation_checker")]
-struct CommutationChecker {
+pub struct CommutationChecker {
     library: CommutationLibrary,
     cache_max_entries: usize,
     cache: HashMap<(String, String), CommutationCacheEntry>,
@@ -227,7 +227,7 @@ impl CommutationChecker {
 
 impl CommutationChecker {
     #[allow(clippy::too_many_arguments)]
-    fn commute_inner(
+    pub fn commute_inner(
         &mut self,
         py: Python,
         op1: &OperationRef,

crates/accelerate/src/lib.rs

@@ -15,6 +15,7 @@ use std::env;
 use pyo3::import_exception;
 
 pub mod circuit_library;
+pub mod commutation_analysis;
 pub mod commutation_checker;
 pub mod convert_2q_block_matrix;
 pub mod dense_layout;

crates/circuit/src/dag_circuit.rs

@@ -231,9 +231,9 @@ pub struct DAGCircuit {
     cregs: Py<PyDict>,
 
     /// The cache used to intern instruction qargs.
-    qargs_interner: Interner<[Qubit]>,
+    pub qargs_interner: Interner<[Qubit]>,
     /// The cache used to intern instruction cargs.
-    cargs_interner: Interner<[Clbit]>,
+    pub cargs_interner: Interner<[Clbit]>,
     /// Qubits registered in the circuit.
     pub qubits: BitData<Qubit>,
     /// Clbits registered in the circuit.
@@ -5523,7 +5523,7 @@ impl DAGCircuit {
     /// Get the nodes on the given wire.
     ///
     /// Note: result is empty if the wire is not in the DAG.
-    fn nodes_on_wire(&self, py: Python, wire: &Wire, only_ops: bool) -> Vec<NodeIndex> {
+    pub fn nodes_on_wire(&self, py: Python, wire: &Wire, only_ops: bool) -> Vec<NodeIndex> {
         let mut nodes = Vec::new();
         let mut current_node = match wire {
             Wire::Qubit(qubit) => self.qubit_io_map.get(qubit.0 as usize).map(|x| x[0]),

crates/pyext/src/lib.rs

@@ -13,15 +13,15 @@
 use pyo3::prelude::*;
 
 use qiskit_accelerate::{
-    circuit_library::circuit_library, commutation_checker::commutation_checker,
-    convert_2q_block_matrix::convert_2q_block_matrix, dense_layout::dense_layout,
-    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, star_prerouting::star_prerouting,
-    stochastic_swap::stochastic_swap, synthesis::synthesis, target_transpiler::target,
-    two_qubit_decompose::two_qubit_decompose, uc_gate::uc_gate, utils::utils,
-    vf2_layout::vf2_layout,
+    circuit_library::circuit_library, commutation_analysis::commutation_analysis,
+    commutation_checker::commutation_checker, convert_2q_block_matrix::convert_2q_block_matrix,
+    dense_layout::dense_layout, 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,
+    star_prerouting::star_prerouting, stochastic_swap::stochastic_swap, synthesis::synthesis,
+    target_transpiler::target, two_qubit_decompose::two_qubit_decompose, uc_gate::uc_gate,
+    utils::utils, vf2_layout::vf2_layout,
 };
 
 #[inline(always)]
@@ -62,5 +62,6 @@ fn _accelerate(m: &Bound<PyModule>) -> PyResult<()> {
     add_submodule(m, utils, "utils")?;
     add_submodule(m, vf2_layout, "vf2_layout")?;
     add_submodule(m, commutation_checker, "commutation_checker")?;
+    add_submodule(m, commutation_analysis, "commutation_analysis")?;
     Ok(())
 }

qiskit/__init__.py

@@ -87,6 +87,7 @@
 sys.modules["qiskit._accelerate.synthesis.linear"] = _accelerate.synthesis.linear
 sys.modules["qiskit._accelerate.synthesis.clifford"] = _accelerate.synthesis.clifford
 sys.modules["qiskit._accelerate.commutation_checker"] = _accelerate.commutation_checker
+sys.modules["qiskit._accelerate.commutation_analysis"] = _accelerate.commutation_analysis
 sys.modules["qiskit._accelerate.synthesis.linear_phase"] = _accelerate.synthesis.linear_phase
 
 from qiskit.exceptions import QiskitError, MissingOptionalLibraryError

qiskit/transpiler/passes/optimization/commutation_analysis.py

@@ -12,11 +12,10 @@
 
 """Analysis pass to find commutation relations between DAG nodes."""
 
-from collections import defaultdict
-
 from qiskit.circuit.commutation_library import SessionCommutationChecker as scc
-from qiskit.dagcircuit import DAGOpNode
+from qiskit.circuit.commutation_checker import CommutationChecker as PyCommutationChecker
 from qiskit.transpiler.basepasses import AnalysisPass
+from qiskit._accelerate.commutation_analysis import analyze_commutations
 
 
 class CommutationAnalysis(AnalysisPass):
@@ -33,6 +32,13 @@ def __init__(self, *, _commutation_checker=None):
         # do not care about commutations of all gates, but just a subset
         if _commutation_checker is None:
             _commutation_checker = scc
+
+        if isinstance(_commutation_checker, PyCommutationChecker):
+            # If an instance of the to-be deprecated Python implementation of the CommutationChecker
+            # is passed, defer to the internally held Rust implementation instead. This is needed
+            # as the ``analyze_commutations`` function expects the Rust implementation.
+            _commutation_checker = _commutation_checker.cc
+
         self.comm_checker = _commutation_checker
 
     def run(self, dag):
@@ -42,49 +48,4 @@ def run(self, dag):
         into the ``property_set``.
         """
         # Initiate the commutation set
-        self.property_set["commutation_set"] = defaultdict(list)
-
-        # Build a dictionary to keep track of the gates on each qubit
-        # The key with format (wire) will store the lists of commutation sets
-        # The key with format (node, wire) will store the index of the commutation set
-        # on the specified wire, thus, for example:
-        # self.property_set['commutation_set'][wire][(node, wire)] will give the
-        # commutation set that contains node.
-
-        for wire in dag.qubits:
-            self.property_set["commutation_set"][wire] = []
-
-        # Add edges to the dictionary for each qubit
-        for node in dag.topological_op_nodes():
-            for _, _, edge_wire in dag.edges(node):
-                self.property_set["commutation_set"][(node, edge_wire)] = -1
-
-        # Construct the commutation set
-        for wire in dag.qubits:
-
-            for current_gate in dag.nodes_on_wire(wire):
-
-                current_comm_set = self.property_set["commutation_set"][wire]
-                if not current_comm_set:
-                    current_comm_set.append([current_gate])
-
-                if current_gate not in current_comm_set[-1]:
-                    does_commute = True
-
-                    # Check if the current gate commutes with all the gates in the current block
-                    for prev_gate in current_comm_set[-1]:
-                        does_commute = (
-                            isinstance(current_gate, DAGOpNode)
-                            and isinstance(prev_gate, DAGOpNode)
-                            and self.comm_checker.commute_nodes(current_gate, prev_gate)
-                        )
-                        if not does_commute:
-                            break
-
-                    if does_commute:
-                        current_comm_set[-1].append(current_gate)
-                    else:
-                        current_comm_set.append([current_gate])
-
-                temp_len = len(current_comm_set)
-                self.property_set["commutation_set"][(current_gate, wire)] = temp_len - 1
+        self.property_set["commutation_set"] = analyze_commutations(dag, self.comm_checker)

releasenotes/notes/oxidize-commutation-analysis-d2fc81feb6ca80aa.yaml

@@ -0,0 +1,4 @@
+---
+features_transpiler:
+  - |
+    Added a Rust implementation of :class:`.CommutationAnalysis` in :func:`.analyze_commutations`.