Remove setRotationKind and getRotationKind from QEC interface

doc/releases/changelog-dev.md

@@ -322,6 +322,10 @@
   [(#2253](https://github.com/PennyLaneAI/catalyst/pull/2253)	
 
 
+  * Removed the `getRotationKind` and `setRotationKind` methods from 
+  the QEC interface `QECOpInterface` to simplify the interface.
+  [(#2250)](https://github.com/PennyLaneAI/catalyst/pull/2250)
+
 <h3>Documentation 📝</h3>
 
 * A typo in the code example for :func:`~.passes.ppr_to_ppm` has been corrected.

mlir/include/QEC/IR/QECOpInterfaces.td

@@ -17,12 +17,16 @@
 
 include "mlir/IR/OpBase.td"
 
+//===----------------------------------------------------------------------===//
+// QEC Operation Interface
+//===----------------------------------------------------------------------===//
+
 def QECOpInterface : OpInterface<"QECOpInterface"> {
 
     let description = [{
         This interface provides a generic way to interact with instructions that are
-        considered QEC Operations. These are characterized by operating on zero
-        or more qubit values, and returning the same amount of qubit values.
+        considered QEC Operations. These are characterized by operating on one or more qubit values,
+        and returning the same amount of qubit values.
     }];
 
     let cppNamespace = "::catalyst::qec";
@@ -40,32 +44,16 @@ def QECOpInterface : OpInterface<"QECOpInterface"> {
         >,
         InterfaceMethod<
             /*desc=*/"Get the Pauli product for this operation.",
-            /*retTy=*/" ::mlir::ArrayAttr",
+            /*retTy=*/"::mlir::ArrayAttr",
             /*methodName=*/"getPauliProduct"
         >,
-        InterfaceMethod<
-            /*desc=*/"Get the Pauli product for this operation.",
-            /*retTy=*/" ::mlir::ArrayAttr",
-            /*methodName=*/"getPauliProductAttr"
-        >,
-        InterfaceMethod<
-            /*desc=*/"Get the rotation kind for this operation.",
-            /*retTy=*/"uint16_t",
-            /*methodName=*/"getRotationKind"
-        >,
         InterfaceMethod<
             /*desc=*/"Set the Pauli product for this operation.",
             /*retTy=*/"void",
-            /*methodName=*/"setPauliProductAttr", (ins " ::mlir::ArrayAttr":$attr)
-        >,
-        InterfaceMethod<
-            /*desc=*/"Set the rotation kind for this operation.",
-            /*retTy=*/"void",
-            /*methodName=*/"setRotationKind", (ins "uint16_t":$attrValue)
+            /*methodName=*/"setPauliProductAttr", (ins "const ::mlir::ArrayAttr&":$propValue)
         >
     ];
 
-
 }
 
 #endif // QECOP_INTERFACES

mlir/lib/QEC/Transforms/CountPPMSpecs.cpp

@@ -157,7 +157,11 @@ struct CountPPMSpecsPass : public impl::CountPPMSpecsPassBase<CountPPMSpecsPass>
             assert(!layer.empty() && "Layer is empty");
 
             auto op = layer.getOps().back();
-            int16_t absRk = std::abs(static_cast<int16_t>(op.getRotationKind()));
+
+            int16_t absRk = 0;
+            if (auto pprOp = dyn_cast<PPRotationOp>(op.getOperation())) {
+                absRk = std::abs(static_cast<int16_t>(pprOp.getRotationKind()));
+            }
             auto parentFuncOp = op->getParentOfType<func::FuncOp>();
             StringRef funcName = parentFuncOp.getName();
             llvm::StringSaver saver(stringAllocator);

mlir/lib/QEC/Transforms/DecomposeNonCliffordPPR.cpp

@@ -34,7 +34,7 @@ using namespace catalyst::quantum;
 namespace {
 
 // Return the magic state or complex conjugate of the magic state
-LogicalInitKind getMagicState(QECOpInterface op)
+LogicalInitKind getMagicState(PPRotationOp op)
 {
     int16_t rotationKind = static_cast<int16_t>(op.getRotationKind());
     if (rotationKind > 0) {

mlir/lib/QEC/Transforms/PPRToMBQC.cpp

@@ -142,11 +142,14 @@ void constructKernelOperation(SmallVector<Value> &qubits, Value &measResult, QEC
         measResult = measOp.getMres();
         qubits[0] = measOp.getOutQubit();
     }
-    else {
-        int16_t signedRk = static_cast<int16_t>(op.getRotationKind());
+    else if (auto pprOp = dyn_cast<PPRotationOp>(op.getOperation())) {
+        int16_t signedRk = static_cast<int16_t>(pprOp.getRotationKind());
         double rk = llvm::numbers::pi / (static_cast<double>(signedRk) / 2);
         qubits[0] = buildSingleQubitGate(qubits[0], "RZ", {rk}, rewriter).getOutQubits().front();
     }
+    else if (isa<PPRotationArbitraryOp>(op)) {
+        op->emitError("Unsupported qec.ppr.arbitrary operation.");
+    }
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/QEC/Transforms/TLayerReduction.cpp

@@ -20,6 +20,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
+#include "QEC/IR/QECOps.h"
 #include "QEC/Utils/PauliStringWrapper.h"
 #include "QEC/Utils/QECLayer.h"
 #include "QEC/Utils/QECOpUtils.h"
@@ -64,6 +65,8 @@ std::pair<bool, QECOpInterface> checkCommutationAndFindMerge(QECOpInterface rhsO
         auto normalizedOps =
             normalizePPROps(lhsOp, rhsOp, lhsOp.getInQubits(), rhsOpInQubitsFromLhsOp);
 
+        // TODO: Handle PPRotationArbitraryOp properly
+
         if (!normalizedOps.first.commutes(normalizedOps.second)) {
             return std::pair(false, nullptr);
         }
@@ -133,8 +136,11 @@ void moveOpToLayer(QECOpInterface rhsOp, QECLayer &rhsLayer, QECOpInterface merg
 // then just remove the `rhsOp` from the rhsLayer.
 void mergePPR(QECOpInterface rhsOp, QECLayer &rhsLayer, QECOpInterface mergeOp, IRRewriter &writer)
 {
-    int16_t signedRk = static_cast<int16_t>(mergeOp.getRotationKind());
-    mergeOp.setRotationKind(static_cast<uint16_t>(signedRk / 2));
+    auto mergeOpPprOp = dyn_cast<PPRotationOp>(mergeOp.getOperation());
+    assert(mergeOpPprOp != nullptr && "Op is not a PPRotationOp");
+
+    int16_t signedRk = static_cast<int16_t>(mergeOpPprOp.getRotationKind());
+    mergeOpPprOp.setRotationKind(static_cast<uint16_t>(signedRk / 2));
 
     rhsLayer.eraseOp(rhsOp);
     writer.replaceOp(rhsOp, rhsOp->getOperands());

mlir/lib/QEC/Utils/PauliStringWrapper.cpp

@@ -151,8 +151,21 @@ PauliWordPair normalizePPROps(QECOpInterface lhs, QECOpInterface rhs, ValueRange
     lhsPSWrapper.op = lhs;
     rhsPSWrapper.op = rhs;
 
-    lhsPSWrapper.updateSign((int16_t)lhs.getRotationKind() < 0);
-    rhsPSWrapper.updateSign((int16_t)rhs.getRotationKind() < 0);
+    auto applySignFromOp = [](PauliStringWrapper &wrapper, QECOpInterface qecOp) {
+        Operation *operation = qecOp.getOperation();
+
+        if (auto pprOp = dyn_cast<PPRotationOp>(operation)) {
+            wrapper.updateSign(static_cast<int16_t>(pprOp.getRotationKind()) < 0);
+            return;
+        }
+
+        if (auto ppmOp = dyn_cast<PPMeasurementOp>(operation)) {
+            wrapper.updateSign(static_cast<int16_t>(ppmOp.getRotationSign()) < 0);
+        }
+    };
+
+    applySignFromOp(lhsPSWrapper, lhs);
+    applySignFromOp(rhsPSWrapper, rhs);
 
     return std::make_pair(std::move(lhsPSWrapper), std::move(rhsPSWrapper));
 }
@@ -211,10 +224,17 @@ void updatePauliWord(QECOpInterface op, const PauliWord &newPauliWord, PatternRe
 
 void updatePauliWordSign(QECOpInterface op, bool isNegated, PatternRewriter &rewriter)
 {
-    int16_t rotationKind = static_cast<int16_t>(op.getRotationKind());
-    int16_t sign = isNegated ? -1 : 1;
-    rotationKind = (rotationKind < 0 ? -rotationKind : rotationKind) * sign;
-    op.setRotationKind(rotationKind);
+    if (auto pprOp = dyn_cast<PPRotationOp>(op.getOperation())) {
+        int16_t rotationKind = static_cast<int16_t>(pprOp.getRotationKind());
+        int16_t sign = isNegated ? -1 : 1;
+        rotationKind = (rotationKind < 0 ? -rotationKind : rotationKind) * sign;
+        pprOp.setRotationKind(rotationKind);
+    }
+    else if (auto ppmOp = dyn_cast<PPMeasurementOp>(op.getOperation())) {
+        int16_t rotationSign = static_cast<int16_t>(ppmOp.getRotationSign());
+        rotationSign = (rotationSign < 0 ? -rotationSign : rotationSign) * (isNegated ? -1 : 1);
+        ppmOp.setRotationSign(rotationSign);
+    }
 }
 
 SmallVector<StringRef> extractPauliString(QECOpInterface op)