✨ Add Dead Gate Elimination Pattern

CHANGELOG.md

@@ -11,6 +11,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 
 ### Added
 
+- ✨ Add Dead Gate Elimination Pattern ([#1755]) ([**DRovara**])
 - 🚸 Add [CMake presets] to provide a standardized and reproducible way to configure builds ([#1660]) ([**@denialhaag**])
 - ✨ Add a `quantum-loop-unroll` pass for unrolling for-loop operations containing quantum operations ([#1718]) ([**@MatthiasReumann**])
 - ✨ Add a `hadamard-lifting` pass for lifting Hadamard gates above Pauli gates ([#1605]) ([**@lirem101**], [**@burgholzer**])
@@ -402,6 +403,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 
 <!-- PR links -->
 
+[#1755]: https://github.com/munich-quantum-toolkit/core/pull/1755
 [#1749]: https://github.com/munich-quantum-toolkit/core/pull/1749
 [#1748]: https://github.com/munich-quantum-toolkit/core/pull/1748
 [#1737]: https://github.com/munich-quantum-toolkit/core/pull/1737

mlir/include/mlir/Dialect/QCO/IR/QCODialect.td

@@ -38,6 +38,8 @@ def QCODialect : Dialect {
   let cppNamespace = "::mlir::qco";
 
   let useDefaultTypePrinterParser = 1;
+
+  let hasCanonicalizer = 1;
 }
 
 #endif // MLIR_DIALECT_QCO_IR_QCODIALECT_TD

mlir/include/mlir/Dialect/QCO/IR/QCOOps.td

@@ -136,6 +136,7 @@ def MeasureOp : QCOOp<"measure"> {
         }]>];
 
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 def ResetOp : QCOOp<"reset", [Idempotent, SameOperandsAndResultType]> {

mlir/include/mlir/Dialect/QCO/QCOUtils.h

@@ -237,4 +237,31 @@ mergeTwoTargetOneParameterWithSwappedTargets(OpType op,
   return success();
 }
 
+/**
+ * @brief Check if given quantum operation is unused (i.e., only used by
+ * deallocations) and remove it if so.
+ *
+ * @param op The operation to check.
+ * @param rewriter The pattern rewriter.
+ * @return LogicalResult Success or failure of the removal.
+ */
+inline LogicalResult checkAndRemoveDeadGate(Operation* op,
+                                            PatternRewriter& rewriter) {
+  if (std::all_of(op->getUsers().begin(), op->getUsers().end(),
+                  [](Operation* user) { return isa<SinkOp>(user); })) {
+    // If the operation is only used by deallocs, we can safely remove it.
+    if (auto u = dyn_cast<UnitaryOpInterface>(op)) {
+      // We specifically have to replace the output *qubits* with the input
+      // *qubits* to ignore parameters.
+      rewriter.replaceOp(op, u.getInputQubits());
+      return success();
+    } else {
+      // This includes the `IfOp` as well as `Reset` and `Measure`.
+      rewriter.replaceOp(op, op->getOperands());
+      return success();
+    }
+  }
+  return failure();
+}
+
 } // namespace mlir::qco

mlir/lib/Dialect/QCO/IR/Operations/MeasureOp.cpp

@@ -9,12 +9,29 @@
  */
 
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/QCOUtils.h"
 
 #include <mlir/Support/LogicalResult.h>
 
 using namespace mlir;
 using namespace mlir::qco;
 
+namespace {
+
+/**
+ * @brief Remove dead measurements.
+ */
+struct DeadMeasurementRemoval final : OpRewritePattern<MeasureOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(MeasureOp op,
+                                PatternRewriter& rewriter) const override {
+    return checkAndRemoveDeadGate(op, rewriter);
+  }
+};
+
+} // namespace
+
 LogicalResult MeasureOp::verify() {
   const auto registerName = getRegisterName();
   const auto registerSize = getRegisterSize();
@@ -37,3 +54,8 @@
   }
   return success();
 }
+
+void MeasureOp::getCanonicalizationPatterns(RewritePatternSet& results,
+                                            MLIRContext* context) {
+  results.add<DeadMeasurementRemoval>(context);
+}

mlir/lib/Dialect/QCO/IR/Operations/ResetOp.cpp

@@ -9,6 +9,7 @@
  */
 
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/QCOUtils.h"
 #include "mlir/Dialect/QTensor/IR/QTensorOps.h"
 #include "mlir/Dialect/QTensor/IR/QTensorUtils.h"
 
@@ -101,6 +102,18 @@ struct RemoveResetAfterExtract final : OpRewritePattern<ResetOp> {
   }
 };
 
+/**
+ * @brief Remove dead resets.
+ */
+struct DeadResetRemoval final : OpRewritePattern<ResetOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ResetOp op,
+                                PatternRewriter& rewriter) const override {
+    return checkAndRemoveDeadGate(op, rewriter);
+  }
+};
+
 } // namespace
 
 OpFoldResult ResetOp::fold(FoldAdaptor /*adaptor*/) {
@@ -114,4 +127,5 @@ OpFoldResult ResetOp::fold(FoldAdaptor /*adaptor*/) {
 void ResetOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                           MLIRContext* context) {
   results.add<RemoveResetAfterExtract>(context);
+  results.add<DeadResetRemoval>(context);
 }

mlir/lib/Dialect/QCO/IR/QCOOps.cpp

@@ -11,6 +11,7 @@
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
 
 #include "mlir/Dialect/QCO/IR/QCODialect.h" // IWYU pragma: associated
+#include "mlir/Dialect/QCO/QCOUtils.h"
 
 #include <llvm/ADT/STLExtras.h>
 #include <mlir/IR/Block.h>
@@ -30,6 +31,29 @@
 using namespace mlir;
 using namespace mlir::qco;
 
+//===----------------------------------------------------------------------===//
+// Dialect-Level Canonicalizers
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct DeadGateElimination final
+    : public OpInterfaceRewritePattern<UnitaryOpInterface> {
+
+  explicit DeadGateElimination(MLIRContext* context)
+      : OpInterfaceRewritePattern(context) {}
+
+  LogicalResult matchAndRewrite(UnitaryOpInterface op,
+                                PatternRewriter& rewriter) const override {
+    if (op->use_empty()) {
+      // This effectively ignores the GPhase operation and variants such as its
+      // inverse, which should never be considered dead.
+      return failure();
+    }
+    return checkAndRemoveDeadGate(op.getOperation(), rewriter);
+  }
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // Custom Parsers
 //===----------------------------------------------------------------------===//
@@ -258,6 +282,10 @@
       >();
 }
 
+void QCODialect::getCanonicalizationPatterns(RewritePatternSet& results) const {
+  results.add<DeadGateElimination>(getContext());
+}
+
 //===----------------------------------------------------------------------===//
 // Types
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp

@@ -9,6 +9,7 @@
  */
 
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/QCOUtils.h"
 
 #include <llvm/ADT/STLExtras.h>
 #include <llvm/ADT/STLFunctionalExtras.h>
@@ -234,11 +235,24 @@ struct ConditionPropagation : public OpRewritePattern<IfOp> {
     return success(changed);
   }
 };
+
+/**
+ * @brief Remove dead resets.
+ */
+struct DeadIfRemoval final : OpRewritePattern<IfOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IfOp op,
+                                PatternRewriter& rewriter) const override {
+    return checkAndRemoveDeadGate(op, rewriter);
+  }
+};
 } // namespace
 
 void IfOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                        MLIRContext* context) {
   results.add<RemoveStaticCondition, ConditionPropagation>(context);
+  results.add<DeadIfRemoval>(context);
   populateRegionBranchOpInterfaceCanonicalizationPatterns(
       results, IfOp::getOperationName());
 }

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

