From 79c76886ad7bf578fee4c2673b17827257a1ac62 Mon Sep 17 00:00:00 2001 From: Simon Hofmann Date: Wed, 19 Aug 2026 16:42:39 +0200 Subject: [PATCH 1/3] =?UTF-8?q?=E2=9C=A8=20Preserve=20classical=20control?= =?UTF-8?q?=20during=20mapping?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scalarize supported qco.if QTensor inputs and keep structured-control boundaries intact while mapping advances independent wires. Assisted-by: Codex Signed-off-by: Simon Hofmann --- CHANGELOG.md | 5 +- mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp | 343 +++++++++++++--- .../QCO/Transforms/Mapping/Mapping.cpp | 58 ++- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 388 ++++++++++++++++++ .../QCO/Transforms/Mapping/test_mapping.cpp | 52 ++- 5 files changed, 780 insertions(+), 66 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a797ebaa65..0b1994b9ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -71,8 +71,8 @@ releases may include breaking changes. - ✨ Add a compiler-target-aware `place-and-route` pass ([#1537], [#1547], [#1568], [#1581], [#1583], [#1588], [#1600], [#1664], [#1709], [#1716], [#1748], [#1805], [#1870], [#1904], [#1911], [#1951], [#1956], [#1997], - [#2016], [#2060], [#2184]) ([**@MatthiasReumann**], [**@burgholzer**], - [**@rturrado**]) + [#2016], [#2060], [#2179], [#2184]) ([**@MatthiasReumann**], + [**@burgholzer**], [**@rturrado**], [**@simon1hofmann**]) - ✨ Add modifier and global-phase normalization passes ([#1986], [#1995], [#2015]) ([**@burgholzer**], [**@denialhaag**]) - ✨ Add single-qubit optimization passes for unitary fusion, Hadamard lifting, @@ -891,6 +891,7 @@ for previous changelogs._ [#2214]: https://github.com/munich-quantum-toolkit/core/pull/2214 [#2193]: https://github.com/munich-quantum-toolkit/core/pull/2193 [#2184]: https://github.com/munich-quantum-toolkit/core/pull/2184 +[#2179]: https://github.com/munich-quantum-toolkit/core/pull/2179 [#2178]: https://github.com/munich-quantum-toolkit/core/pull/2178 [#2176]: https://github.com/munich-quantum-toolkit/core/pull/2176 [#2175]: https://github.com/munich-quantum-toolkit/core/pull/2175 diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp index 89d40d82cc..1091666836 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp @@ -10,13 +10,17 @@ #include "mlir/Dialect/QCO/IR/QCOOps.h" #include "mlir/Dialect/QCO/QCOUtils.h" +#include "mlir/Dialect/QTensor/IR/QTensorOps.h" #include +#include #include #include #include +#include #include #include +#include #include #include #include @@ -31,7 +35,9 @@ #include #include +#include #include +#include using namespace mlir; using namespace mlir::qco; @@ -153,19 +159,9 @@ void IfOp::getRegionInvocationBounds( } } -/** - * @brief Replace operation with the contents of a region - * - * @details - * Replaces the given op with the contents of the given single-block region, - * using the operands of the block terminator to replace operation results. - * - * @param rewriter The used rewriter - * @param op The operation that is replcaed - * @param region The region with the replacement content - * @param blockArgs The block arguments of the region - * - */ +/// Replace an operation with the contents of a single-block region. +/// +/// Use the block terminator operands to replace the operation results. static void replaceOpWithRegion(PatternRewriter& rewriter, Operation* op, Region& region, ValueRange blockArgs = {}) { assert(llvm::hasSingleElement(region) && "expected single-region block"); @@ -179,14 +175,7 @@ static void replaceOpWithRegion(PatternRewriter& rewriter, Operation* op, namespace { -/** - * @brief Remove static conditions - * - * @details - * Removes a qco.if operation with a static condition and replace it with the - * contents of the selected branch. - * - */ +/// Replace an if with a static condition by its selected branch. struct RemoveStaticCondition : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -207,22 +196,7 @@ struct RemoveStaticCondition : public OpRewritePattern { } }; -/** - * @brief Propagate the condition into the branches - * - * @details - * Allow the true region of an if to assume the condition is true - * and vice versa. For example: - * - * qco.if %cmp args(%arg0 = %q0) -> (!qco.qubit) { - * print(true) - * ... - * } else args(%arg = %q0) { - * print(false) - * ... - * } - * - */ +/// Let each branch use the condition value known inside that branch. struct ConditionPropagation : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -271,16 +245,11 @@ struct ConditionPropagation : public OpRewritePattern { } }; -/** - * @brief Forward redundant classical results - * - * @details - * Replaces a classical result with a value yielded by both branches or with an - * earlier classical result whose pair of yielded values is identical. A - * separate pattern removes the result and its yield operands once they become - * unused. Linear results are intentionally excluded because their explicit - * branch threading is part of QCO's quantum dataflow. - */ +/// Forward redundant classical results. +/// +/// Replace a result with a value yielded by both branches or with an earlier +/// result whose pair of yielded values is identical. A separate pattern removes +/// unused results. Linear results retain their explicit quantum dataflow. struct ForwardClassicalResults : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -320,14 +289,9 @@ struct ForwardClassicalResults : public OpRewritePattern { } }; -/** - * @brief Remove unused classical results - * - * @details - * Removes unused classical results and the corresponding operands from both - * branch terminators. The result segment property is updated on the replacement - * operation. The linear result suffix and all quantum dataflow remain intact. - */ +/// Remove unused classical results and their branch yield operands. +/// +/// Update the result segments while preserving the linear result suffix. struct RemoveUnusedClassicalResults : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -369,12 +333,277 @@ struct RemoveUnusedClassicalResults : public OpRewritePattern { return success(); } }; + +} // namespace + +namespace mlir::qco { + +namespace { + +struct QTensorAccess { + qtensor::ExtractOp extract; + qtensor::InsertOp insert; +}; + +struct BranchQTensorAccesses { + DenseMap accesses; + SmallVector qTensorOperations; +}; + +} // namespace + +/// Analyze a QTensor's complete lifetime in one branch. +/// +/// Supported branches extract distinct constant-index qubits, perform +/// QTensor-independent computation, reinsert one qubit at every extracted +/// index, and yield the resulting QTensor. Dynamic indices, repeated accesses, +/// and partial updates do not match. +static std::optional +analyzeQTensorBranch(Block* block, size_t qTensorArgumentIndex, + size_t qTensorYieldIndex) { + BranchQTensorAccesses result; + Value currentQTensor = block->getArgument(qTensorArgumentIndex); + bool reachedInsertPhase = false; + + while (true) { + if (!currentQTensor.hasOneUse()) { + return std::nullopt; + } + Operation* user = *currentQTensor.getUsers().begin(); + if (user->getBlock() != block) { + return std::nullopt; + } + + if (auto extract = dyn_cast(user)) { + auto index = getConstantIntValue(extract.getIndex()); + if (reachedInsertPhase || !index || + !result.accesses + .try_emplace(*index, QTensorAccess{.extract = extract}) + .second) { + return std::nullopt; + } + result.qTensorOperations.push_back(user); + currentQTensor = extract.getOutTensor(); + continue; + } + + if (auto insert = dyn_cast(user)) { + reachedInsertPhase = true; + auto index = getConstantIntValue(insert.getIndex()); + if (!index) { + return std::nullopt; + } + auto access = result.accesses.find(*index); + if (access == result.accesses.end() || access->second.insert) { + return std::nullopt; + } + access->second.insert = insert; + result.qTensorOperations.push_back(user); + currentQTensor = insert.getResult(); + continue; + } + + auto yield = dyn_cast(user); + if (!yield || user != block->getTerminator() || + qTensorYieldIndex >= yield.getTargets().size() || + yield.getTargets()[qTensorYieldIndex] != currentQTensor || + llvm::any_of(result.accesses, [](const auto& access) { + return !access.second.insert; + })) { + return std::nullopt; + } + return result; + } +} + +/// Move a branch while replacing QTensor accesses with scalar qubits. +static void moveScalarizedQTensorBranch(IfOp oldIf, Block* oldBlock, + Block* newBlock, + size_t qTensorArgumentIndex, + BranchQTensorAccesses& accesses, + ArrayRef indices, + PatternRewriter& rewriter) { + auto oldYield = cast(oldBlock->getTerminator()); + auto scalarArguments = newBlock->getArguments().take_back(indices.size()); + auto carriedArguments = newBlock->getArguments().drop_back(indices.size()); + + SmallVector argumentReplacements; + argumentReplacements.reserve(oldBlock->getNumArguments()); + size_t carriedIndex = 0; + for (size_t oldIndex : llvm::seq(oldBlock->getNumArguments())) { + argumentReplacements.push_back(oldIndex == qTensorArgumentIndex + ? oldIf.getQubits()[qTensorArgumentIndex] + : carriedArguments[carriedIndex++]); + } + assert(carriedIndex == carriedArguments.size()); + rewriter.mergeBlocks(oldBlock, newBlock, argumentReplacements); + + for (auto [indexPosition, index] : llvm::enumerate(indices)) { + auto access = accesses.accesses.find(index); + if (access != accesses.accesses.end()) { + rewriter.replaceAllUsesWith(access->second.extract.getResult(), + scalarArguments[indexPosition]); + } + } + + SmallVector scalarYields; + scalarYields.reserve(indices.size()); + for (auto [indexPosition, index] : llvm::enumerate(indices)) { + auto access = accesses.accesses.find(index); + if (access == accesses.accesses.end()) { + scalarYields.push_back(scalarArguments[indexPosition]); + } else { + scalarYields.push_back(access->second.insert.getScalar()); + } + } + + auto oldTargets = oldYield.getTargets(); + size_t classicalResultCount = oldIf.getClassicalResults().size(); + SmallVector newYieldValues; + newYieldValues.reserve(oldTargets.size() - 1 + scalarYields.size()); + llvm::append_range(newYieldValues, + oldTargets.take_front(classicalResultCount)); + for (auto [oldIndex, value] : + llvm::enumerate(oldTargets.drop_front(classicalResultCount))) { + if (oldIndex != qTensorArgumentIndex) { + newYieldValues.push_back(value); + } + } + llvm::append_range(newYieldValues, scalarYields); + + rewriter.setInsertionPoint(oldYield); + rewriter.replaceOpWithNewOp(oldYield, newYieldValues); + + for (Operation* operation : llvm::reverse(accesses.qTensorOperations)) { + rewriter.eraseOp(operation); + } +} + +} // namespace mlir::qco + +namespace { + +/// Replace constant-index QTensor updates in an if with scalar threading. +/// +/// A QTensor carried through an if hides its qubits from target mapping. This +/// pattern extracts the union of constant indices accessed by either branch, +/// threads those qubits through both branches, and reinserts the results. +/// Untouched elements remain in the QTensor outside the if. +struct ScalarizeQTensorInputs final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(IfOp op, + PatternRewriter& rewriter) const override { + size_t classicalResultCount = op.getClassicalResults().size(); + auto oldQubits = op.getQubits(); + + for (auto [qTensorIndex, qTensor] : llvm::enumerate(oldQubits)) { + auto qTensorType = dyn_cast(qTensor.getType()); + if (!qTensorType || !qTensorType.hasStaticShape()) { + continue; + } + + auto thenAccesses = analyzeQTensorBranch( + op.thenBlock(), qTensorIndex, classicalResultCount + qTensorIndex); + auto elseAccesses = analyzeQTensorBranch( + op.elseBlock(), qTensorIndex, classicalResultCount + qTensorIndex); + if (!thenAccesses || !elseAccesses) { + continue; + } + + SmallVector accessedIndices; + accessedIndices.reserve(thenAccesses->accesses.size() + + elseAccesses->accesses.size()); + for (int64_t index : thenAccesses->accesses.keys()) { + accessedIndices.push_back(index); + } + for (int64_t index : elseAccesses->accesses.keys()) { + if (!thenAccesses->accesses.contains(index)) { + accessedIndices.push_back(index); + } + } + llvm::sort(accessedIndices); + ArrayRef indices(accessedIndices); + + rewriter.setInsertionPoint(op); + SmallVector indexValues; + SmallVector scalarInputs; + indexValues.reserve(indices.size()); + scalarInputs.reserve(indices.size()); + Value qTensorWithoutScalars = qTensor; + for (int64_t index : indices) { + auto indexValue = + arith::ConstantIndexOp::create(rewriter, op.getLoc(), index); + auto extract = qtensor::ExtractOp::create(rewriter, op.getLoc(), + qTensorWithoutScalars, + indexValue.getResult()); + indexValues.push_back(indexValue.getResult()); + scalarInputs.push_back(extract.getResult()); + qTensorWithoutScalars = extract.getOutTensor(); + } + + SmallVector newQubits; + newQubits.reserve(oldQubits.size() - 1 + scalarInputs.size()); + for (auto [oldIndex, qubit] : llvm::enumerate(oldQubits)) { + if (oldIndex != qTensorIndex) { + newQubits.push_back(qubit); + } + } + llvm::append_range(newQubits, scalarInputs); + + auto newIf = IfOp::create( + rewriter, op.getLoc(), op.getClassicalResults().getTypes(), + ValueRange(newQubits).getTypes(), op.getCondition(), newQubits); + newIf->setDiscardableAttrs(op->getDiscardableAttrDictionary()); + + SmallVector locations(newQubits.size(), op.getLoc()); + Block* oldThenBlock = op.thenBlock(); + Block* oldElseBlock = op.elseBlock(); + Block* newThenBlock = + rewriter.createBlock(&newIf.getThenRegion(), {}, + ValueRange(newQubits).getTypes(), locations); + Block* newElseBlock = + rewriter.createBlock(&newIf.getElseRegion(), {}, + ValueRange(newQubits).getTypes(), locations); + moveScalarizedQTensorBranch(op, oldThenBlock, newThenBlock, qTensorIndex, + *thenAccesses, indices, rewriter); + moveScalarizedQTensorBranch(op, oldElseBlock, newElseBlock, qTensorIndex, + *elseAccesses, indices, rewriter); + + rewriter.setInsertionPointAfter(newIf); + Value updatedQTensor = qTensorWithoutScalars; + auto scalarResults = newIf.getLinearResults().take_back(indices.size()); + for (auto [scalar, indexValue] : + llvm::zip_equal(scalarResults, indexValues)) { + updatedQTensor = + qtensor::InsertOp::create(rewriter, op.getLoc(), scalar, + updatedQTensor, indexValue) + .getResult(); + } + + SmallVector replacements; + replacements.reserve(op.getNumResults()); + llvm::append_range(replacements, newIf.getClassicalResults()); + size_t newLinearIndex = 0; + for (size_t oldIndex : llvm::seq(oldQubits.size())) { + replacements.push_back( + oldIndex == qTensorIndex + ? updatedQTensor + : newIf.getLinearResults()[newLinearIndex++]); + } + rewriter.replaceOp(op, replacements); + return success(); + } + return failure(); + } +}; } // namespace void IfOp::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { - results.add(context); + results + .add(context); } LogicalResult IfOp::verify() { diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index bd343ec951..bd0477b1b6 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -776,6 +776,39 @@ struct MappingPass : impl::MappingPassBase { return newWhileOp; } + /// Return the value whose wire edge crosses a composite in block order. + static Value valueBeforeBoundary(WireIterator iterator, Operation* boundary) { + assert(boundary != nullptr && boundary->getBlock() != nullptr); + + // Independent wires can advance beyond `boundary`. Rewind to the qubit + // value that crosses it so extending the composite does not move later + // operations before the boundary. + if (iterator == std::default_sentinel) { + --iterator; + } + + while (Operation* operation = iterator.operation()) { + assert(operation->getBlock() == boundary->getBlock()); + if (operation->isBeforeInBlock(boundary)) { + break; + } + --iterator; + } + + Value value = iterator.qubit(); + assert(value && "expected a qubit value before the composite boundary"); + assert(value.hasOneUse() && "expected linear qubit use at boundary"); + Operation* consumer = value.use_begin()->getOwner(); + while (consumer->getBlock() != boundary->getBlock()) { + consumer = consumer->getParentOp(); + assert(consumer != nullptr && "expected consumer in boundary block"); + } + assert((consumer == boundary || boundary->isBeforeInBlock(consumer) || + isa(consumer)) && + "selected qubit value does not cross composite boundary"); + return value; + } + /// Execute `ntrials` many (parallel) initial layout refinement trials and /// return the heuristically best one. /// @@ -1261,6 +1294,26 @@ struct MappingPass : impl::MappingPassBase { return lhs.op->isBeforeInBlock(rhs.op); }); + // Defer a composite while another active wire points to an operation that + // precedes it in the traversal direction. Otherwise, dispatch would remove + // that operation from the routing frontier. + llvm::erase_if(composites, [&](const CompositeUnitary& composite) { + return llvm::any_of(wires, [&](const WireIterator& iterator) { + if (iterator == std::default_sentinel) { + return false; + } + Operation* operation = iterator.operation(); + if (operation == nullptr || operation == composite.op) { + return false; + } + assert(operation->getBlock() == composite.op->getBlock()); + if constexpr (Direction == WireDirection::Forward) { + return operation->isBeforeInBlock(composite.op); + } + return composite.op->isBeforeInBlock(operation); + }); + }); + return composites; } @@ -1291,10 +1344,7 @@ struct MappingPass : impl::MappingPassBase { allIndices, [&](const size_t i) { return !included.contains(i); })); const SmallVector addons(map_range(excluded, [&](const size_t i) { - // Make sure the qubits point to an already processed operation. - const auto& it = std::prev( - parent.wires[i], parent.wires[i] == std::default_sentinel ? 2 : 1); - return it.qubit(); + return valueBeforeBoundary(parent.wires[i], composite.op); })); composite = CompositeUnitary{ diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 2ef1cf24e3..7132fe5414 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -981,6 +981,394 @@ TEST_F(QCOTest, CanonicalizesRedundantClassicalIfResults) { EXPECT_EQ(returnOp.getOperand(2), returnOp.getOperand(1)); } +TEST_F(QCOTest, CanonicalizesConstantIndexQTensorIfToScalarQubits) { + constexpr StringLiteral mlirCode = R"mlir( + module { + func.func @main(%condition: i1) -> i1 { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %tensor0 = qtensor.alloc(%c3) : tensor<3x!qco.qubit> + %flag, %tensor1 = qco.if %condition + args(%arg0 = %tensor0) -> (i1, tensor<3x!qco.qubit>) { + %tensor2, %q0 = qtensor.extract %arg0[%c0] + : tensor<3x!qco.qubit> + %tensor3, %q1 = qtensor.extract %tensor2[%c1] + : tensor<3x!qco.qubit> + %q2, %q3 = qco.swap %q0, %q1 + : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit + %tensor4 = qtensor.insert %q3 into %tensor3[%c1] + : tensor<3x!qco.qubit> + %tensor5 = qtensor.insert %q2 into %tensor4[%c0] + : tensor<3x!qco.qubit> + %true = arith.constant true + qco.yield %true, %tensor5 : i1, tensor<3x!qco.qubit> + } else args(%arg0 = %tensor0) { + %tensor2, %q0 = qtensor.extract %arg0[%c2] + : tensor<3x!qco.qubit> + %q1 = qco.z %q0 : !qco.qubit -> !qco.qubit + %tensor3 = qtensor.insert %q1 into %tensor2[%c2] + : tensor<3x!qco.qubit> + %false = arith.constant false + qco.yield %false, %tensor3 : i1, tensor<3x!qco.qubit> + } {test.marker = "preserved"} + qtensor.dealloc %tensor1 : tensor<3x!qco.qubit> + return %flag : i1 + } + } + )mlir"; + + auto moduleOp = parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(moduleOp); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + + IfOp ifOp; + moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); + ASSERT_TRUE(ifOp); + ASSERT_EQ(ifOp.getClassicalResults().size(), 1); + ASSERT_EQ(ifOp.getQubits().size(), 3); + ASSERT_EQ(ifOp.getLinearResults().size(), 3); + EXPECT_EQ( + cast(ifOp->getDiscardableAttr("test.marker")).getValue(), + "preserved"); + EXPECT_TRUE(llvm::all_of(ifOp.getQubits(), [](Value value) { + return isa(value.getType()); + })); + EXPECT_TRUE(llvm::all_of(ifOp.getLinearResults(), [](Value value) { + return isa(value.getType()); + })); + + size_t nestedExtracts = 0; + size_t nestedInserts = 0; + size_t swaps = 0; + size_t zs = 0; + ifOp->walk([&](Operation* operation) { + nestedExtracts += isa(operation); + nestedInserts += isa(operation); + swaps += isa(operation); + zs += isa(operation); + }); + EXPECT_EQ(nestedExtracts, 0); + EXPECT_EQ(nestedInserts, 0); + EXPECT_EQ(swaps, 1); + EXPECT_EQ(zs, 1); + + size_t extracts = 0; + size_t inserts = 0; + moduleOp->walk([&](qtensor::ExtractOp) { ++extracts; }); + moduleOp->walk([&](qtensor::InsertOp) { ++inserts; }); + EXPECT_EQ(extracts, 3); + EXPECT_EQ(inserts, 3); +} + +TEST_F(QCOTest, ScalarizesOnlyAccessedQTensorElements) { + constexpr StringLiteral mlirCode = R"mlir( + module { + func.func @main(%condition: i1) { + %c1 = arith.constant 1 : index + %c3 = arith.constant 3 : index + %tensor0 = qtensor.alloc(%c3) : tensor<3x!qco.qubit> + %tensor1 = qco.if %condition + args(%arg0 = %tensor0) -> (tensor<3x!qco.qubit>) { + %tensor2, %q0 = qtensor.extract %arg0[%c1] + : tensor<3x!qco.qubit> + %q1 = qco.x %q0 : !qco.qubit -> !qco.qubit + %tensor3 = qtensor.insert %q1 into %tensor2[%c1] + : tensor<3x!qco.qubit> + qco.yield %tensor3 : tensor<3x!qco.qubit> + } else args(%arg0 = %tensor0) { + qco.yield %arg0 : tensor<3x!qco.qubit> + } + qtensor.dealloc %tensor1 : tensor<3x!qco.qubit> + return + } + } + )mlir"; + + auto moduleOp = parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(moduleOp); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + + IfOp ifOp; + moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); + ASSERT_TRUE(ifOp); + ASSERT_EQ(ifOp.getQubits().size(), 1); + ASSERT_EQ(ifOp.getLinearResults().size(), 1); + EXPECT_TRUE(llvm::all_of(ifOp.getQubits(), [](Value value) { + return isa(value.getType()); + })); + + auto thenValues = ifOp.thenYield().getTargets(); + auto elseValues = ifOp.elseYield().getTargets(); + ASSERT_EQ(thenValues.size(), 1); + ASSERT_EQ(elseValues.size(), 1); + EXPECT_TRUE(isa(thenValues[0].getDefiningOp())); + EXPECT_EQ(elseValues[0], ifOp.elseBlock()->getArgument(0)); + + size_t extracts = 0; + size_t inserts = 0; + moduleOp->walk([&](qtensor::ExtractOp) { ++extracts; }); + moduleOp->walk([&](qtensor::InsertOp) { ++inserts; }); + EXPECT_EQ(extracts, 1); + EXPECT_EQ(inserts, 1); +} + +TEST_F(QCOTest, ForwardsUnaccessedQTensorAroundIf) { + constexpr StringLiteral mlirCode = R"mlir( + module { + func.func @main(%condition: i1) { + %c2 = arith.constant 2 : index + %tensor0 = qtensor.alloc(%c2) : tensor<2x!qco.qubit> + %q0 = qco.alloc : !qco.qubit + %tensor1, %q1 = qco.if %condition + args(%tensor = %tensor0, %q = %q0) + -> (tensor<2x!qco.qubit>, !qco.qubit) { + %q2 = qco.h %q : !qco.qubit -> !qco.qubit + qco.yield %tensor, %q2 : tensor<2x!qco.qubit>, !qco.qubit + } else args(%tensor = %tensor0, %q = %q0) { + qco.yield %tensor, %q : tensor<2x!qco.qubit>, !qco.qubit + } + qtensor.dealloc %tensor1 : tensor<2x!qco.qubit> + qco.sink %q1 : !qco.qubit + return + } + } + )mlir"; + + auto moduleOp = parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(moduleOp); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + + IfOp ifOp; + moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); + ASSERT_TRUE(ifOp); + ASSERT_EQ(ifOp.getQubits().size(), 1); + ASSERT_EQ(ifOp.getLinearResults().size(), 1); + EXPECT_TRUE(isa(ifOp.getQubits()[0].getType())); + EXPECT_TRUE(isa(ifOp.getLinearResults()[0].getType())); + + auto thenValues = ifOp.thenYield().getTargets(); + auto elseValues = ifOp.elseYield().getTargets(); + ASSERT_EQ(thenValues.size(), 1); + ASSERT_EQ(elseValues.size(), 1); + EXPECT_TRUE(isa(thenValues[0].getDefiningOp())); + for (auto [index, value] : llvm::enumerate(elseValues)) { + EXPECT_EQ(value, ifOp.elseBlock()->getArgument(index)); + } +} + +TEST_F(QCOTest, PreservesInterleavedResultOrderWhenScalarizingQTensors) { + constexpr StringLiteral mlirCode = R"mlir( + module { + func.func @main(%condition: i1) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %left0 = qco.alloc : !qco.qubit + %tensorA0 = qtensor.alloc(%c1) : tensor<1x!qco.qubit> + %middle0 = qco.alloc : !qco.qubit + %tensorB0 = qtensor.alloc(%c1) : tensor<1x!qco.qubit> + %right0 = qco.alloc : !qco.qubit + %left1, %tensorA1, %middle1, %tensorB1, %right1 = + qco.if %condition + args(%left = %left0, %tensorA = %tensorA0, + %middle = %middle0, %tensorB = %tensorB0, + %right = %right0) + -> (!qco.qubit, tensor<1x!qco.qubit>, !qco.qubit, + tensor<1x!qco.qubit>, !qco.qubit) { + %leftOut = qco.h %left : !qco.qubit -> !qco.qubit + %tensorA2, %tensorAQubit = qtensor.extract %tensorA[%c0] + : tensor<1x!qco.qubit> + %tensorAQubitOut = qco.x %tensorAQubit + : !qco.qubit -> !qco.qubit + %tensorA3 = qtensor.insert %tensorAQubitOut into %tensorA2[%c0] + : tensor<1x!qco.qubit> + %middleOut = qco.y %middle : !qco.qubit -> !qco.qubit + %tensorB2, %tensorBQubit = qtensor.extract %tensorB[%c0] + : tensor<1x!qco.qubit> + %tensorBQubitOut = qco.z %tensorBQubit + : !qco.qubit -> !qco.qubit + %tensorB3 = qtensor.insert %tensorBQubitOut into %tensorB2[%c0] + : tensor<1x!qco.qubit> + %rightOut = qco.sx %right : !qco.qubit -> !qco.qubit + qco.yield %leftOut, %tensorA3, %middleOut, %tensorB3, %rightOut + : !qco.qubit, tensor<1x!qco.qubit>, !qco.qubit, + tensor<1x!qco.qubit>, !qco.qubit + } else args(%left = %left0, %tensorA = %tensorA0, + %middle = %middle0, %tensorB = %tensorB0, + %right = %right0) { + qco.yield %left, %tensorA, %middle, %tensorB, %right + : !qco.qubit, tensor<1x!qco.qubit>, !qco.qubit, + tensor<1x!qco.qubit>, !qco.qubit + } + %left2 = qco.s %left1 : !qco.qubit -> !qco.qubit + %middle2 = qco.t %middle1 : !qco.qubit -> !qco.qubit + %right2 = qco.sdg %right1 : !qco.qubit -> !qco.qubit + qco.sink %left2 : !qco.qubit + qtensor.dealloc %tensorA1 : tensor<1x!qco.qubit> + qco.sink %middle2 : !qco.qubit + qtensor.dealloc %tensorB1 : tensor<1x!qco.qubit> + qco.sink %right2 : !qco.qubit + return + } + } + )mlir"; + + auto moduleOp = parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(moduleOp); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + + IfOp ifOp; + SOp postLeft; + TOp postMiddle; + SdgOp postRight; + SmallVector insertedScalars; + moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); + moduleOp->walk([&](SOp candidate) { postLeft = candidate; }); + moduleOp->walk([&](TOp candidate) { postMiddle = candidate; }); + moduleOp->walk([&](SdgOp candidate) { postRight = candidate; }); + moduleOp->walk([&](qtensor::InsertOp insert) { + insertedScalars.push_back(insert.getScalar()); + }); + + ASSERT_TRUE(ifOp); + ASSERT_EQ(ifOp.getQubits().size(), 5); + ASSERT_EQ(ifOp.getLinearResults().size(), 5); + EXPECT_TRUE(llvm::all_of(ifOp.getQubits(), [](Value value) { + return isa(value.getType()); + })); + EXPECT_TRUE(llvm::all_of(ifOp.getLinearResults(), [](Value value) { + return isa(value.getType()); + })); + + ASSERT_TRUE(postLeft); + ASSERT_TRUE(postMiddle); + ASSERT_TRUE(postRight); + EXPECT_EQ(cast(postLeft.getOperation()) + .getInputQubits() + .front(), + ifOp.getLinearResults()[0]); + EXPECT_EQ(cast(postMiddle.getOperation()) + .getInputQubits() + .front(), + ifOp.getLinearResults()[1]); + EXPECT_EQ(cast(postRight.getOperation()) + .getInputQubits() + .front(), + ifOp.getLinearResults()[2]); + + ASSERT_EQ(insertedScalars.size(), 2); + EXPECT_TRUE(llvm::is_contained(insertedScalars, ifOp.getLinearResults()[3])); + EXPECT_TRUE(llvm::is_contained(insertedScalars, ifOp.getLinearResults()[4])); + + auto thenValues = ifOp.thenYield().getTargets(); + ASSERT_EQ(thenValues.size(), 5); + EXPECT_TRUE(isa(thenValues[0].getDefiningOp())); + EXPECT_TRUE(isa(thenValues[1].getDefiningOp())); + EXPECT_TRUE(isa(thenValues[2].getDefiningOp())); + EXPECT_TRUE(isa(thenValues[3].getDefiningOp())); + EXPECT_TRUE(isa(thenValues[4].getDefiningOp())); +} + +TEST_F(QCOTest, LeavesDynamicIndexQTensorIfUnchanged) { + constexpr StringLiteral mlirCode = R"mlir( + module { + func.func @main(%condition: i1, %index: index) { + %c2 = arith.constant 2 : index + %tensor0 = qtensor.alloc(%c2) : tensor<2x!qco.qubit> + %tensor1 = qco.if %condition + args(%arg0 = %tensor0) -> (tensor<2x!qco.qubit>) { + %tensor2, %q0 = qtensor.extract %arg0[%index] + : tensor<2x!qco.qubit> + %q1 = qco.x %q0 : !qco.qubit -> !qco.qubit + %tensor3 = qtensor.insert %q1 into %tensor2[%index] + : tensor<2x!qco.qubit> + qco.yield %tensor3 : tensor<2x!qco.qubit> + } else args(%arg0 = %tensor0) { + qco.yield %arg0 : tensor<2x!qco.qubit> + } + qtensor.dealloc %tensor1 : tensor<2x!qco.qubit> + return + } + } + )mlir"; + + auto moduleOp = parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(moduleOp); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + + IfOp ifOp; + moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); + ASSERT_TRUE(ifOp); + ASSERT_EQ(ifOp.getQubits().size(), 1); + EXPECT_TRUE(isa(ifOp.getQubits().front().getType())); + size_t nestedExtracts = 0; + size_t nestedInserts = 0; + ifOp->walk([&](Operation* operation) { + nestedExtracts += isa(operation); + nestedInserts += isa(operation); + }); + EXPECT_EQ(nestedExtracts, 1); + EXPECT_EQ(nestedInserts, 1); +} + +TEST_F(QCOTest, LeavesDynamicShapeQTensorIfUnchanged) { + constexpr StringLiteral mlirCode = R"mlir( + module { + func.func @main(%condition: i1, %size: index) { + %c0 = arith.constant 0 : index + %tensor0 = qtensor.alloc(%size) : tensor + %tensor1 = qco.if %condition + args(%arg0 = %tensor0) -> (tensor) { + %tensor2, %q0 = qtensor.extract %arg0[%c0] + : tensor + %q1 = qco.x %q0 : !qco.qubit -> !qco.qubit + %tensor3 = qtensor.insert %q1 into %tensor2[%c0] + : tensor + qco.yield %tensor3 : tensor + } else args(%arg0 = %tensor0) { + qco.yield %arg0 : tensor + } + qtensor.dealloc %tensor1 : tensor + return + } + } + )mlir"; + + auto moduleOp = parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(moduleOp); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + + IfOp ifOp; + moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); + ASSERT_TRUE(ifOp); + ASSERT_EQ(ifOp.getQubits().size(), 1); + auto tensorType = + dyn_cast(ifOp.getQubits().front().getType()); + ASSERT_TRUE(tensorType); + EXPECT_TRUE(tensorType.isDynamicDim(0)); + size_t nestedExtracts = 0; + size_t nestedInserts = 0; + ifOp->walk([&](Operation* operation) { + nestedExtracts += isa(operation); + nestedInserts += isa(operation); + }); + EXPECT_EQ(nestedExtracts, 1); + EXPECT_EQ(nestedInserts, 1); +} + TEST_F(QCOTest, IndexSwitchParser) { // Test IndexSwitch parser const char* mlirCode = R"( diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index 9329d05383..fc32d8c0f5 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -361,7 +361,7 @@ TEST_F(MappingPassFixture, MapTopologyOnlyWithEmptyOperationSet) { std::tie(qubits[1], qubits[2]) = builder.rzx(0.5, qubits[1], qubits[2]); std::tie(qubits[0], qubits[2]) = builder.cx(qubits[0], qubits[2]); - for (int64_t i = 0; i < qubits.size(); ++i) { + for (size_t i = 0; i < qubits.size(); ++i) { std::tie(qubits[i], bits[i]) = builder.measure(qubits[i]); builder.sink(qubits[i]); } @@ -456,7 +456,7 @@ TEST_F(MappingPassFixture, PreserveNoncontiguousTargetSiteIds) { std::tie(qubits[0], qubits[1]) = builder.cx(qubits[0], qubits[1]); std::tie(qubits[1], qubits[2]) = builder.cz(qubits[1], qubits[2]); std::tie(qubits[0], qubits[2]) = builder.cx(qubits[0], qubits[2]); - for (int64_t i = 0; i < qubits.size(); ++i) { + for (size_t i = 0; i < qubits.size(); ++i) { std::tie(qubits[i], bits[i]) = builder.measure(qubits[i]); builder.sink(qubits[i]); } @@ -669,6 +669,52 @@ TEST_P(MappingPassTest, MapScalarAllocation) { EXPECT_EQ(numStatics, 1); } +TEST_F(MappingPassFixture, ExpandNonAdjacentTwoQubitIfOnLineTarget) { + constexpr StringLiteral source = R"mlir( + module { + func.func @main(%condition: i1) + attributes {mqt.entry_point} { + %q0 = qco.alloc : !qco.qubit + %q1 = qco.alloc : !qco.qubit + %q2 = qco.alloc : !qco.qubit + %a0, %a1 = qco.swap %q0, %q1 + : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit + %b1, %b2 = qco.swap %a1, %q2 + : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit + %next0, %next2 = qco.if %condition + args(%arg0 = %a0, %arg2 = %b2) + -> (!qco.qubit, !qco.qubit) { + %then0, %then2 = qco.swap %arg0, %arg2 + : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit + qco.yield %then0, %then2 : !qco.qubit, !qco.qubit + } else args(%arg0 = %a0, %arg2 = %b2) { + qco.yield %arg0, %arg2 : !qco.qubit, !qco.qubit + } + qco.sink %next0 : !qco.qubit + qco.sink %b1 : !qco.qubit + qco.sink %next2 : !qco.qubit + return + } + } + )mlir"; + + auto moduleOp = parseSourceString(source, context.get()); + ASSERT_TRUE(moduleOp); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + + const auto target = llvm::cantFail(CompilerTarget::create( + 3, std::vector{{0, 1}, {1, 2}})); + ASSERT_TRUE(runPass(moduleOp.get(), target, MappingPassOptions{.ntrials = 1}) + .succeeded()); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + EXPECT_TRUE(isExecutable(getEntryPoint(moduleOp.get()), target)); + + IfOp conditional; + moduleOp->walk([&](IfOp candidate) { conditional = candidate; }); + ASSERT_TRUE(conditional); + EXPECT_EQ(conditional.getQubits().size(), 3U); +} + TEST_P(MappingPassTest, MapMixedScalarAndTensorAllocations) { const auto& target = GetParam(); @@ -1866,7 +1912,7 @@ TEST_P(MappingPassTest, MapPaddedCXCZGrid) { qubits[i] = builder.allocQubit(); } cxcz(builder, qubits); - for (int64_t i = 0; i < qubits.size(); ++i) { + for (size_t i = 0; i < qubits.size(); ++i) { std::tie(qubits[i], bits[i]) = builder.measure(qubits[i]); builder.sink(qubits[i]); } From 0f8a6e394cbe1806e4247b0aec3ded83319441bc Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Tue, 1 Sep 2026 23:30:05 +0000 Subject: [PATCH 2/3] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Simplify=20classical-c?= =?UTF-8?q?ontrol=20mapping?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reuse LLVM and MLIR helpers in QTensor scalarization and boundary handling. Keep the same regression coverage with smaller test inputs. Assisted-by: Codex --- mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp | 58 ++--- .../QCO/Transforms/Mapping/Mapping.cpp | 16 +- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 217 +++++++----------- .../QCO/Transforms/Mapping/test_mapping.cpp | 50 ++-- 4 files changed, 122 insertions(+), 219 deletions(-) diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp index 1091666836..cab3c666ed 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp @@ -334,12 +334,6 @@ struct RemoveUnusedClassicalResults : public OpRewritePattern { } }; -} // namespace - -namespace mlir::qco { - -namespace { - struct QTensorAccess { qtensor::ExtractOp extract; qtensor::InsertOp insert; @@ -438,14 +432,6 @@ static void moveScalarizedQTensorBranch(IfOp oldIf, Block* oldBlock, assert(carriedIndex == carriedArguments.size()); rewriter.mergeBlocks(oldBlock, newBlock, argumentReplacements); - for (auto [indexPosition, index] : llvm::enumerate(indices)) { - auto access = accesses.accesses.find(index); - if (access != accesses.accesses.end()) { - rewriter.replaceAllUsesWith(access->second.extract.getResult(), - scalarArguments[indexPosition]); - } - } - SmallVector scalarYields; scalarYields.reserve(indices.size()); for (auto [indexPosition, index] : llvm::enumerate(indices)) { @@ -453,6 +439,8 @@ static void moveScalarizedQTensorBranch(IfOp oldIf, Block* oldBlock, if (access == accesses.accesses.end()) { scalarYields.push_back(scalarArguments[indexPosition]); } else { + rewriter.replaceAllUsesWith(access->second.extract.getResult(), + scalarArguments[indexPosition]); scalarYields.push_back(access->second.insert.getScalar()); } } @@ -479,8 +467,6 @@ static void moveScalarizedQTensorBranch(IfOp oldIf, Block* oldBlock, } } -} // namespace mlir::qco - namespace { /// Replace constant-index QTensor updates in an if with scalar threading. @@ -511,18 +497,11 @@ struct ScalarizeQTensorInputs final : OpRewritePattern { continue; } - SmallVector accessedIndices; - accessedIndices.reserve(thenAccesses->accesses.size() + - elseAccesses->accesses.size()); - for (int64_t index : thenAccesses->accesses.keys()) { - accessedIndices.push_back(index); - } - for (int64_t index : elseAccesses->accesses.keys()) { - if (!thenAccesses->accesses.contains(index)) { - accessedIndices.push_back(index); - } - } + SmallVector accessedIndices(thenAccesses->accesses.keys()); + llvm::append_range(accessedIndices, elseAccesses->accesses.keys()); llvm::sort(accessedIndices); + accessedIndices.erase(llvm::unique(accessedIndices), + accessedIndices.end()); ArrayRef indices(accessedIndices); rewriter.setInsertionPoint(op); @@ -542,13 +521,8 @@ struct ScalarizeQTensorInputs final : OpRewritePattern { qTensorWithoutScalars = extract.getOutTensor(); } - SmallVector newQubits; - newQubits.reserve(oldQubits.size() - 1 + scalarInputs.size()); - for (auto [oldIndex, qubit] : llvm::enumerate(oldQubits)) { - if (oldIndex != qTensorIndex) { - newQubits.push_back(qubit); - } - } + SmallVector newQubits(oldQubits); + newQubits.erase(newQubits.begin() + qTensorIndex); llvm::append_range(newQubits, scalarInputs); auto newIf = IfOp::create( @@ -581,16 +555,12 @@ struct ScalarizeQTensorInputs final : OpRewritePattern { .getResult(); } - SmallVector replacements; - replacements.reserve(op.getNumResults()); - llvm::append_range(replacements, newIf.getClassicalResults()); - size_t newLinearIndex = 0; - for (size_t oldIndex : llvm::seq(oldQubits.size())) { - replacements.push_back( - oldIndex == qTensorIndex - ? updatedQTensor - : newIf.getLinearResults()[newLinearIndex++]); - } + SmallVector replacements( + newIf.getLinearResults().drop_back(indices.size())); + replacements.insert(replacements.begin() + qTensorIndex, updatedQTensor); + replacements.insert(replacements.begin(), + newIf.getClassicalResults().begin(), + newIf.getClassicalResults().end()); rewriter.replaceOp(op, replacements); return success(); } diff --git a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp index bd0477b1b6..fcb8f319c3 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -787,22 +787,18 @@ struct MappingPass : impl::MappingPassBase { --iterator; } - while (Operation* operation = iterator.operation()) { - assert(operation->getBlock() == boundary->getBlock()); - if (operation->isBeforeInBlock(boundary)) { - break; - } + while (iterator.operation() != nullptr && + !iterator.operation()->isBeforeInBlock(boundary)) { + assert(iterator.operation()->getBlock() == boundary->getBlock()); --iterator; } Value value = iterator.qubit(); assert(value && "expected a qubit value before the composite boundary"); assert(value.hasOneUse() && "expected linear qubit use at boundary"); - Operation* consumer = value.use_begin()->getOwner(); - while (consumer->getBlock() != boundary->getBlock()) { - consumer = consumer->getParentOp(); - assert(consumer != nullptr && "expected consumer in boundary block"); - } + Operation* consumer = boundary->getBlock()->findAncestorOpInBlock( + *value.use_begin()->getOwner()); + assert(consumer != nullptr && "expected consumer in boundary block"); assert((consumer == boundary || boundary->isBeforeInBlock(consumer) || isa(consumer)) && "selected qubit value does not cross composite boundary"); diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 7132fe5414..f174874843 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -1170,19 +1170,15 @@ TEST_F(QCOTest, PreservesInterleavedResultOrderWhenScalarizingQTensors) { func.func @main(%condition: i1) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index - %left0 = qco.alloc : !qco.qubit %tensorA0 = qtensor.alloc(%c1) : tensor<1x!qco.qubit> %middle0 = qco.alloc : !qco.qubit %tensorB0 = qtensor.alloc(%c1) : tensor<1x!qco.qubit> - %right0 = qco.alloc : !qco.qubit - %left1, %tensorA1, %middle1, %tensorB1, %right1 = + %tensorA1, %middle1, %tensorB1 = qco.if %condition - args(%left = %left0, %tensorA = %tensorA0, - %middle = %middle0, %tensorB = %tensorB0, - %right = %right0) - -> (!qco.qubit, tensor<1x!qco.qubit>, !qco.qubit, - tensor<1x!qco.qubit>, !qco.qubit) { - %leftOut = qco.h %left : !qco.qubit -> !qco.qubit + args(%tensorA = %tensorA0, %middle = %middle0, + %tensorB = %tensorB0) + -> (tensor<1x!qco.qubit>, !qco.qubit, + tensor<1x!qco.qubit>) { %tensorA2, %tensorAQubit = qtensor.extract %tensorA[%c0] : tensor<1x!qco.qubit> %tensorAQubitOut = qco.x %tensorAQubit @@ -1196,25 +1192,17 @@ TEST_F(QCOTest, PreservesInterleavedResultOrderWhenScalarizingQTensors) { : !qco.qubit -> !qco.qubit %tensorB3 = qtensor.insert %tensorBQubitOut into %tensorB2[%c0] : tensor<1x!qco.qubit> - %rightOut = qco.sx %right : !qco.qubit -> !qco.qubit - qco.yield %leftOut, %tensorA3, %middleOut, %tensorB3, %rightOut - : !qco.qubit, tensor<1x!qco.qubit>, !qco.qubit, - tensor<1x!qco.qubit>, !qco.qubit - } else args(%left = %left0, %tensorA = %tensorA0, - %middle = %middle0, %tensorB = %tensorB0, - %right = %right0) { - qco.yield %left, %tensorA, %middle, %tensorB, %right - : !qco.qubit, tensor<1x!qco.qubit>, !qco.qubit, - tensor<1x!qco.qubit>, !qco.qubit + qco.yield %tensorA3, %middleOut, %tensorB3 + : tensor<1x!qco.qubit>, !qco.qubit, tensor<1x!qco.qubit> + } else args(%tensorA = %tensorA0, %middle = %middle0, + %tensorB = %tensorB0) { + qco.yield %tensorA, %middle, %tensorB + : tensor<1x!qco.qubit>, !qco.qubit, tensor<1x!qco.qubit> } - %left2 = qco.s %left1 : !qco.qubit -> !qco.qubit %middle2 = qco.t %middle1 : !qco.qubit -> !qco.qubit - %right2 = qco.sdg %right1 : !qco.qubit -> !qco.qubit - qco.sink %left2 : !qco.qubit qtensor.dealloc %tensorA1 : tensor<1x!qco.qubit> qco.sink %middle2 : !qco.qubit qtensor.dealloc %tensorB1 : tensor<1x!qco.qubit> - qco.sink %right2 : !qco.qubit return } } @@ -1227,21 +1215,17 @@ TEST_F(QCOTest, PreservesInterleavedResultOrderWhenScalarizingQTensors) { ASSERT_TRUE(succeeded(verify(*moduleOp))); IfOp ifOp; - SOp postLeft; TOp postMiddle; - SdgOp postRight; SmallVector insertedScalars; moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); - moduleOp->walk([&](SOp candidate) { postLeft = candidate; }); moduleOp->walk([&](TOp candidate) { postMiddle = candidate; }); - moduleOp->walk([&](SdgOp candidate) { postRight = candidate; }); moduleOp->walk([&](qtensor::InsertOp insert) { insertedScalars.push_back(insert.getScalar()); }); ASSERT_TRUE(ifOp); - ASSERT_EQ(ifOp.getQubits().size(), 5); - ASSERT_EQ(ifOp.getLinearResults().size(), 5); + ASSERT_EQ(ifOp.getQubits().size(), 3); + ASSERT_EQ(ifOp.getLinearResults().size(), 3); EXPECT_TRUE(llvm::all_of(ifOp.getQubits(), [](Value value) { return isa(value.getType()); })); @@ -1249,124 +1233,89 @@ TEST_F(QCOTest, PreservesInterleavedResultOrderWhenScalarizingQTensors) { return isa(value.getType()); })); - ASSERT_TRUE(postLeft); ASSERT_TRUE(postMiddle); - ASSERT_TRUE(postRight); - EXPECT_EQ(cast(postLeft.getOperation()) - .getInputQubits() - .front(), - ifOp.getLinearResults()[0]); EXPECT_EQ(cast(postMiddle.getOperation()) .getInputQubits() .front(), - ifOp.getLinearResults()[1]); - EXPECT_EQ(cast(postRight.getOperation()) - .getInputQubits() - .front(), - ifOp.getLinearResults()[2]); + ifOp.getLinearResults()[0]); ASSERT_EQ(insertedScalars.size(), 2); - EXPECT_TRUE(llvm::is_contained(insertedScalars, ifOp.getLinearResults()[3])); - EXPECT_TRUE(llvm::is_contained(insertedScalars, ifOp.getLinearResults()[4])); + EXPECT_TRUE(llvm::is_contained(insertedScalars, ifOp.getLinearResults()[1])); + EXPECT_TRUE(llvm::is_contained(insertedScalars, ifOp.getLinearResults()[2])); auto thenValues = ifOp.thenYield().getTargets(); - ASSERT_EQ(thenValues.size(), 5); - EXPECT_TRUE(isa(thenValues[0].getDefiningOp())); - EXPECT_TRUE(isa(thenValues[1].getDefiningOp())); - EXPECT_TRUE(isa(thenValues[2].getDefiningOp())); - EXPECT_TRUE(isa(thenValues[3].getDefiningOp())); - EXPECT_TRUE(isa(thenValues[4].getDefiningOp())); + ASSERT_EQ(thenValues.size(), 3); + EXPECT_TRUE(isa(thenValues[0].getDefiningOp())); + EXPECT_TRUE(isa(thenValues[1].getDefiningOp())); + EXPECT_TRUE(isa(thenValues[2].getDefiningOp())); } -TEST_F(QCOTest, LeavesDynamicIndexQTensorIfUnchanged) { - constexpr StringLiteral mlirCode = R"mlir( - module { - func.func @main(%condition: i1, %index: index) { - %c2 = arith.constant 2 : index - %tensor0 = qtensor.alloc(%c2) : tensor<2x!qco.qubit> - %tensor1 = qco.if %condition - args(%arg0 = %tensor0) -> (tensor<2x!qco.qubit>) { - %tensor2, %q0 = qtensor.extract %arg0[%index] - : tensor<2x!qco.qubit> - %q1 = qco.x %q0 : !qco.qubit -> !qco.qubit - %tensor3 = qtensor.insert %q1 into %tensor2[%index] - : tensor<2x!qco.qubit> - qco.yield %tensor3 : tensor<2x!qco.qubit> - } else args(%arg0 = %tensor0) { - qco.yield %arg0 : tensor<2x!qco.qubit> +TEST_F(QCOTest, LeavesUnsupportedQTensorIfUnchanged) { + constexpr std::array mlirCodes = { + R"mlir( + module { + func.func @main(%condition: i1, %index: index) { + %c2 = arith.constant 2 : index + %tensor0 = qtensor.alloc(%c2) : tensor<2x!qco.qubit> + %tensor1 = qco.if %condition + args(%arg0 = %tensor0) -> (tensor<2x!qco.qubit>) { + %tensor2, %q0 = qtensor.extract %arg0[%index] + : tensor<2x!qco.qubit> + %q1 = qco.x %q0 : !qco.qubit -> !qco.qubit + %tensor3 = qtensor.insert %q1 into %tensor2[%index] + : tensor<2x!qco.qubit> + qco.yield %tensor3 : tensor<2x!qco.qubit> + } else args(%arg0 = %tensor0) { + qco.yield %arg0 : tensor<2x!qco.qubit> + } + qtensor.dealloc %tensor1 : tensor<2x!qco.qubit> + return + } } - qtensor.dealloc %tensor1 : tensor<2x!qco.qubit> - return - } - } - )mlir"; - - auto moduleOp = parseSourceString(mlirCode, context.get()); - ASSERT_TRUE(moduleOp); - ASSERT_TRUE(succeeded(verify(*moduleOp))); - ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); - ASSERT_TRUE(succeeded(verify(*moduleOp))); - - IfOp ifOp; - moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); - ASSERT_TRUE(ifOp); - ASSERT_EQ(ifOp.getQubits().size(), 1); - EXPECT_TRUE(isa(ifOp.getQubits().front().getType())); - size_t nestedExtracts = 0; - size_t nestedInserts = 0; - ifOp->walk([&](Operation* operation) { - nestedExtracts += isa(operation); - nestedInserts += isa(operation); - }); - EXPECT_EQ(nestedExtracts, 1); - EXPECT_EQ(nestedInserts, 1); -} - -TEST_F(QCOTest, LeavesDynamicShapeQTensorIfUnchanged) { - constexpr StringLiteral mlirCode = R"mlir( - module { - func.func @main(%condition: i1, %size: index) { - %c0 = arith.constant 0 : index - %tensor0 = qtensor.alloc(%size) : tensor - %tensor1 = qco.if %condition - args(%arg0 = %tensor0) -> (tensor) { - %tensor2, %q0 = qtensor.extract %arg0[%c0] - : tensor - %q1 = qco.x %q0 : !qco.qubit -> !qco.qubit - %tensor3 = qtensor.insert %q1 into %tensor2[%c0] - : tensor - qco.yield %tensor3 : tensor - } else args(%arg0 = %tensor0) { - qco.yield %arg0 : tensor + )mlir", + R"mlir( + module { + func.func @main(%condition: i1, %size: index) { + %c0 = arith.constant 0 : index + %tensor0 = qtensor.alloc(%size) : tensor + %tensor1 = qco.if %condition + args(%arg0 = %tensor0) -> (tensor) { + %tensor2, %q0 = qtensor.extract %arg0[%c0] + : tensor + %q1 = qco.x %q0 : !qco.qubit -> !qco.qubit + %tensor3 = qtensor.insert %q1 into %tensor2[%c0] + : tensor + qco.yield %tensor3 : tensor + } else args(%arg0 = %tensor0) { + qco.yield %arg0 : tensor + } + qtensor.dealloc %tensor1 : tensor + return + } } - qtensor.dealloc %tensor1 : tensor - return - } - } - )mlir"; - - auto moduleOp = parseSourceString(mlirCode, context.get()); - ASSERT_TRUE(moduleOp); - ASSERT_TRUE(succeeded(verify(*moduleOp))); - ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); - ASSERT_TRUE(succeeded(verify(*moduleOp))); + )mlir"}; - IfOp ifOp; - moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); - ASSERT_TRUE(ifOp); - ASSERT_EQ(ifOp.getQubits().size(), 1); - auto tensorType = - dyn_cast(ifOp.getQubits().front().getType()); - ASSERT_TRUE(tensorType); - EXPECT_TRUE(tensorType.isDynamicDim(0)); - size_t nestedExtracts = 0; - size_t nestedInserts = 0; - ifOp->walk([&](Operation* operation) { - nestedExtracts += isa(operation); - nestedInserts += isa(operation); - }); - EXPECT_EQ(nestedExtracts, 1); - EXPECT_EQ(nestedInserts, 1); + for (StringRef mlirCode : mlirCodes) { + auto moduleOp = parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(moduleOp); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + ASSERT_TRUE(succeeded(runQCOCleanupPipeline(moduleOp.get()))); + ASSERT_TRUE(succeeded(verify(*moduleOp))); + + IfOp ifOp; + moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); + ASSERT_TRUE(ifOp); + ASSERT_EQ(ifOp.getQubits().size(), 1); + EXPECT_TRUE(isa(ifOp.getQubits().front().getType())); + size_t nestedExtracts = 0; + size_t nestedInserts = 0; + ifOp->walk([&](Operation* operation) { + nestedExtracts += isa(operation); + nestedInserts += isa(operation); + }); + EXPECT_EQ(nestedExtracts, 1); + EXPECT_EQ(nestedInserts, 1); + } } TEST_F(QCOTest, IndexSwitchParser) { diff --git a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp index fc32d8c0f5..e4ead777b7 100644 --- a/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp +++ b/mlir/unittests/Dialect/QCO/Transforms/Mapping/test_mapping.cpp @@ -670,37 +670,25 @@ TEST_P(MappingPassTest, MapScalarAllocation) { } TEST_F(MappingPassFixture, ExpandNonAdjacentTwoQubitIfOnLineTarget) { - constexpr StringLiteral source = R"mlir( - module { - func.func @main(%condition: i1) - attributes {mqt.entry_point} { - %q0 = qco.alloc : !qco.qubit - %q1 = qco.alloc : !qco.qubit - %q2 = qco.alloc : !qco.qubit - %a0, %a1 = qco.swap %q0, %q1 - : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit - %b1, %b2 = qco.swap %a1, %q2 - : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit - %next0, %next2 = qco.if %condition - args(%arg0 = %a0, %arg2 = %b2) - -> (!qco.qubit, !qco.qubit) { - %then0, %then2 = qco.swap %arg0, %arg2 - : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit - qco.yield %then0, %then2 : !qco.qubit, !qco.qubit - } else args(%arg0 = %a0, %arg2 = %b2) { - qco.yield %arg0, %arg2 : !qco.qubit, !qco.qubit - } - qco.sink %next0 : !qco.qubit - qco.sink %b1 : !qco.qubit - qco.sink %next2 : !qco.qubit - return - } - } - )mlir"; - - auto moduleOp = parseSourceString(source, context.get()); - ASSERT_TRUE(moduleOp); - ASSERT_TRUE(succeeded(verify(*moduleOp))); + QCOProgramBuilder builder(context.get()); + builder.initialize(); + Value q0 = builder.allocQubit(); + Value q1 = builder.allocQubit(); + Value q2 = builder.allocQubit(); + std::tie(q0, q1) = builder.swap(q0, q1); + std::tie(q1, q2) = builder.swap(q1, q2); + SmallVector conditionalInputs{q0, q2}; + auto conditionalResults = builder.qcoIf( + true, conditionalInputs, + [&](ValueRange args) { + auto [then0, then2] = builder.swap(args[0], args[1]); + return SmallVector{then0, then2}; + }, + [](ValueRange args) { return llvm::to_vector(args); }); + builder.sink(conditionalResults[0]); + builder.sink(q1); + builder.sink(conditionalResults[1]); + auto moduleOp = builder.finalize(); const auto target = llvm::cantFail(CompilerTarget::create( 3, std::vector{{0, 1}, {1, 2}})); From b57bcf8ed402f9191829f047439511655e45c721 Mon Sep 17 00:00:00 2001 From: Simon Hofmann Date: Wed, 2 Sep 2026 13:27:38 +0200 Subject: [PATCH 3/3] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Assert=20verified=20QT?= =?UTF-8?q?ensor=20linearity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace an unreachable match failure with an assertion of the verifier-owned one-use invariant. Assisted-by: Codex Signed-off-by: Simon Hofmann --- mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp index cab3c666ed..8a272fefce 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp @@ -360,9 +360,7 @@ analyzeQTensorBranch(Block* block, size_t qTensorArgumentIndex, bool reachedInsertPhase = false; while (true) { - if (!currentQTensor.hasOneUse()) { - return std::nullopt; - } + assert(currentQTensor.hasOneUse() && "expected linear typing"); Operation* user = *currentQTensor.getUsers().begin(); if (user->getBlock() != block) { return std::nullopt;