Add additional Z and X rotation gate handling to CommutativeCancellation

crates/transpiler/src/passes/commutation_cancellation.rs

@@ -21,6 +21,7 @@ use smallvec::{SmallVec, smallvec};
 
 use super::analyze_commutations;
 use crate::commutation_checker::CommutationChecker;
+use approx::abs_diff_eq;
 use qiskit_circuit::Qubit;
 use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType};
 use qiskit_circuit::operations::{Operation, Param, StandardGate};
@@ -34,15 +35,22 @@ static VAR_Z_MAP: [(&str, StandardGate); 3] = [
     ("p", StandardGate::Phase),
     ("u1", StandardGate::U1),
 ];
-static Z_ROTATIONS: [StandardGate; 6] = [
+static Z_ROTATIONS: [StandardGate; 8] = [
     StandardGate::Phase,
     StandardGate::Z,
     StandardGate::U1,
     StandardGate::RZ,
     StandardGate::T,
+    StandardGate::Tdg,
     StandardGate::S,
+    StandardGate::Sdg,
+];
+static X_ROTATIONS: [StandardGate; 4] = [
+    StandardGate::X,
+    StandardGate::RX,
+    StandardGate::SX,
+    StandardGate::SXdg,
 ];
-static X_ROTATIONS: [StandardGate; 2] = [StandardGate::X, StandardGate::RX];
 static SUPPORTED_GATES: [StandardGate; 5] = [
     StandardGate::CX,
     StandardGate::CY,
@@ -92,6 +100,8 @@ pub fn cancel_commutations(
                     .map(|(_, gate)| gate)
             })
         });
+    let sx_supported = dag.get_op_counts().contains_key("sx") || basis.iter().any(|x| x == "sx");
+    let x_supported = dag.get_op_counts().contains_key("x") || basis.iter().any(|x| x == "x");
 
     // RZ and P/U1 have a phase difference of angle/2, which we need to account for
     let z_phase_shift = match z_var_gate {
@@ -247,9 +257,13 @@ pub fn cancel_commutations(
                     } else {
                         match node_op_name {
                             "t" => Ok((FRAC_PI_4, z_phase_shift("p", FRAC_PI_4))),
+                            "tdg" => Ok((-FRAC_PI_4, -z_phase_shift("p", FRAC_PI_4))),
                             "s" => Ok((FRAC_PI_2, z_phase_shift("p", FRAC_PI_2))),
+                            "sdg" => Ok((-FRAC_PI_2, -z_phase_shift("p", FRAC_PI_2))),
                             "z" => Ok((PI, z_phase_shift("p", PI))),
                             "x" => Ok((PI, FRAC_PI_2)),
+                            "sx" => Ok((FRAC_PI_2, FRAC_PI_4)),
+                            "sxdg" => Ok((-FRAC_PI_2, -FRAC_PI_4)),
                             _ => Err(PyRuntimeError::new_err(format!(
                                 "Angle for operation {node_op_name} is not defined"
                             ))),
@@ -261,22 +275,36 @@ pub fn cancel_commutations(
 
                 let new_op = match cancel_key.gate {
                     GateOrRotation::ZRotation => z_var_gate.unwrap(),
-                    GateOrRotation::XRotation => &StandardGate::RX,
+                    GateOrRotation::XRotation => {
+                        if x_supported && is_multiple_of_pi_k(total_angle, 1.) {
+                            &StandardGate::X
+                        } else if sx_supported && is_multiple_of_pi_k(total_angle, 2.) {
+                            &StandardGate::SX
+                        } else {
+                            &StandardGate::RX
+                        }
+                    }
                     _ => unreachable!(),
                 };
 
-                let pi_multiple = total_angle / PI;
-
-                let mod4 = pi_multiple.rem_euclid(4.);
-                if mod4 < _CUTOFF_PRECISION || (4. - mod4) < _CUTOFF_PRECISION {
+                if is_multiple_of_pi_k(total_angle, 0.25) {
                     // if the angle is close to a 4-pi multiple (from above or below), then the
                     // operator is equal to the identity
-                } else if (mod4 - 2.).abs() < _CUTOFF_PRECISION {
+                } else if is_multiple_of_pi_k(total_angle, 0.5) {
                     // a 2-pi multiple has a phase of pi: RX(2pi) = RZ(2pi) = -I = I exp(i pi)
                     total_phase -= PI;
                 } else {
                     // any other is not the identity and we add the gate
-                    dag.insert_1q_on_incoming_qubit((*new_op, &[total_angle]), cancel_set[0]);
+                    if new_op == &StandardGate::X {
+                        dag.insert_1q_on_incoming_qubit((*new_op, &[]), cancel_set[0]);
+                    } else if new_op == &StandardGate::SX {
+                        let gate_counts = ((total_angle / FRAC_PI_2) as u64) % 4;
+                        for _ in 0..gate_counts {
+                            dag.insert_1q_on_incoming_qubit((*new_op, &[]), cancel_set[0]);
+                        }
+                    } else {
+                        dag.insert_1q_on_incoming_qubit((*new_op, &[total_angle]), cancel_set[0]);
+                    }
                 }
 
                 dag.add_global_phase(&Param::Float(total_phase))?;
@@ -291,6 +319,13 @@ pub fn cancel_commutations(
     Ok(())
 }
 
+fn is_multiple_of_pi_k(angle: f64, k: f64) -> bool {
+    let modulo = angle * k / PI;
+    let remainder = modulo.rem_euclid(1.0);
+    abs_diff_eq!(remainder, 0., epsilon = _CUTOFF_PRECISION)
+        || abs_diff_eq!(remainder, 1., epsilon = _CUTOFF_PRECISION)
+}
+
 pub fn commutation_cancellation_mod(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(cancel_commutations))?;
     Ok(())