✨ 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/QCOUtils.h

@@ -237,4 +237,42 @@ mergeTwoTargetOneParameterWithSwappedTargets(OpType op,
   return success();
 }
 
+/**
+ * @brief Check if given quantum operation is unused (i.e., only used by
+ * sinks) 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 sinks, 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 if (auto m = dyn_cast<MeasureOp>(op)) {
+      // We specifically have to replace the output *qubits* with the input
+      // *qubits* to ignore the classical outcome.
+      rewriter.replaceAllUsesWith(m.getQubitOut(), m.getQubitIn());
+      rewriter.eraseOp(op);
+      return success();
+    } else if (auto i = dyn_cast<IfOp>(op)) {
+      // We specifically have to replace the output *qubits* with the input
+      // *qubits* to ignore the condition.
+      rewriter.replaceOp(op, i.getQubits());
+      return success();
+    } else {
+      // This currently only includes the `Reset` operation.
+      rewriter.replaceOp(op, op->getOperands());
+      return success();
+    }
+  }
+  return failure();
+}
+
 } // namespace mlir::qco

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

@@ -9,12 +9,31 @@
  */
 
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/QCOUtils.h"
 
+#include <mlir/IR/MLIRContext.h>
+#include <mlir/IR/PatternMatch.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 +56,8 @@ LogicalResult MeasureOp::verify() {
   }
   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,14 +11,19 @@
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
 
 #include "mlir/Dialect/QCO/IR/QCODialect.h" // IWYU pragma: associated
+#include "mlir/Dialect/QCO/IR/QCOInterfaces.h"
+#include "mlir/Dialect/QCO/QCOUtils.h"
 
 #include <llvm/ADT/STLExtras.h>
 #include <mlir/IR/Block.h>
+#include <mlir/IR/MLIRContext.h>
 #include <mlir/IR/OpImplementation.h>
 #include <mlir/IR/Operation.h>
 #include <mlir/IR/OperationSupport.h>
+#include <mlir/IR/PatternMatch.h>
 #include <mlir/IR/Region.h>
 #include <mlir/IR/ValueRange.h>
+#include <mlir/Interfaces/ShapedOpInterfaces.h>
 #include <mlir/Support/LLVM.h>
 
 // The following headers are needed for some template instantiations.
@@ -30,6 +35,34 @@
 using namespace mlir;
 using namespace mlir::qco;
 
+//===----------------------------------------------------------------------===//
+// Dialect-Level Canonicalizers
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/**
+ * @brief Remove dead measurements.
+ */
+struct DeadGateElimination final
+    : public OpInterfaceRewritePattern<UnitaryOpInterface> {
+
+  explicit DeadGateElimination(MLIRContext* context)
+      : OpInterfaceRewritePattern(context) {}
+
+  LogicalResult matchAndRewrite(UnitaryOpInterface op,
+                                PatternRewriter& rewriter) const override {
+    if (!isMemoryEffectFree(op)) {
+      // This effectively ignores the GPhase operation and variants such as its
+      // inverse or `if` ops containing it, which should never be considered
+      // dead.
+      return failure();
+    }
+    return checkAndRemoveDeadGate(op.getOperation(), rewriter);
+  }
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // Custom Parsers
 //===----------------------------------------------------------------------===//
@@ -258,6 +291,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 `IfOp` instructions.
+ */
+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,114 @@ TEST_F(QCOTest, BuilderRejectsMixedStaticAndDynamicQubitAllocationModes) {
       "Cannot mix dynamic and static qubit allocation modes");
 }
 
