diff --git a/CHANGELOG.md b/CHANGELOG.md index ec41c4e172..e1255dd91a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,9 +44,9 @@ releases may include breaking changes. [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1710], [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765], [#1780], [#1781], [#1782], [#1787], [#1806], [#1807], [#1808], [#1823], [#1824], - [#1830], [#1869]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**], - [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], - [**@simon1hofmann**]) + [#1830], [#1869], [#1872]) ([**@burgholzer**], [**@denialhaag**], + [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], + [**@MatthiasReumann**], [**@simon1hofmann**]) ### Changed @@ -622,6 +622,7 @@ changelogs._ [#1873]: https://github.com/munich-quantum-toolkit/core/pull/1873 +[#1872]: https://github.com/munich-quantum-toolkit/core/pull/1872 [#1869]: https://github.com/munich-quantum-toolkit/core/pull/1869 [#1850]: https://github.com/munich-quantum-toolkit/core/pull/1850 [#1849]: https://github.com/munich-quantum-toolkit/core/pull/1849 diff --git a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td index 49ab89c5d9..e1b86d79d0 100644 --- a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td +++ b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td @@ -1325,6 +1325,13 @@ def IfOp /// Return the yielded value that corresponds to the given argument /// for the else-block, or `nullptr` on failure. OpOperand* getTiedElseYieldedValue(BlockArgument bbArg); + + /// Append the specified additional "qubit" operands: replace this + /// if-op with a new if-op that has the additional qubit operands. + /// The operands can be of qubit or qtensor type. + /// The branch bodies of this if-op are moved over to the new if-op. + /// The newly added qubits are yielded from each branch. + IfOp replaceWithAdditionalQubits(RewriterBase& rewriter, ValueRange addons); }]; let extraClassDefinition = [{ diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp index c377947153..1cd45a58f6 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp @@ -325,3 +325,43 @@ OpOperand* IfOp::getTiedElseYieldedValue(BlockArgument bbArg) { } return &elseYield().getTargetsMutable()[bbArg.getArgNumber()]; } + +IfOp IfOp::replaceWithAdditionalQubits(RewriterBase& rewriter, + ValueRange addons) { + if (addons.empty()) { + return *this; + } + + SmallVector allQubits; + allQubits.reserve(getQubits().size() + addons.size()); + allQubits.append(getQubits().begin(), getQubits().end()); + allQubits.append(addons.begin(), addons.end()); + const auto allQubitTypes = ValueRange(allQubits).getTypes(); + + auto newIfOp = create(rewriter, getLoc(), getCondition(), allQubits); + + const auto rewriteRegion = [&rewriter, &allQubitTypes, + &addons](Region& oldRegion, Region& newRegion) { + auto* oldBlock = &oldRegion.front(); + const auto numOldArgs = oldBlock->getNumArguments(); + auto* newBlock = rewriter.createBlock(&newRegion, {}, allQubitTypes); + const auto oldArgs = newBlock->getArguments().take_front(numOldArgs); + const auto addonArgs = newBlock->getArguments().drop_front(numOldArgs); + + rewriter.mergeBlocks(oldBlock, newBlock, oldArgs); + + auto yield = cast(newBlock->getTerminator()); + SmallVector yieldedValues; + yieldedValues.reserve(yield.getTargets().size() + addons.size()); + yieldedValues.append(yield.getTargets().begin(), yield.getTargets().end()); + yieldedValues.append(addonArgs.begin(), addonArgs.end()); + rewriter.replaceOpWithNewOp(yield, yieldedValues); + }; + + rewriteRegion(getThenRegion(), newIfOp.getThenRegion()); + rewriteRegion(getElseRegion(), newIfOp.getElseRegion()); + + rewriter.eraseOp(*this); + + return newIfOp; +}