SWAP removal (#36)

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),

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

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)