Move TK2 SWAP decomp into c++, update overall method

pytket/binders/circuit/library.cpp

@@ -230,6 +230,10 @@ void init_library(py::module &m) {
       "_TK2_using_ZZPhase_and_swap", &CircPool::TK2_using_ZZPhase_and_swap,
       "Equivalent to TK2, up to a wire swap that is encoded in the implicit "
       " qubit permutation of the Circuit, using up to 3 ZZPhase gates.");
+  library_m.def(
+      "_TK2_using_TK2_or_swap", &CircPool::TK2_using_TK2_or_swap,
+      "Either the exact TK2, or a wire swap encoded in the implicit qubit "
+      "permutation of the Circuit and single qubit gates.");
   library_m.def(
       "_approx_TK2_using_1xZZPhase", &CircPool::approx_TK2_using_1xZZPhase,
       "Approximate equivalent to TK2, using 1 ZZPhase gate and single-qubit "

pytket/pytket/passes/auto_rebase.py

@@ -37,14 +37,6 @@ def _TK2_using_TK2(a: Param, b: Param, c: Param) -> Circuit:
     return Circuit(2).TK2(a, b, c, 0, 1)
 
 
-def _TK2_using_TK2_or_swap(a: Param, b: Param, c: Param) -> Circuit:
-    if a == b == c == 0.5:
-        c = Circuit(2).SWAP(0, 1)
-        c.replace_SWAPs()
-        return c
-    return Circuit(2).TK2(a, b, c, 0, 1)
-
-
 _TK2_CIRCS: Dict[OpType, Callable[[Param, Param, Param], "Circuit"]] = {
     OpType.TK2: _TK2_using_TK2,
     OpType.ZZPhase: _library._TK2_using_ZZPhase,
@@ -53,8 +45,8 @@ def _TK2_using_TK2_or_swap(a: Param, b: Param, c: Param) -> Circuit:
 }
 
 _TK2_CIRCS_WIRE_SWAP: Dict[OpType, Callable[[Param, Param, Param], "Circuit"]] = {
-    OpType.TK2: _TK2_using_TK2_or_swap,
-    OpType.ZZPhase: _library._TK2_using_ZZPhase,
+    OpType.TK2: _library._TK2_using_TK2_or_swap,
+    OpType.ZZPhase: _library._TK2_using_ZZPhase_and_swap,
     OpType.CX: _library._TK2_using_CX_and_swap,
     OpType.ZZMax: _library._TK2_using_ZZMax_and_swap,
 }

pytket/tests/transform_test.py

@@ -1388,7 +1388,7 @@ def test_auto_rebase_with_swap_zzphase() -> None:
     c_no_swap = Circuit(2).ISWAP(0.3, 0, 1)
     no_swap_pass.apply(c_no_swap)
     assert c_no_swap.n_gates_of_type(OpType.ZZPhase) == 2
-    assert c_no_swap.n_gates == 13
+    assert c_no_swap.n_gates == 14
 
     c_swap = Circuit(2).ISWAPMax(0, 1).ISWAPMax(1, 0)
     swap_pass.apply(c_swap)
@@ -1412,9 +1412,9 @@ def test_auto_rebase_with_swap_zzphase() -> None:
     c_no_swap = Circuit(2).SWAP(0, 1)
     no_swap_pass.apply(c_no_swap)
     assert c_no_swap.n_gates_of_type(OpType.ZZPhase) == 3
-    assert c_no_swap.n_gates == 16
+    assert c_no_swap.n_gates == 13
 
-    c_swap = Circuit(2).ZZPhase(0, 1)
+    c_swap = Circuit(2).ZZPhase(0.4, 0, 1)
     swap_pass.apply(c_swap)
     assert c_swap.n_gates == 1
 
@@ -1428,6 +1428,7 @@ def test_auto_rebase_with_swap_tk2() -> None:
     no_swap_pass = auto_rebase_pass({OpType.TK2, OpType.PhasedX, OpType.Rz}, False)
     c_swap = Circuit(2).SWAP(0, 1)
     swap_pass.apply(c_swap)
+    print(c_swap.get_commands())
     assert c_swap.n_gates == 0
     c_no_swap = Circuit(2).SWAP(0, 1)
     no_swap_pass.apply(c_no_swap)

tket/include/tket/Circuit/CircPool.hpp

@@ -387,6 +387,15 @@ Circuit TK2_using_ZZPhase(
 Circuit TK2_using_ZZPhase_and_swap(
     const Expr &alpha, const Expr &beta, const Expr &gamma);
 
+/**
+ * @brief Either returns TK2(α, β, γ), or a wire swap and single qubit
+ * corrections
+ *
+ * @return Circuit Equivalent circuit, either a wire swap with single
+ * qubit corrections or TK2(α, β, γ).
+ */
+Circuit TK2_using_TK2_or_swap(
+    const Expr &alpha, const Expr &beta, const Expr &gamma);
 /**
  * @brief Equivalent to TK2(α, β, γ), using up to 3 ZZMax gates.
  *

tket/src/Circuit/CircPool.cpp

@@ -1149,6 +1149,17 @@ Circuit TK2_using_ZZPhase_and_swap(
   return TK2_using_ZZPhase(alpha, beta, gamma);
 }
 
+Circuit TK2_using_TK2_or_swap(
+    const Expr &alpha, const Expr &beta, const Expr &gamma) {
+  Circuit c = TK2_using_CX_and_swap(alpha, beta, gamma);
+  if (c.count_gates(OpType::CX) == 0) {
+    return c;
+  }
+  Circuit tk2(2);
+  tk2.add_op<unsigned>(OpType::TK2, {alpha, beta, gamma}, {0, 1});
+  return tk2;
+}
+
 Circuit TK2_using_ZZMax(
     const Expr &alpha, const Expr &beta, const Expr &gamma) {
   Circuit c = TK2_using_CX(alpha, beta, gamma);
@@ -1169,7 +1180,6 @@ Circuit TK2_using_ZZMax(
 Circuit TK2_using_ZZMax_and_swap(
     const Expr &alpha, const Expr &beta, const Expr &gamma) {
   Circuit c = TK2_using_CX_and_swap(alpha, beta, gamma);
-
   if (c.count_gates(OpType::CX) < 3) {
     // Find the CX gates and replace them with ZZMax.
     VertexSet bin;

tket/src/Transformations/Rebase.cpp

@@ -122,7 +122,6 @@ static bool standard_rebase_via_tk2(
         tk2_replacement) {
   bool success = false;
   VertexSet bin;
-  const bool TK2_allowed = allowed_gates.contains(OpType::TK2);
   // 1. Replace all multi-qubit gates outside the target gateset to TK2.
   for (const Vertex& v : circ.all_vertices()) {
     Op_ptr op = circ.get_Op_ptr_from_Vertex(v);
@@ -139,27 +138,24 @@ static bool standard_rebase_via_tk2(
       continue;
     // need to convert
     Circuit replacement = TK2_circ_from_multiq(op);
-    // If TK2 gates are not allowed in the target gateset we find a
-    // replacement circuit by decomposing them
-    if (!TK2_allowed) {
-      VertexSet TK2_bin;
-      for (const Vertex& u : replacement.all_vertices()) {
-        Op_ptr op = circ.get_Op_ptr_from_Vertex(u);
-        TKET_ASSERT(op->get_type() != OpType::Conditional);
-        if (op->get_type() == OpType::TK2) {
-          std::vector<Expr> params = op->get_params();
-          TKET_ASSERT(params.size() == 3);
-          Circuit u_replacement =
-              tk2_replacement(params[0], params[1], params[2]);
-          replacement.substitute(u_replacement, u, Circuit::VertexDeletion::No);
-          TK2_bin.insert(u);
-        }
+    // Find replacement Circuit for all TK2 gates
+    VertexSet TK2_bin;
+    for (const Vertex& u : replacement.all_vertices()) {
+      Op_ptr op = circ.get_Op_ptr_from_Vertex(u);
+      TKET_ASSERT(op->get_type() != OpType::Conditional);
+      if (op->get_type() == OpType::TK2) {
+        std::vector<Expr> params = op->get_params();
+        TKET_ASSERT(params.size() == 3);
+        Circuit u_replacement =
+            tk2_replacement(params[0], params[1], params[2]);
+        replacement.substitute(u_replacement, u, Circuit::VertexDeletion::No);
+        TK2_bin.insert(u);
       }
-      Transforms::squash_1qb_to_tk1().apply(replacement);
-      remove_redundancies().apply(replacement);
-      replacement.remove_vertices(
-          TK2_bin, Circuit::GraphRewiring::No, Circuit::VertexDeletion::Yes);
     }
+    Transforms::squash_1qb_to_tk1().apply(replacement);
+    remove_redundancies().apply(replacement);
+    replacement.remove_vertices(
+        TK2_bin, Circuit::GraphRewiring::No, Circuit::VertexDeletion::Yes);
     if (conditional) {
       circ.substitute_conditional(replacement, v, Circuit::VertexDeletion::No);
     } else {

tket/test/src/test_CompilerPass.cpp

@@ -1751,15 +1751,9 @@ SCENARIO("Custom rebase pass with implicit wire swaps.") {
   //   c.add_op<unsigned>(OpType::Sycamore, {0, 1});
   //   CompilationUnit cu(c);
   //   CHECK(pp_rebase_cx->apply(cu));
-  //   REQUIRE(cu.get_circ_ref().count_gates(OpType::CX) == 1);
+  //   REQUIRE(cu.get_circ_ref().count_gates(OpType::CX) ==2);
   //   auto u1 = tket_sim::get_unitary(c);
   //   auto u2 = tket_sim::get_unitary(cu.get_circ_ref());
-  //   CompilationUnit cu0(c);
-  //   gen_rebase_pass_via_tk2(
-  //     allowed_gates_cx, CircPool::TK2_using_CX,
-  //     CircPool::tk1_to_PhasedXRz)->apply(cu0);
-  //   auto u3 = tket_sim::get_unitary(cu.get_circ_ref());
-  //   REQUIRE(u1.isApprox(u3));
   //   REQUIRE(u1.isApprox(u2));
   // }
   GIVEN("Targeting CX gates, ISWAP gate.") {
@@ -1952,6 +1946,19 @@ SCENARIO("Custom rebase pass with implicit wire swaps.") {
     auto u2 = tket_sim::get_unitary(cu.get_circ_ref());
     REQUIRE(u1.isApprox(u2));
   }
+  GIVEN("Targeting TK2 gates, SWAP gate.") {
+    Circuit c(2);
+    c.add_op<unsigned>(OpType::SWAP, {0, 1});
+    CompilationUnit cu(c);
+    CHECK(gen_rebase_pass_via_tk2(
+              {OpType::PhasedX, OpType::Rz, OpType::TK2},
+              CircPool::TK2_using_TK2_or_swap, CircPool::tk1_to_PhasedXRz)
+              ->apply(cu));
+    REQUIRE(cu.get_circ_ref().n_gates() == 0);
+    auto u1 = tket_sim::get_unitary(c);
+    auto u2 = tket_sim::get_unitary(cu.get_circ_ref());
+    REQUIRE(u1.isApprox(u2));
+  }
 }
 }  // namespace test_CompilerPass
 }  // namespace tket