✨ Add an unroll-modifiers pass for unrolling multi-operation modifiers

mlir/include/mlir/Conversion/QCToQIR/QIRCommon/QIRCommon.h

@@ -73,7 +73,7 @@ struct LoweringState {
   DenseSet<Operation*> returnedStaticResults;
 
   /// Modifier information
-  size_t inCtrlOp = 0;
+  bool inCtrlOp = false;
   SmallVector<Value> controls;
 
   /// Allocator and StringSaver for stable StringRefs

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

@@ -34,4 +34,23 @@ def NormalizeGlobalPhases : Pass<"normalize-global-phases", "mlir::ModuleOp"> {
   }];
 }
 
+def UnrollModifiers : Pass<"unroll-modifiers", "mlir::ModuleOp"> {
+  let dependentDialects = ["mlir::qc::QCDialect", "mlir::qco::QCODialect"];
+  let summary = "Unroll multi-operation modifiers into single-operation ones";
+  let description = [{
+    Splits `ctrl` and `inv` modifiers whose body holds more than one unitary
+    operation into a sequence of modifiers that each hold a single operation.
+    The targets of the new modifiers are narrowed to the qubits of the
+    respective operation. For `inv`, the order of the operations is reversed.
+    Classical operations of the body are moved in front of the new modifiers.
+
+    If a modifier cannot be unrolled, the modifier is left untouched. Such
+    modifiers are skipped silently; the pass never fails.
+
+    In addition, `pow` modifiers are left untouched because `pow(r) { a; b }`
+    is generally not equivalent to `pow(r) { a }; pow(r) { b }`. Modifiers
+    nested in a `pow` body are still unrolled.
+  }];
+}
+
 #endif // MLIR_DIALECT_UTILS_TRANSFORMS_PASSES_TD

mlir/lib/Compiler/Programs.cpp

@@ -26,6 +26,7 @@
 #include "mlir/Dialect/QCO/Transforms/Passes.h"
 #include "mlir/Dialect/QTensor/IR/QTensorDialect.h"
 #include "mlir/Dialect/Utils/Transforms/GlobalPhaseNormalization.h"
+#include "mlir/Dialect/Utils/Transforms/Passes.h"
 #include "mlir/Support/Passes.h"
 
 #include <capnp/common.h>
