diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d87f3b37c..eb1013bee7 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], [#2198]) ([**@DRovara**], [**@burgholzer**]) + [#2197], [#2198], [#2199]) ([**@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 +[#2199]: https://github.com/munich-quantum-toolkit/core/pull/2199 [#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 diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td index a91862c949..42eebceed3 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td @@ -370,6 +370,26 @@ def QuantumArgumentPromotion "mlir::qco::QCODialect"]; } +def AuxiliaryQubitHoisting + : Pass<"quantum-auxiliary-qubit-hoisting", "mlir::ModuleOp"> { + let summary = "Turn callee-internal qubits into arguments"; + let description = [{ + A qubit that a callee allocates and releases itself becomes an extra + argument, and its release point becomes a `qco.reset` handed back as an + extra result. The caller then owns the allocation and can reuse one + qubit across several calls. + + Externally visible functions, declarations, recursive functions, and + functions with more than one block are left alone, as are allocations + nested inside a region. + }]; + + 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/AuxiliaryQubitHoisting.cpp b/mlir/lib/Dialect/QCO/Transforms/Optimizations/AuxiliaryQubitHoisting.cpp new file mode 100644 index 0000000000..18297ed8b6 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/AuxiliaryQubitHoisting.cpp @@ -0,0 +1,417 @@ +/* + * 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/Analysis/CallGraph.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/QCO/Utils/WireIterator.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 // IWYU pragma: keep (Passes.h.inc) +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace mlir::qco { + +#define GEN_PASS_DEF_AUXILIARYQUBITHOISTING +#include "mlir/Dialect/QCO/Transforms/Passes.h.inc" + +/** + * @brief Find the operation that releases the qubit produced by @p alloc. + * + * @details + * Follows the qubit forward along its linear use chain, through gates, + * measurements, resets and calls, and also while it is parked inside a qubit + * tensor. Returns a null op when the qubit escapes the function or when the + * chain cannot be followed, for example because a tensor index is not known at + * compile time. + * + * @param alloc The allocation whose release point is searched. + * @param callMapping Resolves how qubits flow across call boundaries. + * @return The `qco.sink` releasing the qubit, or a null op if there is none. + */ +static SinkOp findDeallocForAlloc(AllocOp alloc, + CallQubitMapping& callMapping) { + Value currentValue = alloc.getResult(); + uint64_t currentIndexInTensor = 0; + bool isInTensor = false; + + while (currentValue) { + // Both qubits and qubit tensors are linear values, so every step of the + // chain has exactly one user. + if (!currentValue.hasOneUse()) { + return nullptr; + } + auto* user = *currentValue.getUsers().begin(); + + if (isInTensor) { + // The qubit currently lives at `currentIndexInTensor` of the tensor in + // `currentValue`. Follow the tensor until it is extracted again. + if (auto extractOp = dyn_cast(user)) { + const auto index = getConstantIntValue(extractOp.getIndex()); + if (!index) { + // Dynamic index, cannot tell whether it is our qubit. + return nullptr; + } + if (std::cmp_equal(*index, currentIndexInTensor)) { + currentValue = extractOp.getResult(); + isInTensor = false; + } else { + currentValue = extractOp.getOutTensor(); + } + continue; + } + if (auto insertOp = dyn_cast(user)) { + const auto index = getConstantIntValue(insertOp.getIndex()); + if (!index || std::cmp_equal(*index, currentIndexInTensor)) { + // Dynamic index, or our slot is overwritten by another qubit. + return nullptr; + } + currentValue = insertOp.getResult(); + continue; + } + // Anything else (a dealloc, a call, ...) takes the qubit out of reach. + return nullptr; + } + + if (auto deallocOp = dyn_cast(user)) { + return deallocOp; + } + 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; + } + if (auto callOp = dyn_cast(user)) { + const auto threadedOr = + callMapping.getResultForOperand(callOp, currentValue); + if (failed(threadedOr)) { + return nullptr; + } + const auto threaded = *threadedOr; + if (!threaded) { + // The callee keeps the qubit, so it is never released here. + return nullptr; + } + currentValue = threaded; + continue; + } + if (auto fromElementsOp = dyn_cast(user)) { + for (auto i = 0ULL; i < user->getNumOperands(); i++) { + if (user->getOperand(i) == currentValue) { + currentIndexInTensor = i; + isInTensor = true; + break; + } + } + currentValue = fromElementsOp.getResult(); + continue; + } + if (auto insertOp = dyn_cast(user)) { + const auto index = getConstantIntValue(insertOp.getIndex()); + if (!index) { + return nullptr; + } + currentIndexInTensor = static_cast(*index); + isInTensor = true; + currentValue = insertOp.getResult(); + continue; + } + // Anything else is not known to thread the qubit. Guessing that a single + // result carries it on would silently follow an unrelated value, so give up + // instead. + return nullptr; + } + return nullptr; +} + +/** + * @brief Check whether the given function takes part in a call cycle. + * + * @details + * Performs an iterative reachability search over the call graph, starting at + * the function's callees rather than at the function itself, so that a + * non-recursive function is not reported as recursive merely because the + * search begins at its own node. A worklist is used instead of recursion so + * that deep call chains cannot exhaust the stack. + * + * @param cg The call graph of the surrounding module. + * @param func The function to check. + * @return True if @p func can reach itself through a chain of calls. + */ +static bool isRecursive(CallGraph& cg, func::FuncOp func) { + CallGraphNode* node = cg.lookupNode(func.getCallableRegion()); + if (node == nullptr) { + return false; + } + + llvm::DenseSet visited; + SmallVector worklist; + + // Seed the search with the function's callees so that the function itself is + // only reported as recursive when a call chain leads back to it. + for (const auto& edge : *node) { + worklist.emplace_back(edge.getTarget()); + } + + while (!worklist.empty()) { + auto* current = worklist.pop_back_val(); + if (current == node) { + return true; + } + if (!visited.insert(current).second) { + continue; + } + for (const auto& edge : *current) { + worklist.emplace_back(edge.getTarget()); + } + } + + return false; +} + +/** + * @brief Turn every auxiliary qubit of the given function into an argument. + * + * @details + * An auxiliary qubit is one that the function allocates and releases itself. + * Hoisting it makes the caller own the allocation, which lets the caller reuse + * one qubit across several calls. The release point becomes a `qco.reset` that + * is handed back as an additional result, so the caller receives the qubit in a + * known state. + * + * @param funcOp The function to transform. + * @param callMapping Resolves how qubits flow across call boundaries. + */ +static void tryAuxiliaryQubitHoisting(func::FuncOp funcOp, + CallQubitMapping& callMapping) { + // The release point is rewritten into a reset whose result is appended to + // every return. That is only sound while there is a single block, because a + // reset in one block need not reach a return in another. + if (!funcOp.getBody().hasOneBlock()) { + return; + } + + // Collect the allocations up front: the loop below erases operations, which + // would invalidate a walk in progress. + SmallVector allocOps; + funcOp.walk([&](AllocOp allocOp) { + if (allocOp->getBlock()->getParentOp() != funcOp) { + // Not directly in the function body, skip. + return; + } + allocOps.emplace_back(allocOp); + }); + + for (auto allocOp : allocOps) { + auto dealloc = findDeallocForAlloc(allocOp, callMapping); + + if (!dealloc) { + // No matching dealloc found, skip. + continue; + } + + // Collect the call sites before touching the signature. Once the signature + // changes the existing calls no longer match it, so if the uses cannot be + // determined it must not have been changed in the first place. + const auto uses = SymbolTable::getSymbolUses(funcOp, funcOp->getParentOp()); + if (!uses) { + continue; + } + + // Every reference has to be a direct call. Anything else, such as the + // symbol captured in an attribute or taken as a function value, has no + // operand list to extend and would be left pointing at the old signature. + SmallVector callOps; + auto onlyDirectCalls = true; + for (const auto use : *uses) { + auto callOp = dyn_cast(use.getUser()); + if (!callOp || callOp.getCallee() != funcOp.getName()) { + onlyDirectCalls = false; + break; + } + callOps.emplace_back(callOp); + } + if (!onlyDirectCalls) { + continue; + } + + // Add a block argument for the auxiliary qubit. + OpBuilder builder(dealloc); + auto* block = allocOp->getBlock(); + auto loc = allocOp.getLoc(); + auto qubitType = allocOp.getType(); + auto newArg = block->addArgument(qubitType, loc); + + // Replace all uses of the alloc with the new block argument. + allocOp.replaceAllUsesWith(newArg); + + // Erase the original alloc operation. + allocOp.erase(); + + // Replace the dealloc with a reset + builder.setInsertionPoint(dealloc); + auto resetOp = + ResetOp::create(builder, dealloc.getLoc(), dealloc.getQubit()); + dealloc.erase(); + + // Add reset outcome to function results and alloc to function arguments + auto funcType = funcOp.getFunctionType(); + SmallVector newArgTypes(funcType.getInputs().begin(), + funcType.getInputs().end()); + SmallVector newResultTypes(funcType.getResults().begin(), + funcType.getResults().end()); + newArgTypes.push_back(newArg.getType()); + newResultTypes.push_back(resetOp.getResult().getType()); + auto newFuncType = + FunctionType::get(funcOp.getContext(), newArgTypes, newResultTypes); + funcOp.setType(newFuncType); + // The cached mapping describes the old signature, so it is stale now. + callMapping.invalidate(); + + // Also add the reset outcome to every return. The operands are updated in + // place so that the terminator stays valid for the ongoing walk. + funcOp.walk([&](func::ReturnOp returnOp) { + SmallVector newReturnValues(returnOp.getOperands().begin(), + returnOp.getOperands().end()); + newReturnValues.emplace_back(resetOp.getResult()); + returnOp->setOperands(newReturnValues); + }); + + // Update the call sites collected above to handle the new return value. + for (auto callOp : callOps) { + builder.setInsertionPoint(callOp); + + // A. Add new alloc + auto newAlloc = AllocOp::create(builder, loc); + + // B. Create New Call + SmallVector newCallOperands = + llvm::to_vector(callOp.getOperands()); + newCallOperands.emplace_back(newAlloc); + auto newCall = + func::CallOp::create(builder, loc, funcOp, newCallOperands); + + // C. Add dealloc after call + SinkOp::create(builder, loc, + newCall.getResult(newCall.getNumResults() - 1)); + for (unsigned i = 0; i < callOp.getNumResults(); ++i) { + callOp.getResult(i).replaceAllUsesWith(newCall.getResult(i)); + } + callOp.erase(); + } + } +} + +/** + * @brief Hoist the auxiliary qubits of every eligible function in the module. + * + * @details + * Externally visible functions and declarations are skipped because their + * signature cannot be changed, and recursive functions are skipped because + * their allocation would have to be threaded through every level of the + * recursion. + * + * @param moduleOp The module to transform. + */ +/** + * @brief Order the hoisting candidates so that callees come before callers. + * + * @details + * Recursive functions are not candidates, so the graph is acyclic here. The + * traversal starts at the external caller node, leaving out functions no entry + * point reaches; those have no call sites to hoist into anyway. + * + * @param cg The call graph of the surrounding module. + * @param candidates The functions to order. + * @return The candidates, callees first. + */ +static SmallVector +orderCalleesFirst(const CallGraph& cg, + const SmallVector& candidates) { + llvm::DenseMap candidateNodes; + for (auto func : candidates) { + if (auto* node = cg.lookupNode(func.getCallableRegion())) { + candidateNodes.try_emplace(node, func); + } + } + + SmallVector ordered; + ordered.reserve(candidates.size()); + for (auto* node : llvm::post_order(&cg)) { + if (const auto it = candidateNodes.find(node); it != candidateNodes.end()) { + ordered.emplace_back(it->second); + } + } + return ordered; +} + +namespace { +/// Turns qubits a callee allocates and releases itself into arguments. +struct AuxiliaryQubitHoisting final + : impl::AuxiliaryQubitHoistingBase { + using impl::AuxiliaryQubitHoistingBase< + AuxiliaryQubitHoisting>::AuxiliaryQubitHoistingBase; + +protected: + void runOnOperation() override { + auto moduleOp = getOperation(); + SmallVector hoistingCandidates; + CallGraph callGraph(moduleOp); + // One shared mapping so that each callee is threaded at most once. + CallQubitMapping callMapping; + + moduleOp.walk([&](func::FuncOp func) { + if (func.isPublic() || func.isDeclaration()) { + return; + } + if (isRecursive(callGraph, func)) { + return; + } + hoistingCandidates.push_back(func); + }); + + // Hoisting out of a callee puts an allocation into each of its callers, + // which may itself be hoistable. Visiting callees first lets such an + // allocation travel all the way up in a single run instead of stopping + // wherever the module happens to declare the functions. + for (auto func : orderCalleesFirst(callGraph, hoistingCandidates)) { + tryAuxiliaryQubitHoisting(func, callMapping); + } + } +}; +} // namespace + +} // namespace mlir::qco diff --git a/mlir/lib/Support/Passes.cpp b/mlir/lib/Support/Passes.cpp index b2de9fc79c..3b2487c292 100644 --- a/mlir/lib/Support/Passes.cpp +++ b/mlir/lib/Support/Passes.cpp @@ -63,6 +63,7 @@ void registerMQTCompilerPasses() { qco::registerReuseQubits(); qco::registerContextSensitiveSpecialization(); qco::registerQuantumArgumentPromotion(); + qco::registerAuxiliaryQubitHoisting(); 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 314ce6486d..7beba5b928 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt @@ -9,6 +9,7 @@ set(target_name mqt-core-mlir-unittest-optimizations) add_executable( ${target_name} + test_qco_auxiliary_qubit_hoisting.cpp test_qco_context_sensitive_specialization.cpp test_qco_hadamard_lifting.cpp test_qco_measurement_lifting.cpp diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_auxiliary_qubit_hoisting.cpp b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_auxiliary_qubit_hoisting.cpp new file mode 100644 index 0000000000..222fb17c38 --- /dev/null +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_auxiliary_qubit_hoisting.cpp @@ -0,0 +1,610 @@ +/* + * 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_auxiliary_qubit_hoisting.cpp + * @brief Tests for the `quantum-auxiliary-qubit-hoisting` pass. + */ + +#include "IPOTestFixture.h" +#include "mlir/Dialect/QCO/Transforms/Passes.h" + +#include +#include +#include +#include + +#include +#include + +namespace { + +using QCOAuxiliaryQubitHoistingTest = ::mqt::test::IPOTestBase; +using namespace mlir; +using namespace mlir::qco; + +// Auxiliary qubit hoisting. +// ========================================================================== + +/** + * @brief A qubit that a callee allocates and releases internally is turned into + * an extra argument, so the caller owns the allocation and can reuse it. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, hoistAuxiliaryQubitIntoCaller) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + auto aux = programBuilder.allocQubit(); + auto target = args[0]; + std::tie(aux, target) = programBuilder.cx(aux, target); + programBuilder.sink(aux); + programBuilder.endFunction({target}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + // The auxiliary qubit becomes a trailing argument and is returned in a reset + // state as a trailing result. + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refAux = refArgs[1]; + auto refTarget = refArgs[0]; + std::tie(refAux, refTarget) = referenceBuilder.cx(refAux, refTarget); + refAux = referenceBuilder.reset(refAux); + referenceBuilder.endFunction({refTarget, refAux}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAuxAlloc = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f", {refQ, refAuxAlloc}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief A qubit that the callee allocates but hands back to the caller is not + * auxiliary and must stay where it is. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, noHoistingForReturnedQubit) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = + programBuilder.startFunction("f", {qubitType}, {qubitType, qubitType}); + auto fresh = programBuilder.allocQubit(); + programBuilder.endFunction({args[0], fresh}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + programBuilder.sink(results[1]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {qubitType}, {qubitType, qubitType}); + auto refFresh = referenceBuilder.allocQubit(); + referenceBuilder.endFunction({refArgs[0], refFresh}); + + auto refQ = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f", {refQ}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief The auxiliary qubit is tracked across a measurement and a reset on its + * way to the release point. + * + * The measurement outcome is handed back to the caller so that the measurement + * is not dead, and the reset sits between two gates so that it is neither + * folded into the allocation nor into the release. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, + hoistAuxiliaryQubitThroughMeasureAndReset) { + const auto qubitType = programBuilder.getQubitType(); + const auto bitType = programBuilder.getI1Type(); + + programBuilder.initialize({bitType}); + auto args = + programBuilder.startFunction("f", {qubitType}, {qubitType, bitType}); + auto aux = programBuilder.h(programBuilder.allocQubit()); + Value bit; + std::tie(aux, bit) = programBuilder.measure(aux); + aux = programBuilder.reset(aux); + auto target = args[0]; + std::tie(aux, target) = programBuilder.cx(aux, target); + programBuilder.sink(aux); + programBuilder.endFunction({target, bit}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize({results[1]}); + + referenceBuilder.initialize({bitType}); + auto refArgs = referenceBuilder.startFunction( + "f", {qubitType, qubitType}, {qubitType, bitType, qubitType}); + auto refAux = referenceBuilder.h(refArgs[1]); + Value refBit; + std::tie(refAux, refBit) = referenceBuilder.measure(refAux); + refAux = referenceBuilder.reset(refAux); + auto refTarget = refArgs[0]; + std::tie(refAux, refTarget) = referenceBuilder.cx(refAux, refTarget); + refAux = referenceBuilder.reset(refAux); + referenceBuilder.endFunction({refTarget, refBit, refAux}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAuxAlloc = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f", {refQ, refAuxAlloc}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[2]); + reference = referenceBuilder.finalize({refResults[1]}); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief The auxiliary qubit is tracked while it is parked in a tensor, past + * an extraction of an unrelated element. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, hoistAuxiliaryQubitThroughTensor) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + auto aux = programBuilder.allocQubit(); + auto target = args[0]; + std::tie(aux, target) = programBuilder.cx(aux, target); + // The auxiliary qubit sits at index 0, the argument qubit at index 1. + auto tensor = programBuilder.qtensorFromElements({aux, target}); + auto [afterOther, other] = programBuilder.qtensorExtract(tensor, 1); + auto [afterAux, auxBack] = programBuilder.qtensorExtract(afterOther, 0); + programBuilder.sink(auxBack); + programBuilder.qtensorDealloc(afterAux); + programBuilder.endFunction({other}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refAux = refArgs[1]; + auto refTarget = refArgs[0]; + std::tie(refAux, refTarget) = referenceBuilder.cx(refAux, refTarget); + auto refTensor = referenceBuilder.qtensorFromElements({refAux, refTarget}); + auto [refAfterOther, refOther] = + referenceBuilder.qtensorExtract(refTensor, 1); + auto [refAfterAux, refAuxBack] = + referenceBuilder.qtensorExtract(refAfterOther, 0); + auto refReset = referenceBuilder.reset(refAuxBack); + referenceBuilder.qtensorDealloc(refAfterAux); + referenceBuilder.endFunction({refOther, refReset}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAuxAlloc = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f", {refQ, refAuxAlloc}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief The auxiliary qubit is tracked across a nested call on its way to the + * release point. + * + * The nested callee returns more than one qubit and the auxiliary one is not + * the first, so the walk has to match the operand position rather than simply + * taking the first result. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, hoistAuxiliaryQubitThroughNestedCall) { + const auto qubitType = programBuilder.getQubitType(); + + const auto buildNestedCallee = [&qubitType](QCOProgramBuilder& b) { + auto innerArgs = + b.startFunction("g", {qubitType, qubitType}, {qubitType, qubitType}); + b.endFunction({b.h(innerArgs[0]), innerArgs[1]}); + }; + + programBuilder.initialize(); + buildNestedCallee(programBuilder); + + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + auto aux = programBuilder.allocQubit(); + // The auxiliary qubit is the second operand and the second result. + auto nested = programBuilder.call("g", {args[0], aux}); + auto target = nested[0]; + aux = nested[1]; + std::tie(aux, target) = programBuilder.cx(aux, target); + programBuilder.sink(aux); + programBuilder.endFunction({target}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + buildNestedCallee(referenceBuilder); + + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refNested = referenceBuilder.call("g", {refArgs[0], refArgs[1]}); + auto refTarget = refNested[0]; + auto refAux = refNested[1]; + std::tie(refAux, refTarget) = referenceBuilder.cx(refAux, refTarget); + refAux = referenceBuilder.reset(refAux); + referenceBuilder.endFunction({refTarget, refAux}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAuxAlloc = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f", {refQ, refAuxAlloc}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief Every call site of a hoisted callee gets its own allocation. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, + hoistAuxiliaryQubitWithMultipleCallSites) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + auto aux = programBuilder.allocQubit(); + auto target = args[0]; + std::tie(aux, target) = programBuilder.cx(aux, target); + programBuilder.sink(aux); + programBuilder.endFunction({target}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto results0 = programBuilder.call("f", {q0}); + auto results1 = programBuilder.call("f", {q1}); + programBuilder.sink(results0[0]); + programBuilder.sink(results1[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refAux = refArgs[1]; + auto refTarget = refArgs[0]; + std::tie(refAux, refTarget) = referenceBuilder.cx(refAux, refTarget); + refAux = referenceBuilder.reset(refAux); + referenceBuilder.endFunction({refTarget, refAux}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refAux0 = referenceBuilder.allocQubit(); + auto refResults0 = referenceBuilder.call("f", {refQ0, refAux0}); + referenceBuilder.sink(refResults0[1]); + auto refAux1 = referenceBuilder.allocQubit(); + auto refResults1 = referenceBuilder.call("f", {refQ1, refAux1}); + referenceBuilder.sink(refResults1[1]); + referenceBuilder.sink(refResults0[0]); + referenceBuilder.sink(refResults1[0]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief An allocation threaded through a recursive callee is not hoisted. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, noHoistingForRecursiveFunction) { + auto module = parseModule(R"mlir( +func.func private @recursive(%q: !qco.qubit) -> !qco.qubit { + %r = func.call @recursive(%q) : (!qco.qubit) -> !qco.qubit + return %r : !qco.qubit +} +func.func private @outer(%q: !qco.qubit) -> !qco.qubit { + %aux = qco.alloc : !qco.qubit + %r = func.call @recursive(%aux) : (!qco.qubit) -> !qco.qubit + qco.sink %r : !qco.qubit + return %q : !qco.qubit +} +func.func @main(%q: !qco.qubit) -> !qco.qubit { + %r = func.call @outer(%q) : (!qco.qubit) -> !qco.qubit + return %r : !qco.qubit +} +)mlir"); + ASSERT_TRUE(module); + ASSERT_TRUE( + runStage(module.get(), createAuxiliaryQubitHoisting()).succeeded()); + EXPECT_EQ(countAllocsIn(module.get(), "outer"), 1U); +} + +/** + * @brief An allocation nested inside a region is not hoisted, because it is not + * executed on every path through the function. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, noHoistingForAllocInsideRegion) { + const auto qubitType = programBuilder.getQubitType(); + + const auto buildProgram = [&qubitType](QCOProgramBuilder& b) { + b.initialize(); + auto args = b.startFunction("f", {qubitType, b.getI1Type()}, {qubitType}); + auto result = b.qcoIf(args[1], args[0], [&](Value qubit) { + auto aux = b.allocQubit(); + auto inner = qubit; + std::tie(aux, inner) = b.cx(aux, inner); + b.sink(aux); + return inner; + }); + b.endFunction({result}); + + auto q = b.allocQubit(); + Value bit; + std::tie(q, bit) = b.measure(q); + auto results = b.call("f", {q, bit}); + b.sink(results[0]); + }; + + buildProgram(programBuilder); + moduleOp = programBuilder.finalize(); + + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief The auxiliary qubit is tracked when it enters a tensor through an + * insertion and while unrelated elements are moved in and out around it. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, hoistAuxiliaryQubitParkedInTensor) { + const auto qubitType = programBuilder.getQubitType(); + + const auto buildBody = [](QCOProgramBuilder& b, Value aux, Value target) { + // Park the auxiliary qubit in a scratch register at index 0. + auto scratch = b.qtensorAlloc(2); + auto [afterPlaceholder, placeholder] = b.qtensorExtract(scratch, 0); + b.sink(placeholder); + auto parked = b.qtensorInsert(aux, afterPlaceholder, 0); + // Move an unrelated element out and back in while the auxiliary qubit + // stays parked at index 0. + auto [afterOther, other] = b.qtensorExtract(parked, 1); + auto restored = b.qtensorInsert(other, afterOther, 1); + auto [afterAux, auxBack] = b.qtensorExtract(restored, 0); + b.qtensorDealloc(afterAux); + return std::pair{auxBack, target}; + }; + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + auto aux = programBuilder.allocQubit(); + auto target = args[0]; + std::tie(aux, target) = programBuilder.cx(aux, target); + auto [auxBack, finalTarget] = buildBody(programBuilder, aux, target); + programBuilder.sink(auxBack); + programBuilder.endFunction({finalTarget}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refAux = refArgs[1]; + auto refTarget = refArgs[0]; + std::tie(refAux, refTarget) = referenceBuilder.cx(refAux, refTarget); + auto [refAuxBack, refFinalTarget] = + buildBody(referenceBuilder, refAux, refTarget); + referenceBuilder.endFunction( + {refFinalTarget, referenceBuilder.reset(refAuxBack)}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAuxAlloc = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f", {refQ, refAuxAlloc}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief A call that only consumes linear values, and one that only produces + * them, keep the builder's tracking consistent. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, callConsumesAndProducesLinearValues) { + const auto qubitType = programBuilder.getQubitType(); + const auto tensorType = programBuilder.getQubitTensorType(2); + + const auto buildProgram = [&](QCOProgramBuilder& b) { + b.initialize(); + auto consumeArgs = b.startFunction("consume", {qubitType, tensorType}, {}); + b.sink(consumeArgs[0]); + b.qtensorDealloc(consumeArgs[1]); + b.endFunction({}); + + b.startFunction("produce", {}, {tensorType}); + b.endFunction({b.qtensorAlloc(2)}); + + auto q = b.allocQubit(); + auto scratch = b.qtensorAlloc(2); + b.call("consume", {q, scratch}); + auto produced = b.call("produce", {}); + b.qtensorDealloc(produced[0]); + }; + + buildProgram(programBuilder); + moduleOp = programBuilder.finalize(); + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief The auxiliary qubit is tracked across a call whose callee also returns + * a classical value ahead of the qubit. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, + hoistAuxiliaryQubitThroughCallWithClassicalResult) { + const auto qubitType = programBuilder.getQubitType(); + const auto bitType = programBuilder.getI1Type(); + + const auto buildCallee = [&](QCOProgramBuilder& b) { + auto innerArgs = b.startFunction("g", {qubitType}, {bitType, qubitType}); + auto [inner, bit] = b.measure(innerArgs[0]); + b.endFunction({bit, inner}); + }; + + programBuilder.initialize({bitType}); + buildCallee(programBuilder); + + auto args = + programBuilder.startFunction("f", {qubitType}, {qubitType, bitType}); + auto nested = programBuilder.call("g", {programBuilder.allocQubit()}); + auto aux = nested[1]; + auto target = args[0]; + std::tie(aux, target) = programBuilder.cx(aux, target); + programBuilder.sink(aux); + programBuilder.endFunction({target, nested[0]}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize({results[1]}); + + referenceBuilder.initialize({bitType}); + buildCallee(referenceBuilder); + + auto refArgs = referenceBuilder.startFunction( + "f", {qubitType, qubitType}, {qubitType, bitType, qubitType}); + auto refNested = referenceBuilder.call("g", {refArgs[1]}); + auto refAux = refNested[1]; + auto refTarget = refArgs[0]; + std::tie(refAux, refTarget) = referenceBuilder.cx(refAux, refTarget); + refAux = referenceBuilder.reset(refAux); + referenceBuilder.endFunction({refTarget, refNested[0], refAux}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAuxAlloc = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f", {refQ, refAuxAlloc}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[2]); + reference = referenceBuilder.finalize({refResults[1]}); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief A qubit handed to a callee that keeps it is not auxiliary, so it is + * left alone rather than indexed past the end of the call's results. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, noHoistingWhenCalleeKeepsAuxiliaryQubit) { + const auto qubitType = programBuilder.getQubitType(); + + const auto buildProgram = [&](QCOProgramBuilder& b) { + b.initialize(); + auto consumeArgs = b.startFunction("consume", {qubitType}, {}); + b.sink(consumeArgs[0]); + b.endFunction({}); + + auto args = b.startFunction("f", {qubitType}, {qubitType}); + b.call("consume", {b.allocQubit()}); + b.endFunction({b.h(args[0])}); + + auto q = b.allocQubit(); + auto results = b.call("f", {q}); + b.sink(results[0]); + }; + + buildProgram(programBuilder); + moduleOp = programBuilder.finalize(); + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createAuxiliaryQubitHoisting()); +} + +/** + * @brief Hoisting reaches the outermost caller regardless of declaration order. + * + * @details + * An allocation hoisted into a caller may be hoistable again. Processing in + * module order used to strand it wherever the declarations happened to sit. + */ +TEST_F(QCOAuxiliaryQubitHoistingTest, hoistingIsIndependentOfDeclarationOrder) { + const auto qubitType = programBuilder.getQubitType(); + + // Builds `main -> mid -> leaf`, where `leaf` owns an auxiliary qubit. + const auto build = [&](QCOProgramBuilder& builder) { + builder.initialize(); + auto leafArgs = builder.startFunction("leaf", {qubitType}, {qubitType}); + auto aux = builder.allocQubit(); + auto target = leafArgs[0]; + std::tie(aux, target) = builder.cx(aux, target); + builder.sink(aux); + builder.endFunction({target}); + + auto midArgs = builder.startFunction("mid", {qubitType}, {qubitType}); + auto midResults = builder.call("leaf", {midArgs[0]}); + builder.endFunction({midResults[0]}); + + auto q = builder.allocQubit(); + auto results = builder.call("mid", {q}); + builder.sink(results[0]); + return builder.finalize(); + }; + + moduleOp = build(programBuilder); + reference = build(referenceBuilder); + + // The builder has to declare a callee before the call, so the second module + // is reordered afterwards. Both now describe the same call graph and differ + // only in the order the module walk visits the two callees. + auto refLeaf = reference->lookupSymbol("leaf"); + auto refMid = reference->lookupSymbol("mid"); + ASSERT_TRUE(refLeaf); + ASSERT_TRUE(refMid); + refLeaf->moveAfter(refMid.getOperation()); + + ASSERT_TRUE( + runStage(moduleOp.get(), createAuxiliaryQubitHoisting()).succeeded()); + ASSERT_TRUE( + runStage(reference.get(), createAuxiliaryQubitHoisting()).succeeded()); + + // The auxiliary allocation belongs in the entry function either way: one + // allocation for the qubit passed in and one for the hoisted auxiliary. + for (auto* module : {&moduleOp, &reference}) { + EXPECT_EQ(countAllocsIn(module->get(), "leaf"), 0U); + EXPECT_EQ(countAllocsIn(module->get(), "mid"), 0U); + EXPECT_EQ(countAllocsIn(module->get(), "main"), 2U); + } +} + +// ========================================================================== + +} // namespace