✨ 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,21 @@ 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 modifier; a
+    modifier is left untouched if one of them depends on the body's qubits.
+
+    `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/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

mlir/lib/Dialect/Utils/Transforms/UnrollModifiers.cpp

@@ -0,0 +1,261 @@
+/*
+ * 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
+ */
+
+#include "mlir/Dialect/QC/IR/QCOps.h"
+#include "mlir/Dialect/QCO/IR/QCOOps.h"
+#include "mlir/Dialect/Utils/Transforms/Passes.h"
+#include "mlir/Dialect/Utils/Utils.h"
+
+#include <llvm/ADT/STLExtras.h>
+#include <llvm/ADT/SmallVector.h>
+#include <llvm/ADT/SmallVectorExtras.h>
+#include <llvm/ADT/TypeSwitch.h>
+#include <mlir/IR/Block.h>
+#include <mlir/IR/IRMapping.h>
+#include <mlir/IR/OpDefinition.h>
+#include <mlir/IR/Operation.h>
+#include <mlir/IR/PatternMatch.h>
+#include <mlir/IR/Value.h>
+#include <mlir/Interfaces/SideEffectInterfaces.h>
+#include <mlir/Support/LLVM.h>
+#include <mlir/Support/LogicalResult.h>
+
+namespace mlir::mqt {
+
+#define GEN_PASS_DEF_UNROLLMODIFIERS
+#include "mlir/Dialect/Utils/Transforms/Passes.h.inc"
+
+namespace {
+
+/// Return the unitary operations in @p body.
+template <typename UnitaryOpInterface>
+SmallVector<UnitaryOpInterface> getBodyUnitaries(Block& body) {
+  return llvm::to_vector(body.getOps<UnitaryOpInterface>());
+}
+
+/// Return the distinct qubit operands of @p op in operand order.
+template <typename QubitType>
+SmallVector<Value> getQubitOperands(Operation* op) {
+  SmallVector<Value> qubits;
+  for (auto operand : op->getOperands()) {
+    if (isa<QubitType>(operand.getType()) &&
+        !llvm::is_contained(qubits, operand)) {
+      qubits.push_back(operand);
+    }
+  }
+  return qubits;
+}
+
+/// Move the classical operations of @p body in front of @p modifier.
+///
+/// Fails if a classical operation is impure or depends on values defined in
+/// @p body.
+template <typename UnitaryOpInterface>
+LogicalResult hoistClassicalOps(Block& body, Operation* modifier,
+                                RewriterBase& rewriter) {
+  const auto isClassical = [](Operation& op) {
+    return !isa<UnitaryOpInterface>(op) &&
+           !op.hasTrait<OpTrait::IsTerminator>();
+  };
+  for (auto& op : body) {
+    if (isClassical(op) &&
+        (!isPure(&op) || llvm::any_of(op.getOperands(), [&](Value operand) {
+          return operand.getParentBlock() == &body;
+        }))) {
+      return failure();
+    }
+  }
+  for (auto& op : llvm::make_early_inc_range(body)) {
+    if (isClassical(op)) {
+      rewriter.moveOpBefore(&op, modifier);
+    }
+  }
+  return success();
+}
+
+/// Clone @p unitary into the body of a new modifier, replacing its qubit
+/// operands @p qubits with the block arguments @p args, and return its results.
+SmallVector<Value> cloneIntoBody(Operation* unitary, ValueRange qubits,
+                                 ValueRange args, RewriterBase& rewriter) {
+  IRMapping mapping;
+  mapping.map(qubits, args);
+  auto results = rewriter.clone(*unitary, mapping)->getResults();
+  return {results.begin(), results.end()};
+}
+
+//===----------------------------------------------------------------------===//
+// QC
+//===----------------------------------------------------------------------===//
+
+/// Unroll a `qc.ctrl` modifier with more than one body unitary.
+LogicalResult unrollModifier(qc::CtrlOp op, RewriterBase& rewriter) {
+  if (op.getNumBodyUnitaries() < 2) {
+    return failure();
+  }
+  auto* body = op.getBody();
+  if (failed(hoistClassicalOps<qc::UnitaryOpInterface>(*body, op, rewriter))) {
+    return failure();
+  }
+
+  rewriter.setInsertionPoint(op);
+  for (auto unitary : getBodyUnitaries<qc::UnitaryOpInterface>(*body)) {
+    const auto qubits = getQubitOperands<qc::QubitType>(unitary);
+    const auto targets = llvm::map_to_vector(qubits, [&](Value qubit) {
+      return utils::getValueFromBlockArgument(qubit, op.getTargets());
+    });
+    qc::CtrlOp::create(rewriter, op.getLoc(), op.getControls(), targets,
+                       [&](ValueRange args) {
+                         cloneIntoBody(unitary, qubits, args, rewriter);
+                       });
+  }
+  rewriter.eraseOp(op);
+  return success();
+}
+
+/// Unroll a `qc.inv` modifier with more than one body unitary.
+LogicalResult unrollModifier(qc::InvOp op, RewriterBase& rewriter) {
+  if (op.getNumBodyUnitaries() < 2) {
+    return failure();
+  }
+  auto* body = op.getBody();
+  if (failed(hoistClassicalOps<qc::UnitaryOpInterface>(*body, op, rewriter))) {
+    return failure();
+  }
+
+  rewriter.setInsertionPoint(op);
+  // (a b)^-1 = b^-1 a^-1, so the operations are inverted in reverse order.
+  auto unitaries = getBodyUnitaries<qc::UnitaryOpInterface>(*body);
+  for (auto unitary : llvm::reverse(unitaries)) {
+    const auto qubits = getQubitOperands<qc::QubitType>(unitary);
+    const auto targets = llvm::map_to_vector(qubits, [&](Value qubit) {
+      return utils::getValueFromBlockArgument(qubit, op.getQubits());
+    });
+    qc::InvOp::create(rewriter, op.getLoc(), targets, [&](ValueRange args) {
+      cloneIntoBody(unitary, qubits, args, rewriter);
+    });
+  }
+  rewriter.eraseOp(op);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// QCO
+//===----------------------------------------------------------------------===//
+
+/// Check that every unitary operation in @p body threads its qubit operands to
+/// its results, which is required to rewire the unrolled modifiers.
+bool hasThreadedBodyUnitaries(Block& body) {
+  return llvm::all_of(body.getOps<qco::UnitaryOpInterface>(),
+                      [](qco::UnitaryOpInterface unitary) {
+                        return unitary->getNumResults() ==
+                               getQubitOperands<qco::QubitType>(unitary).size();
+                      });
+}
+
+/// Unroll a `qco.ctrl` modifier with more than one body unitary.
+LogicalResult unrollModifier(qco::CtrlOp op, RewriterBase& rewriter) {
+  auto* body = op.getBody();
+  if (op.getNumBodyUnitaries() < 2 || !hasThreadedBodyUnitaries(*body)) {
+    return failure();
+  }
+  if (failed(hoistClassicalOps<qco::UnitaryOpInterface>(*body, op, rewriter))) {
+    return failure();
+  }
+
+  rewriter.setInsertionPoint(op);
+  // Maps the qubits of the body to the qubits threaded through the new
+  // modifiers.
+  IRMapping qubits;
+  qubits.map(body->getArguments(), op.getTargetsIn());
+
+  SmallVector<Value> controls(op.getControlsIn());
+  for (auto unitary : getBodyUnitaries<qco::UnitaryOpInterface>(*body)) {
+    const auto operands = getQubitOperands<qco::QubitType>(unitary);
+    const auto targets = llvm::map_to_vector(
+        operands, [&](Value qubit) { return qubits.lookup(qubit); });
+    auto ctrlOp = qco::CtrlOp::create(
+        rewriter, op.getLoc(), controls, targets,
+        [&](ValueRange args) -> SmallVector<Value> {
+          return cloneIntoBody(unitary, operands, args, rewriter);
+        });
+    auto controlsOut = ctrlOp.getControlsOut();
+    controls.assign(controlsOut.begin(), controlsOut.end());
+    qubits.map(unitary->getResults(), ctrlOp.getTargetsOut());
+  }
+
+  SmallVector<Value> results(controls);
+  for (auto yielded : body->getTerminator()->getOperands()) {
+    results.push_back(qubits.lookup(yielded));
+  }
+  rewriter.replaceOp(op, results);
+  return success();
+}
+
+/// Unroll a `qco.inv` modifier with more than one body unitary.
+LogicalResult unrollModifier(qco::InvOp op, RewriterBase& rewriter) {
+  auto* body = op.getBody();
+  if (op.getNumBodyUnitaries() < 2 || !hasThreadedBodyUnitaries(*body)) {
+    return failure();
+  }
+  if (failed(hoistClassicalOps<qco::UnitaryOpInterface>(*body, op, rewriter))) {
+    return failure();
+  }
+
+  rewriter.setInsertionPoint(op);
+  // (a b)^-1 = b^-1 a^-1, so the operations are inverted in reverse order.
+  // Consequently, the inputs of the modifier feed the qubits that its body
+  // yields.
+  IRMapping qubits;
+  qubits.map(body->getTerminator()->getOperands(), op.getQubitsIn());
+
+  auto unitaries = getBodyUnitaries<qco::UnitaryOpInterface>(*body);
+  for (auto unitary : llvm::reverse(unitaries)) {
+    const auto operands = getQubitOperands<qco::QubitType>(unitary);
+    const auto targets =
+        llvm::map_to_vector(unitary->getResults(),
+                            [&](Value qubit) { return qubits.lookup(qubit); });
+    auto invOp = qco::InvOp::create(rewriter, op.getLoc(), targets,
+                                    [&](ValueRange args) -> SmallVector<Value> {
+                                      return cloneIntoBody(unitary, operands,
+                                                           args, rewriter);
+                                    });
+    qubits.map(operands, invOp.getResults());
+  }
+
+  rewriter.replaceOp(
+      op, llvm::map_to_vector(body->getArguments(),
+                              [&](Value arg) { return qubits.lookup(arg); }));
+  return success();
+}
+
+struct UnrollModifiers final : impl::UnrollModifiersBase<UnrollModifiers> {
+protected:
+  void runOnOperation() override {
+    SmallVector<Operation*> modifiers;
+    getOperation()->walk([&](Operation* op) {
+      if (isa<qc::CtrlOp, qc::InvOp, qco::CtrlOp, qco::InvOp>(op)) {
+        modifiers.push_back(op);
+      }
+    });
+
+    // The walk visits nested modifiers before their parents, so unrolling the
+    // collected modifiers in order reaches a fixpoint in a single sweep.
+    IRRewriter rewriter(&getContext());
+    for (auto* modifier : modifiers) {
+      llvm::TypeSwitch<Operation*>(modifier)
+          .Case<qc::CtrlOp, qc::InvOp, qco::CtrlOp, qco::InvOp>([&](auto op) {
+            static_cast<void>(unrollModifier(op, rewriter));
+          });
+    }
+  }
+};
+
+} // namespace
+} // namespace mlir::mqt