Add RemoveSubsequentEquivalentCtrl Canonicalization Pattern

mlir/lib/Dialect/QCO/IR/Modifiers/CtrlOp.cpp

@@ -9,8 +9,10 @@
  */
 
 #include "mlir/Dialect/QCO/IR/QCODialect.h"
+#include "mlir/Dialect/QCO/IR/QCOInterfaces.h"
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
 
+#include <llvm/ADT/DenseSet.h>
 #include <llvm/ADT/STLExtras.h>
 #include <llvm/ADT/STLFunctionalExtras.h>
 #include <llvm/ADT/SmallVector.h>
@@ -34,6 +36,73 @@ using namespace mlir::qco;
 
 namespace {
 
+/**
+ * @brief Remove subsequent equivalent ctrl operation on the same qubits.
+ */
+struct RemoveSubsequentEquivalentCtrl final : OpRewritePattern<CtrlOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CtrlOp op,
+                                PatternRewriter& rewriter) const override {
+    assert(!op.use_empty() && "linear typing assumption violated!");
+
+    // All uses point to the same operation.
+    if (!llvm::all_equal(op->getUsers())) {
+      return failure();
+    }
+
+    // The CtrlOp has a removable body unitary.
+    UnitaryOpInterface bodyU = op.getBodyUnitary();
+    if (!isa<XOp, YOp, ZOp, HOp, SWAPOp>(bodyU)) {
+      return failure();
+    }
+
+    // The next operation is also a CtrlOp.
+    auto nextOp = dyn_cast<CtrlOp>(*(op->getUsers().begin()));
+    if (!nextOp) {
+      return failure();
+    }
+
+    // The next operation applies the equivalent body unitary.
+    if (nextOp.getBodyUnitary()->getName() != bodyU->getName()) {
+      return failure();
+    }
+
+    // Both operations have the same number of controls and targets.
+    if (op.getNumQubits() != nextOp.getNumQubits() ||
+        op.getNumControls() != nextOp.getNumControls()) {
+      return failure();
+    }
+
+    // Make sure that there is a bijective mapping between output (op) and input
+    // (nextOp) control qubits.
+    llvm::SmallDenseSet<Value> outs(op.getControlsOut().begin(),
+                                    op.getControlsOut().end());
+    for (const auto& in : nextOp.getControlsIn()) {
+      outs.erase(in);
+    }
+
+    if (!outs.empty()) {
+      return failure();
+    }
+
+    // Both operations have the same order for the target qubits:
+    // The i-th output is the i-th input of the next op.
+    for (const auto& [opOut, nextOpIn] :
+         llvm::zip_equal(op.getTargetsOut(), nextOp.getTargetsIn())) {
+      if (opOut != nextOpIn) {
+        return failure();
+      }
+    }
+
+    // Remove both ops and rewire accordingly.
+    rewriter.replaceOp(nextOp, op.getInputQubits());
+    rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+
 /**
  * @brief Merge nested control modifiers into a single one.
  */
@@ -337,7 +406,8 @@ LogicalResult CtrlOp::verify() {
 
 void CtrlOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                          MLIRContext* context) {
-  results.add<MergeNestedCtrl, ReduceCtrl>(context);
+  results.add<RemoveSubsequentEquivalentCtrl, MergeNestedCtrl, ReduceCtrl>(
+      context);
 }
 
 std::optional<Eigen::MatrixXcd> CtrlOp::getUnitaryMatrix() {

mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp

@@ -175,19 +175,40 @@ INSTANTIATE_TEST_SUITE_P(
 /// @{
 INSTANTIATE_TEST_SUITE_P(
     QCOCtrlOpTest, QCOTest,
-    testing::Values(QCOTestCase{"TrivialCtrl", MQT_NAMED_BUILDER(trivialCtrl),
-                                MQT_NAMED_BUILDER(rxx)},
-                    QCOTestCase{"NestedCtrl", MQT_NAMED_BUILDER(nestedCtrl),
-                                MQT_NAMED_BUILDER(multipleControlledRxx)},
-                    QCOTestCase{"TripleNestedCtrl",
-                                MQT_NAMED_BUILDER(tripleNestedCtrl),
-                                MQT_NAMED_BUILDER(tripleControlledRxx)},
-                    QCOTestCase{"CtrlInvSandwich",
-                                MQT_NAMED_BUILDER(ctrlInvSandwich),
-                                MQT_NAMED_BUILDER(multipleControlledRxx)},
-                    QCOTestCase{"DoubleNestedCtrlTwoQubits",
-                                MQT_NAMED_BUILDER(doubleNestedCtrlTwoQubits),
-                                MQT_NAMED_BUILDER(fourControlledRxx)}));
+    testing::Values(
+        QCOTestCase{"TrivialCtrl", MQT_NAMED_BUILDER(trivialCtrl),
+                    MQT_NAMED_BUILDER(rxx)},
+        QCOTestCase{"NestedCtrl", MQT_NAMED_BUILDER(nestedCtrl),
+                    MQT_NAMED_BUILDER(multipleControlledRxx)},
+        QCOTestCase{"TripleNestedCtrl", MQT_NAMED_BUILDER(tripleNestedCtrl),
+                    MQT_NAMED_BUILDER(tripleControlledRxx)},
+        QCOTestCase{"CtrlInvSandwich", MQT_NAMED_BUILDER(ctrlInvSandwich),
+                    MQT_NAMED_BUILDER(multipleControlledRxx)},
+        QCOTestCase{"DoubleNestedCtrlTwoQubits",
+                    MQT_NAMED_BUILDER(doubleNestedCtrlTwoQubits),
+                    MQT_NAMED_BUILDER(fourControlledRxx)},
+        QCOTestCase{"EvenConsecutiveCtrlsOneTarget",
+                    MQT_NAMED_BUILDER(evenConsecutiveCtrlsOneTarget),
+                    MQT_NAMED_BUILDER(emptyQCO)},
+        QCOTestCase{"OddConsecutiveCtrlsOneTarget",
+                    MQT_NAMED_BUILDER(oddConsecutiveCtrlsOneTarget),
+                    MQT_NAMED_BUILDER(singleControlledX)},
+        QCOTestCase{"EvenConsecutiveCtrlsTwoTarget",
+                    MQT_NAMED_BUILDER(evenConsecutiveCtrlsTwoTargets),
+                    MQT_NAMED_BUILDER(emptyQCO)},
+        QCOTestCase{"OddConsecutiveCtrlsTwoTarget",
+                    MQT_NAMED_BUILDER(oddConsecutiveCtrlsTwoTargets),
+                    MQT_NAMED_BUILDER(singleControlledSwap)},
+        QCOTestCase{"ConsecutiveCtrlsInequivalentBody",
+                    MQT_NAMED_BUILDER(consecutiveCtrlsInequivalentBody),
+                    MQT_NAMED_BUILDER(consecutiveCtrlsInequivalentBody)},
+        QCOTestCase{"ConsecutiveCtrlsInequivalentControls",
+                    MQT_NAMED_BUILDER(consecutiveCtrlsInequivalentControlCnt),
+                    MQT_NAMED_BUILDER(consecutiveCtrlsInequivalentControlCnt)},
+        QCOTestCase{
+            "ConsecutiveCtrlsInequivalentTargetOrder",
+            MQT_NAMED_BUILDER(consecutiveCtrlsInequivalentTargetOrder),
+            MQT_NAMED_BUILDER(consecutiveCtrlsInequivalentTargetOrder)}));
 /// @}
 
 /// \name QCO/Modifiers/InvOp.cpp

mlir/unittests/programs/qco_programs.cpp

@@ -1994,6 +1994,67 @@ void ctrlInvSandwich(QCOProgramBuilder& b) {
   });
 }
 
