diff --git a/include/mqt-core/qasm3/Types.hpp b/include/mqt-core/qasm3/Types.hpp index 7ccf3def25..5b5549c6b1 100644 --- a/include/mqt-core/qasm3/Types.hpp +++ b/include/mqt-core/qasm3/Types.hpp @@ -154,7 +154,7 @@ template class DesignatedType final : public Type { std::string designatorToString(); }; -enum UnsizedTy : uint8_t { Bool, Duration }; +enum UnsizedTy : uint8_t { Bool, Duration, SingleQubit }; template class UnsizedType final : public Type { public: @@ -178,6 +178,9 @@ template class UnsizedType final : public Type { static std::shared_ptr> getDurationTy() { return std::make_shared(Duration); } + static std::shared_ptr> getSingleQubitTy() { + return std::make_shared(SingleQubit); + } T getDesignator() override { throw std::runtime_error("Unsized types do not have designators"); @@ -195,6 +198,8 @@ template class UnsizedType final : public Type { return "bool"; case Duration: return "duration"; + case SingleQubit: + return "qubit"; } throw std::runtime_error("Unhandled type"); } diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index a1c7e0bfe6..137059fb55 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -10,12 +10,14 @@ #pragma once +#include #include #include #include #include #include +#include #include #include #include @@ -1755,6 +1757,109 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { const std::variant& index, ValueRange yieldedValues); + //===--------------------------------------------------------------------===// + // Additional functions + //===--------------------------------------------------------------------===// + + /** + * @brief Start building an additional private function in the module + * + * @details + * Creates a new private `func.func` at the end of the module and moves the + * insertion point into its entry block. Qubit- and qubit-tensor-typed + * arguments are added to the linear-type tracking so that operations can be + * applied to them; a tensor argument is treated as a register the callee + * owns for the duration of the call. Arguments of other types (e.g., `f64`) + * are returned as-is without tracking. Must be called after `initialize()` + * and must be paired with a call to `endFunction()`; function definitions + * cannot be nested. + * + * @param name The name of the function + * @param argTypes The argument types of the function + * @param resultTypes The result types of the function + * @return The entry block arguments of the new function + * + * @par Example: + * ```c++ + * auto args = builder.startFunction( + * "f", {builder.getQubitType()}, {builder.getQubitType()}); + * auto q = builder.h(args[0]); + * builder.endFunction({q}); + * ``` + * ```mlir + * func.func private @f(%arg0: !qco.qubit) -> !qco.qubit { + * %q = qco.h %arg0 : !qco.qubit -> !qco.qubit + * return %q : !qco.qubit + * } + * ``` + */ + SmallVector startFunction(StringRef name, TypeRange argTypes, + TypeRange resultTypes); + + /** + * @brief Finish the function started with `startFunction()` + * + * @details + * Creates the `func.return` with the given values and restores the insertion + * point to where it was before `startFunction()` was called. All qubits and + * tensors that were created within the function (from arguments or + * operations) must either be consumed (e.g., by `sink()` or + * `qtensorDealloc()`) or returned; otherwise a usage error is reported. + * + * @param returnValues The values to return from the function + */ + void endFunction(ValueRange returnValues); + + /** + * @brief Call a function previously defined in the module + * + * @details + * Creates a `func.call` to the named function. Qubit- and qubit-tensor-typed + * operands are validated and consumed; results of those types are added to + * the tracking. + * + * The i-th qubit operand is paired with the i-th qubit result, and likewise + * for tensors, so a function that threads its linear values through keeps + * the register association intact. Surplus operands are treated as consumed + * and surplus results as freshly created. This positional pairing is a + * calling convention, not something the IR enforces: a callee that returns + * its qubits in a different order than it takes them will be tracked + * incorrectly. + * + * @param callee The name of the function to call + * @param operands The operands to pass to the call + * @return The results of the call operation + * + * @par Example: + * ```c++ + * auto results = builder.call("f", {q0}); + * ``` + * ```mlir + * %q1 = call @f(%q0) : (!qco.qubit) -> !qco.qubit + * ``` + */ + SmallVector call(StringRef callee, ValueRange operands); + + /** + * @brief Get the qubit type + * @return The `!qco.qubit` type + */ + Type getQubitType(); + + /** + * @brief Get a one-dimensional tensor type holding qubits + * @param size The number of qubits in the tensor + * @return The `tensor` type + */ + Type getQubitTensorType(int64_t size); + + /** + * @brief Check whether the given type is a tensor of qubits + * @param type The type to check + * @return True if @p type is a ranked tensor with `!qco.qubit` elements + */ + static bool isQubitTensor(Type type); + //===--------------------------------------------------------------------===// // Finalization //===--------------------------------------------------------------------===// @@ -1943,6 +2048,21 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { /// Ensure static and dynamic qubit allocation modes are not mixed. void ensureAllocationMode(AllocationMode requestedMode); + + /** + * @brief State of an additional function under construction + */ + struct FunctionScope { + /// Insertion point to restore when the function is finished + OpBuilder::InsertPoint savedInsertPoint; + /// Qubit values that were already tracked before the function was started + llvm::DenseSet outerQubits; + /// Tensor values that were already tracked before the function was started + llvm::DenseSet outerTensors; + }; + + /// Active function scope, if a function is currently under construction. + std::optional functionScope; }; } // namespace qco } // namespace mlir diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.h b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.h index de6e78952d..d8c24971ed 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.h @@ -10,6 +10,7 @@ #pragma once +#include #include #include #include @@ -33,6 +34,10 @@ namespace mlir::qco { #define GEN_PASS_REGISTRATION #include "mlir/Dialect/QCO/Transforms/Passes.h.inc" // IWYU pragma: export +void runQuantumArgumentPromotion(ModuleOp module); +void runAuxiliaryQubitHoisting(ModuleOp module); +void runQuantumFunctionBoundaryCommutation(ModuleOp module, + SymbolTable& symbolTable); /** * @brief Create target-independent two-qubit gate fusion. */ diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td index cfa0a5a797..2adda35da0 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td @@ -305,6 +305,18 @@ def ReuseQubits : Pass<"reuse-qubits", "mlir::ModuleOp"> { }]; } +def QuantumIPO : Pass<"quantum-ipo", "mlir::ModuleOp"> { + let summary = "Quantum Interprocedural Optimization Pass"; + let description = [{ + Performs interprocedural optimizations on quantum functions. + }]; + + let dependentDialects = ["::mlir::func::FuncDialect", + "::mlir::arith::ArithDialect", + "::mlir::qtensor::QTensorDialect", + "mlir::qco::QCODialect"]; +} + //===----------------------------------------------------------------------===// // Decomposition Passes //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 46fda77441..4e26a1c93b 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -33,10 +33,12 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -1374,6 +1376,171 @@ QCOProgramBuilder::scfCondition(Value reg, return scfCondition(condition, yieldedValues); } +//===----------------------------------------------------------------------===// +// Additional Functions +//===----------------------------------------------------------------------===// + +Type QCOProgramBuilder::getQubitType() { return QubitType::get(ctx); } + +Type QCOProgramBuilder::getQubitTensorType(const int64_t size) { + return RankedTensorType::get({size}, getQubitType()); +} + +bool QCOProgramBuilder::isQubitTensor(Type type) { + const auto tensorType = dyn_cast(type); + return tensorType && isa(tensorType.getElementType()); +} + +SmallVector QCOProgramBuilder::startFunction(StringRef name, + TypeRange argTypes, + TypeRange resultTypes) { + checkFinalized(); + + if (functionScope.has_value()) { + llvm::reportFatalUsageError( + "Cannot start a function while another one is being built"); + } + + FunctionScope scope{.savedInsertPoint = saveInsertionPoint(), + .outerQubits = {}, + .outerTensors = {}}; + for (const auto& [qubit, info] : validQubits) { + scope.outerQubits.insert(qubit); + } + for (const auto& [tensor, info] : validTensors) { + scope.outerTensors.insert(tensor); + } + + setInsertionPointToEnd(cast(module).getBody()); + auto funcOp = + func::FuncOp::create(*this, name, getFunctionType(argTypes, resultTypes)); + // The interprocedural passes only consider functions that are not externally + // visible, so additional functions are private by default. + funcOp.setPrivate(); + + auto& entryBlock = funcOp.getBody().emplaceBlock(); + const SmallVector locs(argTypes.size(), getLoc()); + entryBlock.addArguments(argTypes, locs); + setInsertionPointToStart(&entryBlock); + + SmallVector args; + for (const auto arg : entryBlock.getArguments()) { + if (isa(arg.getType())) { + validQubits.try_emplace(arg, QubitInfo{}); + } else if (isQubitTensor(arg.getType())) { + // A tensor argument acts like a register the callee owns for the + // duration of the call, so give it its own register id. + validTensors.try_emplace(arg, TensorInfo{tensorCounter++}); + } + args.emplace_back(arg); + } + + functionScope = std::move(scope); + return args; +} + +void QCOProgramBuilder::endFunction(ValueRange returnValues) { + checkFinalized(); + + if (!functionScope.has_value()) { + llvm::reportFatalUsageError( + "endFunction() called without a matching startFunction()"); + } + + for (const auto value : returnValues) { + if (isa(value.getType())) { + validateQubitValue(value); + validQubits.erase(value); + } else if (isQubitTensor(value.getType())) { + validateTensorValue(value); + validTensors.erase(value); + } + } + + for (const auto& [qubit, info] : validQubits) { + if (!functionScope->outerQubits.contains(qubit)) { + llvm::reportFatalUsageError( + "Function body has qubit values that are neither returned nor " + "consumed"); + } + } + for (const auto& [tensor, info] : validTensors) { + if (!functionScope->outerTensors.contains(tensor)) { + llvm::reportFatalUsageError( + "Function body has tensor values that are neither returned nor " + "deallocated"); + } + } + + func::ReturnOp::create(*this, returnValues); + + restoreInsertionPoint(functionScope->savedInsertPoint); + functionScope.reset(); +} + +SmallVector QCOProgramBuilder::call(StringRef callee, + ValueRange operands) { + checkFinalized(); + + auto funcOp = dyn_cast_or_null( + SymbolTable::lookupSymbolIn(module, getStringAttr(callee))); + if (!funcOp) { + llvm::reportFatalUsageError("Callee not found in module"); + } + + SmallVector qubitOperands; + SmallVector tensorOperands; + for (const auto operand : operands) { + if (isa(operand.getType())) { + validateQubitValue(operand); + qubitOperands.emplace_back(operand); + } else if (isQubitTensor(operand.getType())) { + validateTensorValue(operand); + tensorOperands.emplace_back(operand); + } + } + + auto callOp = func::CallOp::create(*this, funcOp, operands); + + SmallVector qubitResults; + SmallVector tensorResults; + for (const auto result : callOp.getResults()) { + if (isa(result.getType())) { + qubitResults.emplace_back(result); + } else if (isQubitTensor(result.getType())) { + tensorResults.emplace_back(result); + } + } + + // Thread the i-th linear operand into the i-th linear result of the same + // kind. Any operand without a matching result is consumed by the call, any + // result without a matching operand is newly created by it. + const auto pairedQubits = std::min(qubitOperands.size(), qubitResults.size()); + for (size_t i = 0; i < pairedQubits; ++i) { + updateQubitTracking(qubitOperands[i], qubitResults[i]); + } + for (size_t i = pairedQubits; i < qubitOperands.size(); ++i) { + validQubits.erase(qubitOperands[i]); + } + for (size_t i = pairedQubits; i < qubitResults.size(); ++i) { + validQubits.try_emplace(qubitResults[i], QubitInfo{}); + } + + const auto pairedTensors = + std::min(tensorOperands.size(), tensorResults.size()); + for (size_t i = 0; i < pairedTensors; ++i) { + updateTensorTracking(tensorOperands[i], tensorResults[i]); + } + for (size_t i = pairedTensors; i < tensorOperands.size(); ++i) { + validTensors.erase(tensorOperands[i]); + } + for (size_t i = pairedTensors; i < tensorResults.size(); ++i) { + validTensors.try_emplace(tensorResults[i], TensorInfo{tensorCounter++}); + } + + return SmallVector(callOp.getResults()); +} + //===----------------------------------------------------------------------===// // Finalization //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp index dae57f52c9..3bf5f36b68 100644 --- a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp +++ b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp @@ -479,6 +479,36 @@ void IndexSwitchOp::print(OpAsmPrinter& p) { //===----------------------------------------------------------------------===// #include "mlir/Dialect/QCO/IR/QCOOpsDialect.cpp.inc" +#include "mlir/Transforms/InliningUtils.h" + +namespace { +// Define the opt-in rules for inlining the qc dialect +struct QCOInlinerInterface : public mlir::DialectInlinerInterface { + using DialectInlinerInterface::DialectInlinerInterface; + + // Tell MLIR that any operation from the qc dialect can be inlined + bool isLegalToInline(mlir::Operation* /*op*/, mlir::Region* /*dest*/, + bool /*wouldBeCloned*/, + mlir::IRMapping& /*valueMapping*/) const override { + return true; + } + + // Tell MLIR that regions (like the inside of loops/ifs) in the qc dialect can + // be inlined + bool isLegalToInline(mlir::Region* /*dest*/, mlir::Region* /*src*/, + bool /*wouldBeCloned*/, + mlir::IRMapping& /*valueMapping*/) const override { + return true; + } + + // Tell MLIR that it's safe to inline calls to functions containing qc + // operations + bool isLegalToInline(mlir::Operation* /*call*/, mlir::Operation* /*callable*/, + bool /*wouldBeCloned*/) const override { + return true; + } +}; +} // namespace void QCODialect::initialize() { // NOLINTNEXTLINE(clang-analyzer-core.StackAddressEscape) @@ -493,6 +523,8 @@ void QCODialect::initialize() { #include "mlir/Dialect/QCO/IR/QCOOps.cpp.inc" >(); + + addInterface(); } //===----------------------------------------------------------------------===// 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..f3ee77da31 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/AuxiliaryQubitHoisting.cpp @@ -0,0 +1,287 @@ +/* + * 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/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 { + +static SinkOp findDeallocForAlloc(AllocOp alloc) { + 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 (isa(user)) { + // Relies on the QCO calling convention that the i-th qubit result of a + // call corresponds to its i-th qubit operand. + // TODO-Damian this only works if the indices are the same. Implement a + // helper function to get the index + for (auto i = 0ULL; i < user->getNumOperands(); i++) { + if (user->getOperand(i) == currentValue) { + currentValue = user->getResult(i); + break; + } + } + 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; + } + if (user->getNumResults() != 1) { + // Multiple results, should not happen. + return nullptr; + } + currentValue = user->getResult(0); + } + return nullptr; +} + +static bool isRecursiveHelper(CallGraphNode* current, CallGraphNode* target, + llvm::DenseSet& visited) { + if (!visited.insert(current).second) { + return false; // Already visited + } + + for (const auto& edge : *current) { + CallGraphNode* callee = edge.getTarget(); + if (callee == target) { + return true; + } + if (isRecursiveHelper(callee, target, visited)) { + return true; + } + } + + return false; +} + +static bool isRecursive(CallGraph& cg, func::FuncOp func) { + CallGraphNode* node = cg.lookupNode(func.getCallableRegion()); + if (node == nullptr) { + return false; + } + + llvm::DenseSet visited; + // Start from the function's callees to avoid immediately returning true + for (const auto& edge : *node) { + if (isRecursiveHelper(edge.getTarget(), node, visited)) { + return true; + } + } + + return false; +} + +static void tryAuxiliaryQubitHoisting(func::FuncOp funcOp) { + funcOp.walk([&](AllocOp allocOp) { + if (allocOp->getBlock()->getParentOp() != funcOp) { + // Not directly in the function body, skip. + return; + } + + auto dealloc = findDeallocForAlloc(allocOp); + + if (!dealloc) { + // No matching dealloc found, skip. + return; + } + + // 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 = + builder.create(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); + + // Also add reset outcome to return + funcOp.walk([&](func::ReturnOp returnOp) { + OpBuilder returnBuilder(returnOp); + SmallVector newReturnValues(returnOp.getOperands().begin(), + returnOp.getOperands().end()); + newReturnValues.push_back(resetOp.getResult()); + returnBuilder.create(returnOp.getLoc(), newReturnValues); + returnOp.erase(); + }); + + // Update all call sites to handle the new return value + // We use the SymbolTable to find all calls to this function + if (auto uses = SymbolTable::getSymbolUses(funcOp, funcOp->getParentOp())) { + for (auto use : *uses) { + if (auto callOp = dyn_cast(use.getUser())) { + builder.setInsertionPoint(callOp); + + // A. Add new alloc + auto newAlloc = builder.create(loc); + + // B. Create New Call + SmallVector newCallOperands = + llvm::to_vector(callOp.getOperands()); + newCallOperands.push_back(newAlloc); + auto newCall = + builder.create(loc, funcOp, newCallOperands); + + // C. Add dealloc after call + builder.create( + loc, newCall.getResult(newCall.getNumResults() - 1)); + for (unsigned i = 0; i < callOp.getNumResults(); ++i) { + callOp.getResult(i).replaceAllUsesWith(newCall.getResult(i)); + } + callOp.erase(); + } + } + } + }); +} + +void runAuxiliaryQubitHoisting(ModuleOp module) { + SmallVector hoistingCandidates; + CallGraph callGraph(module); + + module.walk([&](func::FuncOp func) { + if (func.isPublic() || func.isDeclaration()) { + return; + } + if (isRecursive(callGraph, func)) { + return; + } + hoistingCandidates.push_back(func); + }); + + for (auto& func : hoistingCandidates) { + tryAuxiliaryQubitHoisting(func); + + RewritePatternSet patterns(module.getContext()); + if (!applyPatternsGreedily(module, std::move(patterns)).succeeded()) { + throw std::runtime_error("Failed to apply reuse qubits patterns after " + "auxiliary qubit hoisting."); + } + } +} + +} // namespace mlir::qco 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..d9e1f812e4 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumArgumentPromotion.cpp @@ -0,0 +1,340 @@ +/* + * 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 + */ + +/** + * @brief This pass performs quantum inter-procedural optimizations (IPO). + */ + +#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/QTensorOps.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace mlir::qco { + +namespace { + +/// A tensor slot that crosses the call boundary as a scalar qubit: the qubit +/// is taken out of the tensor at `extractIndex` and put back at `insertIndex`. +struct PromotedSlot { + qtensor::ExtractOp extract; + qtensor::InsertOp insert; + int64_t extractIndex; + int64_t insertIndex; +}; + +} // namespace + +/// Follow the qubit produced by `extract` forward through gate-like operations +/// until it is inserted back into a tensor at a compile-time constant index. +/// Returns 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 (user->getNumResults() != 1) { + return nullptr; + } + currentValue = user->getResult(0); + } + return nullptr; +} + +/// Determine whether the given tensor argument can be replaced by scalar qubit +/// arguments, and if so, which slots have to cross the call boundary. +/// +/// 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. +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 {}; + } + + // Walk the chain of the threaded tensor and collect the accesses on it. + SmallVector extracts; + DenseSet insertsOnChain; + 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)) { + if (!getConstantIntValue(extractOp.getIndex())) { + return {}; + } + extracts.emplace_back(extractOp); + currentTensor = extractOp.getOutTensor(); + continue; + } + if (auto insertOp = dyn_cast(user)) { + if (insertOp.getDest() != currentTensor || + !getConstantIntValue(insertOp.getIndex())) { + 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 {}; + } + + // 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())}); + } + + return slots; +} + +/// Replace the given tensor argument by one scalar qubit argument per slot in +/// `slots`, and update every call site accordingly. +static void promoteArgument(BlockArgument arg, ArrayRef 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 slot : slots) { + returnedQubits.emplace_back(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); + + 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 = + builder.create(loc, slots[i].extractIndex); + auto extractOp = + builder.create(loc, currentTensor, index); + currentTensor = extractOp.getOutTensor(); + newOperands.insert( + std::next(newOperands.begin(), static_cast(argIndex + i)), + extractOp.getResult()); + } + + auto newCall = builder.create(loc, funcOp, newOperands); + + // ... and put them back afterwards. + for (size_t i = 0; i < numSlots; ++i) { + Value index = + builder.create(loc, slots[i].insertIndex); + currentTensor = builder + .create(loc, 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(); + } +} + +void runQuantumArgumentPromotion(ModuleOp module) { + SmallVector>> + argsToPromote; + + module.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) { + promoteArgument(arg, slots); + } +} + +} // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumFunctionBoundaryCommutation.cpp b/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumFunctionBoundaryCommutation.cpp new file mode 100644 index 0000000000..fdfa24fe62 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumFunctionBoundaryCommutation.cpp @@ -0,0 +1,116 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace mlir::qco { + +static func::FuncOp copyFunction(func::FuncOp funcOp, StringRef newName) { + auto newFunc = funcOp.clone(); + newFunc.setName(newName.str()); + + return newFunc; +} + +static bool doOpsCancel(UnitaryOpInterface first, UnitaryOpInterface second) { + // For now, let's just consider self-inverses and single-qubit, non-controlled + // gates. + if (first.getOperation()->getName() != second.getOperation()->getName()) { + return false; + } + if (isa(first)) { + return true; + } + return false; +} + +static void tryBoundaryCommutation( + func::CallOp call, SymbolTable& symbolTable, uint32_t parameter, + std::unordered_map& previousSpecializations) { + auto calleeName = call.getCallee(); + auto funcOp = symbolTable.lookup(calleeName); + + if (!funcOp || funcOp.isExternal()) { + return; + } + + auto argOutside = call.getArgOperands()[parameter]; + auto argInside = funcOp.getArgument(parameter); + + if (!argInside.hasOneUse()) { + return; + } + if (argOutside.getDefiningOp() == nullptr) { + return; + } + + auto lastOp = dyn_cast(argOutside.getDefiningOp()); + auto nextOp = dyn_cast(*argInside.getUsers().begin()); + + if (!lastOp || !nextOp) { + return; + } + + if (!doOpsCancel(lastOp, nextOp)) { + return; + } + argOutside.replaceAllUsesWith(lastOp.getInputQubit(0)); + lastOp.erase(); + + if (previousSpecializations.contains(funcOp.getName().str())) { + call.setCallee(previousSpecializations[funcOp.getName().str()].getName()); + return; + } + + auto newFunc = copyFunction(funcOp, funcOp.getName().str() + + "_spec_boundary_commutation"); + symbolTable.insert(newFunc); + + auto newParameter = newFunc.getArgument(parameter); + auto newUser = dyn_cast(*newParameter.getUsers().begin()); + + for (auto i = 0U; i < newUser.getNumQubits(); ++i) { + newUser.getOutputQubit(i).replaceAllUsesWith(newUser.getInputQubit(i)); + } + newUser.erase(); + previousSpecializations[funcOp.getName().str()] = newFunc; + + call.setCallee(newFunc.getName()); +} + +void runQuantumFunctionBoundaryCommutation(ModuleOp module, + SymbolTable& symbolTable) { + std::unordered_map previousSpecializations; + module.walk([&](func::CallOp call) { + for (uint32_t i = 0; i < call.getArgOperands().size(); ++i) { + const auto arg = call.getArgOperands()[i]; + if (!isa(arg.getType())) { + continue; + } + tryBoundaryCommutation(call, symbolTable, i, previousSpecializations); + } + }); +} + +} // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumIPO.cpp b/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumIPO.cpp new file mode 100644 index 0000000000..deb23ac12f --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/QuantumIPO.cpp @@ -0,0 +1,341 @@ +/* + * 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 + */ + +// +// Created by damian on 1/21/26. +// + +#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" +// IWYU pragma: begin_keep (Passes.h.inc) +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" +// IWYU pragma: end_keep + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mlir::qco { + +static void updateSpecializedCall(func::CallOp callOp, func::FuncOp newCallee, + PatternRewriter& rewriter) { + rewriter.modifyOpInPlace(callOp, + [&] { callOp.setCallee(newCallee.getName()); }); +} + +static func::FuncOp copyFunction(func::FuncOp funcOp, StringRef newName, + PatternRewriter& rewriter) { + const OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointAfter(funcOp); + + auto newFunc = funcOp.clone(); + + rewriter.modifyOpInPlace(newFunc, [&] { newFunc.setName(newName.str()); }); + + return newFunc; +} + +#define GEN_PASS_DEF_QUANTUMIPO +#include "mlir/Dialect/QCO/Transforms/Passes.h.inc" + +namespace { + +struct PreviousSpecializations { + std::map, func::FuncOp> zeroSpecializations; + std::map, func::FuncOp> plusSpecializations; + std::map, func::FuncOp> + rotationSpecializations; +}; + +/** + * @brief This pattern attempts to perform context-sensitive specialization. + */ +struct ContextSensitiveSpecializationPattern final + : OpRewritePattern { + + // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members) + SymbolTable& symbolTable; + // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members) + PreviousSpecializations& previousSpecializations; + + constexpr static const auto ANGLES_TO_SPECIALIZE = + std::array{0.0, std::numbers::pi, std::numbers::pi / 2, + 1.5 * std::numbers::pi, 2 * std::numbers::pi}; + + static bool operationIsNopOnZero(Operation* op, Value zeroArgument) { + if (auto ctrl = dyn_cast(op)) { + return std::find(ctrl.getControlsIn().begin(), ctrl.getControlsIn().end(), + zeroArgument) != ctrl.getControlsIn().end(); + } + return isa(op) || isa(op) || isa(op); // TODO more ops? + } + + static bool operationIsNopOnPlus(Operation* op) { return isa(op); } + + explicit ContextSensitiveSpecializationPattern(MLIRContext* context, + SymbolTable& symbolTable, + PreviousSpecializations& prev) + : OpRewritePattern(context), symbolTable(symbolTable), + previousSpecializations(prev) {} + + LogicalResult matchAndRewrite(func::CallOp callOp, + PatternRewriter& rewriter) const override { + auto found = false; + for (auto i = 0U; i < callOp.getArgOperands().size(); ++i) { + if (trySpecialize(callOp, i, rewriter)) { + found = true; + } + } + return LogicalResult::success(found); + } + + bool trySpecialize(func::CallOp callOp, unsigned operand, + PatternRewriter& rewriter) const { + const auto argValue = callOp.getArgOperands()[operand]; + + auto calleeName = callOp.getCallee(); + auto funcOp = symbolTable.lookup(calleeName); + + if (!funcOp || funcOp.isExternal()) { + return false; + } + + auto* definingOp = argValue.getDefiningOp(); + + if (definingOp == nullptr) { + return false; + } + + if (argValue.getType() == QubitType::get(rewriter.getContext())) { + // CSS for qubit types. + if (isa(definingOp) || isa(definingOp)) { + return trySpecializeZero(callOp, funcOp, operand, rewriter); + } + if (isa(definingOp)) { + const auto* precedingOp = definingOp->getOperand(0).getDefiningOp(); + if (precedingOp != nullptr && + (isa(precedingOp) || isa(precedingOp))) { + return trySpecializePlus(callOp, funcOp, operand, rewriter); + } + } + } + if (argValue.getType() == Float64Type::get(rewriter.getContext())) { + // CSS for double types. + if (isa(definingOp)) { + auto constOp = cast(definingOp); + return trySpecializeRotationArguments( + callOp, funcOp, + cast(constOp.getValue()).getValueAsDouble(), operand, + rewriter); + } + } + + return false; + } + + bool trySpecializeZero(func::CallOp callOp, func::FuncOp funcOp, + unsigned operand, PatternRewriter& rewriter) const { + auto parameter = funcOp.getArgument(operand); + if (!parameter.hasOneUse()) { + return false; + } + if (!operationIsNopOnZero(*parameter.getUsers().begin(), parameter)) { + return false; + } + + auto key = std::make_pair(funcOp.getName().str(), operand); + if (previousSpecializations.zeroSpecializations.contains(key)) { + updateSpecializedCall(callOp, + previousSpecializations.zeroSpecializations.at(key), + rewriter); + return true; + } + + auto newFunc = copyFunction(funcOp, + funcOp.getName().str() + "_spec_zero_arg_" + + std::to_string(operand), + rewriter); + symbolTable.insert(newFunc); + previousSpecializations.zeroSpecializations.insert({key, newFunc}); + + auto newParameter = newFunc.getArgument(operand); + while ( + newParameter.hasOneUse() && + operationIsNopOnZero(*newParameter.getUsers().begin(), newParameter)) { + auto newUser = + dyn_cast(*newParameter.getUsers().begin()); + for (auto i = 0U; i < newUser.getNumQubits(); ++i) { + // TODO-DAMIAN use getOutputQubit/Input again (at current version, this + // seems to use the output of the inner op) + rewriter.replaceAllUsesWith(newUser->getResult(i), + newUser->getOperand(i)); + } + rewriter.eraseOp(newUser); + break; + } + + updateSpecializedCall(callOp, newFunc, rewriter); + return true; + } + + bool trySpecializePlus(func::CallOp callOp, func::FuncOp funcOp, + unsigned operand, PatternRewriter& rewriter) const { + auto parameter = funcOp.getArgument(operand); + if (!parameter.hasOneUse()) { + return false; + } + if (!operationIsNopOnPlus(*parameter.getUsers().begin())) { + return false; + } + + auto key = std::make_pair(funcOp.getName().str(), operand); + if (previousSpecializations.plusSpecializations.contains(key)) { + updateSpecializedCall(callOp, + previousSpecializations.plusSpecializations.at(key), + rewriter); + return true; + } + + auto newFunc = copyFunction(funcOp, + funcOp.getName().str() + "_spec_plus_arg_" + + std::to_string(operand), + rewriter); + symbolTable.insert(newFunc); + previousSpecializations.plusSpecializations.insert({key, newFunc}); + + auto newParameter = newFunc.getArgument(operand); + while (newParameter.hasOneUse() && + operationIsNopOnPlus(*newParameter.getUsers().begin())) { + auto newUser = + dyn_cast(*newParameter.getUsers().begin()); + for (auto i = 0U; i < newUser.getNumQubits(); ++i) { + rewriter.replaceAllUsesWith(newUser.getOutputQubit(i), + newUser.getInputQubit(i)); + } + rewriter.eraseOp(newUser); + } + + updateSpecializedCall(callOp, newFunc, rewriter); + return true; + } + + bool trySpecializeRotationArguments(func::CallOp callOp, func::FuncOp funcOp, + double angle, unsigned operand, + PatternRewriter& rewriter) const { + if (std::ranges::none_of(ANGLES_TO_SPECIALIZE, [angle](double a) { + return std::abs(a - angle) < 1e-9; + })) { + return false; + } + + const std::string suffix = "_spec_fixed_angle_" + std::to_string(operand); + if (funcOp.getName().contains(suffix)) { + // Already specialized + return false; + } + + auto key = std::make_tuple(funcOp.getName().str(), operand, angle); + if (previousSpecializations.rotationSpecializations.contains(key)) { + updateSpecializedCall( + callOp, previousSpecializations.rotationSpecializations.at(key), + rewriter); + return true; + } + + auto newFunc = + copyFunction(funcOp, funcOp.getName().str() + suffix, rewriter); + symbolTable.insert(newFunc); + previousSpecializations.rotationSpecializations.insert({key, newFunc}); + + auto newParameter = newFunc.getArgument(operand); + rewriter.setInsertionPointToStart(&*newFunc.getBody().getBlocks().begin()); + auto constant = rewriter.create( + newFunc.getBody().getLoc(), + rewriter.getFloatAttr(Float64Type::get(rewriter.getContext()), angle)); + rewriter.replaceAllUsesWith(newParameter, constant.getResult()); + + updateSpecializedCall(callOp, newFunc, rewriter); + return true; + } +}; + +} // namespace + +/** + * @brief Populates the given pattern set with the different IPO patterns. + * + * @param patterns The pattern set to populate. + */ +static void +populateQuantumIPOPatterns(RewritePatternSet& patterns, + SymbolTable& symbolTable, + PreviousSpecializations& previousSpecializations) { + patterns.add( + patterns.getContext(), symbolTable, previousSpecializations); +} + +namespace { + +/** + * @brief This pass performs quantum inter-procedural optimizations (IPO). + */ +struct QuantumIPO final : impl::QuantumIPOBase { + using impl::QuantumIPOBase::QuantumIPOBase; + +protected: + void runOnOperation() override { + // Get the current operation being operated on. + auto op = getOperation(); + auto* ctx = &getContext(); + SymbolTable symbolTable(op); + + // Define the set of patterns to use. + RewritePatternSet patterns(ctx); + PreviousSpecializations previousSpecializations; + populateQuantumIPOPatterns(patterns, symbolTable, previousSpecializations); + + // Apply patterns in an iterative and greedy manner. + if (failed(applyPatternsGreedily(op, std::move(patterns)))) { + signalPassFailure(); + } + + runQuantumArgumentPromotion(op); + runAuxiliaryQubitHoisting(op); + runQuantumFunctionBoundaryCommutation(op, symbolTable); + runQuantumFunctionBoundaryCommutation(op, symbolTable); + } +}; + +} // namespace + +} // namespace mlir::qco diff --git a/mlir/lib/Dialect/QTensor/Utils/CMakeLists.txt b/mlir/lib/Dialect/QTensor/Utils/CMakeLists.txt index 2014cafbe6..4e3943c1d0 100644 --- a/mlir/lib/Dialect/QTensor/Utils/CMakeLists.txt +++ b/mlir/lib/Dialect/QTensor/Utils/CMakeLists.txt @@ -19,6 +19,7 @@ add_mlir_dialect_library( PUBLIC MLIRQCODialect PRIVATE + MLIRFuncDialect MLIRSCFDialect) mqt_mlir_target_use_project_options(MLIRQTensorUtils) diff --git a/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp b/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp index 3d984f44ea..83081a6781 100644 --- a/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp +++ b/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -97,8 +98,11 @@ void TensorIterator::forward() { assert(tensor_.hasOneUse() && "expected linear typing"); op_ = *(tensor_.user_begin()); - // The following operations define the end of the tensor's life-chain. - if (isa(op_)) { + // The following operations define the end of the tensor's life-chain. A + // `func.call` ends it because the tensor is handed to the callee; the tensor + // the call returns starts a life-chain of its own. + if (isa(op_)) { isFinal_ = true; return; } @@ -144,8 +148,16 @@ void TensorIterator::backward() { return; } + // A `func.call` sits on both sides of a life-chain: it consumes the caller's + // tensor and produces a fresh one. When the tensor is the call's result, it + // is the start of its chain, just like an allocation. + if (isa(op_) && tensor_.getDefiningOp() == op_) { + return; + } + // For these operations, tensor_ is an OpOperand. Hence, only get the def-op. - if (isa(op_)) { + if (isa(op_)) { op_ = tensor_.getDefiningOp(); isFinal_ = false; return; diff --git a/mlir/lib/Support/IRVerification.cpp b/mlir/lib/Support/IRVerification.cpp index c7cbd6da84..c9efe19e9e 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -11,7 +11,6 @@ #include "mlir/Support/IRVerification.h" #include "mlir/Dialect/QC/IR/QCOps.h" -#include "mlir/Dialect/QCO/IR/QCODialect.h" #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" @@ -61,6 +60,19 @@ struct TensorMapping { const auto i = lhsEquivGroups.at(lhs); return equivGroupMapping.at(i) == rhsEquivGroups.at(rhs); } + + /// Return true if the given lhs value takes part in the equivalence + /// tracking. Only tensors reachable from a `qtensor` allocation are tracked; + /// builtin tensors of qubits are compared through the regular SSA mapping. + [[nodiscard]] bool tracksLhs(Value lhs) const { + return lhsEquivGroups.contains(lhs); + } + + /// Return true if the given rhs value takes part in the equivalence + /// tracking. + [[nodiscard]] bool tracksRhs(Value rhs) const { + return rhsEquivGroups.contains(rhs); + } }; } // namespace @@ -69,16 +81,6 @@ static bool compareRegions(Region& lhs, Region& rhs, SetVector& rhsClosed, IRMapping& m, TensorMapping& tm); -/// Return true, if the given value has the type `tensor`. -static bool hasTypeQubitTensor(Value v) { - auto tensor = dyn_cast(v.getType()); - if (!tensor) { - return false; - } - - return isa(tensor.getElementType()); -} - /// Recursively initialize the equivalence group for a tensor value. static void initEquivGroup(TypedValue v, size_t id, DenseMap& group) { @@ -206,10 +208,10 @@ getPermutation(const LhsRange& lhs, const RhsRange& rhs, const IRMapping& m, const TensorMapping& tm) { SmallVector permutation(lhs.size()); for (const auto& [i, lhsValue] : llvm::enumerate(lhs)) { - const auto it = hasTypeQubitTensor(lhsValue) + const auto it = tm.tracksLhs(lhsValue) ? llvm::find_if(rhs, [&](const auto rhsValue) { - if (!hasTypeQubitTensor(rhsValue)) { + if (!tm.tracksRhs(rhsValue)) { return false; } return tm.equals(lhsValue, rhsValue); @@ -233,9 +235,9 @@ static bool compareValueLists(const LhsRange& lhs, const RhsRange& rhs, for (const auto lhsValue : lhs) { Value mapped; - if (hasTypeQubitTensor(lhsValue)) { + if (tm.tracksLhs(lhsValue)) { const auto it = llvm::find_if(rhs, [&](const auto rhsValue) { - return hasTypeQubitTensor(rhsValue) && tm.equals(lhsValue, rhsValue); + return tm.tracksRhs(rhsValue) && tm.equals(lhsValue, rhsValue); }); if (it == rhs.end()) { return false; @@ -478,8 +480,10 @@ static bool compareOperations(Operation* lhs, Operation* rhs, } else { for (const auto& [lhsOperand, rhsOperand] : llvm::zip_equal(lhs->getOperands(), rhs->getOperands())) { - if (hasTypeQubitTensor(lhsOperand)) { - assert(hasTypeQubitTensor(rhsOperand)); + if (tm.tracksLhs(lhsOperand)) { + if (!tm.tracksRhs(rhsOperand)) { + return false; + } if (!tm.equals(lhsOperand, rhsOperand)) { return false; @@ -745,6 +749,9 @@ static bool compareBlocks(Block& lhs, Block& rhs, auto lhsExtract = cast(lhsOp); auto rhsExtract = cast(rhsOp); m.map(lhsExtract.getResult(), rhsExtract.getResult()); + // The threaded tensor is only covered by the equivalence groups + // when it descends from an allocation, so map it here as well. + m.map(lhsExtract.getOutTensor(), rhsExtract.getOutTensor()); } else { SmallVector permutation(lhsOp->getNumResults()); std::iota(permutation.begin(), permutation.end(), 0); diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 960d2e5675..777e1ebe77 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -174,6 +174,60 @@ TEST_F(QCOTest, BuilderRejectsOutOfBoundsClassicalRegisterIndices) { "Register index is out of bounds"); } +TEST_F(QCOTest, BuilderRejectsMisuseOfAdditionalFunctions) { + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + const auto qubitType = builder.getQubitType(); + builder.startFunction("f", {qubitType}, {qubitType}); + builder.startFunction("g", {qubitType}, {qubitType}); + }, + "Cannot start a function while another one is being built"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + builder.endFunction({}); + }, + "endFunction\\(\\) called without a matching startFunction\\(\\)"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + builder.call("does_not_exist", {}); + }, + "Callee not found in module"); +} + +TEST_F(QCOTest, BuilderRejectsLinearValuesLeakingOutOfFunctions) { + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + const auto qubitType = builder.getQubitType(); + const auto args = builder.startFunction("f", {qubitType}, {qubitType}); + // The freshly allocated qubit is neither returned nor sunk. + builder.allocQubit(); + builder.endFunction({args[0]}); + }, + "neither returned nor consumed"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + const auto qubitType = builder.getQubitType(); + const auto args = builder.startFunction("f", {qubitType}, {qubitType}); + // The freshly allocated register is neither returned nor deallocated. + builder.qtensorAlloc(2); + builder.endFunction({args[0]}); + }, + "neither returned nor deallocated"); +} + TEST_F(QCOTest, DirectSingleQubitPowBuilder) { QCOProgramBuilder builder(context.get()); builder.initialize(); diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt index e269cc8057..7cb5770e21 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt @@ -9,9 +9,13 @@ set(target_name mqt-core-mlir-unittest-optimizations) add_executable( ${target_name} - test_qco_hadamard_lifting.cpp test_qco_measurement_lifting.cpp - test_qco_merge_single_qubit_rotation.cpp test_qco_replace_classical_controls.cpp - test_qco_reuse_qubits.cpp test_quantum_loop_unroll.cpp) + test_qco_hadamard_lifting.cpp + test_qco_measurement_lifting.cpp + test_qco_merge_single_qubit_rotation.cpp + test_qco_quantum_ipo.cpp + test_qco_replace_classical_controls.cpp + test_qco_reuse_qubits.cpp + test_quantum_loop_unroll.cpp) target_link_libraries( ${target_name} diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_quantum_ipo.cpp b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_quantum_ipo.cpp new file mode 100644 index 0000000000..b439c50f7c --- /dev/null +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_quantum_ipo.cpp @@ -0,0 +1,1675 @@ +/* + * 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/Builder/QCOProgramBuilder.h" +#include "mlir/Dialect/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QCO/Transforms/Passes.h" +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" +#include "mlir/Support/IRVerification.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace { + +using namespace mlir; +using namespace mlir::qco; + +class QCOQuantumIPOTest : public testing::Test { + +protected: + MLIRContext context; + QCOProgramBuilder programBuilder; + QCOProgramBuilder referenceBuilder; + OwningOpRef module; + OwningOpRef reference; + + QCOQuantumIPOTest() : programBuilder(&context), referenceBuilder(&context) {} + + void SetUp() override { + // Register all necessary dialects + DialectRegistry registry; + registry.insert(); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + } + + /** + * @brief Adds the quantum IPO pass to the current context and runs it. + * + * @param module The module to run the pass on. + */ + static LogicalResult runQuantumIPOPass(ModuleOp module) { + PassManager pm(module.getContext()); + pm.addPass(createQuantumIPO()); + pm.addPass(createCanonicalizerPass()); + return pm.run(module); + } + + /** + * @brief Adds the canonicalizerPass to the current context and runs it. + */ + static LogicalResult runCanonicalizerPass(ModuleOp module) { + PassManager pm(module.getContext()); + pm.addPass(createCanonicalizerPass()); + return pm.run(module); + } + + /** + * @brief Runs the pass on the constructed module and compares it against the + * constructed reference. + */ + void expectModuleMatchesReference() { + ASSERT_TRUE(runQuantumIPOPass(module.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(module.get(), reference.get())); + } +}; + +} // namespace + +// ========================================================================== +// Context-sensitive specialization for arguments in the |0> state. +// ========================================================================== + +/** + * @brief A gate that acts trivially on |0> is dropped from a specialized copy + * of the callee when the caller passes a freshly allocated qubit. + */ +TEST_F(QCOQuantumIPOTest, specializeZeroArgumentDropsDiagonalGate) { + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {programBuilder.getQubitType()}, + {programBuilder.getQubitType()}); + programBuilder.endFunction({programBuilder.z(args[0])}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + // The original callee is retained, ... + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.z(refArgs[0])}); + // ... while the call is redirected to a specialization without the gate. + auto specArgs = referenceBuilder.startFunction( + "f_spec_zero_arg_0", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({specArgs[0]}); + + auto refQ = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f_spec_zero_arg_0", {refQ}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief A controlled gate whose control is known to be in the |0> state is + * dropped entirely, together with its effect on the target qubit. + */ +TEST_F(QCOQuantumIPOTest, specializeZeroArgumentDropsControlledGate) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto control = args[0]; + auto target = args[1]; + std::tie(control, target) = programBuilder.cx(control, target); + programBuilder.endFunction({control, target}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.h(programBuilder.allocQubit()); + auto results = programBuilder.call("f", {q0, q1}); + programBuilder.sink(results[0]); + programBuilder.sink(results[1]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refControl = refArgs[0]; + auto refTarget = refArgs[1]; + std::tie(refControl, refTarget) = referenceBuilder.cx(refControl, refTarget); + referenceBuilder.endFunction({refControl, refTarget}); + + auto specArgs = referenceBuilder.startFunction( + "f_spec_zero_arg_0", {qubitType, qubitType}, {qubitType, qubitType}); + referenceBuilder.endFunction({specArgs[0], specArgs[1]}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.h(referenceBuilder.allocQubit()); + auto refResults = referenceBuilder.call("f_spec_zero_arg_0", {refQ0, refQ1}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief A gate that does not act trivially on |0> must not be dropped. + */ +TEST_F(QCOQuantumIPOTest, noZeroSpecializationForNonTrivialGate) { + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {programBuilder.getQubitType()}, + {programBuilder.getQubitType()}); + programBuilder.endFunction({programBuilder.x(args[0])}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.x(refArgs[0])}); + + auto refQ = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("f", {refQ}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief If the state of the argument is unknown, no specialization applies. + */ +TEST_F(QCOQuantumIPOTest, noSpecializationForUnknownArgumentState) { + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {programBuilder.getQubitType()}, + {programBuilder.getQubitType()}); + programBuilder.endFunction({programBuilder.z(args[0])}); + + // A `y` gate leaves the qubit in a state the pass cannot reason about. + auto q = programBuilder.y(programBuilder.allocQubit()); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.z(refArgs[0])}); + + auto refQ = referenceBuilder.y(referenceBuilder.allocQubit()); + auto refResults = referenceBuilder.call("f", {refQ}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief Two call sites that qualify for the same specialization share a single + * specialized copy of the callee. + */ +TEST_F(QCOQuantumIPOTest, reuseZeroSpecializationAcrossCallSites) { + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {programBuilder.getQubitType()}, + {programBuilder.getQubitType()}); + programBuilder.endFunction({programBuilder.s(args[0])}); + + 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]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.s(refArgs[0])}); + + auto specArgs = referenceBuilder.startFunction( + "f_spec_zero_arg_0", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({specArgs[0]}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refResults0 = referenceBuilder.call("f_spec_zero_arg_0", {refQ0}); + auto refResults1 = referenceBuilder.call("f_spec_zero_arg_0", {refQ1}); + referenceBuilder.sink(refResults0[0]); + referenceBuilder.sink(refResults1[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +// ========================================================================== +// Context-sensitive specialization for arguments in the |+> state. +// ========================================================================== + +/** + * @brief An `x` gate acting on a qubit known to be in the |+> state is dropped + * from a specialized copy of the callee. + */ +TEST_F(QCOQuantumIPOTest, specializePlusArgumentDropsXGate) { + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {programBuilder.getQubitType()}, + {programBuilder.getQubitType()}); + programBuilder.endFunction({programBuilder.x(args[0])}); + + auto q = programBuilder.h(programBuilder.allocQubit()); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.x(refArgs[0])}); + + auto specArgs = referenceBuilder.startFunction( + "f_spec_plus_arg_0", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({specArgs[0]}); + + auto refQ = referenceBuilder.h(referenceBuilder.allocQubit()); + auto refResults = referenceBuilder.call("f_spec_plus_arg_0", {refQ}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief A gate that does not act trivially on |+> must not be dropped. + */ +TEST_F(QCOQuantumIPOTest, noPlusSpecializationForNonXGate) { + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {programBuilder.getQubitType()}, + {programBuilder.getQubitType()}); + programBuilder.endFunction({programBuilder.z(args[0])}); + + auto q = programBuilder.h(programBuilder.allocQubit()); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.z(refArgs[0])}); + + auto refQ = referenceBuilder.h(referenceBuilder.allocQubit()); + auto refResults = referenceBuilder.call("f", {refQ}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +// ========================================================================== +// Context-sensitive specialization for constant rotation angles. +// ========================================================================== + +/** + * @brief A rotation angle of pi passed at the call site is baked into a + * specialized copy of the callee. + */ +TEST_F(QCOQuantumIPOTest, specializeConstantRotationAngle) { + const auto qubitType = programBuilder.getQubitType(); + const auto floatType = programBuilder.getF64Type(); + + programBuilder.initialize(); + auto args = + programBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + programBuilder.endFunction({programBuilder.rz(args[1], args[0])}); + + auto q = programBuilder.allocQubit(); + auto angle = programBuilder.floatConstant(std::numbers::pi); + auto results = programBuilder.call("f", {q, angle}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction({referenceBuilder.rz(refArgs[1], refArgs[0])}); + + // The specialized copy keeps the parameter in its signature but no longer + // reads it; the angle becomes a constant in the body. + auto specArgs = referenceBuilder.startFunction( + "f_spec_fixed_angle_1", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.rz(std::numbers::pi, specArgs[0])}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAngle = referenceBuilder.floatConstant(std::numbers::pi); + auto refResults = + referenceBuilder.call("f_spec_fixed_angle_1", {refQ, refAngle}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief A rotation angle of pi/2 is likewise specialized. + */ +TEST_F(QCOQuantumIPOTest, specializeHalfPiRotationAngle) { + const auto qubitType = programBuilder.getQubitType(); + const auto floatType = programBuilder.getF64Type(); + + programBuilder.initialize(); + auto args = + programBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + programBuilder.endFunction({programBuilder.rx(args[1], args[0])}); + + auto q = programBuilder.allocQubit(); + auto angle = programBuilder.floatConstant(std::numbers::pi / 2); + auto results = programBuilder.call("f", {q, angle}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction({referenceBuilder.rx(refArgs[1], refArgs[0])}); + + auto specArgs = referenceBuilder.startFunction( + "f_spec_fixed_angle_1", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.rx(std::numbers::pi / 2, specArgs[0])}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAngle = referenceBuilder.floatConstant(std::numbers::pi / 2); + auto refResults = + referenceBuilder.call("f_spec_fixed_angle_1", {refQ, refAngle}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief An angle outside the set of specialized angles leaves the callee + * untouched. + */ +TEST_F(QCOQuantumIPOTest, noSpecializationForArbitraryRotationAngle) { + const auto qubitType = programBuilder.getQubitType(); + const auto floatType = programBuilder.getF64Type(); + + programBuilder.initialize(); + auto args = + programBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + programBuilder.endFunction({programBuilder.rz(args[1], args[0])}); + + auto q = programBuilder.allocQubit(); + auto angle = programBuilder.floatConstant(0.7); + auto results = programBuilder.call("f", {q, angle}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction({referenceBuilder.rz(refArgs[1], refArgs[0])}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAngle = referenceBuilder.floatConstant(0.7); + auto refResults = referenceBuilder.call("f", {refQ, refAngle}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief Two call sites that qualify for the same |+> specialization share a + * single specialized copy of the callee. + */ +TEST_F(QCOQuantumIPOTest, reusePlusSpecializationAcrossCallSites) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + programBuilder.endFunction({programBuilder.x(args[0])}); + + auto q0 = programBuilder.h(programBuilder.allocQubit()); + auto q1 = programBuilder.h(programBuilder.allocQubit()); + auto results0 = programBuilder.call("f", {q0}); + auto results1 = programBuilder.call("f", {q1}); + programBuilder.sink(results0[0]); + programBuilder.sink(results1[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType}, {qubitType}); + referenceBuilder.endFunction({referenceBuilder.x(refArgs[0])}); + auto specArgs = referenceBuilder.startFunction("f_spec_plus_arg_0", + {qubitType}, {qubitType}); + referenceBuilder.endFunction({specArgs[0]}); + + auto refQ0 = referenceBuilder.h(referenceBuilder.allocQubit()); + auto refQ1 = referenceBuilder.h(referenceBuilder.allocQubit()); + auto refResults0 = referenceBuilder.call("f_spec_plus_arg_0", {refQ0}); + auto refResults1 = referenceBuilder.call("f_spec_plus_arg_0", {refQ1}); + referenceBuilder.sink(refResults0[0]); + referenceBuilder.sink(refResults1[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief Two call sites passing the same constant angle share a single + * specialized copy of the callee. + */ +TEST_F(QCOQuantumIPOTest, reuseRotationSpecializationAcrossCallSites) { + const auto qubitType = programBuilder.getQubitType(); + const auto floatType = programBuilder.getF64Type(); + + programBuilder.initialize(); + auto args = + programBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + programBuilder.endFunction({programBuilder.rz(args[1], args[0])}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto angle = programBuilder.floatConstant(std::numbers::pi); + auto results0 = programBuilder.call("f", {q0, angle}); + auto results1 = programBuilder.call("f", {q1, angle}); + programBuilder.sink(results0[0]); + programBuilder.sink(results1[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction({referenceBuilder.rz(refArgs[1], refArgs[0])}); + auto specArgs = referenceBuilder.startFunction( + "f_spec_fixed_angle_1", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.rz(std::numbers::pi, specArgs[0])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refAngle = referenceBuilder.floatConstant(std::numbers::pi); + auto refResults0 = + referenceBuilder.call("f_spec_fixed_angle_1", {refQ0, refAngle}); + auto refResults1 = + referenceBuilder.call("f_spec_fixed_angle_1", {refQ1, refAngle}); + referenceBuilder.sink(refResults0[0]); + referenceBuilder.sink(refResults1[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +// ========================================================================== +// 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(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @brief Only the tensor elements the callee actually touches become scalar + * arguments; untouched elements never cross the call boundary. + */ +TEST_F(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @brief A qubit that is moved to a different slot is promoted with the + * extraction and insertion indices kept apart. + */ +TEST_F(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @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(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @brief A tensor argument that never has an element taken out of it has + * nothing to promote. + */ +TEST_F(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @brief A callee that touches several tensor elements gets one scalar argument + * and one scalar result per element. + */ +TEST_F(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @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(QCOQuantumIPOTest, 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]); + module = 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]}); + + expectModuleMatchesReference(); +} + +/** + * @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(QCOQuantumIPOTest, 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]; + }; + + module = programBuilder.finalize({buildProgram(programBuilder)}); + reference = referenceBuilder.finalize({buildProgram(referenceBuilder)}); + + expectModuleMatchesReference(); +} + +// ========================================================================== +// 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(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @brief A qubit that the callee allocates but hands back to the caller is not + * auxiliary and must stay where it is. + */ +TEST_F(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @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(QCOQuantumIPOTest, 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]); + module = 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]}); + + expectModuleMatchesReference(); +} + +/** + * @brief The auxiliary qubit is tracked while it is parked in a tensor, past + * an extraction of an unrelated element. + */ +TEST_F(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @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(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @brief Every call site of a hoisted callee gets its own allocation. + */ +TEST_F(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @brief A recursive callee is left alone, because its allocation would have to + * be threaded through every level of the recursion. A caller of a recursive + * function is still hoisted. + */ +TEST_F(QCOQuantumIPOTest, noHoistingForRecursiveFunction) { + const auto qubitType = programBuilder.getQubitType(); + + const auto buildRecursiveCallee = [&qubitType](QCOProgramBuilder& b) { + auto innerArgs = b.startFunction("inner", {qubitType}, {qubitType}); + auto innerAux = b.allocQubit(); + auto innerTarget = innerArgs[0]; + std::tie(innerAux, innerTarget) = b.cx(innerAux, innerTarget); + b.sink(innerAux); + b.endFunction({b.call("inner", {innerTarget})[0]}); + }; + + programBuilder.initialize(); + buildRecursiveCallee(programBuilder); + + auto args = programBuilder.startFunction("outer", {qubitType}, {qubitType}); + auto aux = programBuilder.allocQubit(); + auto target = args[0]; + std::tie(aux, target) = programBuilder.cx(aux, target); + programBuilder.sink(aux); + programBuilder.endFunction({programBuilder.call("inner", {target})[0]}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("outer", {q}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + // `inner` is recursive and therefore untouched, ... + buildRecursiveCallee(referenceBuilder); + + // ... while `outer` is hoisted even though it calls a recursive function. + auto refArgs = referenceBuilder.startFunction("outer", {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( + {referenceBuilder.call("inner", {refTarget})[0], refAux}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAuxAlloc = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call("outer", {refQ, refAuxAlloc}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief An allocation nested inside a region is not hoisted, because it is not + * executed on every path through the function. + */ +TEST_F(QCOQuantumIPOTest, 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); + module = programBuilder.finalize(); + + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @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(QCOQuantumIPOTest, 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]); + module = 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(); + + expectModuleMatchesReference(); +} + +/** + * @brief A call that only consumes linear values, and one that only produces + * them, keep the builder's tracking consistent. + */ +TEST_F(QCOQuantumIPOTest, callConsumesAndProducesLinearValues) { + const auto qubitType = programBuilder.getQubitType(); + const auto tensorType = programBuilder.getQubitTensorType(2); + + const auto buildProgram = [&](QCOProgramBuilder& b) { + b.initialize(); + // Allocate before declaring the helpers so that the function scope has to + // remember the already-tracked values of the caller. + auto q = b.allocQubit(); + auto scratch = b.qtensorAlloc(2); + + 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)}); + + b.call("consume", {q, scratch}); + auto produced = b.call("produce", {}); + b.qtensorDealloc(produced[0]); + }; + + buildProgram(programBuilder); + module = programBuilder.finalize(); + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +// ========================================================================== +// Quantum function boundary commutation. +// ========================================================================== + +/** + * @brief A self-inverse gate applied right before a call cancels with the same + * gate at the start of the callee. + */ +TEST_F(QCOQuantumIPOTest, cancelSelfInverseGateAcrossCallBoundary) { + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {programBuilder.getQubitType()}, + {programBuilder.getQubitType()}); + programBuilder.endFunction({programBuilder.h(programBuilder.x(args[0]))}); + + auto q = programBuilder.x(programBuilder.allocQubit()); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction( + {referenceBuilder.h(referenceBuilder.x(refArgs[0]))}); + + // Both the caller-side and the callee-side gate disappear. + auto specArgs = referenceBuilder.startFunction( + "f_spec_boundary_commutation", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.h(specArgs[0])}); + + auto refQ = referenceBuilder.allocQubit(); + auto refResults = + referenceBuilder.call("f_spec_boundary_commutation", {refQ}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief Two different gates across the call boundary do not cancel. + */ +TEST_F(QCOQuantumIPOTest, noCancellationForDifferentGates) { + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {programBuilder.getQubitType()}, + {programBuilder.getQubitType()}); + programBuilder.endFunction({programBuilder.y(args[0])}); + + auto q = programBuilder.x(programBuilder.allocQubit()); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = + referenceBuilder.startFunction("f", {referenceBuilder.getQubitType()}, + {referenceBuilder.getQubitType()}); + referenceBuilder.endFunction({referenceBuilder.y(refArgs[0])}); + + auto refQ = referenceBuilder.x(referenceBuilder.allocQubit()); + auto refResults = referenceBuilder.call("f", {refQ}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief Controlled gates are out of scope for boundary commutation, even when + * the same one appears on both sides of the call. Cancelling them would require + * reasoning about the control qubits as well. + */ +TEST_F(QCOQuantumIPOTest, noCancellationForControlledGates) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto innerControl = args[0]; + auto innerTarget = args[1]; + std::tie(innerControl, innerTarget) = + programBuilder.cx(innerControl, innerTarget); + programBuilder.endFunction({innerControl, innerTarget}); + + auto q0 = programBuilder.y(programBuilder.allocQubit()); + auto q1 = programBuilder.y(programBuilder.allocQubit()); + std::tie(q0, q1) = programBuilder.cx(q0, q1); + auto results = programBuilder.call("f", {q0, q1}); + programBuilder.sink(results[0]); + programBuilder.sink(results[1]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refInnerControl = refArgs[0]; + auto refInnerTarget = refArgs[1]; + std::tie(refInnerControl, refInnerTarget) = + referenceBuilder.cx(refInnerControl, refInnerTarget); + referenceBuilder.endFunction({refInnerControl, refInnerTarget}); + + auto refQ0 = referenceBuilder.y(referenceBuilder.allocQubit()); + auto refQ1 = referenceBuilder.y(referenceBuilder.allocQubit()); + std::tie(refQ0, refQ1) = referenceBuilder.cx(refQ0, refQ1); + auto refResults = referenceBuilder.call("f", {refQ0, refQ1}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief Two call sites that cancel the same gate share a single commuted copy + * of the callee. + */ +TEST_F(QCOQuantumIPOTest, reuseBoundaryCommutationAcrossCallSites) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + programBuilder.endFunction({programBuilder.h(programBuilder.x(args[0]))}); + + auto q0 = programBuilder.x(programBuilder.allocQubit()); + auto q1 = programBuilder.x(programBuilder.allocQubit()); + auto results0 = programBuilder.call("f", {q0}); + auto results1 = programBuilder.call("f", {q1}); + programBuilder.sink(results0[0]); + programBuilder.sink(results1[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.h(referenceBuilder.x(refArgs[0]))}); + auto specArgs = referenceBuilder.startFunction("f_spec_boundary_commutation", + {qubitType}, {qubitType}); + referenceBuilder.endFunction({referenceBuilder.h(specArgs[0])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refResults0 = + referenceBuilder.call("f_spec_boundary_commutation", {refQ0}); + auto refResults1 = + referenceBuilder.call("f_spec_boundary_commutation", {refQ1}); + referenceBuilder.sink(refResults0[0]); + referenceBuilder.sink(refResults1[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +// ========================================================================== +// Integration tests combining several IPO approaches. +// ========================================================================== + +/** + * @brief A callee that both starts with a gate that is trivial on |0> and uses + * an internal auxiliary qubit is first specialized and then hoisted. The + * hoisting applies to the original and the specialized copy alike. + */ +TEST_F(QCOQuantumIPOTest, specializationAndHoistingCombined) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + auto aux = programBuilder.allocQubit(); + auto target = programBuilder.z(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]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + // The original keeps its `z` gate, but is hoisted as well. + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refAux = refArgs[1]; + auto refTarget = referenceBuilder.z(refArgs[0]); + std::tie(refAux, refTarget) = referenceBuilder.cx(refAux, refTarget); + refAux = referenceBuilder.reset(refAux); + referenceBuilder.endFunction({refTarget, refAux}); + + // The specialization drops the `z` gate and is hoisted too. + auto specArgs = referenceBuilder.startFunction( + "f_spec_zero_arg_0", {qubitType, qubitType}, {qubitType, qubitType}); + auto specAux = specArgs[1]; + auto specTarget = specArgs[0]; + std::tie(specAux, specTarget) = referenceBuilder.cx(specAux, specTarget); + specAux = referenceBuilder.reset(specAux); + referenceBuilder.endFunction({specTarget, specAux}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAuxAlloc = referenceBuilder.allocQubit(); + auto refResults = + referenceBuilder.call("f_spec_zero_arg_0", {refQ, refAuxAlloc}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief A callee with two qubit arguments where one argument is specialized + * for the |0> state and the other cancels a gate across the call boundary. + */ +TEST_F(QCOQuantumIPOTest, specializationAndBoundaryCommutationCombined) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto first = programBuilder.z(args[0]); + auto second = programBuilder.x(args[1]); + second = programBuilder.h(second); + programBuilder.endFunction({first, second}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.x(programBuilder.allocQubit()); + auto results = programBuilder.call("f", {q0, q1}); + programBuilder.sink(results[0]); + programBuilder.sink(results[1]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refArgs = referenceBuilder.startFunction("f", {qubitType, qubitType}, + {qubitType, qubitType}); + auto refFirst = referenceBuilder.z(refArgs[0]); + auto refSecond = referenceBuilder.x(refArgs[1]); + refSecond = referenceBuilder.h(refSecond); + referenceBuilder.endFunction({refFirst, refSecond}); + + // The |0> specialization drops the `z` gate on the first argument. + auto specArgs = referenceBuilder.startFunction( + "f_spec_zero_arg_0", {qubitType, qubitType}, {qubitType, qubitType}); + auto specSecond = referenceBuilder.x(specArgs[1]); + specSecond = referenceBuilder.h(specSecond); + referenceBuilder.endFunction({specArgs[0], specSecond}); + + // Boundary commutation then removes the `x` gates around the call. + auto commutedArgs = referenceBuilder.startFunction( + "f_spec_zero_arg_0_spec_boundary_commutation", {qubitType, qubitType}, + {qubitType, qubitType}); + referenceBuilder.endFunction( + {commutedArgs[0], referenceBuilder.h(commutedArgs[1])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refResults = referenceBuilder.call( + "f_spec_zero_arg_0_spec_boundary_commutation", {refQ0, refQ1}); + referenceBuilder.sink(refResults[0]); + referenceBuilder.sink(refResults[1]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} + +/** + * @brief A program with several distinct callees, each hitting a different IPO + * approach: a |0> specialization, a fixed rotation angle, and a cancellation + * across the call boundary. + */ +TEST_F(QCOQuantumIPOTest, multipleFunctionsWithDistinctOptimizations) { + const auto qubitType = programBuilder.getQubitType(); + const auto floatType = programBuilder.getF64Type(); + + programBuilder.initialize(); + auto prepareArgs = + programBuilder.startFunction("prepare", {qubitType}, {qubitType}); + programBuilder.endFunction( + {programBuilder.h(programBuilder.z(prepareArgs[0]))}); + + auto rotateArgs = programBuilder.startFunction( + "rotate", {qubitType, floatType}, {qubitType}); + programBuilder.endFunction({programBuilder.rz(rotateArgs[1], rotateArgs[0])}); + + auto flipArgs = + programBuilder.startFunction("flip", {qubitType}, {qubitType}); + programBuilder.endFunction({programBuilder.y(programBuilder.x(flipArgs[0]))}); + + auto q0 = programBuilder.allocQubit(); + auto q1 = programBuilder.allocQubit(); + auto q2 = programBuilder.x(programBuilder.allocQubit()); + auto prepared = programBuilder.call("prepare", {q0}); + auto angle = programBuilder.floatConstant(std::numbers::pi / 2); + auto rotated = programBuilder.call("rotate", {q1, angle}); + auto flipped = programBuilder.call("flip", {q2}); + programBuilder.sink(prepared[0]); + programBuilder.sink(rotated[0]); + programBuilder.sink(flipped[0]); + module = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto refPrepareArgs = + referenceBuilder.startFunction("prepare", {qubitType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.h(referenceBuilder.z(refPrepareArgs[0]))}); + auto preparedSpecArgs = referenceBuilder.startFunction( + "prepare_spec_zero_arg_0", {qubitType}, {qubitType}); + referenceBuilder.endFunction({referenceBuilder.h(preparedSpecArgs[0])}); + + auto refRotateArgs = referenceBuilder.startFunction( + "rotate", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.rz(refRotateArgs[1], refRotateArgs[0])}); + auto rotateSpecArgs = referenceBuilder.startFunction( + "rotate_spec_fixed_angle_1", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.rz(std::numbers::pi / 2, rotateSpecArgs[0])}); + + auto refFlipArgs = + referenceBuilder.startFunction("flip", {qubitType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.y(referenceBuilder.x(refFlipArgs[0]))}); + auto flipSpecArgs = referenceBuilder.startFunction( + "flip_spec_boundary_commutation", {qubitType}, {qubitType}); + referenceBuilder.endFunction({referenceBuilder.y(flipSpecArgs[0])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refQ2 = referenceBuilder.allocQubit(); + auto refPrepared = referenceBuilder.call("prepare_spec_zero_arg_0", {refQ0}); + auto refAngle = referenceBuilder.floatConstant(std::numbers::pi / 2); + auto refRotated = + referenceBuilder.call("rotate_spec_fixed_angle_1", {refQ1, refAngle}); + auto refFlipped = + referenceBuilder.call("flip_spec_boundary_commutation", {refQ2}); + referenceBuilder.sink(refPrepared[0]); + referenceBuilder.sink(refRotated[0]); + referenceBuilder.sink(refFlipped[0]); + reference = referenceBuilder.finalize(); + + expectModuleMatchesReference(); +} diff --git a/src/qasm3/Importer.cpp b/src/qasm3/Importer.cpp index 6a9fa75818..586fc29b86 100644 --- a/src/qasm3/Importer.cpp +++ b/src/qasm3/Importer.cpp @@ -372,8 +372,16 @@ void Importer::visitDeclarationStatement( throw CompilerError("Angle type is currently not supported.", declarationStatement->debugInfo); } + } else if (const auto sizedTy = + std::dynamic_pointer_cast>(ty)) { + if (sizedTy->type == SingleQubit) { + qc->addQubitRegister(1, identifier); + } else { + throw CompilerError("Only sized types or single qubits are supported.", + declarationStatement->debugInfo); + } } else { - throw CompilerError("Only sized types are supported.", + throw CompilerError("Only sized types or single qubits are supported.", declarationStatement->debugInfo); } declarations.emplace(identifier, declarationStatement); diff --git a/src/qasm3/Parser.cpp b/src/qasm3/Parser.cpp index 66f66439ac..c159a55d7d 100644 --- a/src/qasm3/Parser.cpp +++ b/src/qasm3/Parser.cpp @@ -859,7 +859,8 @@ std::pair, bool> Parser::parseType() { std::shared_ptr type; bool isOldStyleDeclaration = false; - switch (current().kind) { + auto keyword = current().kind; + switch (keyword) { case Token::Kind::CReg: type = DesignatedType>::getBitTy(nullptr); isOldStyleDeclaration = true; @@ -906,6 +907,11 @@ std::pair, bool> Parser::parseType() { type->setDesignator(std::move(designator)); return std::pair{std::move(type), isOldStyleDeclaration}; } + if (keyword == Token::Kind::Qubit) { + return std::pair{ + UnsizedType>::getSingleQubitTy(), + isOldStyleDeclaration}; + } return std::pair{std::move(type), isOldStyleDeclaration}; } diff --git a/src/qasm3/passes/TypeCheckPass.cpp b/src/qasm3/passes/TypeCheckPass.cpp index ab2d736739..9eec373470 100644 --- a/src/qasm3/passes/TypeCheckPass.cpp +++ b/src/qasm3/passes/TypeCheckPass.cpp @@ -375,7 +375,17 @@ InferredType TypeCheckPass::visitMeasureExpression( error("Unknown identifier '" + indexedIdentifier->identifier + "'."); return InferredType::error(); } - const auto width = it->second.type->getDesignator(); + uint64_t width = 0; + const auto type = it->second.type; + if (type->allowsDesignator()) { + width = type->getDesignator(); + } else { + const auto unsized = std::dynamic_pointer_cast>(type); + if (!unsized || unsized->type != qasm3::SingleQubit) { + return error("Cannot measure non-qubit type."); + } + width = 1; + } return InferredType{std::dynamic_pointer_cast( DesignatedType::getBitTy(width))}; } diff --git a/test/ir/test_qasm3_parser.cpp b/test/ir/test_qasm3_parser.cpp index 2dd4ae8e59..a404235202 100644 --- a/test/ir/test_qasm3_parser.cpp +++ b/test/ir/test_qasm3_parser.cpp @@ -206,7 +206,7 @@ include "stdgates.inc"; gate empty q0 { pow(0) @ x q0; } -qubit q; +qubit[1] q; pow(2) @ empty q; x q; )";