✨ Add pass and patterns for Classical Control Replacement

CHANGELOG.md

@@ -50,15 +50,17 @@ releases may include breaking changes.
   restricted topologies ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588],
   [#1600], [#1664], [#1709], [#1716], [#1748], [#1805], [#1870], [#1904])
   ([**@MatthiasReumann**], [**@burgholzer**])
+- ✨ Add a pass for qubit reuse in quantum programs, as well as related
+  auxiliary passes and patterns ([#1705], [#1755], [#1756]) ([**@DRovara**])
 - ✨ 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], [#1710],
-  [#1717], [#1728], [#1730], [#1749], [#1751], [#1755], [#1762], [#1765],
-  [#1780], [#1781], [#1782], [#1787], [#1806], [#1807], [#1815], [#1808],
-  [#1823], [#1824], [#1830], [#1869], [#1872], [#1914]) ([**@burgholzer**],
+  [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765], [#1780],
+  [#1781], [#1782], [#1787], [#1806], [#1807], [#1815], [#1808], [#1823],
+  [#1824], [#1830], [#1869], [#1872], [#1914]) ([**@burgholzer**],
   [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**],
   [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**])
 
@@ -679,6 +681,7 @@ changelogs._
 [#1766]: https://github.com/munich-quantum-toolkit/core/pull/1766
 [#1765]: https://github.com/munich-quantum-toolkit/core/pull/1765
 [#1762]: https://github.com/munich-quantum-toolkit/core/pull/1762
+[#1756]: https://github.com/munich-quantum-toolkit/core/pull/1756
 [#1755]: https://github.com/munich-quantum-toolkit/core/pull/1755
 [#1751]: https://github.com/munich-quantum-toolkit/core/pull/1751
 [#1749]: https://github.com/munich-quantum-toolkit/core/pull/1749
@@ -694,6 +697,7 @@ changelogs._
 [#1710]: https://github.com/munich-quantum-toolkit/core/pull/1710
 [#1709]: https://github.com/munich-quantum-toolkit/core/pull/1709
 [#1706]: https://github.com/munich-quantum-toolkit/core/pull/1706
+[#1705]: https://github.com/munich-quantum-toolkit/core/pull/1705
 [#1702]: https://github.com/munich-quantum-toolkit/core/pull/1702
 [#1700]: https://github.com/munich-quantum-toolkit/core/pull/1700
 [#1694]: https://github.com/munich-quantum-toolkit/core/pull/1694

mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h

@@ -138,6 +138,21 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
    */
   Value floatConstant(double value);
 
+  /**
+   * @brief Create a constant boolean value
+   * @param value The value to store in the constant
+   * @return The value produced by the constant operation
+   *
+   * @par Example:
+   * ```c++
+   * auto c = builder.boolConstant(true);
+   * ```
+   * ```mlir
+   * %c = arith.constant 1 : i1
+   * ```
+   */
+  Value boolConstant(bool value);
+
   //===--------------------------------------------------------------------===//
   // Memory Management
   //===--------------------------------------------------------------------===//

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

@@ -201,4 +201,55 @@ def HadamardLifting : Pass<"hadamard-lifting", "mlir::ModuleOp"> {
   }];
 }
 
+def MeasurementLifting : Pass<"measurement-lifting", "mlir::ModuleOp"> {
+  let dependentDialects = ["mlir::qco::QCODialect",
+                           "::mlir::arith::ArithDialect"];
+  let summary = "This pass attempts to move measurements as far up as "
+                "possible, shifting them above gates that commute with them. "
+                "This is done to enable qubit reuse and other optimizations.";
+  let description = [{
+    This pass applies measurement lifting, moving measurements up the code as far as possible.
+    Measurement lifting is a subroutine of the qubit reuse routine. The goal is to measure qubits earlier in the
+    circuit to reuse them and to potentially remove some quantum gates.
+
+    Measurement lifting uses the following commutation rules:
+                 ┌──────┐             ┌──────┐
+            ──■──┤ Meas │──────      ─┤ Meas ├──■───
+              │  └──────┘             └──────┘  │
+            ┌─┴─┐                 =>          ┌─┴─┐
+            ┤ U ├────────────────    ─────────┤ U ├─
+            └───┘                             └───┘
+            (Where U is any (controlled) unitary gate)
+
+            ┌───┐┌──────┐    ┌──────┐┌───┐
+            ┤ P ├┤ Meas ├ => ┤ Meas ├┤ P ├
+            └───┘└──────┘    └──────┘└───┘
+            (Where P is any diagonal gate, e.g., `z`, `s`, ...)
+
+            ┌───┐┌──────┐    ┌───────┐┌───┐
+            ┤ X ├┤ Meas ├ => ┤ Meas* ├┤ X ├
+            └───┘└──────┘    └───────┘└───┘
+            (Where Meas* is a measurement after which the outcome is classically negated and X can be replaced by Y, which is also anti-diagonal)
+
+  }];
+}
+
+def ReplaceClassicalControls
+    : Pass<"replace-classical-controls", "mlir::ModuleOp"> {
+  let dependentDialects = ["mlir::qco::QCODialect",
+                           "::mlir::arith::ArithDialect",
+  ];
+  let summary =
+      "This pass attempts to replace controlled gates that are controlled by "
+      "values that are in a computational basis state and available as "
+      "measurement outcome by a classical "
+      "`qco.if` instruction. This is done to enable qubit reuse and other "
+      "optimizations.";
+  let description = [{
+    This pass searches for control operations that immediately follow measurements and replaces them
+    with classically controlled operations, represented by `qco.if` operations.
+    This reduces quantum interactions and allows for more optimizations, e.g., by enabling qubit reuse.
+  }];
+}
+
 #endif // MLIR_DIALECT_QCO_TRANSFORMS_PASSES_TD

mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp

@@ -92,6 +92,11 @@ Value QCOProgramBuilder::floatConstant(const double value) {
   return arith::ConstantOp::create(*this, getF64FloatAttr(value)).getResult();
 }
 
+Value QCOProgramBuilder::boolConstant(const bool value) {
+  checkFinalized();
+  return arith::ConstantOp::create(*this, getBoolAttr(value)).getResult();
+}
+
 Value& QCOProgramBuilder::QubitRegister::operator[](const size_t index) {
   if (index >= qubits.size()) {
     llvm::reportFatalUsageError("Qubit index out of bounds");

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

@@ -0,0 +1,225 @@
+/*
+ * 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
+ */
+
+//
+// Created by damian on 5/13/26.
+//
+
+#include "mlir/Dialect/QCO/IR/QCOInterfaces.h"
+#include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/Transforms/Passes.h"
+#include "mlir/Dialect/Utils/Utils.h"
+
+#include <llvm/ADT/STLExtras.h>
+#include <mlir/Dialect/Arith/IR/Arith.h>
+#include <mlir/IR/MLIRContext.h>
+#include <mlir/IR/PatternMatch.h>
+#include <mlir/IR/Value.h>
+#include <mlir/Support/LLVM.h>
+#include <mlir/Transforms/GreedyPatternRewriteDriver.h>
+
+#include <utility>
+
+namespace mlir::qco {
+
+#define GEN_PASS_DEF_MEASUREMENTLIFTING
+#include "mlir/Dialect/QCO/Transforms/Passes.h.inc"
+
+/**
+ * @brief Checks if the given operation is an inverting gate.
+ * @param op The operation to check.
+ * @return True if the operation is an inverting gate, false otherwise.
+ */
+static bool isInverting(Operation* op) { return isa<XOp, YOp>(op); }
+
+/**
+ * @brief Checks if the given operation is a diagonal gate.
+ * @param op The operation to check.
+ * @return True if the operation is a diagonal gate, false otherwise.
+ */
+static bool isDiagonal(Operation* op) {
+  if (op == nullptr) {
+    return false;
+  }
+  if (isa<CtrlOp, InvOp>(op)) {
+    return isDiagonal(utils::getSoleBodyUnitary<UnitaryOpInterface>(
+        *op->getRegion(0).getBlocks().begin()));
+  }
+  return isa<ZOp, SOp, TOp, POp, RZOp, SdgOp, TdgOp, IdOp>(op);
+}
+
+/**
+ * @brief This method swaps a gate with a measurement.
+ * @param gate The gate to swap.
+ * @param measurement The measurement to swap.
+ * @param rewriter The used rewriter.
+ */
+static void swapGateWithMeasurement(UnitaryOpInterface gate,
+                                    MeasureOp measurement,
+                                    mlir::PatternRewriter& rewriter) {
+  auto measurementInput = measurement.getQubitIn();
+  auto gateInput = gate.getInputForOutput(measurementInput);
+  rewriter.replaceUsesWithIf(measurementInput, gateInput,
+                             [&](mlir::OpOperand& operand) {
+                               // We only replace the single use by the
+                               // measure op
+                               return operand.getOwner() == measurement;
+                             });
+  rewriter.replaceUsesWithIf(gateInput, measurement.getQubitOut(),
+                             [&](mlir::OpOperand& operand) {
+                               // We only replace the single use by the
+                               // predecessor
+                               return operand.getOwner() == gate;
+                             });
+  rewriter.replaceUsesWithIf(measurement.getQubitOut(), measurementInput,
+                             [&](mlir::OpOperand& operand) {
+                               // All further uses of the measurement output now
+                               // use the gate output
+                               return operand.getOwner() != gate;
+                             });
+  rewriter.moveOpBefore(measurement, gate);
+}
+
+namespace {
+/**
+ * @brief This pattern is responsible for lifting measurements above any phase
+ * gates.
+ */
+struct LiftMeasurementsAbovePhaseGatesPattern final
+    : mlir::OpRewritePattern<MeasureOp> {
+
+  explicit LiftMeasurementsAbovePhaseGatesPattern(mlir::MLIRContext* context)
+      : OpRewritePattern(context) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(MeasureOp op,
+                  mlir::PatternRewriter& rewriter) const override {
+    const auto qubitVariable = op.getQubitIn();
+    auto* predecessor = qubitVariable.getDefiningOp();
+
+    auto predecessorUnitary = mlir::dyn_cast<UnitaryOpInterface>(predecessor);
+
+    if (!predecessorUnitary) {
+      return mlir::failure();
+    }
+
+    if (isDiagonal(predecessor)) {
+      swapGateWithMeasurement(predecessorUnitary, op, rewriter);
+      return mlir::success();
+    }
+
+    return mlir::failure();
+  }
+};
+
+/**
+ * @brief This pattern is responsible for lifting measurements above any
+ * anti-diagonal gates.
+ */
+struct LiftMeasurementsAboveInvertingGatesPattern final
+    : mlir::OpRewritePattern<MeasureOp> {
+
+  explicit LiftMeasurementsAboveInvertingGatesPattern(
+      mlir::MLIRContext* context)
+      : OpRewritePattern(context) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(MeasureOp op,
+                  mlir::PatternRewriter& rewriter) const override {
+    const auto qubitVariable = op.getQubitIn();
+    auto* predecessor = qubitVariable.getDefiningOp();
+
+    auto predecessorUnitary = mlir::dyn_cast<UnitaryOpInterface>(predecessor);
+
+    if (!predecessorUnitary) {
+      return mlir::failure();
+    }
+
+    if (isInverting(predecessor)) {
+      swapGateWithMeasurement(predecessorUnitary, op, rewriter);
+      rewriter.setInsertionPointAfter(op);
+      const mlir::Value trueConstant = rewriter.create<mlir::arith::ConstantOp>(
+          op.getLoc(), rewriter.getBoolAttr(true));
+      auto inversion = rewriter.create<mlir::arith::XOrIOp>(
+          op.getLoc(), op.getResult(), trueConstant);
+      // We need `replaceUsesWithIf` so that we can replace all uses except for
+      // the one use that defines the inverted bit.
+      rewriter.replaceUsesWithIf(op.getResult(), inversion.getResult(),
+                                 [&](mlir::OpOperand& operand) {
+                                   return operand.getOwner() != inversion;
+                                 });
+      return mlir::success();
+    }
+
+    return mlir::failure();
+  }
+};
+
+/**
+ * @brief This pattern is responsible for applying the "deferred measurement
+ * principle", lifting measurements above controls.
+ */
+struct LiftMeasurementsAboveControlsPattern final
+    : mlir::OpRewritePattern<MeasureOp> {
+
+  explicit LiftMeasurementsAboveControlsPattern(mlir::MLIRContext* context)
+      : OpRewritePattern(context) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(MeasureOp op,
+                  mlir::PatternRewriter& rewriter) const override {
+    const auto qubitVariable = op.getQubitIn();
+    auto* predecessor = qubitVariable.getDefiningOp();
+    auto predecessorCtrl = mlir::dyn_cast<CtrlOp>(predecessor);
+
+    if (!predecessorCtrl) {
+      return mlir::failure();
+    }
+
+    if (llvm::find(predecessorCtrl.getControlsOut(), qubitVariable) ==
+        predecessorCtrl.getControlsOut().end()) {
+      // The measured qubit is a target, not a control of the gate.
+      return mlir::failure();
+    }
+
+    swapGateWithMeasurement(predecessorCtrl, op, rewriter);
+
+    return mlir::success();
+  }
+};
+
+/**
+ * @brief Pass raises Measurements above controlled and uncontrolled gates.
+ */
+struct MeasurementLifting final
+    : impl::MeasurementLiftingBase<MeasurementLifting> {
+  using MeasurementLiftingBase::MeasurementLiftingBase;
+
+protected:
+  void runOnOperation() override {
+    const auto op = getOperation();
+    auto* ctx = &getContext();
+
+    // Define the set of patterns to use.
+    RewritePatternSet patterns(ctx);
+    patterns.add<LiftMeasurementsAboveControlsPattern,
+                 LiftMeasurementsAboveInvertingGatesPattern,
+                 LiftMeasurementsAbovePhaseGatesPattern>(patterns.getContext());
+
+    // Apply patterns in an iterative and greedy manner.
+    if (failed(applyPatternsGreedily(op, std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+
+} // namespace
+
+} // namespace mlir::qco