@@ -114,6 +114,39 @@ TEST_F(QCOTest, BuilderRejectsMixedStaticAndDynamicQubitAllocationModes) {
       "Cannot mix dynamic and static qubit allocation modes");
 }
 
+TEST_F(QCOTest, CheckDeadGateElimination) {
+  QCOProgramBuilder builder(context.get());
+  builder.initialize();
+  auto q0_0 = builder.allocQubit();
+  auto q1_0 = builder.allocQubit();
+  auto q0_1 = builder.h(q0_0);
+  auto [q0_2, q1_1] = builder.cx(q0_1, q1_0);
+  auto q1_2 = builder.h(q1_1);
+  builder.sink(q0_2);
+  builder.sink(q1_2);
+  auto module = builder.finalize();
+
+  QCOProgramBuilder reference(context.get());
+  reference.initialize();
+  auto r0 = reference.allocQubit();
+  auto r1 = reference.allocQubit();
+  reference.sink(r0);
+  reference.sink(r1);
+  auto ref = reference.finalize();
+
+  ASSERT_TRUE(module);
+  EXPECT_TRUE(verify(*module).succeeded());
+  EXPECT_TRUE(runQCOCleanupPipeline(module.get()).succeeded());
+  EXPECT_TRUE(verify(*module).succeeded());
+
+  ASSERT_TRUE(ref);
+  EXPECT_TRUE(verify(*ref).succeeded());
+  EXPECT_TRUE(runQCOCleanupPipeline(ref.get()).succeeded());
+  EXPECT_TRUE(verify(*ref).succeeded());
+
+  EXPECT_TRUE(areModulesEquivalentWithPermutations(module.get(), ref.get()));
+}
+
 TEST_F(QCOTest, DirectIfBuilder) {
   // Test If construction directly
   QCOProgramBuilder builder(context.get());