✨ Add SWAPAbsorption Pass

mlir/include/mlir/Dialect/QCO/Transforms/Passes.td

@@ -101,4 +101,10 @@ def MappingPass : Pass<"place-and-route", "mlir::ModuleOp"> {
                               "The number of inserted SWAPs">];
 }
 
+def SwapAbsorptionPass : Pass<"absorb-swaps", "mlir::ModuleOp"> {
+  let dependentDialects = ["mlir::qco::QCODialect"];
+  let summary = "";
+  let options = [];
+}
+
 #endif // MLIR_DIALECT_QCO_TRANSFORMS_PASSES_TD

mlir/lib/Dialect/QCO/Transforms/Optimizations/SwapAbsorption.cpp

@@ -0,0 +1,82 @@
+/*
+ * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
+ * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH
+ * All rights reserved.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Licensed under the MIT License
+ */
+
+#include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/Transforms/Passes.h"
+#include "mlir/Dialect/QCO/Utils/Drivers.h"
+#include "mlir/Dialect/QCO/Utils/WireIterator.h"
+#include "mlir/Dialect/QTensor/IR/QTensorOps.h"
+
+#include <llvm/ADT/STLExtras.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/PatternMatch.h>
+#include <mlir/IR/Visitors.h>
+#include <mlir/Support/LLVM.h>
+#include <mlir/Support/WalkResult.h>
+
+#include <cstddef>
+#include <utility>
+
+namespace mlir::qco {
+#define GEN_PASS_DEF_SWAPABSORPTIONPASS
+#include "mlir/Dialect/QCO/Transforms/Passes.h.inc"
+
+namespace {
+struct SwapAbsorption : impl::SwapAbsorptionPassBase<SwapAbsorption> {
+public:
+  using SwapAbsorptionPassBase::SwapAbsorptionPassBase;
+
+protected:
+  void runOnOperation() override {
+    ModuleOp anchor = getOperation();
+    IRRewriter rewriter(&getContext());
+
+    for (auto func : anchor.getOps<func::FuncOp>()) {
+      SmallVector<WireIterator> wires;
+      for (auto op : func.getOps<StaticOp>()) {
+        wires.emplace_back(op.getQubit());
+      }
+
+      SmallVector<SWAPOp> readyToAbsorb;
+      readyToAbsorb.reserve((wires.size() + 1) / 2);
+
+      std::ignore = walkProgramGraph<WireDirection::Forward>(
+          wires, [&](const ReadyRange& ready, ReleasedOps& released) {
+            for (const auto& [op, indices] : ready) {
+              if (isa<SWAPOp>(op)) {
+                readyToAbsorb.emplace_back(op);
+              }
+              released.emplace_back(op);
+            }
+            return WalkResult::interrupt();
+          });
+
+      for (auto swapOp : readyToAbsorb) {
+        auto in0 = swapOp.getQubit0In();
+        auto in1 = swapOp.getQubit1In();
+
+        auto out0 = swapOp.getQubit0Out();
+        auto out1 = swapOp.getQubit1Out();
+
+        // TODO: What if single qubit gates are in front of the SWAP input?
+        // Tipp: Use the WireIterator.
+        StaticOp op0 = cast<StaticOp>(in0.getDefiningOp());
+        StaticOp op1 = cast<StaticOp>(in1.getDefiningOp());
+
+        rewriter.replaceAllUsesWith(out0, op1.getQubit());
+        rewriter.replaceAllUsesWith(out1, op0.getQubit());
+        rewriter.eraseOp(swapOp);
+      }
+    }
+  }
+};
+} // namespace
+} // namespace mlir::qco

mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt

@@ -7,7 +7,7 @@
 # Licensed under the MIT License
 
 set(target_name mqt-core-mlir-unittest-optimizations)
-add_executable(${target_name} test_qco_merge_single_qubit_rotation.cpp)
+add_executable(${target_name} test_qco_merge_single_qubit_rotation.cpp test_swapabsorption.cpp)
 
 target_link_libraries(
   ${target_name}

mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_swapabsorption.cpp

@@ -0,0 +1,104 @@
+/*
+ * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
+ * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH
+ * All rights reserved.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Licensed under the MIT License
+ */
+
+#include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h"
+#include "mlir/Dialect/QCO/IR/QCODialect.h"
+#include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/Transforms/Passes.h"
+
+#include <gtest/gtest.h>
+#include <llvm/Support/LogicalResult.h>
+#include <mlir/Dialect/Arith/IR/Arith.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/DialectRegistry.h>
+#include <mlir/IR/OperationSupport.h>
+#include <mlir/IR/OwningOpRef.h>
+#include <mlir/IR/Value.h>
+#include <mlir/Pass/PassManager.h>
+#include <mlir/Support/LLVM.h>
+
+#include <cassert>
+#include <memory>
+
+using namespace mlir;
+using namespace mlir::qco;
+
+namespace {
+
+class SwapAbsorbPassTest : public testing::Test {
+
+protected:
+  void SetUp() override {
+    // Register all necessary dialects
+    DialectRegistry registry;
+    registry.insert<qco::QCODialect, arith::ArithDialect, func::FuncDialect>();
+    context = std::make_unique<MLIRContext>();
+    context->appendDialectRegistry(registry);
+    context->loadAllAvailableDialects();
+  }
+
+  static void applySwapAbsorb(OwningOpRef<ModuleOp>& moduleOp) {
+    PassManager pm(moduleOp->getContext());
+    pm.addPass(qco::createSwapAbsorptionPass());
+    auto res = pm.run(*moduleOp);
+
+    ASSERT_TRUE(succeeded(res));
+  }
+
+  std::unique_ptr<MLIRContext> context;
+};
+}; // namespace
+
+TEST_F(SwapAbsorbPassTest, PassDoesNotChangeSwaplessProgram) {
+
+  qco::QCOProgramBuilder builder(context.get());
+  builder.initialize();
+
+  const auto q00 = builder.staticQubit(0);
+  const auto q10 = builder.staticQubit(1);
+
+  const auto q01 = builder.h(q00);
+  const auto [q02, q11] = builder.cx(q01, q10);
+
+  builder.sink(q02);
+  builder.sink(q11);
+
+  auto moduleThroughPass = builder.finalize();
+  auto originalModule = moduleThroughPass->clone();
+
+  applySwapAbsorb(moduleThroughPass);
+  ASSERT_TRUE(mlir::OperationEquivalence::isEquivalentTo(
+      moduleThroughPass.get(), originalModule,
+      mlir::OperationEquivalence::Flags::IgnoreLocations));
+}
+
+TEST_F(SwapAbsorbPassTest, PassReordersTwoQubitCircuitWithLeadingSwap) {
+
+  qco::QCOProgramBuilder builder(context.get());
+  builder.initialize();
+
+  const auto q00 = builder.staticQubit(0);
+  const auto q10 = builder.staticQubit(1);
+
+  const auto [q01, q11] = builder.swap(q00, q10);
+
+  const auto q02 = builder.id(q01);
+  const auto q12 = builder.id(q11);
+
+  builder.sink(q02);
+  builder.sink(q12);
+
+  auto moduleThroughPass = builder.finalize();
+  applySwapAbsorb(moduleThroughPass);
+
+  ASSERT_EQ(q10, ((IdOp)q02.getDefiningOp()).getInputQubit(0));
+  ASSERT_EQ(q00, ((IdOp)q12.getDefiningOp()).getInputQubit(0));
+}