feat: Add T-count optimization pass (#2049)

Adds a t optimization pass that applies todd with a potential ancilla
budget to pg libs

---------

Co-authored-by: Silas Dilkes <36165522+sjdilkes@users.noreply.github.com>
Co-authored-by: Alec Edgington <54802828+cqc-alec@users.noreply.github.com>
Co-authored-by: yao-cqc <yao.tang@cambridgequantum.com>
Co-authored-by: Alec Edgington <alec.edgington@quantinuum.com>

Author: 4 people

.github/workflows/ci.yml

@@ -199,7 +199,7 @@ jobs:
         run: uv sync --only-group pg-libs --python ${{ env.PYTHON_HIGHEST }}
       - name: Run unstable_simd tests
         working-directory: pg-libs
-        run: uv run --only-group pg-libs cargo +${{ steps.nightly-toolchain.outputs.channel }} test -p tk-pg-bitpacked -p tk-pg-qm-tableau -p tk-pg-greedy-synth --features unstable_simd
+        run: uv run --only-group pg-libs cargo +${{ steps.nightly-toolchain.outputs.channel }} test -p tk-pg-bitpacked -p tk-pg-qm-tableau -p tk-pg-greedy-synth -p tk-pg-t-optimize --features unstable_simd
 
   benches:
     name: Continuous benchmarking 🏋️

Cargo.lock

@@ -19,6 +19,7 @@ members = [
     "pg-libs/pg-utils/tk-pg-utils",
     "pg-libs/pg-passes/tk-pg-canonical-form",
     "pg-libs/pg-passes/tk-pg-optimize",
+    "pg-libs/pg-passes/tk-pg-t-optimize",
     "pg-libs/pg-passes/tk-pg-greedy-synth",
     "pg-libs/pg-passes/tk-pg-rebase",
     "pg-libs/tk-pg-qm-tableau",
@@ -125,6 +126,7 @@ pg-core = { package = "tk-pg-core", path = "pg-libs/tk-pg-core", version = "0.1.
 pg-canonical-form = { package = "tk-pg-canonical-form", path = "pg-libs/pg-passes/tk-pg-canonical-form", version = "0.1.0" }
 pg-greedy-synth = { package = "tk-pg-greedy-synth", path = "pg-libs/pg-passes/tk-pg-greedy-synth", version = "0.1.0" }
 pg-optimise = { package = "tk-pg-optimize", path = "pg-libs/pg-passes/tk-pg-optimize", version = "0.1.0" }
+pg-t-optimize = { package = "tk-pg-t-optimize", path = "pg-libs/pg-passes/tk-pg-t-optimize", version = "0.1.0" }
 pg-rebase = { package = "tk-pg-rebase", path = "pg-libs/pg-passes/tk-pg-rebase", version = "0.1.0" }
 
 [profile.release.package.tket-py]
@@ -146,6 +148,8 @@ opt-level = 1
 opt-level = 1
 [profile.test.package.tk-pg-optimize]
 opt-level = 1
+[profile.test.package.tk-pg-t-optimize]
+opt-level = 1
 [profile.test.package.tk-pg-greedy-synth]
 opt-level = 1
 [profile.test.package.tk-pg-rebase]

Cargo.toml

@@ -1,6 +1,7 @@
 //! Pass to convert a Pauli graph into canonical form.
 use tk_pg_core::{
-    BlackBoxData, ConditionalBoxData, GateData, GateType, Op, PGPass, PauliGraph, TableauData,
+    BlackBoxData, ConditionalBoxData, GateData, GateType, Op, PGPass, Pauli, PauliGraph,
+    RotationData, TableauData,
 };
 use tk_pg_ir_kernels::{PGTableau, get_dagger, is_clifford, is_clifford_gate_type};
 use tk_pg_qm_tableau::Tableau as QubitMajorTableau;
@@ -260,9 +261,7 @@ fn process_op(
                         if compose {
                             tab.precompose_op(&op);
                         } else {
-                            for conjugated_op in tab.conjugate(&op) {
-                                pg.add_op(conjugated_op);
-                            }
+                            process_op(pg, &op, tab, forward, cliff_eval);
                         }
                     }
                     return;
@@ -283,7 +282,25 @@ fn process_op(
     if !forward {
         conjugated_ops.reverse();
     }
-    for conjugated_op in conjugated_ops {
+    for mut conjugated_op in conjugated_ops {
+        if cliff_eval && let Op::Rotation { data } = &mut conjugated_op {
+            let angle = data.get_angle();
+            if angle.is_finite() && angle.abs() >= 0.5 {
+                let remainder = angle % 0.5;
+                let correction = angle - remainder;
+                // The string is already conjugated, so append the correction to the tableau.
+                // Forward traversal maintains the inverse tableau.
+                if data.get_string().iter().any(|p| *p != Pauli::I) {
+                    tab.postcompose_op(&Op::Rotation {
+                        data: RotationData::new(
+                            data.get_string().clone(),
+                            if forward { -correction } else { correction },
+                        ),
+                    });
+                }
+                *data = RotationData::new(data.get_string().clone(), remainder);
+            }
+        }
         pg.add_op(conjugated_op);
     }
 }
@@ -339,7 +356,9 @@ fn to_canonical_form(pg: &PauliGraph, forward: bool, cliff_eval: bool) -> PauliG
 /// tableaux, interleaved with any obstructions and the non-Clifford ops (rotations, measures,
 /// resets, conditional operations) that are conjugated through and left in place.
 /// `cliff_eval` controls whether rotations with a Clifford angle are folded into the tableau
-/// (`true`) or kept as explicit rotations (`false`).
+/// (`true`) or kept as explicit rotations (`false`). When enabled, it also extracts
+/// Clifford corrections from unconditional rotations, leaving their angles strictly
+/// between -0.5 and 0.5 half turns.
 ///
 /// # Panics
 ///
@@ -369,7 +388,9 @@ impl CanonicalFormPass {
         self.forward = forward;
         self
     }
-    /// Set whether to evaluate Clifford angle rotations as Clifford gates or keep them as rotations.
+    /// Set whether to absorb Clifford rotations and the Clifford part of larger angles
+    /// into the tableau. Enabled by default; remaining unconditional rotation angles
+    /// lie strictly between -0.5 and 0.5 half turns.
     pub fn with_cliff_eval(mut self, cliff_eval: bool) -> Self {
         self.cliff_eval = cliff_eval;
         self
@@ -455,6 +476,87 @@ mod tests {
         assert!(compare_unitaries_via_tk(&pg, &transformed));
     }
 
+    #[rstest]
+    #[case(true, true, 0.75, 0.25)]
+    #[case(false, true, 0.75, 0.25)]
+    #[case(true, true, -0.75, -0.25)]
+    #[case(false, true, -0.75, -0.25)]
+    #[case(true, true, 2.1, 0.1)]
+    #[case(false, true, 2.1, 0.1)]
+    #[case(true, true, -2.1, -0.1)]
+    #[case(false, true, -2.1, -0.1)]
+    #[case(true, false, 0.75, 0.75)]
+    #[case(false, false, 0.75, 0.75)]
+    #[case(true, false, -0.75, -0.75)]
+    #[case(false, false, -0.75, -0.75)]
+    #[case(true, false, 2.1, 2.1)]
+    #[case(false, false, 2.1, 2.1)]
+    #[case(true, false, -2.1, -2.1)]
+    #[case(false, false, -2.1, -2.1)]
+    fn test_extract_clifford_corrections(
+        #[case] forward: bool,
+        #[case] cliff_eval: bool,
+        #[case] angle: f64,
+        #[case] expected_angle: f64,
+    ) {
+        let mut pg = PauliGraph::new(2);
+        pg.add_op(Op::Gate {
+            data: GateData::new(GateType::RZ, vec![0]).with_params(vec![angle]),
+        });
+        pg.add_op(Op::Rotation {
+            data: RotationData::new(vec![Pauli::I, Pauli::Z], angle),
+        });
+        let pass = CanonicalFormPass::new()
+            .with_forward(forward)
+            .with_cliff_eval(cliff_eval);
+        let transformed = pass.transform(&pg);
+
+        assert_eq!(transformed.get_ops().len(), 3);
+        for op in transformed.get_ops() {
+            if let Op::Rotation { data } = op {
+                assert!((data.get_angle() - expected_angle).abs() < 1e-10);
+            }
+        }
+        assert!(compare_unitaries_via_tk(&pg, &transformed));
+    }
+
+    #[rstest]
+    #[case(true)]
+    #[case(false)]
+    fn test_clifford_corrections_propagate(#[case] forward: bool) {
+        let mut pg = PauliGraph::new(2);
+        pg.add_op(Op::Gate {
+            data: GateData::new(GateType::H, vec![0]),
+        });
+        pg.add_op(Op::Gate {
+            data: GateData::new(GateType::RX, vec![0]).with_params(vec![0.75]),
+        });
+        pg.add_op(Op::Gate {
+            data: GateData::new(GateType::RY, vec![0]).with_params(vec![-0.75]),
+        });
+        pg.add_op(Op::Rotation {
+            data: RotationData::new(vec![Pauli::X, Pauli::Y], 1.25),
+        });
+        pg.add_op(Op::Gate {
+            data: GateData::new(GateType::PHASEDX, vec![0]).with_params(vec![0.75, 1.25]),
+        });
+        pg.add_op(Op::Gate {
+            data: GateData::new(GateType::ZZPHASE, vec![0, 1]).with_params(vec![-1.25]),
+        });
+        pg.add_op(Op::Gate {
+            data: GateData::new(GateType::H, vec![1]),
+        });
+        let pass = CanonicalFormPass::new().with_forward(forward);
+        let transformed = pass.transform(&pg);
+
+        for op in transformed.get_ops() {
+            if let Op::Rotation { data } = op {
+                assert!(data.get_angle().abs() < 0.5);
+            }
+        }
+        assert!(compare_unitaries_via_tk(&pg, &transformed));
+    }
+
     #[rstest]
     #[case(true, true, vec![0.5, 1.0], 1)]
     #[case(false, true, vec![0.5, 1.0], 1)]

pg-libs/pg-passes/tk-pg-canonical-form/src/lib.rs

@@ -0,0 +1,29 @@
+[package]
+name = "tk-pg-t-optimize"
+version = "0.1.0"
+description = "T-count optimization passes for Pauli graphs"
+edition = { workspace = true }
+rust-version = { workspace = true }
+license = { workspace = true }
+documentation = "https://docs.rs/tk-pg-t-optimize"
+homepage = { workspace = true }
+repository = { workspace = true }
+keywords = { workspace = true }
+categories = { workspace = true }
+
+[features]
+default = []
+unstable_simd = []
+
+[dependencies]
+bitvec = "1.0.1"
+tk-pg-core = { path = "../../tk-pg-core", version = "0.1.0" }
+tk-pg-canonical-form = { path = "../tk-pg-canonical-form", version = "0.1.0" }
+tk-pg-ir-kernels = { path = "../../pg-kernels/tk-pg-ir-kernels", version = "0.1.0" }
+tk-pg-qm-tableau = { path = "../../tk-pg-qm-tableau", version = "0.1.0" }
+
+[dev-dependencies]
+tk-pg-converter = { path = "../../pg-converters/tk-pg-converter", version = "0.1.0" }
+
+[lints]
+workspace = true

pg-libs/pg-passes/tk-pg-t-optimize/Cargo.toml

@@ -0,0 +1,29 @@
+# tk-pg-t-optimize
+
+`tk-pg-t-optimize` provides `TOptimizationPass`, which reduces the T count of
+Clifford + T Pauli graphs using phase polynomial resynthesis by the TODD
+algorithm. Optional ancillas allow Hadamard gadgets to combine more rotations
+into each optimization batch, increasing the performance of the pass.
+
+The pass expects a pauli graph which has been processed by the 
+`CanonicalFormPass` pass, and performs best after also applying phase folding 
+from the `RotationMergingPass`. Input measurements, resets, conditional operations, 
+black boxes and arbitrary Rz gates are not currently supported.
+
+By default, the pass uses no ancillas. `with_ancilla_budget` reserves the last
+qubits of the input graph as ancillas; these must already be present, idle, and
+prepared in zero. The pass does not allocate additional qubits.
+
+Using ancillas hadamard gadgets which contain measurements, resets, and 
+conditional Clifford.
+
+To synthesize the result into gates apply`GreedySynthPass`.
+
+Bit-vector operations are scalar by default. The optional `unstable_simd`
+feature enables portable SIMD and requires a nightly Rust toolchain.
+
+## License
+
+This project is licensed under Apache License, Version 2.0 ([LICENCE][] or <https://www.apache.org/licenses/LICENSE-2.0>).
+
+  [LICENCE]: https://github.com/quantinuum/tket2/blob/main/LICENCE