@@ -333,6 +334,7 @@ std::optional<QIRProgram> QCProgram::intoQIR(const QIRProfile profile) && {
   if (failed(runPasses(
           mod(),
           [profile](OpPassManager& pm) {
+            pm.addPass(mqt::createUnrollModifiers());
             if (profile == QIRProfile::Adaptive) {
               pm.addPass(createQCToQIRAdaptive());
             } else {
@@ -467,7 +469,11 @@ std::optional<QCProgram> QCOProgram::intoQC() && {
 
 std::optional<JeffProgram> QCOProgram::intoJeff() && {
   if (failed(runPasses(
-          mod(), [](OpPassManager& pm) { pm.addPass(createQCOToJeff()); },
+          mod(),
+          [](OpPassManager& pm) {
+            pm.addPass(mqt::createUnrollModifiers());
+            pm.addPass(createQCOToJeff());
+          },
           "failed to convert QCO to jeff"))) {
     return std::nullopt;
   }

mlir/lib/Conversion/QCToQIR/QIRCommon/QIRCommon.cpp

@@ -97,9 +97,8 @@ convertUnitaryToCallOp(QCOpType& op, QCOpAdaptorType& adaptor,
                        LoweringState& state, StringRef fnName,
                        size_t numTargets, size_t numParams) {
   // Query state for modifier information
-  const auto inCtrlOp = state.inCtrlOp;
   const SmallVector<Value> controls =
-      inCtrlOp != 0 ? state.controls : SmallVector<Value>{};
+      state.inCtrlOp ? state.controls : SmallVector<Value>{};
   const size_t numCtrls = controls.size();
 
   // Define argument types
@@ -134,11 +133,9 @@ convertUnitaryToCallOp(QCOpType& op, QCOpAdaptorType& adaptor,
   operands.append(adaptor.getOperands().begin(), adaptor.getOperands().end());
 
   // Clean up modifier information
-  if (inCtrlOp != 0) {
-    state.inCtrlOp--;
-    if (state.inCtrlOp == 0) {
-      state.controls.clear();
-    }
+  if (state.inCtrlOp) {
+    state.inCtrlOp = false;
+    state.controls.clear();
   }
 
   // Replace operation with CallOp
@@ -242,8 +239,7 @@ struct ConvertQCUnitaryOpQIR : StatefulOpConversionPattern<OpType> {
   matchAndRewrite(OpType op, OpType::Adaptor adaptor,
                   ConversionPatternRewriter& rewriter) const override {
     auto& state = this->getState();
-    const auto inCtrlOp = state.inCtrlOp;
-    const size_t numCtrls = inCtrlOp != 0 ? state.controls.size() : 0;
+    const size_t numCtrls = state.inCtrlOp ? state.controls.size() : 0;
     const auto fnName = GetFnName(numCtrls);
     return convertUnitaryToCallOp(op, adaptor, rewriter, this->getContext(),
                                   state, fnName, NumTargets, NumParams);
@@ -325,7 +321,7 @@ struct ConvertQCGPhaseOp final : StatefulOpConversionPattern<GPhaseOp> {
   matchAndRewrite(GPhaseOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter& rewriter) const override {
     auto& state = getState();
-    if (state.inCtrlOp != 0) {
+    if (state.inCtrlOp) {
       return op.emitError("Controlled GPhaseOps cannot be converted to QIR");
     }
     return convertUnitaryToCallOp(op, adaptor, rewriter, getContext(), state,
@@ -360,13 +356,19 @@ struct ConvertQCCtrlOp final : StatefulOpConversionPattern<CtrlOp> {
                   ConversionPatternRewriter& rewriter) const override {
     auto& state = getState();
 
-    if (state.inCtrlOp != 0) {
+    if (state.inCtrlOp) {
       return rewriter.notifyMatchFailure(op,
                                          "Nested CtrlOps are not supported");
     }
 
+    if (op.getNumBodyUnitaries() > 1) {
+      return rewriter.notifyMatchFailure(
+          op, "CtrlOps with multiple body unitaries are not supported. Run the "
+              "unroll-modifiers pass before the conversion");
+    }
+
     // Update modifier information
-    state.inCtrlOp = op.getNumBodyUnitaries();
+    state.inCtrlOp = true;
     state.controls = llvm::to_vector(adaptor.getControls());
 
     // Inline block and remove operation

mlir/lib/Dialect/QC/IR/Modifiers/ModifierUtils.cpp

@@ -28,8 +28,8 @@ namespace mlir::qc::detail {
 LogicalResult verifyModifierBody(Operation* modifierOp, Block& body) {
   const auto hasNonUnitaryOperation =
       body.walk([](Operation* operation) {
-            return isa<AllocOp, DeallocOp, MeasureOp, ResetOp, memref::LoadOp,
-                       memref::StoreOp>(operation)
+            return isa<AllocOp, DeallocOp, StaticOp, MeasureOp, ResetOp,
+                       memref::LoadOp, memref::StoreOp>(operation)
                        ? WalkResult::interrupt()
                        : WalkResult::advance();
           })

mlir/lib/Dialect/QCO/IR/Modifiers/CtrlOp.cpp

@@ -272,8 +272,8 @@ void CtrlOp::build(OpBuilder& odsBuilder, OperationState& odsState,
 LogicalResult CtrlOp::verify() {
   auto& block = *getBody();
   if (llvm::any_of(block, [](Operation& op) {
-        return isa<AllocOp, SinkOp, MeasureOp, ResetOp, qtensor::ExtractOp,
-                   qtensor::InsertOp>(op);
+        return isa<AllocOp, SinkOp, StaticOp, MeasureOp, ResetOp,
+                   qtensor::ExtractOp, qtensor::InsertOp>(op);
       })) {
     return emitOpError("body must not contain non-unitary quantum operations "
                        "or modify a quantum register");

mlir/lib/Dialect/QCO/IR/Modifiers/InvOp.cpp

@@ -442,8 +442,8 @@ void InvOp::build(OpBuilder& odsBuilder, OperationState& odsState, Value qubit,
 LogicalResult InvOp::verify() {
   auto& block = *getBody();
   if (llvm::any_of(block, [](Operation& op) {
-        return isa<AllocOp, SinkOp, MeasureOp, ResetOp, qtensor::ExtractOp,
-                   qtensor::InsertOp>(op);
+        return isa<AllocOp, SinkOp, StaticOp, MeasureOp, ResetOp,
+                   qtensor::ExtractOp, qtensor::InsertOp>(op);
       })) {
     return emitOpError("body must not contain non-unitary quantum operations "
                        "or modify a quantum register");

mlir/lib/Dialect/QCO/IR/Modifiers/PowOp.cpp

@@ -670,8 +670,8 @@ void PowOp::build(OpBuilder& odsBuilder, OperationState& odsState, Value qubit,
 LogicalResult PowOp::verify() {
   auto& block = *getBody();
   if (llvm::any_of(block, [](Operation& op) {
-        return isa<AllocOp, SinkOp, MeasureOp, ResetOp, qtensor::ExtractOp,
-                   qtensor::InsertOp>(op);
+        return isa<AllocOp, SinkOp, StaticOp, MeasureOp, ResetOp,
+                   qtensor::ExtractOp, qtensor::InsertOp>(op);
       })) {
     return emitOpError("body must not contain non-unitary quantum operations "
                        "or modify a quantum register");

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

@@ -9,6 +9,7 @@
 add_mlir_library(
   MLIRMQTTransforms
   NormalizeGlobalPhases.cpp
+  UnrollModifiers.cpp
   ADDITIONAL_HEADER_DIRS
   ${MQT_MLIR_SOURCE_INCLUDE_DIR}/mlir/Dialect/Utils/Transforms
   LINK_LIBS