♻️ Final improvements: type checks, parameter ordering, pipeline test

Author: J4MMlE

mlir/lib/Dialect/QCO/Transforms/CMakeLists.txt

@@ -10,8 +10,6 @@ get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 set(LIBRARIES ${dialect_libs} MQT::CoreIR)
 add_compile_options(-fexceptions)
 
-message(STATUS "MLIR_DIALECT_LIBS contains: ${dialect_libs}")
-
 file(GLOB_RECURSE TRANSFORMS_SOURCES *.cpp)
 
 add_mlir_library(MLIRQCOTransforms ${TRANSFORMS_SOURCES} LINK_LIBS ${LIBRARIES} DEPENDS

mlir/lib/Dialect/QCO/Transforms/QuaternionMergeRotationGates.cpp

@@ -11,11 +11,10 @@
 #include "mlir/Dialect/QCO/IR/QCODialect.h"
 #include "mlir/Dialect/QCO/Transforms/Passes.h"
 
-#include <algorithm>
-#include <array>
 #include <cassert>
 #include <cmath>
 #include <cstdint>
+#include <llvm/ADT/TypeSwitch.h>
 #include <llvm/Support/ErrorHandling.h>
 #include <mlir/Dialect/Arith/IR/Arith.h>
 #include <mlir/Dialect/Math/IR/Math.h>
@@ -26,7 +25,7 @@
 #include <mlir/Support/LLVM.h>
 #include <mlir/Transforms/GreedyPatternRewriteDriver.h>
 #include <numbers>
-#include <string_view>
+#include <optional>
 #include <utility>
 
 namespace mlir::qco {
@@ -52,17 +51,14 @@ struct MergeRotationGatesPattern final
 
   enum class RotationAxis : std::uint8_t { X, Y, Z };
 
-  static constexpr std::array<std::string_view, 4> MERGEABLE_GATES = {
-      "u", "rx", "ry", "rz"};
-
   /**
    * @brief Checks if an operation is a mergeable rotation gate (rx, ry, rz, u).
    *
-   * @param name Name of the operation to check
+   * @param op The operation to check
    * @return True if mergeable, false otherwise
    */
-  static bool isMergeable(std::string_view name) {
-    return std::ranges::find(MERGEABLE_GATES, name) != MERGEABLE_GATES.end();
+  static bool isMergeable(mlir::Operation* op) {
+    return mlir::isa<RXOp, RYOp, RZOp, UOp>(op);
   }
 
   /**
@@ -78,13 +74,28 @@ struct MergeRotationGatesPattern final
    */
   [[nodiscard]] static bool areQuaternionMergeable(mlir::Operation& a,
                                                    mlir::Operation& b) {
-    const auto aName = a.getName().stripDialect().str();
-    const auto bName = b.getName().stripDialect().str();
-
-    if (!(isMergeable(aName) && isMergeable(bName))) {
+    if (!isMergeable(&a) || !isMergeable(&b)) {
       return false;
     }
-    return (aName != bName) || (aName == "u" && bName == "u");
+
+    // Different gate types OR both are U gates
+    return (a.getName() != b.getName()) ||
+           (mlir::isa<UOp>(a) && mlir::isa<UOp>(b));
+  }
+
+  /**
+   * @brief Returns the rotation axis for a single-axis rotation gate.
+   *
+   * @param op The operation to query
+   * @return The rotation axis, or std::nullopt if the operation is not a
+   *         single-axis rotation gate (RX, RY, RZ)
+   */
+  static std::optional<RotationAxis> getRotationAxis(mlir::Operation* op) {
+    return llvm::TypeSwitch<mlir::Operation*, std::optional<RotationAxis>>(op)
+        .Case<RXOp>([](auto) { return RotationAxis::X; })
+        .Case<RYOp>([](auto) { return RotationAxis::Y; })
+        .Case<RZOp>([](auto) { return RotationAxis::Z; })
+        .Default([](auto) { return std::nullopt; });
   }
 
   /**
@@ -130,6 +141,8 @@ struct MergeRotationGatesPattern final
     case RotationAxis::Z:
       return {.w = cos, .x = zero, .y = zero, .z = sin};
     } // NOLINT(bugprone-branch-clone): false positive, branches differ
+
+    llvm_unreachable("Invalid rotation axis");
   }
 
   /**
@@ -142,24 +155,15 @@ struct MergeRotationGatesPattern final
    */
   static Quaternion quaternionFromRotation(UnitaryOpInterface op,
                                            mlir::PatternRewriter& rewriter) {
-    auto const type = op->getName().stripDialect().str();
-
-    if (type == "u") {
+    if (mlir::isa<UOp>(op)) {
       return quaternionFromUGate(op, rewriter);
     }
 
-    auto loc = op->getLoc();
-    auto angle = op.getParameter(0);
-
-    if (type == "rx") {
-      return createAxisQuaternion(angle, RotationAxis::X, loc, rewriter);
-    }
-    if (type == "ry") {
-      return createAxisQuaternion(angle, RotationAxis::Y, loc, rewriter);
-    }
-    if (type == "rz") {
-      return createAxisQuaternion(angle, RotationAxis::Z, loc, rewriter);
+    if (auto axis = getRotationAxis(op.getOperation())) {
+      return createAxisQuaternion(op.getParameter(0), *axis, op->getLoc(),
+                                  rewriter);
     }
+
     llvm_unreachable("Unsupported operation type");
   }
 
@@ -227,14 +231,14 @@ struct MergeRotationGatesPattern final
   /**
    * @brief Converts a u-gate to quaternion representation.
    *
-   * U(alpha, beta, gamma) uses ZYZ decomposition: RZ(alpha) -> RY(beta) ->
-   * RZ(gamma).
+   * U(theta, phi, lambda) uses ZYZ decomposition: RZ(lambda) -> RY(theta) ->
+   * RZ(phi).
    *
    * When composing rotations, quaternion multiplication follows matrix
    * multiplication order (right-to-left), which is the reverse of the
    * application sequence:
-   *   Sequential application: RZ(alpha), then RY(beta), then RZ(gamma)
-   *   Quaternion product:     Qgamma * Qbeta * Qalpha
+   *   Sequential application: RZ(lambda), then RY(theta), then RZ(phi)
+   *   Quaternion product:     qPhi * qTheta * qLambda
    *
    * @param op The u-gate operation to convert
    * @param rewriter Pattern rewriter for creating new operations
@@ -245,26 +249,27 @@ struct MergeRotationGatesPattern final
     auto loc = op->getLoc();
 
     // U gate uses ZYZ decomposition:
-    // U(alpha, beta, gamma) = Rz(alpha) -> Ry(beta) -> Rz(gamma)
-    auto qAlpha = createAxisQuaternion(op.getParameter(0), RotationAxis::Z, loc,
-                                       rewriter);
-    auto qBeta = createAxisQuaternion(op.getParameter(1), RotationAxis::Y, loc,
-                                      rewriter);
-    auto qGamma = createAxisQuaternion(op.getParameter(2), RotationAxis::Z, loc,
+    // U(theta, phi, lambda) uses ZYZ decomposition: RZ(lambda) -> RY(theta) ->
+    // RZ(phi)
+    auto qTheta = createAxisQuaternion(op.getParameter(0), RotationAxis::Y, loc,
                                        rewriter);
-
-    // qGamma * qBeta * qAlpha (multiplication in reverse order!)
-    auto temp = hamiltonProduct(qGamma, qBeta, op, rewriter);
-    return hamiltonProduct(temp, qAlpha, op, rewriter);
+    auto qPhi = createAxisQuaternion(op.getParameter(1), RotationAxis::Z, loc,
+                                     rewriter);
+    auto qLambda = createAxisQuaternion(op.getParameter(2), RotationAxis::Z,
+                                        loc, rewriter);
+
+    // qPhi * qTheta * qLambda (multiplication in reverse order!)
+    auto temp = hamiltonProduct(qPhi, qTheta, op, rewriter);
+    return hamiltonProduct(temp, qLambda, op, rewriter);
   }
 
   /**
    * @brief Converts a quaternion to a u-gate using ZYZ Euler angle extraction.
    *
    * For unit quaternion q = w + x*i + y*j + z*k, extracts u-gate parameters:
-   *   alpha = atan2(z, w) - atan2(-x, y)
+   *   alpha = atan2(z, w) + atan2(-x, y)
    *   beta  = acos(2 * (w^2 + z^2) - 1)
-   *   gamma = atan2(z, w) + atan2(-x, y)
+   *   gamma = atan2(z, w) - atan2(-x, y)
    *
    * Based on Bernardes & Viollet (2022), simplified for unit quaternions and
    * proper ZYZ Euler angles (Chapter 3.3):
@@ -287,6 +292,9 @@ struct MergeRotationGatesPattern final
     auto loc = op->getLoc();
 
     auto floatType = op.getParameter(0).getType();
+    // constant -1.0
+    auto negOneAttr = rewriter.getFloatAttr(floatType, -1.0);
+    auto negOne = rewriter.create<mlir::arith::ConstantOp>(loc, negOneAttr);
     // constant 0.0
     auto zeroAttr = rewriter.getFloatAttr(floatType, 0.0);
     auto zero = rewriter.create<mlir::arith::ConstantOp>(loc, zeroAttr);
@@ -305,12 +313,18 @@ struct MergeRotationGatesPattern final
 
     // calculate angle beta (for y-rotation)
     // beta = acos(2 * (w^2 + z^2) - 1)
+    // NOTE: the term (2 * (w^2 + z^2) - 1) is clamped to [-1, 1],
+    // otherwise acos could produce NaN.
     auto ww = rewriter.create<mlir::arith::MulFOp>(loc, q.w, q.w);
     auto zz = rewriter.create<mlir::arith::MulFOp>(loc, q.z, q.z);
     auto bTemp1 = rewriter.create<mlir::arith::AddFOp>(loc, ww, zz);
     auto bTemp2 = rewriter.create<mlir::arith::MulFOp>(loc, two, bTemp1);
     auto bTemp3 = rewriter.create<mlir::arith::SubFOp>(loc, bTemp2, one);
-    auto beta = rewriter.create<mlir::math::AcosOp>(loc, bTemp3);
+    auto clampedLow =
+        rewriter.create<mlir::arith::MaximumFOp>(loc, bTemp3, negOne);
+    auto clamped =
+        rewriter.create<mlir::arith::MinimumFOp>(loc, clampedLow, one);
+    auto beta = rewriter.create<mlir::math::AcosOp>(loc, clamped);
 
     // intermediates to check for gimbal lock (|beta| and |beta - PI|)
     auto absBeta = rewriter.create<mlir::math::AbsFOp>(loc, beta);
@@ -343,31 +357,31 @@ struct MergeRotationGatesPattern final
         rewriter.create<mlir::arith::MulFOp>(loc, two, thetaMinus);
 
     // Safe Case (no gimbal lock):
-    // alphaSafe = theta+ - theta-
-    // gammaSafe = theta+ + theta-
+    // alphaSafe = theta+ + theta-
+    // gammaSafe = theta+ - theta-
     auto alphaSafe =
-        rewriter.create<mlir::arith::SubFOp>(loc, thetaPlus, thetaMinus);
-    auto gammaSafe =
         rewriter.create<mlir::arith::AddFOp>(loc, thetaPlus, thetaMinus);
+    auto gammaSafe =
+        rewriter.create<mlir::arith::SubFOp>(loc, thetaPlus, thetaMinus);
 
     // Unsafe Case (gimbal lock):
     // when b = 0  then alpha = 2 * (atan2(z,w))
-    // when b = PI then alpha = 2 * (atan2(-z, y))
+    // when b = PI then alpha = 2 * (atan2(-x, y))
     // gamma is set to zero in both cases
     auto alphaUnsafe = rewriter.create<mlir::arith::SelectOp>(
         loc, safe1, twoThetaMinus, twoThetaPlus);
 
     // TODO: could add some normalization here for alpha and gamma otherwise
     // they can be outside of [-PI, PI].
 
-    // choose correct alpha and gamma weather safe or not
+    // choose correct alpha and gamma whether safe or not
     auto alpha = rewriter.create<mlir::arith::SelectOp>(loc, safe, alphaSafe,
                                                         alphaUnsafe);
     auto gamma =
         rewriter.create<mlir::arith::SelectOp>(loc, safe, gammaSafe, zero);
 
-    return rewriter.create<UOp>(loc, op.getInputQubit(0), alpha.getResult(),
-                                beta.getResult(), gamma.getResult());
+    return rewriter.create<UOp>(loc, op.getInputQubit(0), beta.getResult(),
+                                alpha.getResult(), gamma.getResult());
   }
 
   /**

mlir/unittests/Compiler/test_compiler_pipeline.cpp

@@ -3736,4 +3736,44 @@ TEST_F(CompilerPipelineTest, Bell) {
   });
 }
 
+// ##################################################
+// # Rotation Merge Pass Test
+// ##################################################
+
+/**
+ * @brief Test: Rotation merging pass is invoked during optimization stage
+ *
+ * @details
+ * The merged U gate parameters are computed via floating-point arithmetic
+ * that is not bit-identical across platforms, so we cannot use
+ * verifyAllStages with hardcoded expected values. Instead, we compare
+ * the optimization output with and without the pass enabled.
+ * Correctness of the pass is tested in a dedicated test.
+ */
+TEST_F(CompilerPipelineTest, RotationGateMergingPass) {
+  ::qc::QuantumComputation comp;
+  comp.addQubitRegister(1, "q");
+  comp.rz(1.0, 0);
+  comp.rx(1.0, 0);
+
+  // Run with merging enabled
+  config.mergeRotationGates = true;
+
+  auto module = importQuantumCircuit(comp);
+  ASSERT_TRUE(module);
+  ASSERT_TRUE(runPipeline(module.get()).succeeded());
+  const auto withMerging = record.afterOptimization;
+
+  // Run with merging disabled
+  config.mergeRotationGates = false;
+  record = {};
+
+  module = importQuantumCircuit(comp);
+  ASSERT_TRUE(module);
+  ASSERT_TRUE(runPipeline(module.get()).succeeded());
+  const auto withoutMerging = record.afterOptimization;
+
+  // The outputs must differ, proving the pass ran and transformed the IR
+  EXPECT_NE(withMerging, withoutMerging);
+}
 } // namespace

mlir/unittests/Dialect/QCO/Transforms/CMakeLists.txt

@@ -10,12 +10,11 @@ add_executable(mqt-core-mlir-dialect-qco-transforms-test test_qco_quaternion_mer
 
 target_link_libraries(
   mqt-core-mlir-dialect-qco-transforms-test
-  # TODO figure out correct dependencies
   PRIVATE GTest::gtest_main
           MLIRQCOProgramBuilder
           MLIRQCOTransforms
           MLIRIR
-          MLIRPass # for PassManager
+          MLIRPass
           MLIRSupport
           LLVMSupport)