🐛 Fix one-time-use of variables for the ctrl modifier

CHANGELOG.md

@@ -11,7 +11,7 @@ This project adheres to [Semantic Versioning], with the exception that minor rel
 
 ### Added
 
-- ✨ Add initial infrastructure for new QC and QCO MLIR dialects ([#1264], [#1402], [#1428], [#1430]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**])
+- ✨ Add initial infrastructure for new QC and QCO MLIR dialects ([#1264], [#1402], [#1428], [#1430], [#1436]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**])
 
 ### Changed
 
@@ -308,6 +308,7 @@ _📚 Refer to the [GitHub Release Notes](https://github.com/munich-quantum-tool
 <!-- PR links -->
 
 [#1437]: https://github.com/munich-quantum-toolkit/core/pull/1437
+[#1436]: https://github.com/munich-quantum-toolkit/core/pull/1436
 [#1430]: https://github.com/munich-quantum-toolkit/core/pull/1430
 [#1428]: https://github.com/munich-quantum-toolkit/core/pull/1428
 [#1415]: https://github.com/munich-quantum-toolkit/core/pull/1415

mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h

@@ -15,6 +15,7 @@
 #include <cstdint>
 #include <functional>
 #include <llvm/ADT/DenseSet.h>
+#include <llvm/ADT/STLFunctionalExtras.h>
 #include <llvm/ADT/SmallVector.h>
 #include <llvm/Support/ErrorHandling.h>
 #include <mlir/IR/Builders.h>
@@ -988,21 +989,21 @@ class QCOProgramBuilder final : public OpBuilder {
    * @par Example:
    * ```c++
    * {controls_out, targets_out} =
-   *   builder.ctrl(q0_in, q1_in, [&](ValueRange targets) {
-   *     auto q1_res = builder.x(targets[0]);
-   *     return {q1_res};
+   *   builder.ctrl(q0_in, q1_in,
+   *     [&](ValueRange targets) -> llvm::SmallVector<Value> {
+   *       return {builder.x(targets[0])};
    *   });
    * ```
    * ```mlir
-   * %controls_out, %targets_out = qco.ctrl(%q0_in) %q1_in {
-   *   %q1_res = qco.x %q1_in : !qco.qubit -> !qco.qubit
+   * %controls_out, %targets_out = qco.ctrl(%q0_in) targets(%t = %q1_in) {
+   *   %q1_res = qco.x %t : !qco.qubit -> !qco.qubit
    *   qco.yield %q1_res
    * } : ({!qco.qubit}, {!qco.qubit}) -> ({!qco.qubit}, {!qco.qubit})
    * ```
    */
   std::pair<ValueRange, ValueRange>
   ctrl(ValueRange controls, ValueRange targets,
-       const std::function<ValueRange(ValueRange)>& body);
+       llvm::function_ref<llvm::SmallVector<Value>(ValueRange)> body);
 
   //===--------------------------------------------------------------------===//
   // Deallocation

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

@@ -1025,10 +1025,10 @@ def YieldOp : QCOOp<"yield", traits = [Terminator]> {
 
         Example:
         ```mlir
-        %ctrl_q_out, %tgt_q_out = qco.ctrl(%ctrl_q_in) %tgt_q_in {
-            %tgt_q_res = qco.h %tgt_q_in : !qco.qubit -> !qco.qubit
-            qco.yield %tgt_q_res : !qco.qubit
-        } : ({!qco.qubit}, {!qco.qubit}) -> ({!qco.qubit}, {!qco.qubit})
+        %res_ctrl, %res_tgt:2 = qco.ctrl(%ctrl) targets(%a0 = %q0, %a1 = %q1) {
+            %a0_1, %a1_1 = qco.swap %a0, %a1 : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit
+            qco.yield %a0_1, %a1_1
+        } : ({!qco.qubit}, {!qco.qubit, !qco.qubit}) -> ({!qco.qubit}, {!qco.qubit, !qco.qubit})
         ```
     }];
 
@@ -1057,10 +1057,10 @@ def CtrlOp : QCOOp<"ctrl", traits =
 
         Example:
         ```mlir
-        %ctrl_q_out, %tgt_q_out = qco.ctrl(%ctrl_q_in) %tgt_q_in {
-            %tgt_q_res = qco.h %tgt_q_in : !qco.qubit -> !qco.qubit
-            qco.yield %tgt_q_res : !qco.qubit
-        } : ({!qco.qubit}, {!qco.qubit}) -> ({!qco.qubit}, {!qco.qubit})
+        %res_ctrl, %res_tgt:2 = qco.ctrl(%ctrl) targets(%a0 = %q0, %a1 = %q1) {
+            %a0_1, %a1_1 = qco.swap %a0, %a1 : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit
+            qco.yield %a0_1, %a1_1
+        } : ({!qco.qubit}, {!qco.qubit, !qco.qubit}) -> ({!qco.qubit}, {!qco.qubit, !qco.qubit})
         ```
     }];
 
@@ -1069,8 +1069,10 @@ def CtrlOp : QCOOp<"ctrl", traits =
     let results = (outs Variadic<QubitType>:$controls_out, Variadic<QubitType>:$targets_out);
     let regions = (region SizedRegion<1>:$region);
     let assemblyFormat = [{
-        `(` $controls_in `)` $targets_in
-        $region attr-dict `:`
+        `(` $controls_in `)`
+        `targets`
+        custom<TargetAliasing>($region, $targets_in)
+        attr-dict `:`
         `(` `{` type($controls_in) `}` ( `,` `{` type($targets_in)^ `}` )? `)`
         `->`
         `(` `{` type($controls_out) `}` ( `,` `{` type($targets_out)^ `}` )? `)`
@@ -1099,7 +1101,7 @@ def CtrlOp : QCOOp<"ctrl", traits =
            build($_builder, $_state, controls.getTypes(), targets.getTypes(), controls, targets);
         }]>,
         OpBuilder<(ins "ValueRange":$controls, "ValueRange":$targets, "UnitaryOpInterface":$bodyUnitary)>,
-        OpBuilder<(ins "ValueRange":$controls, "ValueRange":$targets, "const std::function<ValueRange(ValueRange)>&":$bodyBuilder)>
+        OpBuilder<(ins "ValueRange":$controls, "ValueRange":$targets, "llvm::function_ref<llvm::SmallVector<Value>(ValueRange)>":$bodyBuilder)>
     ];
 
     let hasCanonicalizer = 1;

mlir/lib/Conversion/QCOToQC/QCOToQC.cpp

@@ -774,12 +774,37 @@ struct ConvertQCOCtrlOp final : OpConversionPattern<qco::CtrlOp> {
     const auto& qcControls = adaptor.getControlsIn();
 
     // Create qc.ctrl operation
-    auto qcoOp = rewriter.create<qc::CtrlOp>(op.getLoc(), qcControls);
+    auto qcOp = qc::CtrlOp::create(rewriter, op.getLoc(), qcControls);
 
     // Clone body region from QCO to QC
-    auto& dstRegion = qcoOp.getRegion();
+    auto& dstRegion = qcOp.getRegion();
     rewriter.cloneRegionBefore(op.getRegion(), dstRegion, dstRegion.end());
 
+    auto& entryBlock = dstRegion.front();
+    const auto numArgs = entryBlock.getNumArguments();
+    if (adaptor.getTargetsIn().size() != numArgs) {
+      return op.emitOpError() << "qco.ctrl: entry block args (" << numArgs
+                              << ") must match number of target operands ("
+                              << adaptor.getTargetsIn().size() << ")";
+    }
+
+    // Remove all block arguments in the cloned region
+    rewriter.modifyOpInPlace(qcOp, [&] {
+      // 1. Replace uses (Must be done BEFORE erasing)
+      // We iterate 0..N using indices since the block args are still stable
+      // here.
+      for (auto i = 0UL; i < numArgs; ++i) {
+        entryBlock.getArgument(i).replaceAllUsesWith(adaptor.getTargetsIn()[i]);
+      }
+
+      // 2. Erase all block arguments
+      // Now that they have no uses, we can safely wipe them.
+      // We use a bulk erase for efficiency (start index 0, count N).
+      if (numArgs > 0) {
+        entryBlock.eraseArguments(0, numArgs);
+      }
+    });
+
     // Replace the output qubits with the same QC references
     rewriter.replaceOp(op, adaptor.getOperands());
 

mlir/lib/Conversion/QCToQCO/QCToQCO.cpp

@@ -1106,7 +1106,7 @@ struct ConvertQCCtrlOp final : StatefulOpConversionPattern<qc::CtrlOp> {
 
     // Create qco.ctrl
     auto qcoOp =
-        rewriter.create<qco::CtrlOp>(op.getLoc(), qcoControls, qcoTargets);
+        qco::CtrlOp::create(rewriter, op.getLoc(), qcoControls, qcoTargets);
 
     // Update the state map if this is a top-level CtrlOp
     // Nested CtrlOps are managed via the targetsIn and targetsOut maps
@@ -1124,12 +1124,27 @@ struct ConvertQCCtrlOp final : StatefulOpConversionPattern<qc::CtrlOp> {
 
     // Update modifier information
     state.inCtrlOp++;
-    state.targetsIn.try_emplace(state.inCtrlOp, qcoTargets);
 
     // Clone body region from QC to QCO
     auto& dstRegion = qcoOp.getRegion();
     rewriter.cloneRegionBefore(op.getRegion(), dstRegion, dstRegion.end());
 
+    // Create block arguments for target qubits and store them in
+    // `state.targetsIn`.
+    auto& entryBlock = dstRegion.front();
+    assert(entryBlock.getNumArguments() == 0 &&
+           "QC ctrl region unexpectedly has entry block arguments");
+    SmallVector<Value> qcoTargetAliases;
+    qcoTargetAliases.reserve(numTargets);
+    const auto qubitType = qco::QubitType::get(qcoOp.getContext());
+    const auto opLoc = op.getLoc();
+    rewriter.modifyOpInPlace(qcoOp, [&] {
+      for (auto i = 0UL; i < numTargets; i++) {
+        qcoTargetAliases.emplace_back(entryBlock.addArgument(qubitType, opLoc));
+      }
+    });
+    state.targetsIn[state.inCtrlOp] = std::move(qcoTargetAliases);
+
     rewriter.eraseOp(op);
     return success();
   }