diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e97e1d76d..0d87f3b37c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -64,7 +64,7 @@ releases may include breaking changes. #### Passes and transformations - ✨ Add passes for quantum-specific interprocedural optimizations ([#2193], - [#2197]) ([**@DRovara**], [**@burgholzer**]) + [#2197], [#2198]) ([**@DRovara**], [**@burgholzer**]) - ✨ Add Pauli twirling, quantum loop unrolling, and qubit reuse passes ([#1705], [#1718], [#1755], [#1756], [#1923], [#1924], [#2039], [#2118], [#2216], [#2224]) ([**@MatthiasReumann**], [**@DRovara**], [**@burgholzer**], @@ -890,6 +890,7 @@ for previous changelogs._ [#2216]: https://github.com/munich-quantum-toolkit/core/pull/2216 [#2203]: https://github.com/munich-quantum-toolkit/core/pull/2203 [#2214]: https://github.com/munich-quantum-toolkit/core/pull/2214 +[#2198]: https://github.com/munich-quantum-toolkit/core/pull/2198 [#2197]: https://github.com/munich-quantum-toolkit/core/pull/2197 [#2196]: https://github.com/munich-quantum-toolkit/core/pull/2196 [#2194]: https://github.com/munich-quantum-toolkit/core/pull/2194 diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td index 676825f506..a91862c949 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td @@ -350,6 +350,26 @@ def ContextSensitiveSpecialization "mlir::qco::QCODialect"]; } +def QuantumArgumentPromotion + : Pass<"quantum-argument-promotion", "mlir::ModuleOp"> { + let summary = "Replace qubit-tensor arguments by the qubits a callee uses"; + let description = [{ + A tensor argument whose elements are taken out and put back at + compile-time constant indices is split into one qubit argument and one + qubit result per touched element, so untouched elements never cross the + call boundary. Call sites extract before and re-insert after the call. + + The callee has to hand the tensor back as its first result, every + reference to it has to be a direct call, and elements taken out but + never put back prevent promotion. + }]; + + let dependentDialects = ["::mlir::func::FuncDialect", + "::mlir::arith::ArithDialect", + "::mlir::qtensor::QTensorDialect", + "mlir::qco::QCODialect"]; +} + def RemoveDeadGates : Pass<"remove-dead-gates", "mlir::ModuleOp"> { let dependentDialects = ["mlir::qco::QCODialect"]; let summary = "Remove quantum gates whose results cannot be observed"; diff --git a/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumArgumentPromotion.cpp b/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumArgumentPromotion.cpp new file mode 100644 index 0000000000..cde85b7e28 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumArgumentPromotion.cpp @@ -0,0 +1,463 @@ +/* + * 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/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QCO/IR/QCOInterfaces.h" +#include "mlir/Dialect/QCO/IR/QCOOps.h" +#include "mlir/Dialect/QCO/Transforms/Passes.h" +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: keep (Passes.h.inc) +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace mlir::qco { + +#define GEN_PASS_DEF_QUANTUMARGUMENTPROMOTION +#include "mlir/Dialect/QCO/Transforms/Passes.h.inc" + +namespace { + +/** + * @brief A tensor slot that crosses the call boundary as a scalar qubit. + * + * @details + * The qubit is taken out of the tensor at `extractIndex` and put back at + * `insertIndex`; the two indices need not agree. + */ +struct PromotedSlot { + /// The extraction taking the qubit out of the tensor. + qtensor::ExtractOp extract; + /// The insertion putting the qubit back. + qtensor::InsertOp insert; + /// The index the qubit is taken from. + int64_t extractIndex; + /// The index the qubit is put back at. + int64_t insertIndex; +}; + +} // namespace + +/** + * @brief Find where an extracted qubit is put back into a tensor. + * + * @details + * Follows the qubit produced by @p extract forward through gate-like + * operations until it is inserted back into a tensor at a compile-time + * constant index. + * + * @param extract The extraction whose qubit is followed. + * @return The matching insertion, or a null op if the qubit never comes back. + */ +static qtensor::InsertOp findInsertForExtract(qtensor::ExtractOp extract) { + Value currentValue = extract.getResult(); + while (currentValue) { + if (!currentValue.hasOneUse()) { + // Qubits are linear, so this should not happen. + return nullptr; + } + auto* user = *currentValue.getUsers().begin(); + + if (auto insertOp = dyn_cast(user)) { + if (insertOp.getScalar() != currentValue || + !getConstantIntValue(insertOp.getIndex())) { + return nullptr; + } + return insertOp; + } + if (auto unitaryOp = dyn_cast(user)) { + currentValue = unitaryOp.getOutputForInput(currentValue); + continue; + } + if (auto measureOp = dyn_cast(user)) { + currentValue = measureOp.getQubitOut(); + continue; + } + if (auto resetOp = dyn_cast(user)) { + currentValue = resetOp.getQubitOut(); + continue; + } + // Anything else is not known to thread the qubit. Guessing that a single + // result carries it on would let a slot be promoted that no longer holds + // the extracted qubit, so give up instead. + return nullptr; + } + return nullptr; +} + +/** + * @brief Determine whether a tensor argument can be replaced by scalar qubits. + * + * @details + * This requires that every operation on the argument's tensor chain is a + * `qtensor.extract` or `qtensor.insert` at a constant index, that every + * extracted qubit is inserted back into the same chain, and that the chain + * ends in the function's first result. + * + * @param arg The tensor argument to analyze. + * @return The slots that have to cross the call boundary, empty if the + * argument cannot be promoted. + */ +static SmallVector canPromoteArgument(BlockArgument arg) { + const auto tensorType = dyn_cast(arg.getType()); + if (!tensorType || !isa(tensorType.getElementType())) { + return {}; + } + + auto funcOp = dyn_cast(arg.getOwner()->getParentOp()); + if (!funcOp || funcOp.getNumResults() == 0 || + funcOp.getResultTypes()[0] != tensorType) { + // The rewrite below turns the first result into the promoted qubits, so + // the tensor has to be handed back there. + return {}; + } + + // Promotion rewrites the signature and every call site, so every reference to + // the function has to be a direct call. A symbol captured anywhere else would + // be left pointing at the old signature, so bail before anything is changed. + const auto uses = SymbolTable::getSymbolUses(funcOp, funcOp->getParentOp()); + if (!uses) { + return {}; + } + for (const auto use : *uses) { + auto callOp = dyn_cast(use.getUser()); + if (!callOp || callOp.getCallee() != funcOp.getName()) { + return {}; + } + } + + // Walk the chain of the threaded tensor and collect the accesses on it. + SmallVector extracts; + DenseSet insertsOnChain; + // Where each slot is read and written along the chain. The rewrite reorders + // the accesses, so their relative order has to be checked below. + DenseMap extractPositions; + DenseMap insertPositions; + unsigned position = 0; + Value currentTensor = arg; + auto reachesReturn = false; + + while (currentTensor) { + if (!currentTensor.hasOneUse()) { + // Qubit tensors are linear, so this should not happen. + return {}; + } + auto* user = *currentTensor.getUsers().begin(); + + if (auto extractOp = dyn_cast(user)) { + const auto index = getConstantIntValue(extractOp.getIndex()); + if (!index) { + return {}; + } + // Reading one slot twice would take a qubit out of a slot the first read + // already emptied, because the call site extracts from a single tensor. + if (!extractPositions.try_emplace(*index, position++).second) { + return {}; + } + extracts.emplace_back(extractOp); + currentTensor = extractOp.getOutTensor(); + continue; + } + if (auto insertOp = dyn_cast(user)) { + const auto index = getConstantIntValue(insertOp.getIndex()); + if (insertOp.getDest() != currentTensor || !index) { + return {}; + } + // Two writes to one slot only differ in their order, which the rewrite + // does not preserve. + if (!insertPositions.try_emplace(*index, position++).second) { + return {}; + } + insertsOnChain.insert(insertOp); + currentTensor = insertOp.getResult(); + continue; + } + if (auto returnOp = dyn_cast(user)) { + if (returnOp.getOperands().front() != currentTensor) { + return {}; + } + reachesReturn = true; + break; + } + // Anything else (a call, a dealloc, ...) keeps the tensor alive. + return {}; + } + + if (!reachesReturn || extracts.empty()) { + return {}; + } + + // The rewrite takes every promoted qubit out of the caller's tensor before + // the call and puts them all back afterwards. A slot that the callee writes + // before it reads it would then be read from the caller's original tensor + // instead of from the value written to it, so reject that ordering. + for (const auto& [index, insertPosition] : insertPositions) { + const auto extractPosition = extractPositions.find(index); + if (extractPosition != extractPositions.end() && + extractPosition->second > insertPosition) { + return {}; + } + } + + // Every extracted qubit has to find its way back into the same chain. + SmallVector slots; + for (auto extractOp : extracts) { + auto insertOp = findInsertForExtract(extractOp); + if (!insertOp || !insertsOnChain.contains(insertOp)) { + return {}; + } + slots.emplace_back( + PromotedSlot{.extract = extractOp, + .insert = insertOp, + .extractIndex = *getConstantIntValue(extractOp.getIndex()), + .insertIndex = *getConstantIntValue(insertOp.getIndex())}); + } + + // Every insertion has to belong to one of the promoted slots. One that does + // not is left behind by the rewrite and keeps using the tensor argument that + // is erased right after, which trips MLIR's `use_empty()` assertion. + DenseSet matchedInserts; + for (const auto& slot : slots) { + matchedInserts.insert(slot.insert); + } + if (matchedInserts.size() != insertsOnChain.size()) { + return {}; + } + + return slots; +} + +/** + * @brief Replace a tensor argument by one scalar qubit argument per slot. + * + * @details + * Rewrites the function signature and body, and updates every call site so + * that the promoted elements are taken out of the tensor before the call and + * put back afterwards. + * + * @param arg The tensor argument to promote. + * @param slots The slots that cross the call boundary, as returned by + * `canPromoteArgument`. + */ +static void promoteArgument(BlockArgument arg, + MutableArrayRef slots) { + Block* entryBlock = arg.getOwner(); + auto funcOp = cast(entryBlock->getParentOp()); + + OpBuilder builder(funcOp); + MLIRContext* ctx = funcOp.getContext(); + const unsigned argIndex = arg.getArgNumber(); + const auto loc = arg.getLoc(); + const auto tensorType = cast(arg.getType()); + const auto qubitType = tensorType.getElementType(); + const auto numSlots = slots.size(); + + // ==================================================== + // 1. Update the function signature + // ==================================================== + + SmallVector newArgTypes = llvm::to_vector(funcOp.getArgumentTypes()); + newArgTypes.erase(std::next(newArgTypes.begin(), argIndex)); + for (size_t i = 0; i < numSlots; ++i) { + newArgTypes.insert( + std::next(newArgTypes.begin(), static_cast(argIndex + i)), + qubitType); + } + + // `canPromoteArgument` guarantees that the first result is the tensor. + SmallVector newResultTypes = llvm::to_vector(funcOp.getResultTypes()); + newResultTypes.erase(newResultTypes.begin()); + for (size_t i = 0; i < numSlots; ++i) { + newResultTypes.insert( + std::next(newResultTypes.begin(), static_cast(i)), + qubitType); + } + + funcOp.setFunctionType(FunctionType::get(ctx, newArgTypes, newResultTypes)); + + // ==================================================== + // 2. Add the scalar block arguments + // ==================================================== + + SmallVector newArgs; + newArgs.reserve(numSlots); + for (size_t i = 0; i < numSlots; ++i) { + // Insert behind the original argument to keep its index stable for now. + newArgs.emplace_back( + entryBlock->insertArgument(argIndex + i + 1, qubitType, loc)); + } + + // ==================================================== + // 3. Drop the tensor accesses from the body + // ==================================================== + + // The qubit reaching the insert is what the function returns for that slot. + SmallVector returnedQubits; + returnedQubits.reserve(numSlots); + for (auto&& [i, slot] : llvm::enumerate(slots)) { + // A pass-through slot hands the new argument straight back. + returnedQubits.emplace_back(slot.insert.getScalar() == + slot.extract.getResult() + ? newArgs[i] + : slot.insert.getScalar()); + } + + for (const auto& [i, constSlot] : llvm::enumerate(slots)) { + auto slot = constSlot; + // Feed the new argument in where the qubit used to be extracted, and let + // the tensor bypass both accesses. All of them collapse onto `arg`. + slot.extract.getResult().replaceAllUsesWith(newArgs[i]); + slot.extract.getOutTensor().replaceAllUsesWith(slot.extract.getTensor()); + slot.insert.getResult().replaceAllUsesWith(slot.insert.getDest()); + } + for (auto slot : slots) { + slot.insert.erase(); + slot.extract.erase(); + } + + // ==================================================== + // 4. Update the terminator + // ==================================================== + + auto returnOp = cast(entryBlock->getTerminator()); + SmallVector newReturns = llvm::to_vector(returnOp.getOperands()); + newReturns.erase(newReturns.begin()); + for (size_t i = 0; i < numSlots; ++i) { + newReturns.insert(std::next(newReturns.begin(), static_cast(i)), + returnedQubits[i]); + } + returnOp->setOperands(newReturns); + + entryBlock->eraseArgument(argIndex); + + // ==================================================== + // 5. Update the call sites + // ==================================================== + + auto uses = SymbolTable::getSymbolUses(funcOp, funcOp->getParentOp()); + if (!uses) { + return; + } + for (auto use : *uses) { + auto callOp = dyn_cast(use.getUser()); + if (!callOp) { + continue; + } + builder.setInsertionPoint(callOp); + const auto callLoc = callOp.getLoc(); + + SmallVector newOperands = llvm::to_vector(callOp.getOperands()); + Value currentTensor = newOperands[argIndex]; + newOperands.erase(std::next(newOperands.begin(), argIndex)); + + // Take the promoted qubits out of the tensor before the call, ... + for (size_t i = 0; i < numSlots; ++i) { + Value index = arith::ConstantIndexOp::create(builder, callLoc, + slots[i].extractIndex); + auto extractOp = + qtensor::ExtractOp::create(builder, callLoc, currentTensor, index); + currentTensor = extractOp.getOutTensor(); + newOperands.insert( + std::next(newOperands.begin(), static_cast(argIndex + i)), + extractOp.getResult()); + } + + auto newCall = func::CallOp::create(builder, callLoc, funcOp, newOperands); + + // ... and put them back afterwards. + for (size_t i = 0; i < numSlots; ++i) { + Value index = arith::ConstantIndexOp::create(builder, callLoc, + slots[i].insertIndex); + currentTensor = + qtensor::InsertOp::create(builder, callLoc, newCall.getResult(i), + currentTensor, index) + .getResult(); + } + + callOp.getResult(0).replaceAllUsesWith(currentTensor); + for (unsigned r = 1; r < callOp.getNumResults(); ++r) { + callOp.getResult(r).replaceAllUsesWith( + newCall.getResult(numSlots + r - 1)); + } + callOp.erase(); + } +} + +/** + * @brief Promote tensor arguments to scalar qubits across the whole module. + * + * @details + * Externally visible functions and declarations are skipped because their + * signature cannot be changed. At most one argument per function is promoted + * per run, because promoting shifts the indices of the remaining arguments. + * + * @param moduleOp The module to transform. + */ +namespace { +/// Replaces qubit-tensor arguments by the scalar qubits a callee uses. +struct QuantumArgumentPromotion final + : impl::QuantumArgumentPromotionBase { + using impl::QuantumArgumentPromotionBase< + QuantumArgumentPromotion>::QuantumArgumentPromotionBase; + +protected: + void runOnOperation() override { + SmallVector>> + argsToPromote; + + getOperation().walk([&](func::FuncOp func) { + if (func.isPublic() || func.isDeclaration()) { + return; + } + for (auto arg : func.getArguments()) { + auto slots = canPromoteArgument(arg); + if (!slots.empty()) { + argsToPromote.emplace_back(arg, slots); + // Promoting shifts the indices of the remaining arguments, so handle + // at most one argument per function. Any further tensor argument is + // picked up the next time the pass runs. + break; + } + } + }); + + for (auto& [arg, slots] : argsToPromote) { + // Promoting one function rewrites the call sites in others, which can + // invalidate a chain recorded during the walk. Re-derive it right before + // use rather than trusting the recorded slots. + auto current = canPromoteArgument(arg); + if (current.empty()) { + continue; + } + promoteArgument(arg, current); + } + } +}; +} // namespace + +} // namespace mlir::qco diff --git a/mlir/lib/Support/Passes.cpp b/mlir/lib/Support/Passes.cpp index f6aca3927e..b2de9fc79c 100644 --- a/mlir/lib/Support/Passes.cpp +++ b/mlir/lib/Support/Passes.cpp @@ -62,6 +62,7 @@ void registerMQTCompilerPasses() { qco::registerReplaceClassicalControls(); qco::registerReuseQubits(); qco::registerContextSensitiveSpecialization(); + qco::registerQuantumArgumentPromotion(); mqt::registerNormalizeGlobalPhases(); mqt::registerUnrollModifiers(); PassPipelineRegistration<>("mqt-qco-default", diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt index 0fb76a393a..314ce6486d 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt @@ -14,6 +14,7 @@ add_executable( test_qco_measurement_lifting.cpp test_qco_merge_single_qubit_rotation.cpp test_qco_pauli_twirling.cpp + test_qco_quantum_argument_promotion.cpp test_qco_remove_dead_gates.cpp test_qco_replace_classical_controls.cpp test_qco_reuse_qubits.cpp diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_quantum_argument_promotion.cpp b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_quantum_argument_promotion.cpp new file mode 100644 index 0000000000..4c613cbc38 --- /dev/null +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_quantum_argument_promotion.cpp @@ -0,0 +1,513 @@ +/* + * 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 + */ + +/** + * @file test_qco_quantum_argument_promotion.cpp + * @brief Tests for the `quantum-argument-promotion` pass. + */ + +#include "IPOTestFixture.h" +#include "mlir/Dialect/QCO/Transforms/Passes.h" + +#include +#include +#include +#include + +#include + +namespace { + +using QCOQuantumArgumentPromotionTest = ::mqt::test::IPOTestBase; +using namespace mlir; +using namespace mlir::qco; + +// Quantum argument promotion. +// ========================================================================== + +/** + * @brief A tensor argument whose elements are extracted and re-inserted at + * compile-time constant indices is replaced by scalar qubit arguments. + */ +TEST_F(QCOQuantumArgumentPromotionTest, promoteTensorArgumentToQubitArgument) { + const auto tensorType = programBuilder.getQubitTensorType(2); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {tensorType}, {tensorType}); + auto [tensorIn, inner] = programBuilder.qtensorExtract(args[0], 0); + inner = programBuilder.h(inner); + programBuilder.endFunction( + {programBuilder.qtensorInsert(inner, tensorIn, 0)}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto tensor = programBuilder.qtensorFromElements({q0, q1}); + auto results = programBuilder.call("f", {tensor}); + programBuilder.qtensorDealloc(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.h(refArgs[0])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refTensor = referenceBuilder.qtensorFromElements({refQ0, refQ1}); + // The caller extracts the promoted element, calls, and re-inserts it. + auto [refTensorIn, refExtracted] = + referenceBuilder.qtensorExtract(refTensor, 0); + auto refResults = referenceBuilder.call("f", {refExtracted}); + auto refInserted = + referenceBuilder.qtensorInsert(refResults[0], refTensorIn, 0); + referenceBuilder.qtensorDealloc(refInserted); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief An element that is taken out and put straight back at the same index, + * without any gate in between, leaves nothing to promote once it is folded. + * + * @details + * The folder collapses such an extract/insert pair back into the original + * tensor, which leaves the callee as an identity function. That fold is the + * precondition here, so it is applied explicitly: this pass does not run a + * folder of its own, and promoting an unfolded pass-through is wasted work + * rather than a miscompile. + */ +TEST_F(QCOQuantumArgumentPromotionTest, noPromotionForFoldedPassThrough) { + const auto tensorType = programBuilder.getQubitTensorType(2); + + const auto buildProgram = [&tensorType](QCOProgramBuilder& b) { + b.initialize(); + auto args = b.startFunction("f", {tensorType}, {tensorType}); + auto [rest, inner] = b.qtensorExtract(args[0], 0); + b.endFunction({b.qtensorInsert(inner, rest, 0)}); + + auto q0 = b.allocQubit(); + auto q1 = b.allocQubit(); + auto tensor = b.qtensorFromElements({q0, q1}); + auto results = b.call("f", {tensor}); + b.qtensorDealloc(results[0]); + }; + + buildProgram(programBuilder); + moduleOp = programBuilder.finalize(); + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + ASSERT_TRUE(runCanonicalizerPass(moduleOp.get()).succeeded()); + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief Only the tensor elements the callee actually touches become scalar + * arguments; untouched elements never cross the call boundary. + */ +TEST_F(QCOQuantumArgumentPromotionTest, promoteOnlyUsedTensorElements) { + const auto tensorType = programBuilder.getQubitTensorType(3); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {tensorType}, {tensorType}); + auto [tensorIn, inner] = programBuilder.qtensorExtract(args[0], 1); + inner = programBuilder.x(inner); + programBuilder.endFunction( + {programBuilder.qtensorInsert(inner, tensorIn, 1)}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto q2 = programBuilder.allocQubit(); + auto tensor = programBuilder.qtensorFromElements({q0, q1, q2}); + auto results = programBuilder.call("f", {tensor}); + programBuilder.qtensorDealloc(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.x(refArgs[0])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refQ2 = referenceBuilder.allocQubit(); + auto refTensor = referenceBuilder.qtensorFromElements({refQ0, refQ1, refQ2}); + auto [refTensorIn, refExtracted] = + referenceBuilder.qtensorExtract(refTensor, 1); + auto refResults = referenceBuilder.call("f", {refExtracted}); + auto refInserted = + referenceBuilder.qtensorInsert(refResults[0], refTensorIn, 1); + referenceBuilder.qtensorDealloc(refInserted); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief A qubit that is moved to a different slot is promoted with the + * extraction and insertion indices kept apart. + */ +TEST_F(QCOQuantumArgumentPromotionTest, promoteTensorElementIntoDifferentSlot) { + const auto tensorType = programBuilder.getQubitTensorType(2); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {tensorType}, {tensorType}); + auto [tensorIn, inner] = programBuilder.qtensorExtract(args[0], 0); + inner = programBuilder.h(inner); + programBuilder.endFunction( + {programBuilder.qtensorInsert(inner, tensorIn, 1)}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto tensor = programBuilder.qtensorFromElements({q0, q1}); + auto results = programBuilder.call("f", {tensor}); + programBuilder.qtensorDealloc(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.h(refArgs[0])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refTensor = referenceBuilder.qtensorFromElements({refQ0, refQ1}); + auto [refTensorIn, refExtracted] = + referenceBuilder.qtensorExtract(refTensor, 0); + auto refResults = referenceBuilder.call("f", {refExtracted}); + auto refInserted = + referenceBuilder.qtensorInsert(refResults[0], refTensorIn, 1); + referenceBuilder.qtensorDealloc(refInserted); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief An element that is extracted but never re-inserted cannot be promoted, + * because the promoted callee would have nothing to hand back for that slot. + */ +TEST_F(QCOQuantumArgumentPromotionTest, noPromotionWithoutMatchingInsert) { + const auto tensorType = programBuilder.getQubitTensorType(2); + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = + programBuilder.startFunction("f", {tensorType}, {tensorType, qubitType}); + // The element at index 0 leaves the tensor for good. + auto [tensorIn, escaping] = programBuilder.qtensorExtract(args[0], 0); + escaping = programBuilder.h(escaping); + programBuilder.endFunction({tensorIn, escaping}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto tensor = programBuilder.qtensorFromElements({q0, q1}); + auto results = programBuilder.call("f", {tensor}); + programBuilder.sink(results[1]); + programBuilder.qtensorDealloc(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {tensorType}, + {tensorType, qubitType}); + auto [refTensorIn, refEscaping] = + referenceBuilder.qtensorExtract(refArgs[0], 0); + refEscaping = referenceBuilder.h(refEscaping); + referenceBuilder.endFunction({refTensorIn, refEscaping}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refTensor = referenceBuilder.qtensorFromElements({refQ0, refQ1}); + auto refResults = referenceBuilder.call("f", {refTensor}); + referenceBuilder.sink(refResults[1]); + referenceBuilder.qtensorDealloc(refResults[0]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief An element whose path from extraction to re-insertion runs through a + * call is not promoted. + * + * @details + * The walk only recognises the operations it knows to thread a qubit. A call is + * not one of them, and guessing that its first result carries the qubit on + * would let a slot be promoted that no longer holds the extracted qubit, so the + * callee is left alone. + */ +TEST_F(QCOQuantumArgumentPromotionTest, + noPromotionWhenCallSitsOnExtractedPath) { + const auto tensorType = programBuilder.getQubitTensorType(2); + + const auto buildProgram = [&](QCOProgramBuilder& b) { + const auto qubitType = b.getQubitType(); + b.initialize(); + + // The helper takes the qubit as its second operand and returns one qubit. + auto helperArgs = + b.startFunction("helper", {qubitType, qubitType}, {qubitType}); + b.sink(helperArgs[0]); + b.endFunction({b.h(helperArgs[1])}); + + auto args = b.startFunction("f", {tensorType}, {tensorType}); + auto [rest, inner] = b.qtensorExtract(args[0], 0); + auto helped = b.call("helper", {b.allocQubit(), inner})[0]; + b.endFunction({b.qtensorInsert(helped, rest, 0)}); + + auto q0 = b.allocQubit(); + auto q1 = b.allocQubit(); + auto tensor = b.qtensorFromElements({q0, q1}); + auto results = b.call("f", {tensor}); + b.qtensorDealloc(results[0]); + }; + + buildProgram(programBuilder); + moduleOp = programBuilder.finalize(); + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief A tensor argument that never has an element taken out of it has + * nothing to promote. + */ +TEST_F(QCOQuantumArgumentPromotionTest, noPromotionWithoutElementAccess) { + const auto tensorType = programBuilder.getQubitTensorType(2); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {tensorType}, {tensorType}); + programBuilder.endFunction({args[0]}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto tensor = programBuilder.qtensorFromElements({q0, q1}); + auto results = programBuilder.call("f", {tensor}); + programBuilder.qtensorDealloc(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {tensorType}, {tensorType}); + referenceBuilder.endFunction({refArgs[0]}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refTensor = referenceBuilder.qtensorFromElements({refQ0, refQ1}); + auto refResults = referenceBuilder.call("f", {refTensor}); + referenceBuilder.qtensorDealloc(refResults[0]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief A callee that touches several tensor elements gets one scalar argument + * and one scalar result per element. + */ +TEST_F(QCOQuantumArgumentPromotionTest, promoteMultipleTensorElements) { + const auto tensorType = programBuilder.getQubitTensorType(2); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {tensorType}, {tensorType}); + auto [afterFirst, first] = programBuilder.qtensorExtract(args[0], 0); + auto firstTensor = + programBuilder.qtensorInsert(programBuilder.h(first), afterFirst, 0); + auto [afterSecond, second] = programBuilder.qtensorExtract(firstTensor, 1); + programBuilder.endFunction( + {programBuilder.qtensorInsert(programBuilder.x(second), afterSecond, 1)}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto tensor = programBuilder.qtensorFromElements({q0, q1}); + auto results = programBuilder.call("f", {tensor}); + programBuilder.qtensorDealloc(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + const auto qubitType = referenceBuilder.getQubitType(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.h(refArgs[0]), referenceBuilder.x(refArgs[1])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refTensor = referenceBuilder.qtensorFromElements({refQ0, refQ1}); + // The caller takes every promoted element out before the call and puts them + // all back afterwards. + auto [refAfterFirst, refFirst] = + referenceBuilder.qtensorExtract(refTensor, 0); + auto [refAfterSecond, refSecond] = + referenceBuilder.qtensorExtract(refAfterFirst, 1); + auto refResults = referenceBuilder.call("f", {refFirst, refSecond}); + auto refFirstBack = + referenceBuilder.qtensorInsert(refResults[0], refAfterSecond, 0); + referenceBuilder.qtensorDealloc( + referenceBuilder.qtensorInsert(refResults[1], refFirstBack, 1)); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief A promoted element may be measured inside the callee; the measurement + * outcome stays a separate result and the caller keeps reading it. + */ +TEST_F(QCOQuantumArgumentPromotionTest, promoteTensorElementWithMeasurement) { + const auto tensorType = programBuilder.getQubitTensorType(2); + const auto bitType = programBuilder.getI1Type(); + + programBuilder.initialize({bitType}); + auto args = + programBuilder.startFunction("f", {tensorType}, {tensorType, bitType}); + auto [rest, inner] = programBuilder.qtensorExtract(args[0], 0); + Value bit; + std::tie(inner, bit) = programBuilder.measure(inner); + programBuilder.endFunction( + {programBuilder.qtensorInsert(inner, rest, 0), bit}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto tensor = programBuilder.qtensorFromElements({q0, q1}); + auto results = programBuilder.call("f", {tensor}); + programBuilder.qtensorDealloc(results[0]); + moduleOp = programBuilder.finalize({results[1]}); + + referenceBuilder.initialize({bitType}); + auto refArgs = referenceBuilder.startFunction( + "f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType(), bitType}); + Value refBit; + auto refInner = refArgs[0]; + std::tie(refInner, refBit) = referenceBuilder.measure(refInner); + referenceBuilder.endFunction({refInner, refBit}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refTensor = referenceBuilder.qtensorFromElements({refQ0, refQ1}); + auto [refRest, refExtracted] = referenceBuilder.qtensorExtract(refTensor, 0); + auto refResults = referenceBuilder.call("f", {refExtracted}); + referenceBuilder.qtensorDealloc( + referenceBuilder.qtensorInsert(refResults[0], refRest, 0)); + reference = referenceBuilder.finalize({refResults[1]}); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief The tensor has to be handed back as the first result, because that is + * the result the promoted qubits take the place of. + */ +TEST_F(QCOQuantumArgumentPromotionTest, noPromotionWhenTensorIsNotFirstResult) { + const auto tensorType = programBuilder.getQubitTensorType(2); + const auto bitType = programBuilder.getI1Type(); + + const auto buildProgram = [&](QCOProgramBuilder& b) { + b.initialize({bitType}); + auto args = b.startFunction("f", {tensorType}, {bitType, tensorType}); + auto [rest, inner] = b.qtensorExtract(args[0], 0); + Value bit; + std::tie(inner, bit) = b.measure(inner); + b.endFunction({bit, b.qtensorInsert(inner, rest, 0)}); + + auto q0 = b.allocQubit(); + auto q1 = b.allocQubit(); + auto tensor = b.qtensorFromElements({q0, q1}); + auto results = b.call("f", {tensor}); + b.qtensorDealloc(results[1]); + return results[0]; + }; + + moduleOp = programBuilder.finalize({buildProgram(programBuilder)}); + reference = referenceBuilder.finalize({buildProgram(referenceBuilder)}); + + expectSingleStageMatchesReference(createQuantumArgumentPromotion()); +} + +/** + * @brief A slot the callee writes before reading again must not be promoted. + * + * @details + * Extractions move in front of the call and insertions behind it, so such a + * read would be served from the caller's original tensor. Here the callee + * computes `x(h(slot 0))`, which promotion would turn into `x(slot 1)`. + */ +TEST_F(QCOQuantumArgumentPromotionTest, + noPromotionWhenSlotIsWrittenBeforeItIsRead) { + auto module = parseModule(R"mlir( +func.func private @callee(%t: tensor<2x!qco.qubit>) -> tensor<2x!qco.qubit> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %t1, %q0 = qtensor.extract %t[%c0] : tensor<2x!qco.qubit> + %q0h = qco.h %q0 : !qco.qubit -> !qco.qubit + %t2 = qtensor.insert %q0h into %t1[%c1] : tensor<2x!qco.qubit> + %t3, %q1 = qtensor.extract %t2[%c1] : tensor<2x!qco.qubit> + %q1x = qco.x %q1 : !qco.qubit -> !qco.qubit + %t4 = qtensor.insert %q1x into %t3[%c0] : tensor<2x!qco.qubit> + return %t4 : tensor<2x!qco.qubit> +} +func.func @main(%t: tensor<2x!qco.qubit>) -> tensor<2x!qco.qubit> { + %r = func.call @callee(%t) : (tensor<2x!qco.qubit>) -> tensor<2x!qco.qubit> + return %r : tensor<2x!qco.qubit> +} +)mlir"); + ASSERT_TRUE(module); + ASSERT_TRUE( + runStage(module.get(), createQuantumArgumentPromotion()).succeeded()); + + auto callee = module->lookupSymbol("callee"); + ASSERT_TRUE(callee); + EXPECT_TRUE(isa(callee.getArgumentTypes()[0])) + << "the tensor argument must survive, the accesses depend on each other"; +} + +/** + * @brief An insertion that does not belong to a promoted slot blocks promotion. + * + * @details + * It survives the rewrite still using the tensor argument that is erased right + * afterwards, which used to abort on MLIR's `use_empty()` assertion. + */ +TEST_F(QCOQuantumArgumentPromotionTest, noPromotionForUnmatchedInsertOnChain) { + auto module = parseModule(R"mlir( +func.func private @callee(%t: tensor<2x!qco.qubit>, %extra: !qco.qubit) -> tensor<2x!qco.qubit> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %t1, %q0 = qtensor.extract %t[%c0] : tensor<2x!qco.qubit> + %q0h = qco.h %q0 : !qco.qubit -> !qco.qubit + %t2 = qtensor.insert %q0h into %t1[%c0] : tensor<2x!qco.qubit> + %t3 = qtensor.insert %extra into %t2[%c1] : tensor<2x!qco.qubit> + return %t3 : tensor<2x!qco.qubit> +} +func.func @main(%t: tensor<2x!qco.qubit>, %e: !qco.qubit) -> tensor<2x!qco.qubit> { + %r = func.call @callee(%t, %e) : (tensor<2x!qco.qubit>, !qco.qubit) -> tensor<2x!qco.qubit> + return %r : tensor<2x!qco.qubit> +} +)mlir"); + ASSERT_TRUE(module); + ASSERT_TRUE( + runStage(module.get(), createQuantumArgumentPromotion()).succeeded()); + + auto callee = module->lookupSymbol("callee"); + ASSERT_TRUE(callee); + EXPECT_TRUE(isa(callee.getArgumentTypes()[0])) + << "the tensor argument must survive, one insertion is unmatched"; +} + +// ========================================================================== + +} // namespace