Improve commutation checking of Pauli product rotations and measurements

crates/transpiler/src/commutation_checker.rs

@@ -240,7 +240,8 @@ fn try_extract_op_from_ppr(
     Some(out.compose_map(&local, |i| qubits[i as usize].0))
 }
 
-fn try_pauli_generator(
+/// Attempt to extract generator of a Pauli-based gate in the form of a sparse observable.
+fn try_sparse_observable_generator_for_pauli_based(
     operation: &OperationRef,
     qubits: &[Qubit],
     num_qubits: u32,
@@ -254,7 +255,8 @@ fn try_pauli_generator(
     }
 }
 
-fn try_standard_gate_generator(
+/// Attemp to extract generator of a standard gate in the form of a sparse observable.
+fn try_sparse_observable_generator_for_standard_gate(
     operation: &OperationRef,
     params: &[Param],
     qubits: &[Qubit],
@@ -269,6 +271,18 @@ fn try_standard_gate_generator(
     None
 }
 
+/// Attempt to extract generator of a Pauli-based gate in the form of a single Pauli.
+/// When successful, return the generator in the (Z, X) form.
+pub fn try_pauli_generator_for_pauli_based<'a>(
+    operation: &'a OperationRef,
+) -> Option<(&'a Vec<bool>, &'a Vec<bool>)> {
+    match operation {
+        OperationRef::PauliProductRotation(ppr) => Some((&ppr.z, &ppr.x)),
+        OperationRef::PauliProductMeasurement(ppm) => Some((&ppm.z, &ppm.x)),
+        _ => None,
+    }
+}
+
 fn get_bits_from_py<T>(
     py_bits1: &Bound<'_, PyTuple>,
     py_bits2: &Bound<'_, PyTuple>,
@@ -307,6 +321,7 @@ pub struct CommutationChecker {
     current_cache_entries: usize,
     #[pyo3(get)]
     gates: Option<HashSet<String>>,
+    scratch_map: HashMap<usize, usize>,
 }
 
 #[pymethods]
@@ -460,6 +475,7 @@ impl CommutationChecker {
             cache_max_entries,
             current_cache_entries: 0,
             gates,
+            scratch_map: HashMap::new(),
         }
     }
 
@@ -542,21 +558,53 @@ impl CommutationChecker {
             _ => (),
         };
 
+        // Sort the arguments, such that `op2` always is the larger one.
+        let reversed = (op1.num_qubits(), op1.name().len(), op1.name())
+            > (op2.num_qubits(), op2.name().len(), op2.name());
+
+        let (op1, op2, params1, params2, qargs1, qargs2) = if reversed {
+            (op2, op1, params2, params1, qargs2, qargs1)
+        } else {
+            (op1, op2, params1, params2, qargs1, qargs2)
+        };
+
+        // Handle commutations using Pauli-based generators.
+        if let (Some((z1, x1)), Some((z2, x2))) = (
+            try_pauli_generator_for_pauli_based(op1),
+            try_pauli_generator_for_pauli_based(op2),
+        ) {
+            self.scratch_map.clear();
+            for (i, &q) in qargs1.iter().enumerate() {
+                self.scratch_map.insert(q.index(), i);
+            }
+            let mut parity = false;
+            for (j, &q) in qargs2.iter().enumerate() {
+                if let Some(&i) = self.scratch_map.get(&q.index()) {
+                    parity ^= (x1[i] && z2[j]) ^ (z1[i] && x2[j]);
+                }
+            }
+            return Ok(!parity);
+        }
+
         // Handle commutations between Pauli-based gates among themselves, and with standard gates
         // TODO Support trivial commutations of standard gates with identities in the Paulis
         let size = qargs1.iter().chain(qargs2.iter()).max().unwrap().0 + 1;
-        let maybe_pauli1 = try_pauli_generator(op1, qargs1, size);
-        let maybe_pauli2 = try_pauli_generator(op2, qargs2, size);
+        let maybe_pauli1 = try_sparse_observable_generator_for_pauli_based(op1, qargs1, size);
+        let maybe_pauli2 = try_sparse_observable_generator_for_pauli_based(op2, qargs2, size);
 
         match (maybe_pauli1, maybe_pauli2) {
             (None, None) => (), // No gate is Pauli-based, continue
             (None, Some(pauli2)) => {
-                if let Some(pauli1) = try_standard_gate_generator(op1, params1, qargs1, size) {
+                if let Some(pauli1) =
+                    try_sparse_observable_generator_for_standard_gate(op1, params1, qargs1, size)
+                {
                     return Ok(pauli1.commutes(&pauli2, tol));
                 }
             }
             (Some(pauli1), None) => {
-                if let Some(pauli2) = try_standard_gate_generator(op2, params2, qargs2, size) {
+                if let Some(pauli2) =
+                    try_sparse_observable_generator_for_standard_gate(op2, params2, qargs2, size)
+                {
                     return Ok(pauli1.commutes(&pauli2, tol));
                 }
             }
@@ -568,24 +616,6 @@ impl CommutationChecker {
             return Ok(false);
         }
 
-        // Sort the arguments, such that `second_op` always is the larger one.
-        let reversed = if op1.num_qubits() != op2.num_qubits() {
-            op1.num_qubits() > op2.num_qubits()
-        } else {
-            (op1.name().len(), op1.name()) >= (op2.name().len(), op2.name())
-        };
-        let (first_params, second_params) = if reversed {
-            (params2, params1)
-        } else {
-            (params1, params2)
-        };
-        let (first_op, second_op) = if reversed { (op2, op1) } else { (op1, op2) };
-        let (first_qargs, second_qargs) = if reversed {
-            (qargs2, qargs1)
-        } else {
-            (qargs1, qargs2)
-        };
-
         // For our cache to work correctly, we require the gate's definition to only depend on the
         // ``params`` attribute. This cannot be guaranteed for custom gates, so we only check
         // the cache for
@@ -599,70 +629,53 @@ impl CommutationChecker {
                 false
             }
         };
-        let check_cache =
-            is_cachable(first_op, first_params) && is_cachable(second_op, second_params);
+        let check_cache = is_cachable(op1, params1) && is_cachable(op2, params2);
 
         if !check_cache {
-            // The arguments are sorted, so if first_qargs.len() > matrix_max_num_qubits, then
-            // second_qargs.len() > matrix_max_num_qubits as well.
-            if second_qargs.len() > matrix_max_num_qubits as usize {
+            // The arguments are sorted, so if qargs1.len() > matrix_max_num_qubits, then
+            // qargs2.len() > matrix_max_num_qubits as well.
+            if qargs2.len() > matrix_max_num_qubits as usize {
                 return Ok(false);
             }
-            return self.commute_matmul(
-                first_op,
-                first_params,
-                first_qargs,
-                second_op,
-                second_params,
-                second_qargs,
-                tol,
-            );
+            return self.commute_matmul(op1, params1, qargs1, op2, params2, qargs2, tol);
         }
 
         // Query commutation library
-        let relative_placement = get_relative_placement(first_qargs, second_qargs);
+        let relative_placement = get_relative_placement(qargs1, qargs2);
         if let Some(is_commuting) =
             self.library
-                .check_commutation_entries(first_op, second_op, &relative_placement)
+                .check_commutation_entries(op1, op2, &relative_placement)
         {
             return Ok(is_commuting);
         }
 
         // Query cache
-        let key1 = hashable_params(first_params)?;
-        let key2 = hashable_params(second_params)?;
+        let key1 = hashable_params(params1)?;
+        let key2 = hashable_params(params2)?;
         if let Some(commutation_dict) = self
             .cache
-            .get(&(first_op.name().to_string(), second_op.name().to_string()))
+            .get(&(op1.name().to_string(), op2.name().to_string()))
         {
             let hashes = (key1.clone(), key2.clone());
             if let Some(commutation) = commutation_dict.get(&(relative_placement.clone(), hashes)) {
                 return Ok(*commutation);
             }
         }
 
-        if second_qargs.len() > matrix_max_num_qubits as usize {
+        if qargs2.len() > matrix_max_num_qubits as usize {
             return Ok(false);
         }
 
         // Perform matrix multiplication to determine commutation
-        let is_commuting = self.commute_matmul(
-            first_op,
-            first_params,
-            first_qargs,
-            second_op,
-            second_params,
-            second_qargs,
-            tol,
-        )?;
+        let is_commuting = self.commute_matmul(op1, params1, qargs1, op2, params2, qargs2, tol)?;
 
         // TODO: implement a LRU cache for this
         if self.current_cache_entries >= self.cache_max_entries {
             self.clear_cache();
         }
         // Cache results from is_commuting
         self.cache
-            .entry((first_op.name().to_string(), second_op.name().to_string()))
+            .entry((op1.name().to_string(), op2.name().to_string()))
             .and_modify(|entries| {
                 let key = (relative_placement.clone(), (key1.clone(), key2.clone()));
                 entries.insert(key, is_commuting);
@@ -1184,6 +1197,7 @@ pub fn get_standard_commutation_checker() -> CommutationChecker {
         cache: HashMap::new(),
         current_cache_entries: 0,
         gates: None,
+        scratch_map: HashMap::new(),
     }
 }
 

crates/transpiler/src/passes/commutative_optimization.rs

@@ -20,13 +20,15 @@ use pyo3::{Bound, PyResult, pyfunction, wrap_pyfunction};
 use qiskit_circuit::instruction::Parameters;
 use smallvec::smallvec;
 
-use crate::commutation_checker::{CommutationChecker, try_matrix_with_definition};
+use crate::commutation_checker::{
+    CommutationChecker, try_matrix_with_definition, try_pauli_generator_for_pauli_based,
+};
 use crate::passes::remove_identity_equiv::{average_gate_fidelity_below_tol, is_identity_equiv};
 use qiskit_circuit::circuit_instruction::OperationFromPython;
 use qiskit_circuit::dag_circuit::DAGCircuit;
 use qiskit_circuit::operations::{
-    Operation, OperationRef, Param, PauliBased, PauliProductRotation, StandardGate, multiply_param,
-    radd_param,
+    Operation, OperationRef, Param, PauliBased, PauliProductMeasurement, PauliProductRotation,
+    StandardGate, multiply_param, radd_param,
 };
 use qiskit_circuit::{BlocksMode, Clbit, NoBlocks, Qubit, imports};
 
@@ -104,7 +106,8 @@ static MERGEABLE_ROTATION_GATES: [StandardGate; 12] = [
 /// Computes the canonical representative of a packed instruction, and in particular:
 /// * replaces all types of Z-rotations by RZ-gates,
 /// * replaces all types of X-rotations by RX-gates,
-/// * sorts the qubits for symmetric gates.
+/// * sorts qubits for symmetric gates,
+/// * sorts qubits for [PauliProductRotation] and [PauliProductMeasurement].
 ///
 /// # Arguments:
 ///
@@ -196,6 +199,8 @@ fn canonicalize(
         }
     }
 
+    // Sort qubits for PauliProductRotations: this allows to merge scrambled pauli rotations
+    // and allows a faster commutativity check with other Pauli-based gates.
     if let OperationRef::PauliProductRotation(ppr) = inst.op.view() {
         let qargs = dag.get_qargs(inst.qubits);
         let mut paired = qargs
@@ -229,6 +234,41 @@ fn canonicalize(
         return Some((canonical_instruction, Param::Float(0.)));
     }
 
+    // Sort qubits for PauliProductMeasurements: this allows a faster commutativity check with
+    // other Pauli-based gates.
+    if let OperationRef::PauliProductMeasurement(ppm) = inst.op.view() {
+        let qargs = dag.get_qargs(inst.qubits);
+        let mut paired = qargs
+            .iter()
+            .zip(ppm.z.iter())
+            .zip(ppm.x.iter())
+            .map(|((q, z), x)| (q, z, x))
+            .collect::<Vec<_>>();
+        paired.sort_by_key(|(q, _, _)| **q);
+        let (sorted_qargs, sorted_z, sorted_x) =
+            paired
+                .into_iter()
+                .multiunzip::<(Vec<Qubit>, Vec<bool>, Vec<bool>)>();
+        let sorted_ppm = PauliProductMeasurement {
+            z: sorted_z,
+            x: sorted_x,
+            neg: ppm.neg,
+        };
+
+        let sorted_qubits = dag.add_qargs(&sorted_qargs);
+
+        let canonical_instruction = PackedInstruction {
+            op: PauliBased::PauliProductMeasurement(sorted_ppm).into(),
+            qubits: sorted_qubits,
+            clbits: inst.clbits,
+            params: inst.params.clone(),
+            label: None,
+            #[cfg(feature = "cache_pygates")]
+            py_op: std::sync::OnceLock::new(),
+        };
+        return Some((canonical_instruction, Param::Float(0.)));
+    }
+
     None
 }
 
@@ -256,6 +296,40 @@ fn commute(
     let op1 = inst1.op.view();
     let op2 = inst2.op.view();
 
+    // To check commutation of two Pauli-based gates, we extract their Pauli generators
+    // and check whether they commute.
+    if matches!(
+        op1,
+        OperationRef::PauliProductRotation(_) | OperationRef::PauliProductMeasurement(_)
+    ) && matches!(
+        op2,
+        OperationRef::PauliProductRotation(_) | OperationRef::PauliProductMeasurement(_)
+    ) {
+        if let (Some((z1, x1)), Some((z2, x2))) = (
+            try_pauli_generator_for_pauli_based(&op1),
+            try_pauli_generator_for_pauli_based(&op2),
+        ) {
+            let mut parity = false;
+
+            // Due to canonicalization, the qubits in PPRs and PPMs are known to be
+            // sorted by qubit index.
+            let (n1, n2) = (qargs1.len(), qargs2.len());
+            let (mut i1, mut i2) = (0, 0);
+            while i1 < n1 && i2 < n2 {
+                match qargs1[i1].cmp(&qargs2[i2]) {
+                    std::cmp::Ordering::Less => i1 += 1,
+                    std::cmp::Ordering::Greater => i2 += 1,
+                    std::cmp::Ordering::Equal => {
+                        parity ^= (x1[i1] && z2[i2]) ^ (z1[i1] && x2[i2]);
+                        i1 += 1;
+                        i2 += 1;
+                    }
+                }
+            }
+            return Ok(!parity);
+        }
+    }
+
     Ok(commutation_checker.commute(
         &op1,
         inst1.params.as_deref(),

releasenotes/notes/speedup-pauli-rotations-commutation-checking-c9b65ad1e3d0c453.yaml

@@ -0,0 +1,11 @@
+---
+features_transpiler:
+  - |
+    Improved performance of :class:`.CommutationChecker` when evaluating commutations
+    between :class:`.PauliProductRotationGate` and :class:`.PauliProductMeasurement` objects.
+    This is achieved by representing their generators as single-terms Paulis and performing
+    commutativity checks using these representations.
+  - |
+    Improved performance of the :class:`.CommutativeOptimization` transpiler pass on
+    circuits containing :class:`.PauliProductRotationGate` and :class:`.PauliProductMeasurement`
+    objects.