⚡️ Improve canonicalizations for gate fusion and inverse cancellation

CHANGELOG.md

@@ -42,8 +42,9 @@ with the exception that minor releases may include breaking changes.
   [#1569], [#1570], [#1572], [#1573], [#1580], [#1602], [#1620], [#1623],
   [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700],
   [#1710], [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765],
-  [#1774], [#1781]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**],
-  [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**],
+  [#1774], [#1781], [#1782])
+  ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**],
+  [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**],
   [**@simon1hofmann**])
 
 ### Changed
@@ -597,6 +598,7 @@ changelogs._
 
 <!-- PR links -->
 
+[#1782]: https://github.com/munich-quantum-toolkit/core/pull/1782
 [#1781]: https://github.com/munich-quantum-toolkit/core/pull/1781
 [#1776]: https://github.com/munich-quantum-toolkit/core/pull/1776
 [#1774]: https://github.com/munich-quantum-toolkit/core/pull/1774

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

@@ -10,12 +10,36 @@
 
 #pragma once
 
+#include "mlir/Dialect/QCO/IR/QCOOps.h"
+
 #include <mlir/Dialect/Arith/IR/Arith.h>
+#include <mlir/Dialect/Utils/Utils.h>
 #include <mlir/IR/PatternMatch.h>
 #include <mlir/Support/LLVM.h>
+#include <mlir/Support/LogicalResult.h>
 
 namespace mlir::qco {
 
+/**
+ * @brief Check whether two parameter values match.
+ *
+ * @details
+ * Identical SSA values always match. Otherwise, if both are constants, they
+ * are compared with @ref utils::TOLERANCE.
+ *
+ * @param lhs The first parameter value.
+ * @param rhs The second parameter value.
+ * @return true if the values match.
+ */
+static bool valuesMatchWithinTolerance(Value lhs, Value rhs) {
+  if (lhs == rhs) {
+    return true;
+  }
+  const auto lhsVal = utils::valueToDouble(lhs);
+  const auto rhsVal = utils::valueToDouble(rhs);
+  return lhsVal && rhsVal && std::abs(*lhsVal - *rhsVal) <= utils::TOLERANCE;
+}
+
 /**
  * @brief Remove a pair of inverse one-target, zero-parameter operations
  *
@@ -35,77 +59,49 @@ removeInversePairOneTargetZeroParameter(OpType op, PatternRewriter& rewriter) {
   }
 
   // Erase both operations
-  rewriter.replaceAllUsesWith(nextOp->getResult(0), op.getInputQubit(0));
-  rewriter.eraseOp(nextOp);
-  rewriter.eraseOp(op);
-
+  rewriter.replaceOp(op, op.getInputQubits());
+  rewriter.replaceOp(nextOp, nextOp.getInputQubits());
   return success();
 }
 
 /**
- * @brief Remove a pair of inverse two-target, zero-parameter operations
+ * @brief Remove a pair of inverse two-target, zero-parameter operations.
  *
  * @tparam InverseOpType The type of the inverse operation.
  * @tparam OpType The type of the operation to be checked.
  * @param op The operation instance.
  * @param rewriter The pattern rewriter.
+ * @param symmetric Whether the two-target gate is symmetric (order of the
+ * qubits does not matter)
+ * @param swappedTargets Whether the successor consumes swapped target wires.
  * @return LogicalResult Success or failure of the removal.
  */
 template <typename InverseOpType, typename OpType>
 LogicalResult
-removeInversePairTwoTargetZeroParameter(OpType op, PatternRewriter& rewriter) {
+removeInversePairTwoTargetZeroParameter(OpType op, PatternRewriter& rewriter,
+                                        bool symmetric = false,
+                                        bool swappedTargets = false) {
+  auto output0 = op.getOutputQubit(0);
+
   // Check if the successor is the inverse operation
-  auto nextOp = dyn_cast<InverseOpType>(*op.getOutputQubit(0).user_begin());
+  auto nextOp = dyn_cast<InverseOpType>(*output0.user_begin());
   if (!nextOp) {
     return failure();
   }
 
-  // Confirm operations act on the same qubits
-  if (op.getOutputQubit(1) != nextOp.getInputQubit(1)) {
-    return failure();
-  }
-
-  // Erase both operations
-  rewriter.replaceAllUsesWith(nextOp->getResults(),
-                              {op.getInputQubit(0), op.getInputQubit(1)});
-  rewriter.eraseOp(nextOp);
-  rewriter.eraseOp(op);
-
-  return success();
-}
-
-/**
- * @brief Remove a pair of two-target, zero-parameter operations where
- *        the second operation is the same gate with swapped targets.
- *
- * @tparam OpType The type of the (self-inverse) operation.
- * @param op The operation instance.
- * @param rewriter The pattern rewriter.
- * @return LogicalResult Success or failure of the removal.
- */
-template <typename OpType>
-LogicalResult
-removeTwoTargetZeroParameterPairWithSwappedTargets(OpType op,
-                                                   PatternRewriter& rewriter) {
-  // Check if the successor is the same operation
-  auto nextOp = dyn_cast<OpType>(*op.getOutputQubit(0).user_begin());
-  if (!nextOp) {
+  // Both qubits have to point to the same successor
+  auto nextOp2 = *op.getOutputQubit(1).user_begin();
+  if (nextOp2 != nextOp) {
     return failure();
   }
 
-  // Confirm operations act on the same qubits but with swapped targets
-  if (op.getOutputQubit(0) != nextOp.getInputQubit(1) ||
-      op.getOutputQubit(1) != nextOp.getInputQubit(0)) {
-    return failure();
+  if (symmetric || (swappedTargets && output0 == nextOp.getInputQubit(1)) ||
+      (!swappedTargets && output0 == nextOp.getInputQubit(0))) {
+    rewriter.replaceOp(op, op.getInputQubits());
+    rewriter.replaceOp(nextOp, nextOp.getInputQubits());
+    return success();
   }
-
-  // Erase both operations
-  rewriter.replaceAllUsesWith(nextOp->getResults(),
-                              {op.getInputQubit(1), op.getInputQubit(0)});
-  rewriter.eraseOp(nextOp);
-  rewriter.eraseOp(op);
-
-  return success();
+  return failure();
 }
 
 /**
@@ -166,84 +162,89 @@ LogicalResult mergeOneTargetOneParameter(OpType op, PatternRewriter& rewriter) {
 
   // Replace the second operation with the result of the first operation
   rewriter.replaceOp(nextOp, op.getResult());
-
   return success();
 }
 
 /**
- * @brief Merge two compatible two-target, one-parameter operations
+ * @brief Shared implementation for merging two-target, one-parameter
+ * operations.
  *
- * @details
- * The new parameter is computed as the sum of the two original parameters.
+ * @tparam OpType The type of the operation to be merged.
+ * @param op The first operation instance.
+ * @param nextOp The successor operation instance.
+ * @param rewriter The pattern rewriter.
+ * @param symmetric Whether the two-target gate is symmetric (order of the
+ * qubits does not matter)
+ * @return LogicalResult Success or failure of the merge.
+ */
+template <typename OpType>
+static LogicalResult mergeTwoTargetOneParameterImpl(OpType op, OpType nextOp,
+                                                    PatternRewriter& rewriter,
+                                                    bool symmetric = false) {
+
+  // Both qubits have to point to the same successor
+  auto nextOp2 = *op.getOutputQubit(1).user_begin();
+  if (nextOp2 != nextOp) {
+    return failure();
+  }
+
+  auto output0 = op.getOutputQubit(0);
+  if (symmetric || output0 == nextOp.getInputQubit(0)) {
+    // Compute and set the new parameter
+    auto newParameter = arith::AddFOp::create(
+        rewriter, op.getLoc(), op.getOperand(2), nextOp.getOperand(2));
+    op->setOperand(2, newParameter.getResult());
+    rewriter.replaceOp(nextOp, nextOp.getInputQubits());
+    return success();
+  }
+  return failure();
+}
+
+/**
+ * @brief Merge two compatible two-target, one-parameter operations.
  *
  * @tparam OpType The type of the operation to be merged.
  * @param op The operation instance.
  * @param rewriter The pattern rewriter.
+ * @param symmetric Whether the two-target gate is symmetric (order of the
+ * qubits does not matter)
  * @return LogicalResult Success or failure of the merge.
  */
 template <typename OpType>
-LogicalResult mergeTwoTargetOneParameter(OpType op, PatternRewriter& rewriter) {
+LogicalResult mergeTwoTargetOneParameter(OpType op, PatternRewriter& rewriter,
+                                         bool symmetric = false) {
   // Check if the successor is the same operation
   auto nextOp = dyn_cast<OpType>(*op.getOutputQubit(0).user_begin());
   if (!nextOp) {
     return failure();
   }
-
-  // Confirm operations act on the same qubits
-  if (op.getOutputQubit(1) != nextOp.getInputQubit(1)) {
-    return failure();
-  }
-
-  // Compute and set the new parameter
-  auto newParameter = arith::AddFOp::create(
-      rewriter, op.getLoc(), op.getOperand(2), nextOp.getOperand(2));
-  op->setOperand(2, newParameter.getResult());
-
-  // Replace the second operation with the result of the first operation
-  rewriter.replaceOp(nextOp, op.getResults());
-
-  return success();
+  return mergeTwoTargetOneParameterImpl(op, nextOp, rewriter, symmetric);
 }
 
 /**
- * @brief Merge two compatible two-target, one-parameter operations where the
- *        second operation consumes the outputs with swapped targets.
+ * @brief Merge consecutive XXPlusYY or XXMinusYY operations.
  *
  * @details
- * This is analogous to mergeTwoTargetOneParameter, but it additionally handles
- * the case where the second operation swaps its target qubits. The new
- * parameter is computed as the sum of the two original parameters.
+ * Sums `theta` when `beta` matches within tolerance.
  *
  * @tparam OpType The type of the operation to be merged.
  * @param op The operation instance.
  * @param rewriter The pattern rewriter.
  * @return LogicalResult Success or failure of the merge.
  */
 template <typename OpType>
-LogicalResult
-mergeTwoTargetOneParameterWithSwappedTargets(OpType op,
-                                             PatternRewriter& rewriter) {
+LogicalResult mergeXXPlusMinusYY(OpType op, PatternRewriter& rewriter) {
   // Check if the successor is the same operation
   auto nextOp = dyn_cast<OpType>(*op.getOutputQubit(0).user_begin());
   if (!nextOp) {
     return failure();
   }
 
-  // Confirm operations act on the same qubits but with swapped targets
-  if (op.getOutputQubit(0) != nextOp.getInputQubit(1) ||
-      op.getOutputQubit(1) != nextOp.getInputQubit(0)) {
+  // Confirm matching beta before summing theta
+  if (!valuesMatchWithinTolerance(op.getBeta(), nextOp.getBeta())) {
     return failure();
   }
-
-  // Compute and set the new parameter on the first operation
-  auto newParameter = arith::AddFOp::create(
-      rewriter, op.getLoc(), op.getOperand(2), nextOp.getOperand(2));
-  op->setOperand(2, newParameter.getResult());
-
-  // nextOp results correspond to swapped operands, so swap replacements too
-  rewriter.replaceOp(nextOp, {op.getOutputQubit(1), op.getOutputQubit(0)});
-
-  return success();
+  return mergeTwoTargetOneParameterImpl(op, nextOp, rewriter, true);
 }
 
 } // namespace mlir::qco

mlir/lib/Dialect/QCO/IR/Operations/StandardGates/DCXOp.cpp

@@ -31,8 +31,8 @@ struct RemoveInversePairDCX final : OpRewritePattern<DCXOp> {
 
   LogicalResult matchAndRewrite(DCXOp op,
                                 PatternRewriter& rewriter) const override {
-    return removeTwoTargetZeroParameterPairWithSwappedTargets<DCXOp>(op,
-                                                                     rewriter);
+    return removeInversePairTwoTargetZeroParameter<DCXOp>(op, rewriter, false,
+                                                          true);
   }
 };
 

mlir/lib/Dialect/QCO/IR/Operations/StandardGates/ROp.cpp

@@ -9,13 +9,16 @@
  */
 
 #include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/QCO/QCOUtils.h"
 #include "mlir/Dialect/QCO/Utils/Matrix.h"
 #include "mlir/Dialect/Utils/Utils.h"
 
+#include <mlir/Dialect/Arith/IR/Arith.h>
 #include <mlir/IR/Builders.h>
 #include <mlir/IR/MLIRContext.h>
 #include <mlir/IR/OperationSupport.h>
 #include <mlir/IR/PatternMatch.h>
+#include <mlir/Support/LLVM.h>
 #include <mlir/Support/LogicalResult.h>
 
 #include <cmath>
@@ -64,6 +67,31 @@ struct ReplaceRWithRY final : OpRewritePattern<ROp> {
   }
 };
 
+/**
+ * @brief Merge subsequent R operations on the same qubit with matching `phi`.
+ */
+struct MergeSubsequentR final : OpRewritePattern<ROp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ROp op,
+                                PatternRewriter& rewriter) const override {
+    auto nextOp = dyn_cast<ROp>(*op.getOutputQubit(0).user_begin());
+    if (!nextOp) {
+      return failure();
+    }
+
+    if (!valuesMatchWithinTolerance(op.getPhi(), nextOp.getPhi())) {
+      return failure();
+    }
+
+    auto newParameter = arith::AddFOp::create(rewriter, op.getLoc(),
+                                              op.getTheta(), nextOp.getTheta());
+    op->setOperand(1, newParameter.getResult());
+    rewriter.replaceOp(nextOp, op.getResult());
+    return success();
+  }
+};
+
 } // namespace
 
 void ROp::build(OpBuilder& odsBuilder, OperationState& odsState, Value qubitIn,
@@ -77,7 +105,7 @@ void ROp::build(OpBuilder& odsBuilder, OperationState& odsState, Value qubitIn,
 
 void ROp::getCanonicalizationPatterns(RewritePatternSet& results,
                                       MLIRContext* context) {
-  results.add<ReplaceRWithRX, ReplaceRWithRY>(context);
+  results.add<ReplaceRWithRX, ReplaceRWithRY, MergeSubsequentR>(context);
 }
 
 std::optional<Matrix2x2> ROp::getUnitaryMatrix() {

mlir/lib/Dialect/QCO/IR/Operations/StandardGates/RXXOp.cpp

@@ -39,20 +39,7 @@ struct MergeSubsequentRXX final : OpRewritePattern<RXXOp> {
 
   LogicalResult matchAndRewrite(RXXOp op,
                                 PatternRewriter& rewriter) const override {
-    return mergeTwoTargetOneParameter(op, rewriter);
-  }
-};
-
-/**
- * @brief Merge subsequent RXX operations with swapped targets by adding their
- * angles.
- */
-struct MergeSwappedTargetsRXX final : OpRewritePattern<RXXOp> {
-  using OpRewritePattern::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(RXXOp op,
-                                PatternRewriter& rewriter) const override {
-    return mergeTwoTargetOneParameterWithSwappedTargets(op, rewriter);
+    return mergeTwoTargetOneParameter(op, rewriter, true);
   }
 };
 
@@ -79,7 +66,7 @@ LogicalResult RXXOp::fold(FoldAdaptor /*adaptor*/,
 
 void RXXOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                         MLIRContext* context) {
-  results.add<MergeSubsequentRXX, MergeSwappedTargetsRXX>(context);
+  results.add<MergeSubsequentRXX>(context);
 }
 
 std::optional<Matrix4x4> RXXOp::getUnitaryMatrix() {