From f7b7ff06f8a9b30996006320086806d8248f4d0b Mon Sep 17 00:00:00 2001
From: "Daniel J. Egger" <38065505+eggerdj@users.noreply.github.com>
Date: Thu, 10 Oct 2024 19:30:29 +0200
Subject: [PATCH] SWAP removal (#36)

---
 .../transpilation/preset_qaoa_passmanager.py  |  2 +
 .../transpilation/swap_cancellation_pass.py   | 44 +++++++++++++++++++
 test/test_qaoa_construction.py                | 41 ++++++++++++++++-
 3 files changed, 86 insertions(+), 1 deletion(-)
 create mode 100644 qopt_best_practices/transpilation/swap_cancellation_pass.py

diff --git a/qopt_best_practices/transpilation/preset_qaoa_passmanager.py b/qopt_best_practices/transpilation/preset_qaoa_passmanager.py
index fd8b6d7..279e198 100644
--- a/qopt_best_practices/transpilation/preset_qaoa_passmanager.py
+++ b/qopt_best_practices/transpilation/preset_qaoa_passmanager.py
@@ -11,6 +11,7 @@
 from qiskit.circuit.library import CXGate
 
 from qopt_best_practices.transpilation.qaoa_construction_pass import QAOAConstructionPass
+from qopt_best_practices.transpilation.swap_cancellation_pass import SwapToFinalMapping
 
 
 def qaoa_swap_strategy_pm(config: Dict[str, Any]):
@@ -39,6 +40,7 @@ def qaoa_swap_strategy_pm(config: Dict[str, Any]):
                 swap_strategy,
                 edge_coloring,
             ),
+            SwapToFinalMapping(),
             HighLevelSynthesis(basis_gates=basis_gates),
             InverseCancellation(gates_to_cancel=[CXGate()]),
             QAOAConstructionPass(num_layers),
diff --git a/qopt_best_practices/transpilation/swap_cancellation_pass.py b/qopt_best_practices/transpilation/swap_cancellation_pass.py
new file mode 100644
index 0000000..431a4bf
--- /dev/null
+++ b/qopt_best_practices/transpilation/swap_cancellation_pass.py
@@ -0,0 +1,44 @@
+"""Pass to remove SWAP gates that are not needed."""
+
+from qiskit.dagcircuit import DAGOutNode, DAGCircuit
+from qiskit.transpiler import TransformationPass
+
+
+class SwapToFinalMapping(TransformationPass):
+    """Absorb any redundent SWAPs in the final layout.
+
+    This pass should be executed after a SWAPStrategy has been applied to a block
+    of commuting gates. It will remove any final redundent SWAP gates and absorb
+    them into the virtual layout. This effectively undoes any possibly redundent
+    SWAP gates that the SWAPStrategy may have inserted.
+    """
+
+    def run(self, dag: DAGCircuit):
+        """run the pass."""
+
+        qmap = self.property_set["virtual_permutation_layout"]
+
+        qreg = dag.qregs[next(iter(dag.qregs))]
+
+        # This will remove SWAP gates that are applied before anything else
+        # This remove is executed multiple times until there are no more SWAP
+        # gates left to remove. Note: a more inteligent DAG traversal could
+        # be implemented here.
+
+        done = False
+
+        while not done:
+            permuted = False
+            for node in dag.topological_op_nodes():
+                if node.op.name == "swap":
+                    successors = list(dag.successors(node))
+                    if len(successors) == 2:
+                        if all(isinstance(successors[idx], DAGOutNode) for idx in [0, 1]):
+                            bits = [qreg.index(qubit) for qubit in node.qargs]
+                            qmap[bits[0]], qmap[bits[1]] = qmap[bits[1]], qmap[bits[0]]
+                            dag.remove_op_node(node)
+                            permuted = True
+
+            done = not permuted
+
+        return dag
diff --git a/test/test_qaoa_construction.py b/test/test_qaoa_construction.py
index e5ab4f1..d0bd085 100644
--- a/test/test_qaoa_construction.py
+++ b/test/test_qaoa_construction.py
@@ -9,10 +9,16 @@
 from qiskit.primitives import StatevectorEstimator
 from qiskit.quantum_info import SparsePauliOp
 from qiskit.transpiler import PassManager
-from qiskit.transpiler.passes.routing.commuting_2q_gate_routing import SwapStrategy
+from qiskit.transpiler.passes import HighLevelSynthesis
+from qiskit.transpiler.passes.routing.commuting_2q_gate_routing import (
+    SwapStrategy,
+    FindCommutingPauliEvolutions,
+    Commuting2qGateRouter,
+)
 
 from qopt_best_practices.transpilation.qaoa_construction_pass import QAOAConstructionPass
 from qopt_best_practices.transpilation.preset_qaoa_passmanager import qaoa_swap_strategy_pm
+from qopt_best_practices.transpilation.swap_cancellation_pass import SwapToFinalMapping
 
 
 class TestQAOAConstruction(TestCase):
@@ -95,3 +101,36 @@ def test_depth_two_qaoa_pass(self):
         )
 
         self.assertAlmostEqual(value, expected)
+
+    def test_swap_construction(self):
+        """Test that redundent SWAP gates are removed."""
+        cost_op = SparsePauliOp.from_list(
+            [("IIIIZZ", 1), ("IIZZII", 1), ("ZZIIII", 1), ("IIZIIZ", 1)],
+        )
+
+        ansatz = QAOAAnsatz(
+            cost_op, reps=1, initial_state=QuantumCircuit(6), mixer_operator=QuantumCircuit(6)
+        )
+
+        # Test with the SWAP removal
+        qaoa_pm = PassManager(
+            [
+                HighLevelSynthesis(basis_gates=["PauliEvolution"]),
+                FindCommutingPauliEvolutions(),
+                Commuting2qGateRouter(SwapStrategy.from_line(range(6))),
+                SwapToFinalMapping(),
+            ]
+        )
+
+        self.assertEqual(qaoa_pm.run(ansatz).count_ops()["swap"], 2)
+
+        # Test without the SWAP removal
+        qaoa_pm = PassManager(
+            [
+                HighLevelSynthesis(basis_gates=["PauliEvolution"]),
+                FindCommutingPauliEvolutions(),
+                Commuting2qGateRouter(SwapStrategy.from_line(range(6))),
+            ]
+        )
+
+        self.assertEqual(qaoa_pm.run(ansatz).count_ops()["swap"], 3)