+void evenConsecutiveCtrlsOneTarget(QCOProgramBuilder& b) {
+  auto q = b.allocQubitRegister(2);
+  const auto& [q01, q11] = b.cx(q[0], q[1]);
+  std::ignore = b.cx(q01, q11);
+}
+
+void oddConsecutiveCtrlsOneTarget(QCOProgramBuilder& b) {
+  auto q = b.allocQubitRegister(2);
+  const auto& [q01, q11] = b.cx(q[0], q[1]);
+  const auto& [q02, q12] = b.cx(q01, q11);
+  std::ignore = b.cx(q02, q12);
+}
+
+void evenConsecutiveCtrlsTwoTargets(QCOProgramBuilder& b) {
+  auto q = b.allocQubitRegister(3);
+  const auto out0 = b.cswap(q[0], q[1], q[2]);
+  const auto q01 = out0.first;
+  const auto [q11, q21] = out0.second;
+
+  std::ignore = b.cswap(q01, q11, q21);
+}
+
+void oddConsecutiveCtrlsTwoTargets(QCOProgramBuilder& b) {
+  auto q = b.allocQubitRegister(3);
+  const auto out0 = b.cswap(q[0], q[1], q[2]);
+  const auto q01 = out0.first;
+  const auto [q11, q21] = out0.second;
+
+  const auto out1 = b.cswap(q01, q11, q21);
+  const auto q02 = out1.first;
+  const auto [q12, q22] = out1.second;
+
+  std::ignore = b.cswap(q02, q12, q22);
+}
+
+void consecutiveCtrlsInequivalentBody(QCOProgramBuilder& b) {
+  auto q = b.allocQubitRegister(2);
+  const auto& [q01, q11] = b.cx(q[0], q[1]);
+  std::ignore = b.cy(q01, q11);
+}
+
+void consecutiveCtrlsInequivalentControlCnt(QCOProgramBuilder& b) {
+  auto q = b.allocQubitRegister(3);
+  const auto& [q01, q11] = b.cx(q[0], q[1]);
+  std::ignore = b.ctrl({q01, q11}, {q[2]}, [&](ValueRange targets) {
+    const auto q21 = b.x(targets[0]);
+    return SmallVector{q21};
+  });
+}
+
+void consecutiveCtrlsInequivalentTargetOrder(QCOProgramBuilder& b) {
+  auto q = b.allocQubitRegister(3);
+  const auto out0 = b.cswap(q[0], q[1], q[2]);
+  const auto q01 = out0.first;
+  const auto [q11, q21] = out0.second;
+
+  const auto out1 = b.cswap(q11, q01, q21);
+  const auto q02 = out1.first;
+  const auto [q12, q22] = out1.second;
+}
+
 void nestedInv(QCOProgramBuilder& b) {
   auto q = b.allocQubitRegister(2);
   b.inv({q[0], q[1]}, [&](ValueRange qubits) {

mlir/unittests/programs/qco_programs.h

@@ -969,6 +969,34 @@ void doubleNestedCtrlTwoQubits(QCOProgramBuilder& b);
 /// Creates a circuit with control modifiers interleaved by an inverse modifier.
 void ctrlInvSandwich(QCOProgramBuilder& b);
 
+/// Creates a circuit with two consecutive equivalent ctrl modifiers with one
+/// target qubit.
+void evenConsecutiveCtrlsOneTarget(QCOProgramBuilder& b);
+
+/// Creates a circuit with three consecutive equivalent ctrl modifiers with one
+/// target qubit.
+void oddConsecutiveCtrlsOneTarget(QCOProgramBuilder& b);
+
+/// Creates a circuit with two consecutive equivalent ctrl modifiers with two
+/// target qubits.
+void evenConsecutiveCtrlsTwoTargets(QCOProgramBuilder& b);
+
+/// Creates a circuit with three consecutive equivalent ctrl modifiers with two
+/// target qubits.
+void oddConsecutiveCtrlsTwoTargets(QCOProgramBuilder& b);
+
+/// Creates a circuit with two consecutive ctrl modifiers with inequivalent body
+/// unitaries.
+void consecutiveCtrlsInequivalentBody(QCOProgramBuilder& b);
+
+/// Creates a circuit with two consecutive ctrl modifiers with an inequivalent
+/// number of controls.
+void consecutiveCtrlsInequivalentControlCnt(QCOProgramBuilder& b);
+
+/// Creates a circuit with two consecutive ctrl modifiers with an inequivalent
+/// target qubit order.
+void consecutiveCtrlsInequivalentTargetOrder(QCOProgramBuilder& b);
+
 // --- InvOp ---------------------------------------------------------------- //
 
 /// Creates a circuit with nested inverse modifiers.