diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fcb2eecb4..3dbd6409e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -66,7 +66,7 @@ releases may include breaking changes. [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1710], [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765], [#1780], [#1781], [#1782], [#1787], [#1806], [#1807], [#1815], [#1808], - [#1823], [#1824], [#1830], [#1869], [#1872], [#1886], [#1914]) + [#1823], [#1824], [#1830], [#1869], [#1872], [#1886], [#1914], [#1925]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**], [**@J4MMlE**]) @@ -646,6 +646,7 @@ changelogs._ +[#1925]: https://github.com/munich-quantum-toolkit/core/pull/1925 [#1924]: https://github.com/munich-quantum-toolkit/core/pull/1924 [#1915]: https://github.com/munich-quantum-toolkit/core/pull/1915 [#1914]: https://github.com/munich-quantum-toolkit/core/pull/1914 diff --git a/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h b/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h index 95632380c6..57702de649 100644 --- a/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h @@ -1239,6 +1239,37 @@ class QCProgramBuilder final : public ImplicitLocOpBuilder { const function_ref& thenBody, const function_ref& elseBody = nullptr); + /** + * @brief Construct an scf.index_switch operation + * + * @param arg Index argument. + * @param cases The individual switch cases. + * @param caseBodies An array of functions that build the case bodies. + * @param defaultBody Function that builds the default body. + * @return Reference to this builder for method chaining + * + * @par Example: + * ```c++ + * builder.scfIndexSwitch(index, + * SmallVector{0}, + * SmallVector>{[&] { b.x(q0); }}, + * [&] { b.z(q0); }); + * ``` + * ```mlir + * scf.index_switch %condition + * case 0 { + * qc.x %q0 : !qc.qubit + * } + * default { + * qc.z %q0 : !qc.qubit + * } + * ``` + */ + QCProgramBuilder& scfIndexSwitch(const std::variant& arg, + ArrayRef cases, + ArrayRef> caseBodies, + const function_ref& defaultBody); + /** * @brief Construct an scf.condition operation * diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index f3853a7d54..ec59cfa058 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1556,6 +1556,87 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { function_ref thenBody, function_ref elseBody = nullptr); + /** + * @brief Construct an index switch operation for qubits or tensors of qubits + * with linear typing. + * + * @details + * Constructs an index switch operation that takes an index Value and a range + * of qubit and qtensor values that are used in the case regions of this + * operation. The values are passed down as block arguments to each region. + * Qubits that were extracted from a tensor that is used as an argument for + * this operation are automatically inserted before the operation is + * constructed. + * + * @param arg Index argument. + * @param targets Initial arguments for the index switch branches. + * @param cases The individual switch cases. + * @param caseBodies An array of functions that build the case bodies. + * @param defaultBody Function that builds the default body. + * @return ValueRange of the results. + * + * @par Example: + * ```c++ + * result = b.qcoIndexSwitch(arg, initTargets, + * SmallVector{0}, + * SmallVector(ValueRange)>>{ + * [&](ValueRange args) { + * auto q1 = builder.x(args[0]); + * return {q1}; + * } + * }, + * [&](ValueRange args) { + * auto q2 = builder.x(args[0]); + * return {q2}; + * }); + * ``` + * ```mlir + * %result = qco.index_switch %arg -> !qco.qubit + * case 0 args(%arg0 = %q0) { + * %q1 = qco.x %arg0 : !qco.qubit -> !qco.qubit + * qco.yield %q1 : !qco.qubit + * } + * default args(%arg0 = %q0) { + * %q2 = qco.z %arg0 : !qco.qubit -> !qco.qubit + * qco.yield %q2 : !qco.qubit + * } + * ``` + */ + ValueRange qcoIndexSwitch( + const std::variant& arg, ValueRange targets, + ArrayRef cases, + ArrayRef(ValueRange)>> caseBodies, + function_ref(ValueRange)> defaultBody); + + /** + * @brief Construct an index switch operation with a single linear target. + * + * @details + * Constructs an index switch operation for one qubit or qtensor value. + * Each branch callback receives and returns a single value, avoiding + * one-element ranges and vectors. + * + * @param arg Index argument. + * @param target Initial argument for every index switch branch. + * @param cases The individual switch cases. + * @param caseBodies Functions that build the case bodies. + * @param defaultBody Function that builds the default body. + * @return The single result value. + * + * @par Example: + * ```c++ + * result = builder.qcoIndexSwitch( + * arg, target, SmallVector{0}, + * SmallVector>{ + * [&](Value value) { return builder.x(value); }}, + * [&](Value value) { return builder.z(value); }); + * ``` + */ + Value qcoIndexSwitch(const std::variant& arg, Value target, + ArrayRef cases, + ArrayRef> caseBodies, + function_ref defaultBody); + /** * @brief Construct an scf.for operation * diff --git a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td index 3f34e821f7..87ef2a65f0 100644 --- a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td +++ b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td @@ -1480,4 +1480,132 @@ def IfOp let hasVerifier = 1; } +//===----------------------------------------------------------------------===// +// Index Switch Operation +//===----------------------------------------------------------------------===// + +def IndexSwitchOp + : QCOOp< + "index_switch", + traits = [DeclareOpInterfaceMethods< + RegionBranchOpInterface, ["getRegionInvocationBounds", + "getEntrySuccessorRegions"]>, + SingleBlockImplicitTerminator<"YieldOp">, Pure, + RecursiveMemoryEffects]> { + + let summary = "Index-based switch operation for linear (qubit) types"; + let description = [{ + The `qco.index_switch` operation provides multi-way branching based on an index + value, analogous to `scf.index_switch`. In addition to the index, the operation + takes a variadic number of qubits and qtensors as inputs that are required in all + case branches. These values are passed down to each case region as block arguments. + The number of results and the type of the results must be equivalent to the number + and types of the input qubit and qtensor values. + + The operation has one region for each case value plus a default region. Each region + must terminate with a `qco.yield` operation that yields values matching the operation's + result types. + + Example: + ```mlir + %result = qco.index_switch %index -> !qco.qubit + case 0 args(%arg0 = %q0) { + %q1 = qco.h %arg0 : !qco.qubit -> !qco.qubit + qco.yield %q1 : !qco.qubit + } + case 1 args(%arg0 = %q0) { + %q1 = qco.x %arg0 : !qco.qubit -> !qco.qubit + qco.yield %q1 : !qco.qubit + } + default args(%arg0 = %q0) { + qco.yield %arg0 : !qco.qubit + } + ``` + }]; + + let arguments = (ins Index:$arg, DenseI64ArrayAttr:$cases, + Variadic:$targets); + let results = (outs Variadic:$results); + let regions = (region SizedRegion<1>:$defaultRegion, + VariadicRegion>:$caseRegions); + let hasCustomAssemblyFormat = 1; + + let builders = [OpBuilder<(ins "Value":$arg, "Value":$target, + "ArrayRef":$cases, + "ArrayRef>":$caseBuilders, + "function_ref":$defaultBuilder)>]; + + let extraClassDeclaration = [{ + /// Return the number of case regions + size_t getNumCases(); + + /// Return the block for a specific case + Block* getCaseBlock(size_t index); + + /// Return the yield operation for a specific case + YieldOp getCaseYield(size_t index); + + /// Return the default block + Block* getDefaultBlock(); + + /// Return the default yield operation + YieldOp getDefaultYield(); + + /// Return the result that corresponds to the given target operand, + /// or "empty" OpResult on failure. + OpResult getTiedResult(OpOperand* target); + + /// Return the target operand that corresponds to the given result, + /// or `nullptr` on failure. + OpOperand* getTiedTarget(OpResult result); + + /// Return the argument that corresponds to the given target operand + /// for the i-th case, or "empty" BlockArgument on failure. + BlockArgument getTiedCaseBlockArgument(OpOperand* target, size_t i); + + /// Return the yielded value that corresponds to the given argument + /// for the i-th case, or `nullptr` on failure. + OpOperand* getTiedCaseYieldedValue(BlockArgument bbArg, size_t i); + + /// Return the argument that corresponds to the given target operand + /// for the default case, or "empty" BlockArgument on failure. + BlockArgument getTiedDefaultBlockArgument(OpOperand* target); + + /// Return the yielded value that corresponds to the given argument + /// for the default case, or `nullptr` on failure. + OpOperand* getTiedDefaultYieldedValue(BlockArgument bbArg); + + /// Append the specified additional linear operands: replace this + /// index_switch with a new index_switch that has the additional linear operands. + /// The operands can be of qubit or qtensor type. + /// The branch bodies of this index_switch are moved over to the new index_switch. + /// The newly added qubits are yielded from each branch. + IndexSwitchOp replaceWithAdditionalTargets(RewriterBase& rewriter, ValueRange addons); + }]; + + let extraClassDefinition = [{ + size_t $cppClass::getNumCases() { + return getCaseRegions().size(); + } + + Block* $cppClass::getCaseBlock(size_t index) { + return &getCaseRegions()[index].front(); + } + + YieldOp $cppClass::getCaseYield(size_t index) { + return cast(getCaseBlock(index)->getTerminator()); + } + + Block* $cppClass::getDefaultBlock() { + return &getDefaultRegion().front(); + } + + YieldOp $cppClass::getDefaultYield() { + return cast(getDefaultBlock()->getTerminator()); + } + }]; + + let hasVerifier = 1; +} + #endif // MLIR_DIALECT_QCO_IR_QCOOPS_TD diff --git a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h index 6fe2730cbc..7187931965 100644 --- a/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h +++ b/mlir/include/mlir/Dialect/QCO/Utils/Drivers.h @@ -53,6 +53,11 @@ inline size_t getNumQubitArgs(Operation* op) { return isa(v.getType()); }); }) + .Case([&](qco::IndexSwitchOp op) { + return llvm::count_if(op.getTargets(), [](Value v) { + return isa(v.getType()); + }); + }) .Default([&](Operation* op) { const auto name = op->getName().getStringRef(); reportFatalInternalError("unknown op: " + name); diff --git a/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp b/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp index 7c0815542c..c13f98d16e 100644 --- a/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp +++ b/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp @@ -737,7 +737,7 @@ struct ConvertQCOPowOp final : OpConversionPattern { /** * @brief Converts qco.yield to qc.yield or to scf.yield if the parent is a - * scf::IfOp + * scf::IfOp or scf::IndexSwitchOp. * * @par Example: * ```mlir @@ -754,7 +754,7 @@ struct ConvertQCOYieldOp final : OpConversionPattern { LogicalResult matchAndRewrite(qco::YieldOp op, OpAdaptor /*adaptor*/, ConversionPatternRewriter& rewriter) const override { - if (isa(op->getParentOp())) { + if (isa(op->getParentOp())) { rewriter.replaceOpWithNewOp(op); } else { rewriter.replaceOpWithNewOp(op); @@ -907,6 +907,59 @@ struct ConvertQCOIfOp final : OpConversionPattern { } }; +/** + * @brief Converts qco.index_switch to scf.index_switch + * + * @par Example: + * ```mlir + * %result = qco.index_switch %condition -> !qco.qubit + * case 0 args(%arg0 = %q0) { + * %q1 = qco.x %arg0 : !qco.qubit -> !qco.qubit + * qco.yield %q1 : !qco.qubit + * } + * default args(%arg0 = %q0) { + * %q2 = qco.z %arg0 : !qco.qubit -> !qco.qubit + * qco.yield %q2 : !qco.qubit + * } + * ``` + * is converted to + * ```mlir + * scf.index_switch %condition + * case 0 { + * qc.x %q0 : !qc.qubit + * } + * default { + * qc.z %q0 : !qc.qubit + * } + * ``` + */ +struct ConvertQCOIndexSwitchOp final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(IndexSwitchOp op, OpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + auto newOp = + scf::IndexSwitchOp::create(rewriter, op.getLoc(), {}, adaptor.getArg(), + adaptor.getCases(), op.getNumCases()); + + const auto oldRegions = op.getCaseRegions(); + const auto newCaseRegions = newOp.getCaseRegions(); + for (size_t i = 0; i < op.getNumCases(); ++i) { + inlineRegion(oldRegions[i], newCaseRegions[i], 0, + adaptor.getOperands().drop_front(1), rewriter); + } + + inlineRegion(op.getDefaultRegion(), newOp.getDefaultRegion(), 0, + adaptor.getOperands().drop_front(1), rewriter); + + // Replace the qco results with the input qc values except the condition + rewriter.replaceOp(op, adaptor.getOperands().drop_front(1)); + + return success(); + } +}; + /** * @brief Converts scf.yield with value semantics to scf.yield with memory * semantics for qubit values. This currently assumes no mixed types as yielded @@ -1036,9 +1089,9 @@ struct QCOToQC final : impl::QCOToQCBase { patterns.add(typeConverter, - context); + ConvertQCOIndexSwitchOp, ConvertQCOSCFWhileOp, + ConvertQCOSCFConditionOp, ConvertQCOSCFYieldOp, + ConvertQCOSCFForOp>(typeConverter, context); // Register operation conversion patterns that need state tracking patterns.add { // Create the new IfOp auto newIfOp = - IfOp::create(rewriter, op->getLoc(), op.getCondition(), qcoTargets); + IfOp::create(rewriter, op.getLoc(), op.getCondition(), qcoTargets); assignMappedTensors(state, op.getOperation(), registerMap, newIfOp.getResults().take_front(numRegisters)); assignMappedQubits(state, op.getOperation(), qubitMap, - newIfOp->getResults().take_back(numQubits)); + newIfOp.getResults().take_back(numQubits)); // Extract all the previously inserted qubits again rewriter.setInsertionPointAfter(newIfOp); @@ -1506,9 +1506,100 @@ struct ConvertSCFIfOp final : StatefulOpConversionPattern { } }; +/** + * @brief Converts scf.index_switch to qco.index_switch + * + * @par Example: + * ```mlir + * scf.index_switch %condition + * case 0 { + * qc.x %q0 : !qc.qubit + * } + * default { + * qc.z %q0 : !qc.qubit + * } + * ``` + * is converted to + * ```mlir + * %result = qco.index_switch %condition -> !qco.qubit + * case 0 args(%arg0 = %q0) { + * %q1 = qco.x %arg0 : !qco.qubit -> !qco.qubit + * qco.yield %q1 : !qco.qubit + * } + * default args(%arg0 = %q0) { + * %q2 = qco.z %arg0 : !qco.qubit -> !qco.qubit + * qco.yield %q2 : !qco.qubit + * } + * ``` + */ +struct ConvertSCFIndexSwitchOp final + : StatefulOpConversionPattern { + using StatefulOpConversionPattern::StatefulOpConversionPattern; + + LogicalResult + matchAndRewrite(scf::IndexSwitchOp op, OpAdaptor /*adaptor*/, + ConversionPatternRewriter& rewriter) const override { + auto& state = getState(); + auto* operation = op.getOperation(); + auto& registerMap = state.regionRegisterMap[op]; + auto& qubitMap = state.regionQubitMap[op]; + + const auto qubits = qubitMap.getArrayRef(); + const auto nqubits = qubits.size(); + const auto registers = registerMap.getArrayRef(); + const auto nregisters = registers.size(); + + insertQubitsBeforeOp(state, operation, rewriter); + + const auto targets = resolveAllValues(state, operation); + const auto results = ValueRange(targets).getTypes(); + const SmallVector locs(targets.size(), op.getLoc()); + + auto newOp = + IndexSwitchOp::create(rewriter, op.getLoc(), results, op.getArg(), + op.getCases(), targets, op.getNumCases()); + + assignMappedTensors(state, op.getOperation(), registerMap, + newOp.getResults().take_front(nregisters)); + assignMappedQubits(state, op.getOperation(), qubitMap, + newOp.getResults().take_back(nqubits)); + + rewriter.setInsertionPointAfter(newOp); + extractQubitsAfterOp(state, operation, rewriter); + + const auto newCaseRegions = newOp.getCaseRegions(); + const auto oldCaseRegions = op.getCaseRegions(); + const auto buildRegion = [&](Region& oldRegion, Region& newRegion) { + Block* const oldBlock = &(*oldRegion.begin()); + Block* const newBlock = + rewriter.createBlock(&newRegion, {}, newOp.getResultTypes(), locs); + rewriter.inlineBlockBefore(oldBlock, newBlock, newBlock->begin()); + + const auto args = newBlock->getArguments(); + pushModifierFrameWithRegisters(state, qubits, registers, + args.take_back(nqubits), + args.take_front(nregisters)); + }; + + // Process regions and push frames in the correct order. The top-most stack + // element must be the default region (0th index) followed by the case + // regions. + + for (size_t i = op.getNumCases(); i > 0; --i) { + buildRegion(oldCaseRegions[i - 1], newCaseRegions[i - 1]); + } + buildRegion(op.getDefaultRegion(), newOp.getDefaultRegion()); + + rewriter.eraseOp(op); + + return success(); + } +}; + /** * @brief Converts scf.yield with memory semantics to scf.yield with value - * semantics for qubit values or to qco.scf if the parentOp is a qco::IfOp + * semantics for qubit values or to qco.yield if the parentOp is a qco::IfOp or + * qco::IndexSwitchOp. * * @par Example: * ```mlir @@ -1531,7 +1622,7 @@ struct ConvertSCFYieldOp final : StatefulOpConversionPattern { insertQubitsBeforeOp(state, op.getOperation(), rewriter); SmallVector targets = resolveAllValues(state, operation); - if (isa(op->getParentOp())) { + if (isa(op->getParentOp())) { rewriter.replaceOpWithNewOp(op, targets); } else { rewriter.replaceOpWithNewOp(op, targets); @@ -1652,13 +1743,14 @@ struct QCToQCO final : impl::QCToQCOBase { }); // Register operation conversion patterns with state tracking. - patterns.add( - typeConverter, context, &state); + patterns + .add( + typeConverter, context, &state); // Not part of the central gate table. patterns.add>( diff --git a/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp b/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp index 83ce01febb..c29a93f0bf 100644 --- a/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp +++ b/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp @@ -14,6 +14,7 @@ #include "mlir/Dialect/QC/IR/QCOps.h" #include "mlir/Dialect/Utils/Utils.h" +#include #include #include #include @@ -642,6 +643,42 @@ QCProgramBuilder::scfIf(const std::variant& cond, return *this; } +QCProgramBuilder& +QCProgramBuilder::scfIndexSwitch(const std::variant& arg, + ArrayRef cases, + ArrayRef> caseBodies, + const function_ref& defaultBody) { + checkFinalized(); + + if (cases.size() != caseBodies.size()) { + const char* msg = "Each case must have a corresponding case body function"; + llvm::reportFatalUsageError(msg); + llvm_unreachable(msg); + } + + auto argValue = variantToValue(*this, getLoc(), arg); + auto switchOp = + scf::IndexSwitchOp::create(*this, {}, argValue, cases, cases.size()); + + const InsertionGuard guard(*this); + const auto buildRegion = [&](Region& region, const function_ref& f) { + Block* block = createBlock(®ion); // Implicitly sets the insertion point. + regionStack.emplace_back(®ion); + f(); + scf::YieldOp::create(*this, getLoc()); + regionStack.pop_back(); + }; + + for (auto [region, f] : + llvm::zip_equal(switchOp.getCaseRegions(), caseBodies)) { + buildRegion(region, f); + } + + buildRegion(switchOp.getDefaultRegion(), defaultBody); + + return *this; +} + QCProgramBuilder& QCProgramBuilder::scfCondition(Value condition) { checkFinalized(); scf::ConditionOp::create(*this, condition, ValueRange{}); diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index ad8920e5bd..6bf4f66902 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -19,12 +19,14 @@ #include #include #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -1204,13 +1206,104 @@ ValueRange QCOProgramBuilder::qcoIf( "Else body must return exactly one value per input value"); } YieldOp::create(*this, elseResult); - updateQubitValueTracking(elseResult, ifOp->getResults()); + updateQubitValueTracking(elseResult, ifOp.getResults()); } else { YieldOp::create(*this, elseArgs); - updateQubitValueTracking(thenResult, ifOp->getResults()); + updateQubitValueTracking(thenResult, ifOp.getResults()); } - return ifOp->getResults(); + return ifOp.getResults(); +} + +ValueRange QCOProgramBuilder::qcoIndexSwitch( + const std::variant& arg, ValueRange targets, + ArrayRef cases, + ArrayRef(ValueRange)>> caseBodies, + const function_ref(ValueRange)> defaultBody) { + checkFinalized(); + + if (cases.size() != caseBodies.size()) { + const char* msg = "Each case must have a corresponding case body function"; + llvm::reportFatalUsageError(msg); + llvm_unreachable(msg); + } + + const auto ntargets = targets.size(); + const auto types = targets.getTypes(); + const auto updatedTargets = prepareInitArgs(targets); + const auto argValue = variantToValue(*this, getLoc(), arg); + + auto switchOp = IndexSwitchOp::create(*this, types, argValue, cases, + updatedTargets, cases.size()); + + const InsertionGuard guard(*this); + const SmallVector locs(ntargets, getLoc()); + + const auto buildRegion = [&](Region& region, SmallVector& prev, + function_ref(ValueRange)> f) { + Block* const block = createBlock(®ion, {}, types, locs); + updateQubitValueTracking(prev, block->getArguments()); + + const auto result = f(block->getArguments()); + if (result.size() != ntargets) { + const char* msg = + "Case body must return exactly one value per input value"; + llvm::reportFatalUsageError(msg); + llvm_unreachable(msg); + } + + YieldOp::create(*this, result); + prev = result; + }; + + SmallVector prev(updatedTargets); + for (const auto [region, f] : + llvm::zip_equal(switchOp.getCaseRegions(), caseBodies)) { + buildRegion(region, prev, f); + } + + buildRegion(switchOp.getDefaultRegion(), prev, defaultBody); + updateQubitValueTracking(prev, switchOp.getResults()); + + return switchOp.getResults(); +} + +Value QCOProgramBuilder::qcoIndexSwitch( + const std::variant& arg, Value target, + ArrayRef cases, ArrayRef> caseBodies, + function_ref defaultBody) { + checkFinalized(); + + if (cases.size() != caseBodies.size()) { + llvm::reportFatalUsageError( + "Each case must have a corresponding case body function"); + } + + const auto updatedTarget = prepareInitArg(target); + const auto argValue = variantToValue(*this, getLoc(), arg); + auto switchOp = IndexSwitchOp::create(*this, target.getType(), argValue, + cases, updatedTarget, cases.size()); + + const InsertionGuard guard(*this); + const auto buildRegion = [&](Region& region, Value previous, + function_ref body) -> Value { + auto& block = region.emplaceBlock(); + const auto blockArgument = block.addArgument(target.getType(), getLoc()); + updateQubitValueTracking(previous, blockArgument); + setInsertionPointToStart(&block); + const auto result = body(blockArgument); + YieldOp::create(*this, result); + return result; + }; + + Value previous = updatedTarget; + for (const auto [region, body] : + llvm::zip_equal(switchOp.getCaseRegions(), caseBodies)) { + previous = buildRegion(region, previous, body); + } + previous = buildRegion(switchOp.getDefaultRegion(), previous, defaultBody); + updateQubitValueTracking(previous, switchOp.getResult(0)); + return switchOp.getResult(0); } Value QCOProgramBuilder::qcoIf(const std::variant& condition, diff --git a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp index f1cb23a849..0fb150ee4e 100644 --- a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp +++ b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp @@ -15,6 +15,7 @@ #include #include +#include #include #include #include @@ -22,6 +23,9 @@ #include #include +#include +#include + // The following headers are needed for some template instantiations. // IWYU pragma: begin_keep #include @@ -169,6 +173,198 @@ void IfOp::print(OpAsmPrinter& p) { p.printOptionalAttrDict((*this)->getAttrs()); } +ParseResult IndexSwitchOp::parse(::mlir::OpAsmParser& parser, + ::mlir::OperationState& result) { + OpAsmParser::UnresolvedOperand index; + if (parser.parseOperand(index) || + parser.resolveOperand(index, parser.getBuilder().getIndexType(), + result.operands)) { + return failure(); + } + + if (parser.parseOptionalArrowTypeList(result.types)) { + return failure(); + } + + if (parser.parseOptionalAttrDict(result.attributes)) { + return failure(); + } + + // Create default region here to ensure regions(0) = default. + Region* defaultRegion = result.addRegion(); + + SmallVector operands; + SmallVector caseValues; + SmallVector regionArgs; + SmallVector regionOperands; + + while (succeeded(parser.parseOptionalKeyword("case"))) { + int64_t caseValue = 0; + if (parser.parseInteger(caseValue)) { + return failure(); + } + + caseValues.push_back(caseValue); + + if (parser.parseKeyword("args")) { + return failure(); + } + + if (parser.parseAssignmentList(regionArgs, regionOperands)) { + return failure(); + } + + if (caseValues.size() == 1) { + + // Resolve the operands into the result for the very first case. + + if (parser.resolveOperands(regionOperands, result.types, + parser.getCurrentLocation(), + result.operands)) { + return failure(); + } + } else { + + // Otherwise, verify if the other cases use the equivalent operands (minus + // the case-value) as the first one. + + SmallVector operands; + if (parser.resolveOperands(regionOperands, result.types, + parser.getCurrentLocation(), operands)) { + return failure(); + } + + for (auto [v0, v1] : + llvm::zip_equal(operands, llvm::drop_begin(result.operands, 1))) { + if (v0 != v1) { + return parser.emitError( + parser.getCurrentLocation(), + "else qubits must reference the same SSA values as then qubits"); + } + } + } + + for (auto [arg, type] : llvm::zip_equal(regionArgs, result.types)) { + arg.type = type; + } + + if (parser.parseRegion(*result.addRegion(), regionArgs)) { + return failure(); + } + + operands.clear(); + regionArgs.clear(); + regionOperands.clear(); + } + + result.addAttribute("cases", + DenseI64ArrayAttr::get(parser.getContext(), caseValues)); + + // Parse the default regions and again verify if the default case uses the + // equivalent operands (minus the case-value) as all other cases. + + if (parser.parseKeyword("default")) { + return failure(); + } + + if (parser.parseKeyword("args")) { + return failure(); + } + + if (parser.parseAssignmentList(regionArgs, regionOperands)) { + return failure(); + } + + // Otherwise, verify if the other cases use the equivalent operands (minus + // the case-value) for all other cases. + + if (parser.resolveOperands(regionOperands, result.types, + parser.getCurrentLocation(), operands)) { + return failure(); + } + + if (caseValues.empty()) { + result.operands.append(operands); + } else { + for (auto [v0, v1] : + llvm::zip_equal(operands, llvm::drop_begin(result.operands, 1))) { + if (v0 != v1) { + return parser.emitError( + parser.getCurrentLocation(), + "else qubits must reference the same SSA values as then qubits"); + } + } + } + + for (auto [args, type] : llvm::zip_equal(regionArgs, result.types)) { + args.type = type; + } + + if (parser.parseRegion(*defaultRegion, regionArgs)) { + return failure(); + } + + return success(); +} + +void IndexSwitchOp::print(OpAsmPrinter& p) { + p << " "; + p.printOperand(getArg()); + + // Print result types if present + if (!getResults().empty()) { + p << " -> "; + llvm::interleaveComma(getResultTypes(), p); + } + + // Print attributes (excluding cases which we handle specially) + p.printOptionalAttrDictWithKeyword(getOperation()->getAttrs(), + /*elidedAttrs=*/{"cases"}); + + // Print case regions + const auto cases = getCases(); + for (size_t i = 0; i < getNumCases(); ++i) { + p.printNewline(); + p << "case "; + p << cases[i]; + p << " args("; + + auto& region = getCaseRegions()[i]; + auto& block = region.front(); + + // Print block arguments with their corresponding target operands + for (size_t j = 0; j < block.getNumArguments(); ++j) { + if (j > 0) { + p << ", "; + } + p.printOperand(block.getArgument(j)); + p << " = "; + p.printOperand(getTargets()[j]); + } + p << ") "; + p.printRegion(region, /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/true); + } + + p.printNewline(); + p << "default args("; + auto& defaultRegion = getDefaultRegion(); + auto& defaultBlock = defaultRegion.front(); + + // Print block arguments with their corresponding target operands + for (size_t j = 0; j < defaultBlock.getNumArguments(); ++j) { + if (j > 0) { + p << ", "; + } + p.printOperand(defaultBlock.getArgument(j)); + p << " = "; + p.printOperand(getTargets()[j]); + } + p << ") "; + p.printRegion(defaultRegion, /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/true); +} + //===----------------------------------------------------------------------===// // Dialect //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp index 2b515b858b..2e1a500ad2 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp @@ -353,20 +353,22 @@ IfOp IfOp::replaceWithAdditionalQubits(RewriterBase& rewriter, return *this; } - SmallVector allQubits; - allQubits.reserve(getQubits().size() + addons.size()); - allQubits.append(getQubits().begin(), getQubits().end()); - allQubits.append(addons.begin(), addons.end()); - const auto allQubitTypes = ValueRange(allQubits).getTypes(); + const auto qubits = getQubits(); - auto newIfOp = create(rewriter, getLoc(), getCondition(), allQubits); + SmallVector newQubits; + newQubits.reserve(qubits.size() + addons.size()); + newQubits.append(qubits.begin(), qubits.end()); + newQubits.append(addons.begin(), addons.end()); + const auto allQubitTypes = ValueRange(newQubits).getTypes(); + + auto newIfOp = create(rewriter, getLoc(), getCondition(), newQubits); const auto rewriteRegion = [&](Region& oldRegion, Region& newRegion) { auto* oldBlock = &oldRegion.front(); const auto numOldArgs = oldBlock->getNumArguments(); auto* newBlock = rewriter.createBlock(&newRegion, {}, allQubitTypes, - SmallVector(allQubits.size(), getLoc())); + SmallVector(newQubits.size(), getLoc())); const auto oldArgs = newBlock->getArguments().take_front(numOldArgs); const auto addonArgs = newBlock->getArguments().drop_front(numOldArgs); diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp new file mode 100644 index 0000000000..f57c1e7cb4 --- /dev/null +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -0,0 +1,272 @@ +/* + * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM + * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH + * All rights reserved. + * + * SPDX-License-Identifier: MIT + * + * Licensed under the MIT License + */ + +#include "mlir/Dialect/QCO/IR/QCOOps.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::qco; + +void IndexSwitchOp::build(OpBuilder& odsBuilder, OperationState& odsState, + Value arg, Value target, ArrayRef cases, + ArrayRef> caseBuilders, + function_ref defaultBuilder) { + if (cases.size() != caseBuilders.size()) { + llvm::reportFatalUsageError( + "Each case must have a corresponding case body function"); + } + + build(odsBuilder, odsState, target.getType(), arg, cases, target, + cases.size()); + + const OpBuilder::InsertionGuard guard(odsBuilder); + const auto buildRegion = [&](Region& region, + function_ref bodyBuilder) { + auto& block = region.emplaceBlock(); + const auto blockArgument = + block.addArgument(target.getType(), odsState.location); + odsBuilder.setInsertionPointToStart(&block); + YieldOp::create(odsBuilder, odsState.location, bodyBuilder(blockArgument)); + }; + + for (const auto [index, caseBuilder] : llvm::enumerate(caseBuilders)) { + buildRegion(*odsState.regions[index + 1], caseBuilder); + } + buildRegion(*odsState.regions.front(), defaultBuilder); +} + +// Adapted from +// https://github.com/llvm/llvm-project/blob/llvmorg-22.1.1/mlir/lib/Dialect/SCF/IR/SCF.cpp + +void IndexSwitchOp::getSuccessorRegions( + RegionBranchPoint point, SmallVectorImpl& regions) { + if (!point.isParent()) { + regions.emplace_back(getOperation(), getResults()); + return; + } + + llvm::append_range(regions, getRegions()); +} + +void IndexSwitchOp::getRegionInvocationBounds( + ArrayRef operands, SmallVectorImpl& bounds) { + FoldAdaptor adaptor(operands, *this); + + // If the constant "arg" operand is not provided, we can't reason about the + // invocation bounds and thus assume that all regions are invoked at most + // once. + + auto arg = llvm::dyn_cast_or_null(adaptor.getArg()); + if (!arg) { + bounds.append(getNumRegions(), InvocationBounds(/*lb=*/0, /*ub=*/1)); + return; + } + + // Otherwise, we can reason that all but the "live" case (can be the default + // case) are invoked zero times. + + const auto nregions = getNumRegions(); + const auto* it = llvm::find(getCases(), arg.getInt()); + const auto liveIndex = it != getCases().end() + ? std::distance(getCases().begin(), it) + 1 + : 0; // Default region. + + for (size_t i = 0; i < nregions; ++i) { + bounds.emplace_back(/*lb=*/0, /*ub=*/std::cmp_equal(i, liveIndex) ? 1 : 0); + } +} + +void IndexSwitchOp::getEntrySuccessorRegions( + ArrayRef operands, SmallVectorImpl& regions) { + FoldAdaptor adaptor(operands, *this); + + // If a constant was not provided, all regions are possible successors. + auto arg = dyn_cast_or_null(adaptor.getArg()); + if (!arg) { + llvm::append_range(regions, getRegions()); + return; + } + + // Otherwise, try to find a case with a matching value. If not, the + // default region is the only successor. + + const auto* it = llvm::find(getCases(), arg.getInt()); + if (it == getCases().end()) { + regions.emplace_back(&getDefaultRegion()); + return; + } + + const auto caseIndex = std::distance(getCases().begin(), it); + regions.emplace_back(&getCaseRegions()[caseIndex]); +} + +LogicalResult IndexSwitchOp::verify() { + if (getCases().size() != getNumCases()) { + return emitOpError( + "must have the same number of case values and case regions"); + } + + llvm::SmallDenseSet uniqueCases; + for (const int64_t caseValue : getCases()) { + if (!uniqueCases.insert(caseValue).second) { + return emitOpError("case values must be unique"); + } + } + + const auto targets = getTargets(); + const auto ntargets = targets.size(); + const auto results = getResults(); + const auto nresults = results.size(); + + for (Region* region : getRegions()) { + if (region->getNumArguments() != ntargets) { + return emitOpError( + "Region " + Twine(region->getRegionNumber()) + + " must have the same number of arguments as the number of targets"); + } + } + + SmallPtrSet visited; + for (const auto target : targets) { + if (!visited.insert(target).second) { + return emitOpError("The operation requires unique values as targets."); + } + } + + if (nresults != ntargets) { + return emitOpError( + "The operation must consume and produce the same number of values."); + } + + for (auto [resType, targetType] : + llvm::zip_equal(results.getTypes(), targets.getTypes())) { + if (resType != targetType) { + return emitOpError( + "The operation must consume and produce the same types."); + } + } + + return success(); +} + +OpResult IndexSwitchOp::getTiedResult(OpOperand* target) { + if (target->getOwner() != getOperation()) { + return {}; + } + // Because the first operand is the index, subtract one. + return getResults()[target->getOperandNumber() - 1]; +} + +OpOperand* IndexSwitchOp::getTiedTarget(OpResult result) { + if (result.getDefiningOp() != getOperation()) { + return nullptr; + } + return &getTargetsMutable()[result.getResultNumber()]; +} + +BlockArgument IndexSwitchOp::getTiedCaseBlockArgument(OpOperand* target, + size_t i) { + if (target->getOwner() != getOperation() || i >= getNumCases()) { + return {}; + } + + return getCaseBlock(i)->getArgument(target->getOperandNumber() - 1); +} + +OpOperand* IndexSwitchOp::getTiedCaseYieldedValue(BlockArgument bbArg, + size_t i) { + if (i >= getNumCases() || bbArg.getOwner() != getCaseBlock(i)) { + return nullptr; + } + + return &getCaseYield(i).getTargetsMutable()[bbArg.getArgNumber()]; +} + +BlockArgument IndexSwitchOp::getTiedDefaultBlockArgument(OpOperand* target) { + if (target->getOwner() != getOperation()) { + return {}; + } + + return getDefaultBlock()->getArgument(target->getOperandNumber() - 1); +} + +OpOperand* IndexSwitchOp::getTiedDefaultYieldedValue(BlockArgument bbArg) { + if (bbArg.getOwner() != getDefaultBlock()) { + return nullptr; + } + + return &getDefaultYield().getTargetsMutable()[bbArg.getArgNumber()]; +} + +IndexSwitchOp +IndexSwitchOp::replaceWithAdditionalTargets(RewriterBase& rewriter, + ValueRange addons) { + if (addons.empty()) { + return *this; + } + + const auto targets = getTargets(); + const auto nregions = getNumRegions(); + + SmallVector newTargets; + newTargets.reserve(targets.size() + addons.size()); + newTargets.append(targets.begin(), targets.end()); + newTargets.append(addons.begin(), addons.end()); + const auto newTargetTypes = ValueRange(newTargets).getTypes(); + + auto newSwitchOp = create(rewriter, getLoc(), newTargetTypes, getArg(), + getCases(), newTargets, getNumCases()); + + const auto rewriteRegion = [&](Region& oldRegion, Region& newRegion) { + auto* oldBlock = &oldRegion.front(); + const auto numOldArgs = oldBlock->getNumArguments(); + auto* newBlock = rewriter.createBlock( + &newRegion, {}, newTargetTypes, + SmallVector(newTargets.size(), getLoc())); + const auto oldArgs = newBlock->getArguments().take_front(numOldArgs); + const auto addonArgs = newBlock->getArguments().drop_front(numOldArgs); + + rewriter.mergeBlocks(oldBlock, newBlock, oldArgs); + + auto yield = cast(newBlock->getTerminator()); + SmallVector yieldedValues; + yieldedValues.reserve(yield.getTargets().size() + addons.size()); + yieldedValues.append(yield.getTargets().begin(), yield.getTargets().end()); + yieldedValues.append(addonArgs.begin(), addonArgs.end()); + rewriter.replaceOpWithNewOp(yield, yieldedValues); + }; + + for (size_t i = 0; i < nregions; ++i) { + rewriteRegion(getRegion(i), newSwitchOp.getRegion(i)); + } + + rewriter.replaceOp(*this, + newSwitchOp.getResults().take_front(getNumResults())); + + return newSwitchOp; +} diff --git a/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp b/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp index 3b97c9ff01..34ae2f5aef 100644 --- a/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp +++ b/mlir/lib/Dialect/QCO/Utils/WireIterator.cpp @@ -86,6 +86,9 @@ void WireIterator::forward() { }) .Case( [&](IfOp op) { qubit_ = op.getTiedResult(&(*qubit_.use_begin())); }) + .Case([&](IndexSwitchOp op) { + qubit_ = op.getTiedResult(&(*qubit_.use_begin())); + }) .Default([&](Operation* op) { llvm::reportFatalInternalError("unknown op in def-use chain: " + op->getName().getStringRef()); @@ -152,6 +155,13 @@ void WireIterator::backward() { } llvm::reportFatalInternalError("expected result lookup"); }) + .Case([&](IndexSwitchOp op) { + if (auto result = dyn_cast(qubit_)) { + qubit_ = op.getTiedTarget(result)->get(); + return; + } + llvm::reportFatalInternalError("expected result lookup"); + }) .Default([&](Operation* op) { llvm::reportFatalInternalError("unknown op in def-use chain: " + op->getName().getStringRef()); diff --git a/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp b/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp index 6af8c5bff7..da550a5781 100644 --- a/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp +++ b/mlir/lib/Dialect/QTensor/Utils/TensorIterator.cpp @@ -71,6 +71,10 @@ void TensorIterator::forward() { const auto idx = std::distance(op.getQubits().begin(), it); tensor_ = cast>(op.getResults()[idx]); }) + .Case([&](qco::IndexSwitchOp op) { + tensor_ = cast>( + op.getTiedResult(&(*tensor_.use_begin()))); + }) .Default([&](Operation* op) { report_fatal_error("unknown op in def-use chain: " + op->getName().getStringRef()); @@ -131,6 +135,16 @@ void TensorIterator::backward() { llvm::reportFatalInternalError( "expected scf.for result for tied init lookup"); }) + .Case([&](qco::IndexSwitchOp op) { + if (auto result = dyn_cast(tensor_)) { + tensor_ = cast>( + op.getTiedTarget(result)->get()); + return; + } + + llvm::reportFatalInternalError( + "expected qco.index_switch result for tied target lookup"); + }) .Default([&](Operation* op) { llvm::reportFatalInternalError("unknown op in def-use chain: " + op->getName().getStringRef()); diff --git a/mlir/lib/Support/IRVerification.cpp b/mlir/lib/Support/IRVerification.cpp index 79ec785ba7..edf44be05e 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -16,7 +16,6 @@ #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" #include -#include #include #include #include @@ -95,9 +94,10 @@ static void initEquivGroup(TypedValue v, size_t id, if (auto op = dyn_cast(it.operation())) { const auto prev = std::prev(it); - const auto qIt = llvm::find(op.getQubits(), prev.tensor()); + const auto qubits = op.getQubits(); + const auto qIt = llvm::find(qubits, prev.tensor()); assert(qIt != op.getQubits().end()); - const auto idx = std::distance(op.getQubits().begin(), qIt); + const auto idx = std::distance(qubits.begin(), qIt); auto& thenRegion = op.getThenRegion(); auto& elseRegion = op.getElseRegion(); @@ -107,6 +107,18 @@ static void initEquivGroup(TypedValue v, size_t id, initEquivGroup(cast>(thenArg), id, group); initEquivGroup(cast>(elseArg), id, group); + } else if (auto op = dyn_cast(it.operation())) { + const auto prev = std::prev(it); + const auto targets = op.getTargets(); + const auto targetIt = llvm::find(targets, prev.tensor()); + assert(targetIt != targets.end()); + const auto idx = std::distance(targets.begin(), targetIt); + + for (Region* region : op.getRegions()) { + initEquivGroup( + cast>(region->getArgument(idx)), id, + group); + } } else if (auto forOp = dyn_cast(it.operation())) { const auto& arg = forOp.getTiedLoopRegionIterArg(cast(it.tensor())); @@ -154,99 +166,56 @@ static void mapArguments(Block& lhs, Block& rhs, ArrayRef permutation, } } -/// Return a permutation vector, where permutation[i] maps the i-th target of -/// the lhs to the j-th target of the rhs. -static SmallVector getTargetPermutation(qc::CtrlOp lhs, qc::CtrlOp rhs, - const IRMapping& m) { - SmallVector permutation(lhs.getNumTargets()); - for (const auto& [i, trgt] : llvm::enumerate(lhs.getTargets())) { - const auto it = llvm::find(rhs.getTargets(), m.lookup(trgt)); - const auto j = std::distance(rhs.getTargets().begin(), it); - permutation[i] = j; - } - return permutation; -} - -/// Return a permutation vector, where permutation[i] maps the i-th input -/// target of the lhs to the j-th input target of the rhs. -static SmallVector -getTargetPermutation(qco::CtrlOp lhs, qco::CtrlOp rhs, const IRMapping& m) { - SmallVector permutation(lhs.getNumTargets()); - for (const auto& [i, trgt] : llvm::enumerate(lhs.getInputTargets())) { - const auto it = llvm::find(rhs.getInputTargets(), m.lookup(trgt)); - const auto j = std::distance(rhs.getInputTargets().begin(), it); - permutation[i] = j; - } - return permutation; -} - -/// Return a permutation vector, where permutation[i] maps the i-th input -/// target of the lhs to the j-th input target of the rhs. -static SmallVector -getControlPermutation(qco::CtrlOp lhs, qco::CtrlOp rhs, const IRMapping& m) { - SmallVector permutation(lhs.getNumControls()); - for (const auto& [i, trgt] : llvm::enumerate(lhs.getInputControls())) { - const auto it = llvm::find(rhs.getInputControls(), m.lookup(trgt)); - const auto j = std::distance(rhs.getInputControls().begin(), it); +/// Return a permutation vector, where permutation[i] maps the i-th value of the +/// lhs range to the j-th value of the rhs range. +template +static FailureOr> +getPermutation(const LhsRange& lhs, const RhsRange& rhs, const IRMapping& m, + const TensorMapping& tm) { + SmallVector permutation(lhs.size()); + for (const auto& [i, lhsValue] : llvm::enumerate(lhs)) { + const auto it = hasTypeQubitTensor(lhsValue) + ? llvm::find_if(rhs, + [&](const auto rhsValue) { + if (!hasTypeQubitTensor(rhsValue)) { + return false; + } + return tm.equals(lhsValue, rhsValue); + }) + : llvm::find(rhs, m.lookup(lhsValue)); + if (it == rhs.end()) { + return failure(); + } + const auto j = std::distance(rhs.begin(), it); permutation[i] = j; } return permutation; } -/// Compare two ctrl operations, allowing permutations of control and target -/// qubits. -static bool compareCtrlOps(qc::CtrlOp lhs, qc::CtrlOp rhs, const IRMapping& m) { - DenseSet workset; - workset.insert_range(rhs.getControls()); - for (const auto& ctrl : lhs.getControls()) { - const auto& v = m.lookup(ctrl); - if (!workset.contains(v)) { - return false; - } - workset.erase(v); - } - - if (!workset.empty()) { - return false; - } - - workset.insert_range(rhs.getTargets()); - for (const auto& trgt : lhs.getTargets()) { - const auto& v = m.lookup(trgt); - if (!workset.contains(v)) { - return false; - } - workset.erase(v); - } - - return workset.empty(); -} - -/// Compare two ctrl operations, allowing permutations of input control and -/// input target qubits. -static bool compareCtrlOps(qco::CtrlOp lhs, qco::CtrlOp rhs, - const IRMapping& m) { +/// Compare two value lists, allowing permutations. +template +static bool compareValueLists(const LhsRange& lhs, const RhsRange& rhs, + const IRMapping& m, const TensorMapping& tm) { DenseSet workset; - workset.insert_range(rhs.getInputControls()); - for (const auto& ctrl : lhs.getInputControls()) { - const auto& v = m.lookup(ctrl); - if (!workset.contains(v)) { - return false; + workset.insert_range(rhs); + + for (const auto lhsValue : lhs) { + Value mapped; + if (hasTypeQubitTensor(lhsValue)) { + const auto it = llvm::find_if(rhs, [&](const auto rhsValue) { + return hasTypeQubitTensor(rhsValue) && tm.equals(lhsValue, rhsValue); + }); + if (it == rhs.end()) { + return false; + } + mapped = *it; + } else { + mapped = m.lookup(lhsValue); } - workset.erase(v); - } - - if (!workset.empty()) { - return false; - } - - workset.insert_range(rhs.getInputTargets()); - for (const auto& trgt : lhs.getInputTargets()) { - const auto& v = m.lookup(trgt); - if (!workset.contains(v)) { + if (!workset.contains(mapped)) { return false; } - workset.erase(v); + workset.erase(mapped); } return workset.empty(); @@ -413,12 +382,42 @@ static bool compareOperations(Operation* lhs, Operation* rhs, if (isa(lhs)) { assert(isa(rhs)); - if (!compareCtrlOps(cast(lhs), cast(rhs), m)) { + auto lhsCtrl = cast(lhs); + auto rhsCtrl = cast(rhs); + if (!compareValueLists(lhsCtrl.getControls(), rhsCtrl.getControls(), m, + tm) || + !compareValueLists(lhsCtrl.getTargets(), rhsCtrl.getTargets(), m, tm)) { return false; } } else if (isa(lhs)) { assert(isa(rhs)); - if (!compareCtrlOps(cast(lhs), cast(rhs), m)) { + auto lhsCtrl = cast(lhs); + auto rhsCtrl = cast(rhs); + if (!compareValueLists(lhsCtrl.getInputControls(), + rhsCtrl.getInputControls(), m, tm) || + !compareValueLists(lhsCtrl.getInputTargets(), rhsCtrl.getInputTargets(), + m, tm)) { + return false; + } + } else if (isa(lhs)) { + assert(isa(rhs)); + if (!compareValueLists(cast(lhs).getQubits(), + cast(rhs).getQubits(), m, tm)) { + return false; + } + } else if (isa(lhs)) { + assert(isa(rhs)); + auto lhsSwitch = cast(lhs); + auto rhsSwitch = cast(rhs); + if (lhsSwitch.getCases() != rhsSwitch.getCases() || + !compareValueLists(lhsSwitch.getTargets(), rhsSwitch.getTargets(), m, + tm)) { + return false; + } + } else if (isa(lhs)) { + assert(isa(rhs)); + if (!compareValueLists(cast(lhs).getTargets(), + cast(rhs).getTargets(), m, tm)) { return false; } } else { @@ -550,12 +549,42 @@ static bool compareBlocks(Block& lhs, Block& rhs, assert(isa(rhs.getParentOp())); auto lhsCtrl = cast(lhs.getParentOp()); auto rhsCtrl = cast(rhs.getParentOp()); - mapArguments(lhs, rhs, getTargetPermutation(lhsCtrl, rhsCtrl, m), m); + const auto permutation = + getPermutation(lhsCtrl.getTargets(), rhsCtrl.getTargets(), m, tm); + if (failed(permutation)) { + return false; + } + mapArguments(lhs, rhs, *permutation, m); } else if (isa(lhs.getParentOp())) { assert(isa(rhs.getParentOp())); auto lhsCtrl = cast(lhs.getParentOp()); auto rhsCtrl = cast(rhs.getParentOp()); - mapArguments(lhs, rhs, getTargetPermutation(lhsCtrl, rhsCtrl, m), m); + const auto permutation = getPermutation(lhsCtrl.getInputTargets(), + rhsCtrl.getInputTargets(), m, tm); + if (failed(permutation)) { + return false; + } + mapArguments(lhs, rhs, *permutation, m); + } else if (isa(lhs.getParentOp())) { + assert(isa(rhs.getParentOp())); + auto lhsIf = cast(lhs.getParentOp()); + auto rhsIf = cast(rhs.getParentOp()); + const auto permutation = + getPermutation(lhsIf.getQubits(), rhsIf.getQubits(), m, tm); + if (failed(permutation)) { + return false; + } + mapArguments(lhs, rhs, *permutation, m); + } else if (isa(lhs.getParentOp())) { + assert(isa(rhs.getParentOp())); + auto lhsSwitch = cast(lhs.getParentOp()); + auto rhsSwitch = cast(rhs.getParentOp()); + const auto permutation = + getPermutation(lhsSwitch.getTargets(), rhsSwitch.getTargets(), m, tm); + if (failed(permutation)) { + return false; + } + mapArguments(lhs, rhs, *permutation, m); } else { SmallVector permutation(lhs.getNumArguments()); std::iota(permutation.begin(), permutation.end(), 0); @@ -606,13 +635,40 @@ static bool compareBlocks(Block& lhs, Block& rhs, auto lhsCtrl = cast(lhsOp); auto rhsCtrl = cast(rhsOp); - SmallVector permutation; + const auto controlPermutation = getPermutation( + lhsCtrl.getInputControls(), rhsCtrl.getInputControls(), m, tm); + const auto targetPermutation = getPermutation( + lhsCtrl.getInputTargets(), rhsCtrl.getInputTargets(), m, tm); + if (failed(controlPermutation) || failed(targetPermutation)) { + return false; + } + + SmallVector permutation(*controlPermutation); permutation.reserve(lhsCtrl.getNumQubits()); - permutation.append(getControlPermutation(lhsCtrl, rhsCtrl, m)); - for (const auto i : getTargetPermutation(lhsCtrl, rhsCtrl, m)) { + for (const auto i : *targetPermutation) { permutation.emplace_back(lhsCtrl.getNumControls() + i); } mapResults(lhsCtrl, rhsCtrl, permutation, m); + } else if (isa(lhsOp)) { + assert(isa(rhsOp)); + auto lhsIf = cast(lhsOp); + auto rhsIf = cast(rhsOp); + const auto permutation = + getPermutation(lhsIf.getQubits(), rhsIf.getQubits(), m, tm); + if (failed(permutation)) { + return false; + } + mapResults(lhsIf, rhsIf, *permutation, m); + } else if (isa(lhsOp)) { + assert(isa(rhsOp)); + auto lhsSwitch = cast(lhsOp); + auto rhsSwitch = cast(rhsOp); + const auto permutation = getPermutation( + lhsSwitch.getTargets(), rhsSwitch.getTargets(), m, tm); + if (failed(permutation)) { + return false; + } + mapResults(lhsSwitch, rhsSwitch, *permutation, m); } else if (isa(lhsOp)) { assert(isa(rhsOp)); auto lhsAlloc = cast(lhsOp); diff --git a/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp b/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp index b655df2f0b..c3a49d5ee5 100644 --- a/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp +++ b/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp @@ -696,6 +696,19 @@ INSTANTIATE_TEST_SUITE_P( MQT_NAMED_BUILDER(qc::nestedIfOpForLoop)})); /// @} +/// \name QCOToQC/Operations/IndexSwitchOp.cpp +/// @{ +INSTANTIATE_TEST_SUITE_P( + QCOIndexSwitchOpTest, QCOToQCTest, + testing::Values(QCOToQCTestCase{"SimpleIndexSwitchOp", + MQT_NAMED_BUILDER(qco::simpleIndexSwitch), + MQT_NAMED_BUILDER(qc::simpleIndexSwitch)}, + QCOToQCTestCase{ + "IndexSwitchMultiCase", + MQT_NAMED_BUILDER(qco::indexSwitchMultiCase), + MQT_NAMED_BUILDER(qc::indexSwitchMultiCase)})); +/// @} + /// \name QCOToQC/Operations/WhileOp.cpp /// @{ INSTANTIATE_TEST_SUITE_P( @@ -721,6 +734,9 @@ INSTANTIATE_TEST_SUITE_P( QCOToQCTestCase{"NestedForLoopWhileOp", MQT_NAMED_BUILDER(qco::nestedForLoopWhileOp), MQT_NAMED_BUILDER(qc::nestedForLoopWhileOp)}, + QCOToQCTestCase{"NestedForLoopSwitchOp", + MQT_NAMED_BUILDER(qco::nestedForLoopSwitchOp), + MQT_NAMED_BUILDER(qc::nestedForLoopSwitchOp)}, QCOToQCTestCase{ "nestedForLoopCtrlOpWithSeparateQubit", MQT_NAMED_BUILDER(qco::nestedForLoopCtrlOpWithSeparateQubit), diff --git a/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp b/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp index 082312f45a..67a32dab55 100644 --- a/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp +++ b/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp @@ -698,6 +698,19 @@ INSTANTIATE_TEST_SUITE_P( MQT_NAMED_BUILDER(qco::nestedIfOpForLoop)})); /// @} +/// \name QCToQCO/Operations/IndexSwitchOp.cpp +/// @{ +INSTANTIATE_TEST_SUITE_P( + QCOIndexSwitchOpTest, QCToQCOTest, + testing::Values(QCToQCOTestCase{"SimpleIndexSwitchOp", + MQT_NAMED_BUILDER(qc::simpleIndexSwitch), + MQT_NAMED_BUILDER(qco::simpleIndexSwitch)}, + QCToQCOTestCase{ + "IndexSwitchMultiCase", + MQT_NAMED_BUILDER(qc::indexSwitchMultiCase), + MQT_NAMED_BUILDER(qco::indexSwitchMultiCase)})); +/// @} + /// \name QCToQCO/Operations/WhileOp.cpp /// @{ INSTANTIATE_TEST_SUITE_P( @@ -723,6 +736,9 @@ INSTANTIATE_TEST_SUITE_P( QCToQCOTestCase{"NestedForLoopWhileOp", MQT_NAMED_BUILDER(qc::nestedForLoopWhileOp), MQT_NAMED_BUILDER(qco::nestedForLoopWhileOp)}, + QCToQCOTestCase{"NestedForLoopSwitchOp", + MQT_NAMED_BUILDER(qc::nestedForLoopSwitchOp), + MQT_NAMED_BUILDER(qco::nestedForLoopSwitchOp)}, QCToQCOTestCase{ "nestedForLoopCtrlOpWithSeparateQubit", MQT_NAMED_BUILDER(qc::nestedForLoopCtrlOpWithSeparateQubit), diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 1080dff78a..01a7a37f41 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -20,16 +20,26 @@ #include "qco_programs.h" #include +#include +#include +#include #include #include +#include +#include +#include +#include #include #include +#include #include #include +#include #include #include #include +#include #include #include #include @@ -64,7 +74,7 @@ class QCOTest : public testing::TestWithParam { // Register all necessary dialects DialectRegistry registry; registry.insert(); + scf::SCFDialect, qtensor::QTensorDialect>(); context = std::make_unique(); context->appendDialectRegistry(registry); context->loadAllAvailableDialects(); @@ -176,6 +186,115 @@ TEST_F(QCOTest, DirectIfBuilder) { refBuilder.get())); } +TEST_F(QCOTest, DirectSingleTargetIndexSwitchBuilder) { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + auto index = arith::ConstantIndexOp::create(builder, 0); + const auto target = builder.allocQubit(); + const auto identity = [](Value value) { return value; }; + const SmallVector> caseBuilders{identity}; + + auto switchOp = + IndexSwitchOp::create(builder, index.getResult(), target, + ArrayRef{0}, caseBuilders, identity); + + ASSERT_EQ(switchOp.getTargets().size(), 1); + ASSERT_EQ(switchOp.getResults().size(), 1); + ASSERT_EQ(switchOp.getNumCases(), 1); + EXPECT_EQ(switchOp.getCaseBlock(0)->getArgument(0).getType(), + target.getType()); + EXPECT_EQ(switchOp.getDefaultBlock()->getArgument(0).getType(), + target.getType()); + EXPECT_TRUE(switchOp.verify().succeeded()); +} + +TEST_F(QCOTest, IndexSwitchBuildersRejectMismatchedCasesAndBodies) { + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + auto index = arith::ConstantIndexOp::create(builder, 0); + const auto target = builder.allocQubit(); + const auto identity = [](Value value) { return value; }; + const SmallVector> noCaseBuilders; + IndexSwitchOp::create(builder, index.getResult(), target, + ArrayRef{0}, noCaseBuilders, identity); + }, + "Each case must have a corresponding case body function"); + + EXPECT_DEATH( + { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + const auto target = builder.allocQubit(); + const auto identity = [](Value value) { return value; }; + const SmallVector> noCaseBodies; + builder.qcoIndexSwitch(0, target, ArrayRef{0}, noCaseBodies, + identity); + }, + "Each case must have a corresponding case body function"); +} + +TEST_F(QCOTest, IndexSwitchTiedValuesAndTargetExtension) { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + auto index = arith::ConstantIndexOp::create(builder, 0); + const auto target = builder.allocQubit(); + const auto addon = builder.allocQubit(); + const auto identity = [](Value value) { return value; }; + const SmallVector> caseBuilders{identity}; + auto switchOp = + IndexSwitchOp::create(builder, index.getResult(), target, + ArrayRef{0}, caseBuilders, identity); + + auto* targetOperand = &switchOp->getOpOperand(1); + const auto caseArgument = switchOp.getCaseBlock(0)->getArgument(0); + const auto defaultArgument = switchOp.getDefaultBlock()->getArgument(0); + auto* caseYieldOperand = &switchOp.getCaseYield(0)->getOpOperand(0); + auto* defaultYieldOperand = &switchOp.getDefaultYield()->getOpOperand(0); + + EXPECT_EQ(switchOp.getTiedResult(targetOperand), switchOp.getResult(0)); + EXPECT_EQ(switchOp.getTiedTarget(cast(switchOp.getResult(0))), + targetOperand); + EXPECT_EQ(switchOp.getTiedCaseBlockArgument(targetOperand, 0), caseArgument); + EXPECT_EQ(switchOp.getTiedCaseYieldedValue(caseArgument, 0), + caseYieldOperand); + EXPECT_EQ(switchOp.getTiedDefaultBlockArgument(targetOperand), + defaultArgument); + EXPECT_EQ(switchOp.getTiedDefaultYieldedValue(defaultArgument), + defaultYieldOperand); + + EXPECT_FALSE(switchOp.getTiedResult(caseYieldOperand)); + EXPECT_EQ(switchOp.getTiedTarget(cast(addon)), nullptr); + EXPECT_FALSE(switchOp.getTiedCaseBlockArgument(caseYieldOperand, 0)); + EXPECT_FALSE(switchOp.getTiedCaseBlockArgument(targetOperand, 1)); + EXPECT_EQ(switchOp.getTiedCaseYieldedValue(defaultArgument, 0), nullptr); + EXPECT_EQ(switchOp.getTiedCaseYieldedValue(caseArgument, 1), nullptr); + EXPECT_FALSE(switchOp.getTiedDefaultBlockArgument(defaultYieldOperand)); + EXPECT_EQ(switchOp.getTiedDefaultYieldedValue(caseArgument), nullptr); + + IRRewriter rewriter(context.get()); + rewriter.setInsertionPoint(switchOp); + EXPECT_EQ(switchOp.replaceWithAdditionalTargets(rewriter, ValueRange{}), + switchOp); + + auto expanded = switchOp.replaceWithAdditionalTargets(rewriter, addon); + ASSERT_EQ(expanded.getTargets().size(), 2); + ASSERT_EQ(expanded.getResults().size(), 2); + EXPECT_EQ(expanded.getTargets()[0], target); + EXPECT_EQ(expanded.getTargets()[1], addon); + for (Region* region : expanded.getRegions()) { + ASSERT_EQ(region->getNumArguments(), 2); + auto yield = cast(region->front().getTerminator()); + ASSERT_EQ(yield.getTargets().size(), 2); + EXPECT_EQ(yield.getTargets()[0], region->getArgument(0)); + EXPECT_EQ(yield.getTargets()[1], region->getArgument(1)); + } + EXPECT_TRUE(expanded.verify().succeeded()); +} + TEST_F(QCOTest, IfOpParser) { // Test IfOp parser const char* mlirCode = R"( @@ -221,6 +340,241 @@ TEST_F(QCOTest, IfOpParser) { refBuilder.get())); } +TEST_F(QCOTest, IndexSwitchParser) { + // Test IndexSwitch parser + const char* mlirCode = R"( + module { + func.func @main() -> (i1, i1, i1) attributes {passthrough = ["entry_point"]} { + %c2 = arith.constant 2 : index + %c1 = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c3 = arith.constant 3 : index + %0 = qtensor.alloc(%c3) : tensor<3x!qco.qubit> + %1 = scf.for %arg0 = %c0 to %c3 step %c1 iter_args(%arg1 = %0) -> (tensor<3x!qco.qubit>) { + %5 = arith.remui %arg0, %c3 : index + %out_tensor_9, %result_10 = qtensor.extract %arg1[%arg0] : tensor<3x!qco.qubit> + %6 = qco.index_switch %5 -> !qco.qubit + case 0 args(%arg2 = %result_10) { + %8 = qco.x %arg2 : !qco.qubit -> !qco.qubit + qco.yield %8 : !qco.qubit + } + case 1 args(%arg2 = %result_10) { + %8 = qco.y %arg2 : !qco.qubit -> !qco.qubit + qco.yield %8 : !qco.qubit + } + case 2 args(%arg2 = %result_10) { + %8 = qco.x %arg2 : !qco.qubit -> !qco.qubit + %9 = qco.y %8 : !qco.qubit -> !qco.qubit + qco.yield %9 : !qco.qubit + } + default args(%arg2 = %result_10) { + qco.yield %arg2 : !qco.qubit + } + %7 = qtensor.insert %6 into %out_tensor_9[%arg0] : tensor<3x!qco.qubit> + scf.yield %7 : tensor<3x!qco.qubit> + } + %out_tensor, %result = qtensor.extract %1[%c0] : tensor<3x!qco.qubit> + %qubit_out, %result_0 = qco.measure %result : !qco.qubit + %out_tensor_1, %result_2 = qtensor.extract %out_tensor[%c1] : tensor<3x!qco.qubit> + %qubit_out_3, %result_4 = qco.measure %result_2 : !qco.qubit + %out_tensor_5, %result_6 = qtensor.extract %out_tensor_1[%c2] : tensor<3x!qco.qubit> + %2 = qtensor.insert %qubit_out into %out_tensor_5[%c0] : tensor<3x!qco.qubit> + %3 = qtensor.insert %qubit_out_3 into %2[%c1] : tensor<3x!qco.qubit> + %qubit_out_7, %result_8 = qco.measure %result_6 : !qco.qubit + %4 = qtensor.insert %qubit_out_7 into %3[%c2] : tensor<3x!qco.qubit> + qtensor.dealloc %4 : tensor<3x!qco.qubit> + return %result_0, %result_4, %result_8 : i1, i1, i1 + } + })"; + + auto parsedSourceModule = + parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(parsedSourceModule); + EXPECT_TRUE(verify(*parsedSourceModule).succeeded()); + EXPECT_TRUE(runQCOCleanupPipeline(parsedSourceModule.get()).succeeded()); + EXPECT_TRUE(verify(*parsedSourceModule).succeeded()); + + auto refBuilder = mqt::test::buildMLIRProgram( + context.get(), MQT_NAMED_BUILDER(nestedForLoopSwitchOp)); + ASSERT_TRUE(refBuilder); + EXPECT_TRUE(verify(*refBuilder).succeeded()); + EXPECT_TRUE(runQCOCleanupPipeline(refBuilder.get()).succeeded()); + EXPECT_TRUE(verify(*refBuilder).succeeded()); + + EXPECT_TRUE(areModulesEquivalentWithPermutations(parsedSourceModule.get(), + refBuilder.get())); +} + +TEST_F(QCOTest, DefaultOnlyIndexSwitchParser) { + const char* mlirCode = R"( + module { + func.func @main(%index: index) -> i1 { + %q0 = qco.alloc : !qco.qubit + %q1 = qco.index_switch %index -> !qco.qubit + default args(%arg0 = %q0) { + qco.yield %arg0 : !qco.qubit + } + %q2, %result = qco.measure %q1 : !qco.qubit + qco.sink %q2 : !qco.qubit + return %result : i1 + } + })"; + + auto parsedModule = parseSourceString(mlirCode, context.get()); + ASSERT_TRUE(parsedModule); + EXPECT_TRUE(verify(*parsedModule).succeeded()); + + std::string printed; + llvm::raw_string_ostream stream(printed); + parsedModule->print(stream); + stream.flush(); + + auto reparsedModule = parseSourceString(printed, context.get()); + ASSERT_TRUE(reparsedModule); + EXPECT_TRUE(verify(*reparsedModule).succeeded()); +} + +TEST_F(QCOTest, EquivalentTensorIndexSwitches) { + const auto build = [&]() { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; + const SmallVector(ValueRange)>> caseBodies{ + identity}; + + const auto tensor = builder.qtensorAlloc(1); + const auto result = builder.qcoIndexSwitch( + 0, tensor, SmallVector{0}, caseBodies, identity); + builder.qtensorDealloc(result.front()); + return builder.finalize(); + }; + + const auto lhs = build(); + const auto rhs = build(); + ASSERT_TRUE(lhs); + ASSERT_TRUE(rhs); + EXPECT_TRUE(verify(*lhs).succeeded()); + EXPECT_TRUE(verify(*rhs).succeeded()); + EXPECT_TRUE(areModulesEquivalentWithPermutations(lhs.get(), rhs.get())); +} + +TEST_F(QCOTest, IndexSwitchCaseValuesAffectEquivalence) { + const auto build = [&](const int64_t caseValue) { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; + const SmallVector(ValueRange)>> caseBodies{ + identity}; + + const auto q0 = builder.allocQubit(); + const auto result = builder.qcoIndexSwitch( + 0, q0, SmallVector{caseValue}, caseBodies, identity); + builder.sink(result.front()); + return builder.finalize(); + }; + + const auto lhs = build(0); + const auto rhs = build(1); + ASSERT_TRUE(lhs); + ASSERT_TRUE(rhs); + EXPECT_TRUE(verify(*lhs).succeeded()); + EXPECT_TRUE(verify(*rhs).succeeded()); + EXPECT_FALSE(areModulesEquivalentWithPermutations(lhs.get(), rhs.get())); +} + +TEST_F(QCOTest, NonEquivalentTensorIndexSwitches) { + const auto build = [&](const bool switchFirstTensor) { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; + const SmallVector(ValueRange)>> caseBodies{ + identity}; + + const auto first = builder.qtensorAlloc(1); + const auto second = builder.qtensorAlloc(1); + const auto target = switchFirstTensor ? first : second; + const auto untouched = switchFirstTensor ? second : first; + const auto result = builder.qcoIndexSwitch( + 0, target, SmallVector{0}, caseBodies, identity); + builder.qtensorDealloc(result.front()); + builder.qtensorDealloc(untouched); + return builder.finalize(); + }; + + const auto lhs = build(true); + const auto rhs = build(false); + ASSERT_TRUE(lhs); + ASSERT_TRUE(rhs); + EXPECT_TRUE(verify(*lhs).succeeded()); + EXPECT_TRUE(verify(*rhs).succeeded()); + EXPECT_FALSE(areModulesEquivalentWithPermutations(lhs.get(), rhs.get())); +} + +TEST_F(QCOTest, IndexSwitchConstantSuccessor) { + QCOProgramBuilder builder(context.get()); + builder.initialize(); + + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; + const SmallVector(ValueRange)>> caseBodies{ + identity, identity}; + + const auto q0 = builder.allocQubit(); + auto result = builder.qcoIndexSwitch(1, q0, SmallVector{0, 1}, + caseBodies, identity); + builder.sink(result.front()); + [[maybe_unused]] auto module = builder.finalize(); + + auto switchOp = result.front().getDefiningOp(); + ASSERT_TRUE(switchOp); + + SmallVector unknownOperands(switchOp->getNumOperands()); + SmallVector unknownBounds; + switchOp.getRegionInvocationBounds(unknownOperands, unknownBounds); + ASSERT_EQ(unknownBounds.size(), switchOp->getNumRegions()); + for (const auto& bound : unknownBounds) { + EXPECT_EQ(bound.getLowerBound(), 0); + ASSERT_TRUE(bound.getUpperBound().has_value()); + EXPECT_EQ(*bound.getUpperBound(), 1); + } + + const auto checkConstant = [&](const int64_t value, Region* expectedRegion, + const size_t expectedRegionIndex) { + SmallVector operands(switchOp->getNumOperands()); + operands.front() = builder.getIndexAttr(value); + + SmallVector successors; + switchOp.getEntrySuccessorRegions(operands, successors); + ASSERT_EQ(successors.size(), 1); + EXPECT_EQ(successors.front().getSuccessor(), expectedRegion); + + SmallVector bounds; + switchOp.getRegionInvocationBounds(operands, bounds); + ASSERT_EQ(bounds.size(), switchOp->getNumRegions()); + for (const auto [index, bound] : llvm::enumerate(bounds)) { + EXPECT_EQ(bound.getLowerBound(), 0); + ASSERT_TRUE(bound.getUpperBound().has_value()); + EXPECT_EQ(*bound.getUpperBound(), index == expectedRegionIndex ? 1 : 0); + } + }; + + checkConstant(0, switchOp.getCaseRegions().data(), 1); + checkConstant(1, &switchOp.getCaseRegions()[1], 2); + checkConstant(2, &switchOp.getDefaultRegion(), 0); + + switchOp->setAttr( + "cases", DenseI64ArrayAttr::get(context.get(), ArrayRef{0})); + ScopedDiagnosticHandler handler(context.get(), + [](Diagnostic&) { return success(); }); + EXPECT_TRUE(switchOp.verify().failed()); + + switchOp->setAttr( + "cases", DenseI64ArrayAttr::get(context.get(), ArrayRef{0, 0})); + EXPECT_TRUE(switchOp.verify().failed()); +} + /// \name QCO/SCF/IfOp.cpp /// @{ INSTANTIATE_TEST_SUITE_P( diff --git a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp index 7bbc2437dd..c44f381695 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp @@ -27,6 +27,7 @@ #include #include +#include #include #include @@ -109,7 +110,13 @@ TEST_F(DriversTest, ProgramGraphWalk) { return SmallVector{builder.id(args[0])}; })[0]; - builder.measure(q05); + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; + const SmallVector(ValueRange)>> caseBodies{ + identity}; + const auto q06 = builder.qcoIndexSwitch(0, q05, SmallVector{0}, + caseBodies, identity)[0]; + + builder.measure(q06); builder.measure(forResults[1]); builder.measure(forResults[2]); builder.measure(forResults[3]); diff --git a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp index 295ca3c016..c466671c58 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp @@ -13,6 +13,7 @@ #include "mlir/Dialect/QCO/Utils/WireIterator.h" #include +#include #include #include #include @@ -20,6 +21,7 @@ #include #include +#include #include #include #include @@ -80,8 +82,15 @@ TEST_P(WireIteratorTest, Traversal) { [&](ValueRange args) { return SmallVector{args[0], args[1]}; }); const auto q06 = ifOut[0]; const auto q13 = ifOut[1]; - builder.sink(q06); - builder.sink(q13); + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; + const SmallVector(ValueRange)>> caseBodies{ + identity}; + const auto switchOut = builder.qcoIndexSwitch( + 0, {q06, q13}, SmallVector{0}, caseBodies, identity); + const auto q07 = switchOut[0]; + const auto q14 = switchOut[1]; + builder.sink(q07); + builder.sink(q14); [[maybe_unused]] auto module = builder.finalize(); // Setup WireIterator. @@ -119,7 +128,11 @@ TEST_P(WireIteratorTest, Traversal) { ASSERT_EQ(it.qubit(), q06); ++it; - ASSERT_EQ(it.operation(), *(q06.getUsers().begin())); // qco.sink + ASSERT_EQ(it.operation(), q07.getDefiningOp()); // qco.index_switch + ASSERT_EQ(it.qubit(), q07); + + ++it; + ASSERT_EQ(it.operation(), *(q07.getUsers().begin())); // qco.sink ASSERT_EQ(it.qubit(), nullptr); ++it; @@ -133,9 +146,13 @@ TEST_P(WireIteratorTest, Traversal) { // --it; - ASSERT_EQ(it.operation(), *(q06.getUsers().begin())); // qco.sink + ASSERT_EQ(it.operation(), *(q07.getUsers().begin())); // qco.sink ASSERT_EQ(it.qubit(), nullptr); + --it; + ASSERT_EQ(it.operation(), q07.getDefiningOp()); // qco.index_switch + ASSERT_EQ(it.qubit(), q07); + --it; ASSERT_EQ(it.operation(), q06.getDefiningOp()); // qco.if ASSERT_EQ(it.qubit(), q06); diff --git a/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp b/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp index 358f095cd2..0d673a6efd 100644 --- a/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp +++ b/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp @@ -14,6 +14,7 @@ #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" #include +#include #include #include #include @@ -101,7 +102,12 @@ TEST_F(TensorIteratorTest, Traversal) { tensorElse2 = builder.qtensorInsert(q, tensorElse1, 0); return SmallVector{tensorElse2}; })[0]; - builder.qtensorDealloc(tensor8); + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; + const SmallVector(ValueRange)>> caseBodies{ + identity}; + const auto tensor9 = builder.qcoIndexSwitch( + 0, tensor8, SmallVector{0}, caseBodies, identity)[0]; + builder.qtensorDealloc(tensor9); [[maybe_unused]] auto m = builder.finalize(); TensorIterator it(cast>(tensor0)); @@ -142,7 +148,11 @@ TEST_F(TensorIteratorTest, Traversal) { ASSERT_EQ(it.tensor(), tensor8); ++it; - ASSERT_EQ(it.operation(), *(tensor8.user_begin())); // qtensor.dealloc + ASSERT_EQ(it.operation(), tensor9.getDefiningOp()); // qco.index_switch + ASSERT_EQ(it.tensor(), tensor9); + + ++it; + ASSERT_EQ(it.operation(), *(tensor9.user_begin())); // qtensor.dealloc ASSERT_EQ(it.tensor(), nullptr); ++it; @@ -152,9 +162,13 @@ TEST_F(TensorIteratorTest, Traversal) { ASSERT_EQ(it, std::default_sentinel); --it; - ASSERT_EQ(it.operation(), *(tensor8.user_begin())); // qtensor.dealloc + ASSERT_EQ(it.operation(), *(tensor9.user_begin())); // qtensor.dealloc ASSERT_EQ(it.tensor(), nullptr); + --it; + ASSERT_EQ(it.operation(), tensor9.getDefiningOp()); // qco.index_switch + ASSERT_EQ(it.tensor(), tensor9); + --it; ASSERT_EQ(it.operation(), tensor8.getDefiningOp()); // qco.if ASSERT_EQ(it.tensor(), tensor8); diff --git a/mlir/unittests/programs/qc_programs.cpp b/mlir/unittests/programs/qc_programs.cpp index 1646d8ab62..b1a7f3ea90 100644 --- a/mlir/unittests/programs/qc_programs.cpp +++ b/mlir/unittests/programs/qc_programs.cpp @@ -14,9 +14,11 @@ #include #include +#include #include #include +#include #include namespace mlir::qc { @@ -2498,6 +2500,48 @@ Value nestedIfOpForLoop(QCProgramBuilder& b) { return b.measure(q0); } +SmallVector simpleIndexSwitch(QCProgramBuilder& b) { + auto reg = b.allocQubitRegister(1); + b.h(reg[0]); + auto bit0 = b.measure(reg[0]); + auto i0 = arith::IndexCastUIOp::create(b, b.getIndexType(), bit0).getOut(); + b.scfIndexSwitch(i0, SmallVector{0}, + SmallVector>{[&] { b.x(reg[0]); }}, + [&] { b.z(reg[0]); }); + auto bit1 = b.measure(reg[0]); + return {bit0, bit1}; +} + +SmallVector indexSwitchMultiCase(QCProgramBuilder& b) { + constexpr int64_t size = 2; + + auto reg = b.allocQubitRegister(size); + auto c1 = arith::ConstantOp::create(b, b.getIndexType(), b.getIndexAttr(1)) + .getResult(); + auto condition = + arith::ConstantOp::create(b, b.getIndexType(), b.getIndexAttr(0)) + .getResult(); + for (int64_t i = 0; i < size; ++i) { + b.h(reg[i]); + const auto bit = b.measure(reg[i]); + const auto index = + arith::IndexCastUIOp::create(b, b.getIndexType(), bit).getOut(); + condition = arith::OrIOp::create(b, {condition, index}).getResult(); + condition = arith::ShLIOp::create(b, {condition, c1}); + } + + b.scfIndexSwitch(condition, SmallVector{1, 2, 3}, + SmallVector>{[&] { b.x(reg[1]); }, + [&] { b.x(reg[0]); }, + [&] { + b.x(reg[0]); + b.x(reg[1]); + }}, + [&] { /* no-op */ }); + + return measureAndReturn(b, reg.qubits); +} + Value simpleWhileReset(QCProgramBuilder& b) { auto q = b.allocQubit(); b.h(q); @@ -2563,6 +2607,25 @@ SmallVector nestedForLoopWhileOp(QCProgramBuilder& b) { return measureAndReturn(b, reg.qubits); } +SmallVector nestedForLoopSwitchOp(QCProgramBuilder& b) { + constexpr int64_t n = 3; + auto reg = b.allocQubitRegister(n); + auto c3 = arith::ConstantOp::create(b, b.getIndexAttr(3)); + b.scfFor(0, n, 1, [&](Value iv) { + auto rem = arith::RemUIOp::create(b, {iv, c3}).getResult(); + auto q = b.memrefLoad(reg.value, iv); + b.scfIndexSwitch(rem, SmallVector{0, 1, 2}, + SmallVector>{[&] { b.x(q); }, + [&] { b.y(q); }, + [&] { + b.x(q); + b.y(q); + }}, + [&] { /* error */ }); + }); + return measureAndReturn(b, reg.qubits); +} + Value nestedForLoopCtrlOpWithSeparateQubit(QCProgramBuilder& b) { auto reg = b.allocQubitRegister(3); auto control = b.allocQubit(); diff --git a/mlir/unittests/programs/qc_programs.h b/mlir/unittests/programs/qc_programs.h index aed137e8cd..f58a9b6c19 100644 --- a/mlir/unittests/programs/qc_programs.h +++ b/mlir/unittests/programs/qc_programs.h @@ -1225,6 +1225,14 @@ SmallVector ifTwoQubits(QCProgramBuilder& b); /// a register. Value nestedIfOpForLoop(QCProgramBuilder& b); +// --- IndexSwitchOp -------------------------------------------------------- // + +/// Creates a circuit with an index switch operation with one qubit. +SmallVector simpleIndexSwitch(QCProgramBuilder& b); + +/// Creates a circuit with an index switch operation with multiple cases. +SmallVector indexSwitchMultiCase(QCProgramBuilder& b); + // --- WhileOp -------------------------------------------------------------- // /// Creates a circuit with a while operation using a while loop. @@ -1246,6 +1254,10 @@ Value nestedForLoopIfOp(QCProgramBuilder& b); /// operation. SmallVector nestedForLoopWhileOp(QCProgramBuilder& b); +/// Creates a circuit with a for operation with a register and a nested index +/// switch operation. +SmallVector nestedForLoopSwitchOp(QCProgramBuilder& b); + /// Creates a circuit with a for operation with a register and a qubit and a /// nested ctrl operation where the qubit is separately allocated from the /// register. diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 38a54c87b3..6189628162 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -4049,6 +4050,71 @@ SmallVector nestedFalseIf(QCOProgramBuilder& b) { return {measureResult, c}; } +SmallVector simpleIndexSwitch(QCOProgramBuilder& b) { + Value q; + Value bit0; + Value c0; + Value bit1; + + auto reg = b.allocQubitRegister(1); + + q = b.h(reg[0]); + std::tie(q, bit0) = b.measure(q); + c0 = arith::IndexCastUIOp::create(b, b.getIndexType(), bit0).getOut(); + q = b.qcoIndexSwitch(c0, q, SmallVector{0}, + SmallVector>{ + [&](Value arg) { return b.x(arg); }}, + [&](Value arg) { return b.z(arg); }); + + std::tie(q, bit1) = b.measure(q); + + return {bit0, bit1}; +} + +SmallVector indexSwitchMultiCase(QCOProgramBuilder& b) { + constexpr int64_t size = 2; + + auto reg = b.allocQubitRegister(size); + auto c1 = arith::ConstantOp::create(b, b.getIndexType(), b.getIndexAttr(1)) + .getResult(); + auto condition = + arith::ConstantOp::create(b, b.getIndexType(), b.getIndexAttr(0)) + .getResult(); + for (int64_t i = 0; i < size; ++i) { + Value bit; + + reg[i] = b.h(reg[i]); + std::tie(reg[i], bit) = b.measure(reg[i]); + const auto index = + arith::IndexCastUIOp::create(b, b.getIndexType(), bit).getOut(); + condition = arith::OrIOp::create(b, {condition, index}).getResult(); + condition = arith::ShLIOp::create(b, {condition, c1}); + } + + reg.qubits = b.qcoIndexSwitch( + condition, reg.qubits, SmallVector{1, 2, 3}, + SmallVector(ValueRange)>>{ + [&](ValueRange args) { + SmallVector qs(args); + qs[1] = b.x(qs[1]); + return qs; + }, + [&](ValueRange args) { + SmallVector qs(args); + qs[0] = b.x(qs[0]); + return qs; + }, + [&](ValueRange args) { + SmallVector qs(args); + qs[0] = b.x(qs[0]); + qs[1] = b.x(qs[1]); + return qs; + }}, + [&](ValueRange args) { return args; }); + + return measureAndReturn(b, reg.qubits); +} + SmallVector qtensorAlloc(QCOProgramBuilder& b) { (void)b.qtensorAlloc(3); return measureAndReturn(b, {}); @@ -4246,6 +4312,41 @@ SmallVector nestedForLoopWhileOp(QCOProgramBuilder& b) { return measureAndReturnQTensor(b, scfFor[0], 2); } +SmallVector nestedForLoopSwitchOp(QCOProgramBuilder& b) { + constexpr int64_t n = 3; + auto reg = b.allocQubitRegister(n); + auto c3 = arith::ConstantOp::create(b, b.getIndexAttr(3)); + + reg.value = b.scfFor(0, n, 1, reg.value, [&](Value iv, ValueRange iterArgs) { + auto rem = arith::RemUIOp::create(b, {iv, c3}).getResult(); + auto [t, q] = b.qtensorExtract(iterArgs[0], iv); + q = b.qcoIndexSwitch( + rem, {q}, SmallVector{0, 1, 2}, + SmallVector(ValueRange)>>{ + [&](ValueRange args) { + SmallVector qs(args); + qs[0] = b.x(qs[0]); + return qs; + }, + [&](ValueRange args) { + SmallVector qs(args); + qs[0] = b.y(qs[0]); + return qs; + }, + [&](ValueRange args) { + SmallVector qs(args); + qs[0] = b.x(qs[0]); + qs[0] = b.y(qs[0]); + return qs; + }}, + [&](ValueRange args) { return args; })[0]; + auto insert = b.qtensorInsert(q, t, iv); + return SmallVector{insert}; + })[0]; + + return measureAndReturnQTensor(b, reg.value, n); +} + Value nestedForLoopCtrlOpWithSeparateQubit(QCOProgramBuilder& b) { auto reg = b.allocQubitRegister(3); auto control0 = b.allocQubit(); diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 5cb0e7c224..4a4e8d2eee 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -1419,6 +1419,14 @@ SmallVector nestedFalseIf(QCOProgramBuilder& b); /// a register. Value nestedIfOpForLoop(QCOProgramBuilder& b); +// --- IndexSwitchOp ------------------------------------------------------- // + +/// Creates a circuit with an index switch operation with one qubit. +SmallVector simpleIndexSwitch(QCOProgramBuilder& b); + +/// Creates a circuit with an index switch operation with multiple cases. +SmallVector indexSwitchMultiCase(QCOProgramBuilder& b); + // --- WhileOp -------------------------------------------------------------- // /// Creates a circuit with a while operation using a while loop. @@ -1440,6 +1448,10 @@ Value nestedForLoopIfOp(QCOProgramBuilder& b); /// operation. SmallVector nestedForLoopWhileOp(QCOProgramBuilder& b); +/// Creates a circuit with a for operation with a register and a nested index +/// switch operation. +SmallVector nestedForLoopSwitchOp(QCOProgramBuilder& b); + /// Creates a circuit with a for operation with a register and a qubit and a /// nested ctrl operation where the qubit is separately allocated from the /// register.