From 2ecc1d33c30d06efd03290e72d72fe985d26a499 Mon Sep 17 00:00:00 2001 From: Damian Rovara Date: Fri, 28 Aug 2026 12:24:10 +0000 Subject: [PATCH 1/2] =?UTF-8?q?=E2=9C=A8=20Add=20context-sensitive=20speci?= =?UTF-8?q?alization?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Specialize callees from values known at their call sites and cover the interprocedural behavior. Assisted-by: Claude Opus 5 Assisted-by: Codex --- CHANGELOG.md | 4 +- .../mlir/Dialect/QCO/Transforms/Passes.td | 25 + .../ContextSensitiveSpecialization.cpp | 548 ++++++++++++++ .../QCO/Transforms/Optimizations/IPOUtils.cpp | 62 ++ .../QCO/Transforms/Optimizations/IPOUtils.h | 62 ++ mlir/lib/Support/Passes.cpp | 1 + .../Transforms/Optimizations/CMakeLists.txt | 1 + .../Transforms/Optimizations/IPOTestFixture.h | 139 ++++ ...t_qco_context_sensitive_specialization.cpp | 701 ++++++++++++++++++ 9 files changed, 1542 insertions(+), 1 deletion(-) create mode 100644 mlir/lib/Dialect/QCO/Transforms/Optimizations/ContextSensitiveSpecialization.cpp create mode 100644 mlir/lib/Dialect/QCO/Transforms/Optimizations/IPOUtils.cpp create mode 100644 mlir/lib/Dialect/QCO/Transforms/Optimizations/IPOUtils.h create mode 100644 mlir/unittests/Dialect/QCO/Transforms/Optimizations/IPOTestFixture.h create mode 100644 mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_context_sensitive_specialization.cpp diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a3b9e9988..1d451e05e5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -63,7 +63,8 @@ releases may include breaking changes. #### Passes and transformations -- ✨ Add passes for quantum-specific interprocedural optimizations ([#2193]) +- ✨ Add passes for quantum-specific interprocedural optimizations ([#2193], + [#2197]) ([**@DRovara**], [**@burgholzer**]) - ✨ Add Pauli twirling, quantum loop unrolling, and qubit reuse passes ([#1705], [#1718], [#1755], [#1756], [#1923], [#1924], [#2039], [#2118], @@ -890,6 +891,7 @@ for previous changelogs._ [#2216]: https://github.com/munich-quantum-toolkit/core/pull/2216 [#2203]: https://github.com/munich-quantum-toolkit/core/pull/2203 [#2214]: https://github.com/munich-quantum-toolkit/core/pull/2214 +[#2197]: https://github.com/munich-quantum-toolkit/core/pull/2197 [#2196]: https://github.com/munich-quantum-toolkit/core/pull/2196 [#2194]: https://github.com/munich-quantum-toolkit/core/pull/2194 [#2193]: https://github.com/munich-quantum-toolkit/core/pull/2193 diff --git a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td index ff1e049157..676825f506 100644 --- a/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/QCO/Transforms/Passes.td @@ -325,6 +325,31 @@ def ReuseQubits : Pass<"reuse-qubits", "mlir::ModuleOp"> { }]; } +def ContextSensitiveSpecialization + : Pass<"quantum-context-sensitive-specialization", "mlir::ModuleOp"> { + let summary = "Specialize callees for what is known at their call sites"; + let description = [{ + Redirects a `func.call` to a specialized copy of its callee whenever the + call site pins down something the callee can exploit: + + - a qubit argument known to be in the |0> state, which lets operations + that fix |0> be dropped; + - a qubit argument known to be in the |+> state, which lets an `x` be + dropped; + - a compile-time constant rotation angle from a small distinguished set, + which is folded into the copy. + + Copies are shared between call sites with the same context. A callee + left without callers is erased, but only when this pass created the + situation; unrelated unused functions are left alone. + }]; + + let dependentDialects = ["::mlir::func::FuncDialect", + "::mlir::arith::ArithDialect", + "::mlir::qtensor::QTensorDialect", + "mlir::qco::QCODialect"]; +} + def RemoveDeadGates : Pass<"remove-dead-gates", "mlir::ModuleOp"> { let dependentDialects = ["mlir::qco::QCODialect"]; let summary = "Remove quantum gates whose results cannot be observed"; diff --git a/mlir/lib/Dialect/QCO/Transforms/Optimizations/ContextSensitiveSpecialization.cpp b/mlir/lib/Dialect/QCO/Transforms/Optimizations/ContextSensitiveSpecialization.cpp new file mode 100644 index 0000000000..05259418b8 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/ContextSensitiveSpecialization.cpp @@ -0,0 +1,548 @@ +/* + * 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 "IPOUtils.h" +#include "mlir/Dialect/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QCO/IR/QCOInterfaces.h" +#include "mlir/Dialect/QCO/IR/QCOOps.h" +#include "mlir/Dialect/QCO/Transforms/Passes.h" +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" // IWYU pragma: keep (Passes.h.inc) + +#include +#include +#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 { + +#define GEN_PASS_DEF_CONTEXTSENSITIVESPECIALIZATION +#include "mlir/Dialect/QCO/Transforms/Passes.h.inc" + +/** + * @brief Redirect a call to a specialized copy of its callee. + * + * @param callOp The call to redirect. + * @param newCallee The specialization the call should target. + * @param rewriter The rewriter driving the pattern application. + */ +static void updateSpecializedCall(func::CallOp callOp, func::FuncOp newCallee, + PatternRewriter& rewriter) { + rewriter.modifyOpInPlace(callOp, + [&] { callOp.setCallee(newCallee.getName()); }); +} + +namespace { + +/// One cached rotation specialization: the parameter it was created for, the +/// angle baked into it, and the resulting copy of the callee. +struct RotationSpecialization { + uint32_t operand; + double angle; + func::FuncOp func; +}; + +/// Caches the specializations already created for a callee, so that call sites +/// sharing the same context reuse one copy instead of cloning it repeatedly. +/// +/// All three are keyed by callee name first, which lets lookups take a +/// `StringRef` without building a temporary `std::string`. +struct PreviousSpecializations { + /// Keyed by callee name, then by the specialized parameter index. + llvm::StringMap> zeroSpecializations; + llvm::StringMap> plusSpecializations; + /// Keyed by callee name; angles are compared with the same tolerance the + /// pass uses elsewhere rather than for exact equality. + llvm::StringMap> rotationSpecializations; + + /// Every callee a call was redirected away from, and every specialization + /// created. These are exactly the functions the pass may have left without + /// callers, so they are the only ones it is entitled to clean up. + SmallVector touchedFunctions; +}; + +/** + * @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}; + + /// Tolerance used when comparing rotation angles. + constexpr static double ANGLE_TOLERANCE = 1e-9; + + /// Discardable attribute recording which rotation specializations a function + /// already represents. Each entry carries the specialized operand index and + /// the angle that was baked into the body, which mirrors the identity used by + /// `PreviousSpecializations::rotationSpecializations`. + constexpr static llvm::StringLiteral ROTATION_SPECIALIZATION_ATTR = + "qco.rotation_specializations"; + /// Key of the operand index inside one specialization entry. + constexpr static llvm::StringLiteral ROTATION_SPECIALIZATION_OPERAND = + "operand"; + /// Key of the baked-in angle inside one specialization entry. + constexpr static llvm::StringLiteral ROTATION_SPECIALIZATION_ANGLE = "angle"; + + /** + * @brief Compare two rotation angles up to the specialization tolerance. + * + * @param lhs The first angle. + * @param rhs The second angle. + * @return True if the angles are considered equal. + */ + static bool anglesAreEqual(double lhs, double rhs) { + return std::abs(lhs - rhs) < ANGLE_TOLERANCE; + } + + /** + * @brief Check whether a function already is the rotation specialization for + * the given operand and angle. + * + * @details + * Uses a discardable attribute rather than the function name, so that a + * specialization for one operand or angle does not block specialization for + * another, and so that user-chosen names cannot be mistaken for markers. + * + * @param funcOp The candidate callee. + * @param operand The index of the angle argument. + * @param angle The angle passed at the call site. + * @return True if @p funcOp already bakes in that angle for that operand. + */ + static bool hasRotationSpecialization(func::FuncOp funcOp, unsigned operand, + double angle) { + const auto marker = + funcOp->getAttrOfType(ROTATION_SPECIALIZATION_ATTR); + if (!marker) { + return false; + } + return llvm::any_of( + marker.getAsRange(), [&](DictionaryAttr entry) { + const auto operandAttr = + entry.getAs(ROTATION_SPECIALIZATION_OPERAND); + const auto angleAttr = + entry.getAs(ROTATION_SPECIALIZATION_ANGLE); + return operandAttr && angleAttr && operandAttr.getInt() == operand && + anglesAreEqual(angleAttr.getValueAsDouble(), angle); + }); + } + + /** + * @brief Record on a cloned function which rotation specialization it is. + * + * @details + * Existing markers are kept, so a function specialized for several operands + * accumulates one entry per operand. + * + * @param funcOp The clone to mark. + * @param operand The index of the angle argument. + * @param angle The angle baked into the body. + * @param builder Builder used to create the attributes. + */ + static void markRotationSpecialization(func::FuncOp funcOp, unsigned operand, + double angle, OpBuilder& builder) { + SmallVector entries; + if (const auto existing = + funcOp->getAttrOfType(ROTATION_SPECIALIZATION_ATTR)) { + entries.assign(existing.begin(), existing.end()); + } + entries.emplace_back(builder.getDictionaryAttr( + {builder.getNamedAttr( + ROTATION_SPECIALIZATION_OPERAND, + builder.getI32IntegerAttr(static_cast(operand))), + builder.getNamedAttr(ROTATION_SPECIALIZATION_ANGLE, + builder.getF64FloatAttr(angle))})); + funcOp->setAttr(ROTATION_SPECIALIZATION_ATTR, + builder.getArrayAttr(entries)); + } + + /** + * @brief Check whether an operation leaves a qubit in the |0> state alone. + * + * @details + * Only operations that fix |0> *exactly* qualify. That covers a controlled + * operation whose control is the |0> qubit, a reset, and the diagonal gates + * whose first diagonal entry is one: `id`, `z`, `s`, `sdg`, `t`, `tdg`, and + * `p` for any angle, since `p = diag(1, exp(i * theta))`. + * + * `rz` is deliberately excluded: `rz = diag(exp(-i * theta / 2), + * exp(i * theta / 2))` maps |0> to a phase multiple of itself, so removing it + * would silently change the global phase of the program. Admitting it would + * require emitting a compensating `qco.gphase`. + * + * @param op The operation applied to the argument. + * @param zeroArgument The argument known to be in the |0> state. + * @return True if @p op has no effect given that state. + */ + static bool operationIsNopOnZero(Operation* op, Value zeroArgument) { + if (auto ctrl = dyn_cast(op)) { + return llvm::is_contained(ctrl.getControlsIn(), zeroArgument); + } + return isa(op); + } + + /** + * @brief Check whether an operation leaves a qubit in the |+> state alone. + * + * @param op The operation applied to the argument. + * @return True if @p op has no effect given that state. + */ + 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); + } + + /** + * @brief Try to specialize the callee for what is known about one argument. + * + * @param callOp The call whose callee may be specialized. + * @param operand The index of the argument to reason about. + * @param rewriter The rewriter driving the pattern application. + * @return True if a specialization was applied. + */ + 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; + } + + /** + * @brief Specialize a callee for an argument known to be in the |0> state. + * + * @param callOp The call to redirect. + * @param funcOp The current callee. + * @param operand The index of the |0> argument. + * @param rewriter The rewriter driving the pattern application. + * @return True if a specialization was applied. + */ + 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; + } + + const auto calleeName = funcOp.getName(); + if (const auto it = + previousSpecializations.zeroSpecializations.find(calleeName); + it != previousSpecializations.zeroSpecializations.end()) { + if (const auto cached = it->second.find(operand); + cached != it->second.end()) { + previousSpecializations.touchedFunctions.emplace_back(funcOp); + updateSpecializedCall(callOp, cached->second, rewriter); + return true; + } + } + + auto newFunc = + copyFunction(funcOp, funcOp.getName().str() + "_spec_zero_arg_" + + std::to_string(operand)); + symbolTable.insert(newFunc); + previousSpecializations.zeroSpecializations[calleeName][operand] = newFunc; + previousSpecializations.touchedFunctions.emplace_back(funcOp); + previousSpecializations.touchedFunctions.emplace_back(newFunc); + + // Drop the whole run of operations that leave the |0> state alone, not just + // the first one. Every iteration erases an operation, so this terminates. + auto newParameter = newFunc.getArgument(operand); + while ( + newParameter.hasOneUse() && + operationIsNopOnZero(*newParameter.getUsers().begin(), newParameter)) { + auto* newUser = *newParameter.getUsers().begin(); + + // A reset does not implement `UnitaryOpInterface`, so it has to be + // forwarded explicitly instead of going through the qubit accessors. + if (auto resetOp = dyn_cast(newUser)) { + rewriter.replaceAllUsesWith(resetOp.getQubitOut(), + resetOp.getQubitIn()); + rewriter.eraseOp(resetOp); + continue; + } + + auto unitaryOp = dyn_cast(newUser); + if (!unitaryOp) { + break; + } + // Use the qubit accessors rather than raw operand and result indices. + // The gates currently accepted above happen to list their qubits first, + // but the modifier operations do not: `qco.pow` takes its exponent as + // operand 0, so raw indexing would forward that exponent into a qubit + // result. The accessors skip parameters wherever they sit. + for (auto i = 0U; i < unitaryOp.getNumQubits(); ++i) { + rewriter.replaceAllUsesWith(unitaryOp.getOutputQubit(i), + unitaryOp.getInputQubit(i)); + } + rewriter.eraseOp(unitaryOp); + } + + updateSpecializedCall(callOp, newFunc, rewriter); + return true; + } + + /** + * @brief Specialize a callee for an argument known to be in the |+> state. + * + * @param callOp The call to redirect. + * @param funcOp The current callee. + * @param operand The index of the |+> argument. + * @param rewriter The rewriter driving the pattern application. + * @return True if a specialization was applied. + */ + 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; + } + + const auto calleeName = funcOp.getName(); + if (const auto it = + previousSpecializations.plusSpecializations.find(calleeName); + it != previousSpecializations.plusSpecializations.end()) { + if (const auto cached = it->second.find(operand); + cached != it->second.end()) { + previousSpecializations.touchedFunctions.emplace_back(funcOp); + updateSpecializedCall(callOp, cached->second, rewriter); + return true; + } + } + + auto newFunc = + copyFunction(funcOp, funcOp.getName().str() + "_spec_plus_arg_" + + std::to_string(operand)); + symbolTable.insert(newFunc); + previousSpecializations.plusSpecializations[calleeName][operand] = newFunc; + previousSpecializations.touchedFunctions.emplace_back(funcOp); + previousSpecializations.touchedFunctions.emplace_back(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; + } + + /** + * @brief Specialize a callee for a rotation angle known at compile time. + * + * @details + * Only a small set of distinguished angles is specialized, because every + * specialization costs a copy of the callee. The parameter stays in the + * signature; only its uses inside the copy are replaced by a constant. + * + * @param callOp The call to redirect. + * @param funcOp The current callee. + * @param angle The constant angle passed at the call site. + * @param operand The index of the angle argument. + * @param rewriter The rewriter driving the pattern application. + * @return True if a specialization was applied. + */ + 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 anglesAreEqual(a, angle); + })) { + return false; + } + + if (funcOp.getArgument(operand).use_empty()) { + // Nothing inside the callee reads the angle, so a specialization would be + // an exact copy of what it was cloned from. + return false; + } + + if (hasRotationSpecialization(funcOp, operand, angle)) { + // This callee already is the specialization for that operand and angle, + // so specializing it again would clone it forever. + return false; + } + + const auto calleeName = funcOp.getName(); + if (const auto it = + previousSpecializations.rotationSpecializations.find(calleeName); + it != previousSpecializations.rotationSpecializations.end()) { + const auto* const cached = + llvm::find_if(it->second, [&](const RotationSpecialization& entry) { + return entry.operand == operand && + anglesAreEqual(entry.angle, angle); + }); + if (cached != it->second.end()) { + previousSpecializations.touchedFunctions.emplace_back(funcOp); + updateSpecializedCall(callOp, cached->func, rewriter); + return true; + } + } + + auto newFunc = + copyFunction(funcOp, funcOp.getName().str() + "_spec_fixed_angle_" + + std::to_string(operand)); + markRotationSpecialization(newFunc, operand, angle, rewriter); + symbolTable.insert(newFunc); + previousSpecializations.touchedFunctions.emplace_back(funcOp); + previousSpecializations.touchedFunctions.emplace_back(newFunc); + previousSpecializations.rotationSpecializations[calleeName].emplace_back( + RotationSpecialization{ + .operand = operand, .angle = angle, .func = newFunc}); + + auto newParameter = newFunc.getArgument(operand); + const OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(&*newFunc.getBody().getBlocks().begin()); + auto constant = arith::ConstantOp::create( + rewriter, newFunc.getBody().getLoc(), + rewriter.getFloatAttr(Float64Type::get(rewriter.getContext()), angle)); + rewriter.replaceAllUsesWith(newParameter, constant.getResult()); + + updateSpecializedCall(callOp, newFunc, rewriter); + return true; + } +}; + +} // namespace + +/** + * @brief Populate the pattern set with the specialization pattern. + * + * @param patterns The pattern set to populate. + * @param symbolTable The symbol table specializations are inserted into. + * @param previousSpecializations Cache of already-created specializations. + */ +static void populateSpecializationPatterns( + RewritePatternSet& patterns, SymbolTable& symbolTable, + PreviousSpecializations& previousSpecializations) { + patterns.add( + patterns.getContext(), symbolTable, previousSpecializations); +} + +namespace { +/// Specializes callees for what is known at their call sites. +struct ContextSensitiveSpecialization final + : impl::ContextSensitiveSpecializationBase { + using impl::ContextSensitiveSpecializationBase< + ContextSensitiveSpecialization>::ContextSensitiveSpecializationBase; + +protected: + void runOnOperation() override { + auto moduleOp = getOperation(); + SymbolTable symbolTable(moduleOp); + + RewritePatternSet patterns(moduleOp.getContext()); + PreviousSpecializations previousSpecializations; + populateSpecializationPatterns(patterns, symbolTable, + previousSpecializations); + + if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { + llvm::reportFatalInternalError( + "failed to apply context-sensitive specialization patterns"); + } + + // Drop the callees this pass left without callers. + eraseOrphanedSpecializations(symbolTable, + previousSpecializations.touchedFunctions); + } +}; +} // namespace + +} // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Transforms/Optimizations/IPOUtils.cpp b/mlir/lib/Dialect/QCO/Transforms/Optimizations/IPOUtils.cpp new file mode 100644 index 0000000000..5fe3076ae1 --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/IPOUtils.cpp @@ -0,0 +1,62 @@ +/* + * 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 "IPOUtils.h" + +#include +#include +#include +#include +#include + +#include + +namespace mlir::qco { + +func::FuncOp copyFunction(func::FuncOp funcOp, StringRef newName) { + auto newFunc = funcOp.clone(); + newFunc.setName(newName.str()); + // Cloning carries the original's visibility over. A specialization is + // internal to the stage that made it, and orphan cleanup only erases private + // functions, so a public copy would be exported and never reclaimed. + newFunc.setPrivate(); + return newFunc; +} + +void eraseOrphanedSpecializations(SymbolTable& symbolTable, + SmallVector& candidates) { + // Duplicates would leave dangling handles once the first copy is erased. + SmallVector unique; + llvm::DenseSet seen; + for (auto candidate : candidates) { + if (seen.insert(candidate.getOperation()).second) { + unique.emplace_back(candidate); + } + } + candidates = std::move(unique); + + auto erasedAny = true; + while (erasedAny) { + erasedAny = false; + SmallVector remaining; + for (auto candidate : candidates) { + if (candidate.isPrivate() && SymbolTable::symbolKnownUseEmpty( + candidate, candidate->getParentOp())) { + symbolTable.erase(candidate); + erasedAny = true; + continue; + } + remaining.emplace_back(candidate); + } + candidates = std::move(remaining); + } +} + +} // namespace mlir::qco diff --git a/mlir/lib/Dialect/QCO/Transforms/Optimizations/IPOUtils.h b/mlir/lib/Dialect/QCO/Transforms/Optimizations/IPOUtils.h new file mode 100644 index 0000000000..42e211a4dd --- /dev/null +++ b/mlir/lib/Dialect/QCO/Transforms/Optimizations/IPOUtils.h @@ -0,0 +1,62 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +/** + * @file + * @brief Internals shared between the interprocedural passes. + * + * @details + * These are implementation details of the interprocedural passes rather than + * public API, which is why they live next to the sources instead of in the + * dialect's include directory. Nothing outside this directory should need them. + */ + +#pragma once + +#include +#include +#include + +namespace mlir::qco { + +/** + * @brief Create a detached copy of a function under a new name. + * + * @details + * The copy is not inserted into a symbol table; the caller is responsible for + * that, which is also what makes the name unique should @p newName already be + * taken. + * + * @param funcOp The function to copy. + * @param newName The name of the copy. + * @return The detached copy. + */ +[[nodiscard]] func::FuncOp copyFunction(func::FuncOp funcOp, StringRef newName); + +/** + * @brief Erase the functions a stage left without callers. + * + * @details + * Only the functions in @p candidates are considered, which are the callees a + * stage redirected calls away from and the specializations it created. A + * private function no stage touched is left alone even when it is unused, + * because removing it is the user's decision rather than ours. + * + * Erasing one function can orphan another, for example when a specialization is + * itself specialized further, so this repeats until nothing more is removed. + * + * @param symbolTable The symbol table to erase from. + * @param candidates The functions that may have been orphaned. Erased entries + * are removed from it, so the remaining handles stay valid. + */ +void eraseOrphanedSpecializations(SymbolTable& symbolTable, + SmallVector& candidates); + +} // namespace mlir::qco diff --git a/mlir/lib/Support/Passes.cpp b/mlir/lib/Support/Passes.cpp index 8c1302a361..f6aca3927e 100644 --- a/mlir/lib/Support/Passes.cpp +++ b/mlir/lib/Support/Passes.cpp @@ -61,6 +61,7 @@ void registerMQTCompilerPasses() { qco::registerRemoveDeadGates(); qco::registerReplaceClassicalControls(); qco::registerReuseQubits(); + qco::registerContextSensitiveSpecialization(); mqt::registerNormalizeGlobalPhases(); mqt::registerUnrollModifiers(); PassPipelineRegistration<>("mqt-qco-default", diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt index 5cc0548e3e..0fb76a393a 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/CMakeLists.txt @@ -9,6 +9,7 @@ set(target_name mqt-core-mlir-unittest-optimizations) add_executable( ${target_name} + test_qco_context_sensitive_specialization.cpp test_qco_hadamard_lifting.cpp test_qco_measurement_lifting.cpp test_qco_merge_single_qubit_rotation.cpp diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/IPOTestFixture.h b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/IPOTestFixture.h new file mode 100644 index 0000000000..b47e16f838 --- /dev/null +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/IPOTestFixture.h @@ -0,0 +1,139 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +/** + * @file + * @brief Shared fixture for the interprocedural optimization test suites. + * + * @details + * Each interprocedural pass has its own test file so that a case is scheduled + * on the pass it is about, rather than on a pipeline where a later pass could + * mask a regression in an earlier one. + */ + +#pragma once + +#include "Support/IRVerification.h" +#include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h" +#include "mlir/Dialect/QCO/IR/QCODialect.h" +#include "mlir/Dialect/QCO/IR/QCOOps.h" +#include "mlir/Dialect/QTensor/IR/QTensorDialect.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace mqt::test { + +/// Every name is spelled out: this is a header, and a `using namespace` in one +/// leaks into every other test file sharing the translation unit under a unity +/// build. +class IPOTestBase : public testing::Test { + +protected: + mlir::MLIRContext context; + mlir::qco::QCOProgramBuilder programBuilder; + mlir::qco::QCOProgramBuilder referenceBuilder; + mlir::OwningOpRef moduleOp; + mlir::OwningOpRef reference; + + IPOTestBase() : programBuilder(&context), referenceBuilder(&context) {} + + void SetUp() override { + // Register all necessary dialects + mlir::DialectRegistry registry; + registry.insert(); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + } + + /** + * @brief Runs a single interprocedural stage and compares against the + * reference. + * + * @param stage The one stage to schedule. + */ + void expectSingleStageMatchesReference(std::unique_ptr stage) { + mlir::PassManager pm(moduleOp->getContext()); + pm.addPass(std::move(stage)); + pm.addPass(mlir::createCanonicalizerPass()); + ASSERT_TRUE(pm.run(moduleOp.get()).succeeded()); + ASSERT_TRUE(runCanonicalizerPass(reference.get()).succeeded()); + + EXPECT_TRUE( + areModulesEquivalentWithPermutations(moduleOp.get(), reference.get())); + } + + /** + * @brief Parses a module from MLIR source. + * + * @details + * Used by the few cases describing IR `QCOProgramBuilder` cannot build. + * + * @param source The MLIR source to parse. + * @return The parsed module. + */ + mlir::OwningOpRef parseModule(const char* source) { + return mlir::parseSourceString(source, &context); + } + + /** + * @brief Runs one stage on a module without comparing against a reference. + * + * @param module The module to transform. + * @param stage The stage to schedule. + */ + static mlir::LogicalResult runStage(mlir::ModuleOp module, + std::unique_ptr stage) { + mlir::PassManager pm(module.getContext()); + pm.addPass(std::move(stage)); + return pm.run(module); + } + + /** + * @brief Counts the qubit allocations inside a named function. + * + * @param module The module to look in. + * @param name The name of the function to count in. + */ + static unsigned countAllocsIn(mlir::ModuleOp module, mlir::StringRef name) { + unsigned count = 0; + module.walk([&](mlir::func::FuncOp func) { + if (func.getName() == name) { + func.walk([&](mlir::qco::AllocOp) { ++count; }); + } + }); + return count; + } + + /** + * @brief Adds the canonicalizerPass to the current context and runs it. + */ + static mlir::LogicalResult runCanonicalizerPass(mlir::ModuleOp moduleOp) { + mlir::PassManager pm(moduleOp.getContext()); + pm.addPass(mlir::createCanonicalizerPass()); + return pm.run(moduleOp); + } +}; +} // namespace mqt::test diff --git a/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_context_sensitive_specialization.cpp b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_context_sensitive_specialization.cpp new file mode 100644 index 0000000000..655434ce30 --- /dev/null +++ b/mlir/unittests/Dialect/QCO/Transforms/Optimizations/test_qco_context_sensitive_specialization.cpp @@ -0,0 +1,701 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +/** + * @file test_qco_context_sensitive_specialization.cpp + * @brief Tests for the `quantum-context-sensitive-specialization` pass. + */ + +#include "IPOTestFixture.h" +#include "mlir/Dialect/QCO/Transforms/Passes.h" + +#include +#include +#include +#include + +#include +#include + +namespace { + +using QCOContextSensitiveSpecializationTest = ::mqt::test::IPOTestBase; +using namespace mlir; +using namespace mlir::qco; + +// ========================================================================== +// 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(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + // ... 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief A reset applied to an argument that is already in the |0> state is + * dropped. A reset does not implement the unitary interface, so this exercises + * a different removal path than the gate case above. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + specializeZeroArgumentDropsReset) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + programBuilder.endFunction({programBuilder.reset(args[0])}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto specArgs = referenceBuilder.startFunction("f_spec_zero_arg_0", + {qubitType}, {qubitType}); + 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @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(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + + 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief A gate that does not act trivially on |0> must not be dropped. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief If the state of the argument is unknown, no specialization applies. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief Two call sites that qualify for the same specialization share a single + * specialized copy of the callee. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + + 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief A run of operations that all leave |0> alone is dropped in one go. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + specializeZeroArgumentDropsNoOpRun) { + const auto qubitType = programBuilder.getQubitType(); + + programBuilder.initialize(); + auto args = programBuilder.startFunction("f", {qubitType}, {qubitType}); + programBuilder.endFunction( + {programBuilder.t(programBuilder.s(programBuilder.z(args[0])))}); + + auto q = programBuilder.allocQubit(); + auto results = programBuilder.call("f", {q}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + auto specArgs = referenceBuilder.startFunction("f_spec_zero_arg_0", + {qubitType}, {qubitType}); + 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief A phase gate fixes |0> for any angle and is dropped even when the + * angle is only known at run time. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + specializeZeroArgumentDropsPhaseGate) { + const auto qubitType = programBuilder.getQubitType(); + const auto floatType = programBuilder.getF64Type(); + + programBuilder.initialize(); + auto args = + programBuilder.startFunction("f", {qubitType, floatType}, {qubitType}); + programBuilder.endFunction({programBuilder.p(args[1], args[0])}); + + auto q = programBuilder.allocQubit(); + // An angle outside the specialized set, so only the |0> specialization fires. + auto angle = programBuilder.floatConstant(0.7); + auto results = programBuilder.call("f", {q, angle}); + programBuilder.sink(results[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + // The angle stays in the signature, it is simply no longer read. + auto specArgs = referenceBuilder.startFunction( + "f_spec_zero_arg_0", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction({specArgs[0]}); + + auto refQ = referenceBuilder.allocQubit(); + auto refAngle = referenceBuilder.floatConstant(0.7); + auto refResults = + referenceBuilder.call("f_spec_zero_arg_0", {refQ, refAngle}); + referenceBuilder.sink(refResults[0]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief A z-rotation must not be dropped: it maps |0> to a phase multiple of + * itself, so removing it would change the global phase of the program. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + noZeroSpecializationForZRotation) { + const auto qubitType = programBuilder.getQubitType(); + const auto floatType = programBuilder.getF64Type(); + + const auto buildProgram = [&](QCOProgramBuilder& b) { + b.initialize(); + auto args = b.startFunction("f", {qubitType, floatType}, {qubitType}); + b.endFunction({b.rz(args[1], args[0])}); + + auto q = b.allocQubit(); + auto results = b.call("f", {q, b.floatConstant(0.7)}); + b.sink(results[0]); + }; + + buildProgram(programBuilder); + moduleOp = programBuilder.finalize(); + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief A specialization of a public callee is private. + * + * @details + * Cloning carries visibility over, so this used to export the generated symbol + * and, since orphan cleanup skips public functions, never reclaim it. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + specializationOfPublicCalleeIsPrivate) { + auto module = parseModule(R"mlir( +func.func @callee(%q: !qco.qubit) -> !qco.qubit { + %0 = qco.z %q : !qco.qubit -> !qco.qubit + return %0 : !qco.qubit +} +func.func @main() { + %q = qco.alloc : !qco.qubit + %r = func.call @callee(%q) : (!qco.qubit) -> !qco.qubit + qco.sink %r : !qco.qubit + return +} +)mlir"); + ASSERT_TRUE(module); + ASSERT_TRUE(runStage(module.get(), createContextSensitiveSpecialization()) + .succeeded()); + + auto specializations = 0; + module->walk([&](func::FuncOp func) { + if (func.getName() == "callee" || func.getName() == "main") { + return; + } + ++specializations; + EXPECT_TRUE(func.isPrivate()) + << "specialization " << func.getName().str() << " must not be exported"; + }); + EXPECT_EQ(specializations, 1) << "the public callee should be specialized"; +} + +// ========================================================================== +// 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(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + + 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief A gate that does not act trivially on |+> must not be dropped. + */ +TEST_F(QCOContextSensitiveSpecializationTest, 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]); + moduleOp = 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +// ========================================================================== +// 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(QCOContextSensitiveSpecializationTest, 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]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + + // 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief A rotation angle of pi/2 is likewise specialized. + */ +TEST_F(QCOContextSensitiveSpecializationTest, 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]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + + 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief Two call sites passing different constant angles to the same callee + * each get their own specialization. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + separateRotationSpecializationPerAngle) { + 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 results0 = programBuilder.call( + "f", {q0, programBuilder.floatConstant(std::numbers::pi)}); + auto results1 = programBuilder.call( + "f", {q1, programBuilder.floatConstant(std::numbers::pi / 2)}); + programBuilder.sink(results0[0]); + programBuilder.sink(results1[0]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + + // The rewriter reaches the pi/2 call first, so that specialization keeps the + // plain name and the pi one gets the uniqued name. + auto specHalfPiArgs = referenceBuilder.startFunction( + "f_spec_fixed_angle_1", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.rz(std::numbers::pi / 2, specHalfPiArgs[0])}); + + auto specPiArgs = referenceBuilder.startFunction( + "f_spec_fixed_angle_1_0", {qubitType, floatType}, {qubitType}); + referenceBuilder.endFunction( + {referenceBuilder.rz(std::numbers::pi, specPiArgs[0])}); + + auto refQ0 = referenceBuilder.allocQubit(); + auto refQ1 = referenceBuilder.allocQubit(); + auto refResults0 = referenceBuilder.call( + "f_spec_fixed_angle_1_0", + {refQ0, referenceBuilder.floatConstant(std::numbers::pi)}); + auto refResults1 = referenceBuilder.call( + "f_spec_fixed_angle_1", + {refQ1, referenceBuilder.floatConstant(std::numbers::pi / 2)}); + referenceBuilder.sink(refResults0[0]); + referenceBuilder.sink(refResults1[0]); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief An angle the callee never reads is not specialized, because the copy + * would be identical to the callee it was cloned from. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + noSpecializationForUnusedRotationAngle) { + const auto qubitType = programBuilder.getQubitType(); + const auto floatType = programBuilder.getF64Type(); + + const auto buildProgram = [&](QCOProgramBuilder& b) { + b.initialize(); + // The angle is part of the signature but nothing in the body uses it. + auto args = b.startFunction("f", {qubitType, floatType}, {qubitType}); + b.endFunction({b.h(args[0])}); + + auto q = b.allocQubit(); + auto results = b.call("f", {q, b.floatConstant(std::numbers::pi)}); + b.sink(results[0]); + }; + + buildProgram(programBuilder); + moduleOp = programBuilder.finalize(); + buildProgram(referenceBuilder); + reference = referenceBuilder.finalize(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief An angle outside the set of specialized angles leaves the callee + * untouched. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief Two call sites that qualify for the same |+> specialization share a + * single specialized copy of the callee. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +/** + * @brief Two call sites passing the same constant angle share a single + * specialized copy of the callee. + */ +TEST_F(QCOContextSensitiveSpecializationTest, + 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]); + moduleOp = programBuilder.finalize(); + + referenceBuilder.initialize(); + 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(); + + expectSingleStageMatchesReference(createContextSensitiveSpecialization()); +} + +// ========================================================================== + +} // namespace From 91511ae392caa8158b4b1e2968f1f1ccf111d5fa Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 30 Aug 2026 14:23:33 +0000 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=8E=A8=20pre-commit=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CHANGELOG.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d451e05e5..0e97e1d76d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -64,8 +64,7 @@ releases may include breaking changes. #### Passes and transformations - ✨ Add passes for quantum-specific interprocedural optimizations ([#2193], - [#2197]) - ([**@DRovara**], [**@burgholzer**]) + [#2197]) ([**@DRovara**], [**@burgholzer**]) - ✨ Add Pauli twirling, quantum loop unrolling, and qubit reuse passes ([#1705], [#1718], [#1755], [#1756], [#1923], [#1924], [#2039], [#2118], [#2216], [#2224]) ([**@MatthiasReumann**], [**@DRovara**], [**@burgholzer**],