✨ Add Dead Gate Elimination Pattern

CHANGELOG.md

@@ -18,7 +18,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 - ✨ Add conversions between `jeff` and QCO ([#1479], [#1548], [#1565], [#1637], [#1676], [#1706]) ([**@denialhaag**], [**@burgholzer**])
 - ✨ Add a `place-and-route` pass for mapping circuits to architectures with restricted topologies ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588], [#1600], [#1664], [#1709], [#1716], [#1748]) ([**@MatthiasReumann**], [**@burgholzer**])
 - ✨ Add initial infrastructure for new QC and QCO MLIR dialects
-  ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521], [#1542], [#1548], [#1550], [#1554], [#1567], [#1569], [#1570], [#1572], [#1573], [#1580], [#1602], [#1620], [#1623], [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1717], [#1728], [#1730], [#1749])
+  ([#1264], [#1330], [#1402], [#1428], [#1430], [#1436], [#1443], [#1446], [#1464], [#1465], [#1470], [#1471], [#1472], [#1474], [#1475], [#1506], [#1510], [#1513], [#1521], [#1542], [#1548], [#1550], [#1554], [#1567], [#1569], [#1570], [#1572], [#1573], [#1580], [#1602], [#1620], [#1623], [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1717], [#1728], [#1730], [#1749], [#1755])
   ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**])
 
 ### Changed
@@ -402,6 +402,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

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <llvm/ADT/TypeSwitch.h>
 #include <mlir/Dialect/Arith/IR/Arith.h>
 #include <mlir/IR/PatternMatch.h>
 #include <mlir/Support/LLVM.h>
@@ -237,4 +238,45 @@ mergeTwoTargetOneParameterWithSwappedTargets(OpType op,
   return success();
 }
 
+/**
+ * @brief Check if a 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 (!llvm::all_of(op->getUsers(),
+                    [](Operation* user) { return isa<SinkOp>(user); })) {
+    return failure();
+  }
+
+  // If the operation is only used by sinks, we can safely remove it.
+  return TypeSwitch<Operation*, LogicalResult>(op)
+      .Case<UnitaryOpInterface>([&](auto u) {
+        // Replace output *qubits* with input *qubits* to ignore parameters.
+        rewriter.replaceOp(op, u.getInputQubits());
+        return success();
+      })
+      .Case<MeasureOp>([&](auto m) {
+        // Replace output *qubits* with input *qubits* to ignore classical
+        // outcome.
+        rewriter.replaceAllUsesWith(m.getQubitOut(), m.getQubitIn());
+        rewriter.eraseOp(op);
+        return success();
+      })
+      .Case<IfOp>([&](auto i) {
+        // Replace output *qubits* with input *qubits* to ignore the condition.
+        rewriter.replaceOp(op, i.getQubits());
+        return success();
+      })
+      .Default([&](auto*) {
+        // This currently only includes the `Reset` operation.
+        rewriter.replaceOp(op, op->getOperands());
+        return success();
+      });
+}
+
 } // 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/SideEffectInterfaces.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 gates.
+ */
+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 `ctrl` 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::initialize() {
       >();
 }
 
+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>
@@ -25,6 +26,7 @@
 #include <mlir/IR/Value.h>
 #include <mlir/IR/ValueRange.h>
 #include <mlir/Interfaces/ControlFlowInterfaces.h>
+#include <mlir/Interfaces/SideEffectInterfaces.h>
 #include <mlir/Support/LLVM.h>
 
 #include <cassert>
@@ -234,11 +236,28 @@ 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 {
+    if (!isMemoryEffectFree(op)) {
+      // This effectively ignores `IfOp`s with memory effects.
+      return failure();
+    }
+    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());