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..8a272fefce 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,245 @@ struct RemoveUnusedClassicalResults : public OpRewritePattern { return success(); } }; + +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) { + assert(currentQTensor.hasOneUse() && "expected linear typing"); + 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); + + 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 { + rewriter.replaceAllUsesWith(access->second.extract.getResult(), + scalarArguments[indexPosition]); + 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 { + +/// 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(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); + 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(oldQubits); + newQubits.erase(newQubits.begin() + qTensorIndex); + 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( + 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(); + } + 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..fcb8f319c3 100644 --- a/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp +++ b/mlir/lib/Dialect/QCO/Transforms/Mapping/Mapping.cpp @@ -776,6 +776,35 @@ 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 (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 = 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"); + return value; + } + /// Execute `ntrials` many (parallel) initial layout refinement trials and /// return the heuristically best one. /// @@ -1261,6 +1290,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 +1340,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..f174874843 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -981,6 +981,343 @@ 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 + %tensorA0 = qtensor.alloc(%c1) : tensor<1x!qco.qubit> + %middle0 = qco.alloc : !qco.qubit + %tensorB0 = qtensor.alloc(%c1) : tensor<1x!qco.qubit> + %tensorA1, %middle1, %tensorB1 = + qco.if %condition + 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 + : !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> + 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> + } + %middle2 = qco.t %middle1 : !qco.qubit -> !qco.qubit + qtensor.dealloc %tensorA1 : tensor<1x!qco.qubit> + qco.sink %middle2 : !qco.qubit + qtensor.dealloc %tensorB1 : tensor<1x!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; + TOp postMiddle; + SmallVector insertedScalars; + moduleOp->walk([&](IfOp candidate) { ifOp = candidate; }); + moduleOp->walk([&](TOp candidate) { postMiddle = candidate; }); + moduleOp->walk([&](qtensor::InsertOp insert) { + insertedScalars.push_back(insert.getScalar()); + }); + + ASSERT_TRUE(ifOp); + 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()); + })); + EXPECT_TRUE(llvm::all_of(ifOp.getLinearResults(), [](Value value) { + return isa(value.getType()); + })); + + ASSERT_TRUE(postMiddle); + EXPECT_EQ(cast(postMiddle.getOperation()) + .getInputQubits() + .front(), + ifOp.getLinearResults()[0]); + + ASSERT_EQ(insertedScalars.size(), 2); + 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(), 3); + EXPECT_TRUE(isa(thenValues[0].getDefiningOp())); + EXPECT_TRUE(isa(thenValues[1].getDefiningOp())); + EXPECT_TRUE(isa(thenValues[2].getDefiningOp())); +} + +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 + } + } + )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 + } + } + )mlir"}; + + 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) { // 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..e4ead777b7 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,40 @@ TEST_P(MappingPassTest, MapScalarAllocation) { EXPECT_EQ(numStatics, 1); } +TEST_F(MappingPassFixture, ExpandNonAdjacentTwoQubitIfOnLineTarget) { + 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}})); + 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 +1900,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]); }