+TEST_F(QCOTest, CheckDeadGateElimination) {
+  QCOProgramBuilder builder(context.get());
+  builder.initialize();
+  auto q0S0 = builder.allocQubit();
+  auto q1S0 = builder.allocQubit();
+  auto q0S1 = builder.h(q0S0);
+  auto [q0S2, q1S1] = builder.cx(q0S1, q1S0);
+  auto [q1S2, c1] = builder.measure(q1S1);
+  builder.sink(q0S2);
+  builder.sink(q1S2);
+  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 refModule = reference.finalize();
+
+  ASSERT_TRUE(module);
+  EXPECT_TRUE(verify(*module).succeeded());
+  EXPECT_TRUE(runQCOCleanupPipeline(module.get()).succeeded());
+  EXPECT_TRUE(verify(*module).succeeded());
+
+  ASSERT_TRUE(refModule);
+  EXPECT_TRUE(verify(*refModule).succeeded());
+  EXPECT_TRUE(runQCOCleanupPipeline(refModule.get()).succeeded());
+  EXPECT_TRUE(verify(*refModule).succeeded());
+
+  EXPECT_TRUE(
+      areModulesEquivalentWithPermutations(module.get(), refModule.get()));
+}
+
+TEST_F(QCOTest, CheckIfOpDeadGateElimination) {
+  QCOProgramBuilder builder(context.get());
+  builder.initialize();
+  auto q0S0 = builder.allocQubit();
+  auto q1S0 = builder.allocQubit();
+  auto q0S1 = builder.h(q0S0);
+  auto [q0S2, c0] = builder.measure(q0S1);
+
+  // This is an `if` with memory effects - it can't be removed.
+  auto q1S1 = builder.qcoIf(
+      c0, {q1S0},
+      [&](ValueRange qubits) -> SmallVector<Value> {
+        auto q1Then = builder.x(qubits[0]);
+        builder.gphase(0.5);
+        return SmallVector<Value>{q1Then};
+      },
+      [&](ValueRange qubits) -> SmallVector<Value> {
+        auto q1Else = builder.h(qubits[0]);
+        return SmallVector<Value>{q1Else};
+      })[0];
+
+  // This is an `if` without memory effects - it can be removed.
+  auto q1S2 = builder.qcoIf(
+      c0, {q1S1},
+      [&](ValueRange qubits) -> SmallVector<Value> {
+        auto q1Then = builder.x(qubits[0]);
+        return SmallVector<Value>{q1Then};
+      },
+      [&](ValueRange qubits) -> SmallVector<Value> {
+        auto q1Else = builder.h(qubits[0]);
+        return SmallVector<Value>{q1Else};
+      })[0];
+  builder.sink(q0S2);
+  builder.sink(q1S2);
+  auto module = builder.finalize();
+
+  QCOProgramBuilder reference(context.get());
+  reference.initialize();
+  auto r0S0 = reference.allocQubit();
+  auto r1S0 = reference.allocQubit();
+  auto r0S1 = reference.h(r0S0);
+  auto [r0S2, cr0] = reference.measure(r0S1);
+
+  // This is an `if` with memory effects - it can't be removed.
+  auto r1S1 = reference.qcoIf(
+      cr0, {r1S0},
+      [&](ValueRange qubits) -> SmallVector<Value> {
+        auto q1Then = reference.x(qubits[0]);
+        reference.gphase(0.5);
+        return SmallVector<Value>{q1Then};
+      },
+      [&](ValueRange qubits) -> SmallVector<Value> {
+        auto q1Else = reference.h(qubits[0]);
+        return SmallVector<Value>{q1Else};
+      })[0];
+
+  reference.sink(r0S2);
+  reference.sink(r1S1);
+  auto refModule = reference.finalize();
+
+  ASSERT_TRUE(module);
+  EXPECT_TRUE(verify(*module).succeeded());
+  EXPECT_TRUE(runQCOCleanupPipeline(module.get()).succeeded());
+  EXPECT_TRUE(verify(*module).succeeded());
+
+  ASSERT_TRUE(refModule);
+  EXPECT_TRUE(verify(*refModule).succeeded());
+  EXPECT_TRUE(runQCOCleanupPipeline(refModule.get()).succeeded());
+  EXPECT_TRUE(verify(*refModule).succeeded());
+
+  EXPECT_TRUE(
+      areModulesEquivalentWithPermutations(module.get(), refModule.get()));
+}
+
 TEST_F(QCOTest, DirectIfBuilder) {
   // Test If construction directly
   QCOProgramBuilder builder(context.get());