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

CHANGELOG.md

@@ -12,6 +12,8 @@ releases may include breaking changes.
 
 ### Added
 
+- ✨ Add an `unroll-modifiers` pass for unrolling multi-operation modifiers
+  ([#2015]) ([**@denialhaag**])
 - ✨ Add generic C++ and Python FoMaC support for custom device properties that
   contain operation handles ([#2042]) ([**@burgholzer**])
 - ✨ Support retrieving existing jobs by ID through the QDMI client API and C++
@@ -742,6 +744,7 @@ for previous changelogs._
 [#2028]: https://github.com/munich-quantum-toolkit/core/pull/2028
 [#2026]: https://github.com/munich-quantum-toolkit/core/pull/2026
 [#2017]: https://github.com/munich-quantum-toolkit/core/pull/2017
+[#2015]: https://github.com/munich-quantum-toolkit/core/pull/2015
 [#2016]: https://github.com/munich-quantum-toolkit/core/pull/2016
 [#2014]: https://github.com/munich-quantum-toolkit/core/pull/2014
 [#2011]: https://github.com/munich-quantum-toolkit/core/pull/2011

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/QC/IR/QCInterfaces.h

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <mlir/IR/OpDefinition.h>
 #include <mlir/Support/LLVM.h>
 
 #include <cstddef>

mlir/include/mlir/Dialect/QC/IR/QCOps.td

@@ -1021,8 +1021,7 @@ def InvOp : QCOp<"inv",
   let regions = (region SizedRegion<1>:$region);
   let assemblyFormat = [{
         custom<TargetAliasing>($region, $qubits)
-        attr-dict `:`
-        type($qubits)
+        attr-dict ( `:` type($qubits)^ )?
     }];
 
   let extraClassDeclaration = [{

mlir/include/mlir/Dialect/QCO/IR/QCOOps.td

@@ -1120,6 +1120,9 @@ def CtrlOp : QCOOp<"ctrl",
         Control qubits are not modified by the operation and simply pass through to the outputs.
 
         The body region may contain an arbitrary amount of unitary and classical operations.
+        Classical SSA values may be captured from above, but every qubit used
+        in the body must be passed through a modifier operand and accessed via
+        its aliased block argument.
         Non-unitary operations, such as `AllocOp` and `MeasureOp`, are not allowed.
 
         Example:
@@ -1202,6 +1205,9 @@ def InvOp : QCOOp<"inv", traits = [UnitaryOpInterface,
         The operation takes a variadic number of qubits as inputs and produces corresponding output qubits.
 
         The body region may contain an arbitrary amount of unitary and classical operations.
+        Classical SSA values may be captured from above, but every qubit used
+        in the body must be passed through a modifier operand and accessed via
+        its aliased block argument.
         Non-unitary operations, such as `AllocOp` and `MeasureOp`, are not allowed.
 
         Example:
@@ -1219,10 +1225,8 @@ def InvOp : QCOOp<"inv", traits = [UnitaryOpInterface,
   let regions = (region SizedRegion<1>:$region);
   let assemblyFormat = [{
         custom<TargetAliasing>($region, $qubits_in)
-        attr-dict `:`
-        `{` type($qubits_in) `}`
-        `->`
-        `{` type($qubits_out) `}`
+        attr-dict
+        ( `:` `{` type($qubits_in)^ `}` `->` `{` type($qubits_out) `}` )?
     }];
 
   let extraClassDeclaration = [{
@@ -1284,6 +1288,9 @@ def PowOp
         - r < 0: equivalent to inv @ pow(-r) @ g.
 
         The body region may contain an arbitrary amount of unitary and classical operations.
+        Classical SSA values may be captured from above, but every qubit used
+        in the body must be passed through a modifier operand and accessed via
+        its aliased block argument.
         Non-unitary operations, such as `AllocOp` and `MeasureOp`, are not allowed.
 
         Example:
@@ -1302,10 +1309,8 @@ def PowOp
   let regions = (region SizedRegion<1>:$region);
   let assemblyFormat = [{
         `(` $exponent `)` custom<TargetAliasing>($region, $qubits_in)
-        attr-dict `:`
-        `{` type($qubits_in) `}`
-        `->`
-        `{` type($qubits_out) `}`
+        attr-dict
+        ( `:` `{` type($qubits_in)^ `}` `->` `{` type($qubits_out) `}` )?
     }];
 
   let extraClassDeclaration = [{

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

@@ -34,4 +34,25 @@ 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.
+    For `inv`, the order of the operations is reversed. Classical operations of
+    the body are moved in front of the new modifiers.
+
+    A `pow` modifier is only unrolled if its exponent is a compile-time known
+    integer and its operations act on disjoint qubits. Otherwise, `pow(r) { a;
+    b }` is generally not equivalent to `pow(r) { a }; pow(r) { b }`: the
+    operations do not commute if they share a qubit, and for a non-integer
+    exponent, the principal branch of the matrix power does not distribute over
+    them. Modifiers nested in a `pow` body are unrolled either way.
+
+    If a modifier cannot be unrolled, the modifier is left untouched. Such
+    modifiers are skipped silently; the pass never fails.
+  }];
+}
+
 #endif // MLIR_DIALECT_UTILS_TRANSFORMS_PASSES_TD

mlir/include/mlir/Dialect/Utils/Utils.h

@@ -13,6 +13,7 @@
 #include <llvm/ADT/DenseMap.h>
 #include <llvm/ADT/STLExtras.h>
 #include <llvm/ADT/SmallVector.h>
+#include <llvm/ADT/SmallVectorExtras.h>
 #include <llvm/ADT/StringRef.h>
 #include <llvm/Support/ErrorHandling.h>
 #include <mlir/Dialect/Arith/IR/Arith.h>
@@ -447,7 +448,12 @@ inline void inlineModifierBody(Operation* op, Block& body,
                                ValueRange blockArgReplacements,
                                RewriterBase& rewriter) {
   auto* terminator = body.getTerminator();
-  const SmallVector<Value> results(terminator->getOperands());
+  // Yielded block arguments are substituted when the body is inlined, so
+  // resolve them to their replacements before they are erased.
+  const auto results =
+      llvm::map_to_vector(terminator->getOperands(), [&](Value yielded) {
+        return getValueFromBlockArgument(yielded, blockArgReplacements);
+      });
   rewriter.inlineBlockBefore(&body, op, blockArgReplacements);
   rewriter.eraseOp(terminator);
   rewriter.replaceOp(op, results);

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,20 +97,17 @@ convertUnitaryToCallOp(QCOpType& op, QCOpAdaptorType& adaptor,
                        LoweringState& state, StringRef fnName,
                        const size_t numTargets, const 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 auto convertedOperands = adaptor.getOperands();
   const auto targets = convertedOperands.take_front(numTargets);
   const auto parameters = convertedOperands.drop_front(numTargets);
   assert(parameters.size() == numParams && "unexpected gate parameter count");
 
   // 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();
   }
 
   qir::emitQISCall(rewriter, op, op.getLoc(), parameters, controls, targets,
@@ -216,8 +213,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, state, fnName,
                                   NumTargets, NumParams);
@@ -299,7 +295,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, state, QIR_GPHASE, 0,
@@ -334,14 +330,27 @@ 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");
     }
 
-    // Update modifier information
-    state.inCtrlOp = op.getNumBodyUnitaries();
-    state.controls = llvm::to_vector(adaptor.getControls());
+    if (op.getNumBodyUnitaries() > 1) {
+      return rewriter.notifyMatchFailure(
+          op, "CtrlOps with multiple body unitaries are not supported. Run the "
+              "unroll-modifiers pass before the conversion");
+    }
+
+    // Empty control bodies and controls around no-op unitaries do not need
+    // lowering state. In particular, barrier lowering erases the operation
+    // without consuming that state, which would otherwise control the next
+    // gate.
+    auto bodyUnitary = op.getNumBodyUnitaries() == 1 ? op.getBodyUnitary(0)
+                                                     : UnitaryOpInterface{};
+    if (bodyUnitary && !isa<BarrierOp, IdOp>(bodyUnitary.getOperation())) {
+      state.inCtrlOp = true;
+      state.controls = llvm::to_vector(adaptor.getControls());
+    }
 
     // Inline block and remove operation
     rewriter.inlineBlockBefore(&op.getRegion().front(), op,

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

@@ -166,6 +166,32 @@ struct EraseEmptyCtrl final : OpRewritePattern<CtrlOp> {
   }
 };
 
+/**
+ * @brief Drop the target qubits that the body does not use.
+ */
+struct DropUnusedTargets final : OpRewritePattern<CtrlOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CtrlOp op,
+                                PatternRewriter& rewriter) const override {
+    auto* body = op.getBody();
+    const auto used = qc::detail::getUsedQubitIndices(*body);
+    if (used.size() == op.getNumTargets()) {
+      return failure();
+    }
+
+    const auto targets = llvm::map_to_vector(
+        used, [&](const size_t index) { return op.getTargets()[index]; });
+    CtrlOp::create(rewriter, op.getLoc(), op.getControls(), targets,
+                   [&](ValueRange args) {
+                     qc::detail::inlineNarrowedBody(*body, op.getTargets(),
+                                                    used, args, rewriter);
+                   });
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 } // namespace
 
 size_t CtrlOp::getNumBodyUnitaries() {
@@ -228,5 +254,6 @@ LogicalResult CtrlOp::verify() {
 
 void CtrlOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                          MLIRContext* context) {
-  results.add<MergeNestedCtrl, ReduceCtrl, EraseEmptyCtrl>(context);
+  results.add<MergeNestedCtrl, ReduceCtrl, EraseEmptyCtrl, DropUnusedTargets>(
+      context);
 }

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

@@ -346,6 +346,29 @@ struct EraseEmptyInv final : OpRewritePattern<InvOp> {
   }
 };
 
+/**
+ * @brief Drop the qubits that the body does not use.
+ */
+struct DropUnusedInvQubits final : OpRewritePattern<InvOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(InvOp op,
+                                PatternRewriter& rewriter) const override {
+    auto* body = op.getBody();
+    const auto qubits = op.getQubits();
+    return qc::detail::dropUnusedQubits(
+        op, *body, qubits,
+        [&](ValueRange narrowedQubits, ArrayRef<size_t> used) {
+          InvOp::create(rewriter, op.getLoc(), narrowedQubits,
+                        [&](ValueRange args) {
+                          qc::detail::inlineNarrowedBody(*body, qubits, used,
+                                                         args, rewriter);
+                        });
+        },
+        rewriter);
+  }
+};
+
 } // namespace
 
 size_t InvOp::getNumBodyUnitaries() {
@@ -386,5 +409,6 @@ LogicalResult InvOp::verify() {
 void InvOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                         MLIRContext* context) {
   results.add<CancelNestedInv, MoveCtrlOutsideInv, InvPowToNegPow,
-              InlineSelfAdjoint, ReplaceWithKnownGates, EraseEmptyInv>(context);
+              InlineSelfAdjoint, ReplaceWithKnownGates, EraseEmptyInv,
+              DropUnusedInvQubits>(context);
 }