From f00135ff9b087a2e42eb7b8e89bdfad4322ee9d1 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 17 Jul 2026 13:30:12 +0200 Subject: [PATCH 01/22] Add initial implementation --- mlir/include/mlir/Dialect/QCO/IR/QCOOps.td | 127 ++++++++++ mlir/lib/Dialect/QCO/IR/QCOOps.cpp | 151 ++++++++++++ mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp | 16 +- mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 229 ++++++++++++++++++ 4 files changed, 516 insertions(+), 7 deletions(-) create mode 100644 mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp diff --git a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td index a3ee7b405c..9921f32387 100644 --- a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td +++ b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td @@ -1353,4 +1353,131 @@ 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. + + The case values are stored as a `DenseI64ArrayAttr`, matching the SCF dialect's + `scf.index_switch` operation for seamless interoperability. + + Example: + ```mlir + %result = qco.index_switch %index : 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 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/lib/Dialect/QCO/IR/QCOOps.cpp b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp index f1cb23a849..3f1d889105 100644 --- a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp +++ b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp @@ -169,6 +169,157 @@ void IfOp::print(OpAsmPrinter& p) { p.printOptionalAttrDict((*this)->getAttrs()); } +ParseResult IndexSwitchOp::parse(::mlir::OpAsmParser& parser, + ::mlir::OperationState& result) { + auto& builder = parser.getBuilder(); + OpAsmParser::UnresolvedOperand index; + + // Parse the index operand + if (parser.parseOperand(index) || + parser.resolveOperand(index, builder.getIndexType(), result.operands)) { + return failure(); + } + + // Parse optional result type list + if (parser.parseOptionalArrowTypeList(result.types)) { + return failure(); + } + + // Parse optional attributes + if (parser.parseOptionalAttrDict(result.attributes)) { + return failure(); + } + + // Parse the case regions and default region + SmallVector caseValues; + SmallVector regionArgs; + SmallVector regionOperands; + + // Parse case regions + 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(); + } + + regionArgs.clear(); + regionOperands.clear(); + + // Parse assignment list for this case + if (parser.parseAssignmentList(regionArgs, regionOperands)) { + return failure(); + } + + // Set argument types + for (auto [iterArg, type] : llvm::zip_equal(regionArgs, result.types)) { + iterArg.type = type; + } + + // Add case region + Region* caseRegion = result.addRegion(); + if (parser.parseRegion(*caseRegion, regionArgs)) { + return failure(); + } + } + + // Parse default region + if (parser.parseKeyword("default")) { + return failure(); + } + + if (parser.parseKeyword("args")) { + return failure(); + } + + regionArgs.clear(); + regionOperands.clear(); + + // Parse assignment list for default + if (parser.parseAssignmentList(regionArgs, regionOperands)) { + return failure(); + } + + // Set argument types + for (auto [iterArg, type] : llvm::zip_equal(regionArgs, result.types)) { + iterArg.type = type; + } + + // Add default region + Region* defaultRegion = result.addRegion(); + if (parser.parseRegion(*defaultRegion, regionArgs)) { + return failure(); + } + + // Set the cases attribute + auto casesAttr = DenseI64ArrayAttr::get(parser.getContext(), caseValues); + result.addAttribute("cases", casesAttr); + + 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 + for (size_t i = 0; i < getNumCases(); ++i) { + p << "\ncase "; + p << getCases()[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 << "}"; + } + + // Print default region + p << "\ndefault 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); + p << "}"; +} + //===----------------------------------------------------------------------===// // Dialect //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp index f22e963584..aab7c78932 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp @@ -332,20 +332,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..821daef4bc --- /dev/null +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -0,0 +1,229 @@ +/* + * 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 + +using namespace mlir; +using namespace mlir::qco; + +// 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) + : nregions - 1; // Default region. + + for (size_t i = 0; i < nregions; ++i) { + bounds.emplace_back(/*lb=*/0, /*ub=*/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 nregions = getNumRegions(); + const auto* it = llvm::find(getCases(), arg.getInt()); + const auto liveIndex = it != getCases().end() + ? std::distance(getCases().begin(), it) + : nregions - 1; // Default region. + + regions.emplace_back(&getRegion(liveIndex)); +} + +LogicalResult IndexSwitchOp::verify() { + 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 (bbArg.getOwner()->getParentOp() != getOperation() || i >= getNumCases()) { + 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()->getParentOp() != getOperation()) { + 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(getTargets().size() + addons.size()); + newTargets.append(getTargets().begin(), getTargets().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; +} \ No newline at end of file From d23d182d3fdf86f2dffa9e9d7aacfe209f8c9262 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Fri, 17 Jul 2026 13:37:07 +0200 Subject: [PATCH 02/22] Remove unused includes --- mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 6 ------ 1 file changed, 6 deletions(-) diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index 821daef4bc..e73212d3fe 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -11,13 +11,9 @@ #include "mlir/Dialect/QCO/IR/QCOOps.h" #include -#include -#include #include #include -#include #include -#include #include #include #include @@ -25,8 +21,6 @@ #include #include -#include - using namespace mlir; using namespace mlir::qco; From c086c283eb4fe2484dba8d130e71c0ef136adc8d Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Mon, 20 Jul 2026 10:37:28 +0200 Subject: [PATCH 03/22] Add builders --- .../Dialect/QC/Builder/QCProgramBuilder.h | 21 +++++++ .../Dialect/QCO/Builder/QCOProgramBuilder.h | 32 ++++++++++ mlir/include/mlir/Dialect/QCO/IR/QCOOps.td | 3 - .../Dialect/QC/Builder/QCProgramBuilder.cpp | 34 +++++++++++ .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 61 ++++++++++++++++++- mlir/lib/Dialect/QCO/IR/QCOOps.cpp | 19 +++--- .../Conversion/QCOToQC/test_qco_to_qc.cpp | 9 +++ mlir/unittests/programs/qc_programs.cpp | 14 +++++ mlir/unittests/programs/qc_programs.h | 5 ++ mlir/unittests/programs/qco_programs.cpp | 29 +++++++++ mlir/unittests/programs/qco_programs.h | 5 ++ 11 files changed, 217 insertions(+), 15 deletions(-) diff --git a/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h b/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h index dee7248db0..754d2e7464 100644 --- a/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h @@ -1129,6 +1129,27 @@ 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 for the index switch operation + * @param thenBody Function that builds the then body of the if operation + * @param elseBody Function that builds the else body of the if operation + * @return Reference to this builder for method chaining + * + * @par Example: + * ```c++ + * TODO + * ``` + * ```mlir + * TODO + * ``` + */ + 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 d88e12dfa7..09b5a39f98 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1352,6 +1352,38 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { function_ref(ValueRange)> thenBody, function_ref(ValueRange)> 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 Bool condition. + * @param targets Initial arguments for the index switch branches. + * @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++ + * TODO + * ``` + * ```mlir + * TODO + * ``` + */ + ValueRange qcoIndexSwitch( + const std::variant& arg, ValueRange targets, + ArrayRef cases, + ArrayRef(ValueRange)>> caseBodies, + function_ref(ValueRange)> 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 9921f32387..0480dc5039 100644 --- a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td +++ b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td @@ -1380,9 +1380,6 @@ def IndexSwitchOp must terminate with a `qco.yield` operation that yields values matching the operation's result types. - The case values are stored as a `DenseI64ArrayAttr`, matching the SCF dialect's - `scf.index_switch` operation for seamless interoperability. - Example: ```mlir %result = qco.index_switch %index : index -> !qco.qubit diff --git a/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp b/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp index 3a06d4b593..86c80bb4b2 100644 --- a/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp +++ b/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp @@ -599,6 +599,40 @@ 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. + f(); + scf::YieldOp::create(*this, getLoc()); + }; + + 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 f766602a24..db511e7b48 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -20,12 +20,14 @@ #include #include #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -1094,13 +1096,66 @@ 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(); } QCOProgramBuilder& QCOProgramBuilder::scfCondition(Value condition, diff --git a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp index 3f1d889105..4723f73fd0 100644 --- a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp +++ b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp @@ -278,9 +278,11 @@ void IndexSwitchOp::print(OpAsmPrinter& p) { /*elidedAttrs=*/{"cases"}); // Print case regions + const auto cases = getCases(); for (size_t i = 0; i < getNumCases(); ++i) { - p << "\ncase "; - p << getCases()[i]; + p.printNewline(); + p << "case "; + p << cases[i]; p << " args("; auto& region = getCaseRegions()[i]; @@ -288,20 +290,20 @@ void IndexSwitchOp::print(OpAsmPrinter& p) { // Print block arguments with their corresponding target operands for (size_t j = 0; j < block.getNumArguments(); ++j) { - if (j > 0) + if (j > 0) { p << ", "; + } p.printOperand(block.getArgument(j)); p << " = "; p.printOperand(getTargets()[j]); } - p << ") {"; + p << ") "; p.printRegion(region, /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true); - p << "}"; } - // Print default region - p << "\ndefault args("; + p.printNewline(); + p << "default args("; auto& defaultRegion = getDefaultRegion(); auto& defaultBlock = defaultRegion.front(); @@ -314,10 +316,9 @@ void IndexSwitchOp::print(OpAsmPrinter& p) { p << " = "; p.printOperand(getTargets()[j]); } - p << ") {"; + p << ") "; p.printRegion(defaultRegion, /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true); - p << "}"; } //===----------------------------------------------------------------------===// diff --git a/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp b/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp index 3cab63449f..16d875c474 100644 --- a/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp +++ b/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp @@ -669,6 +669,15 @@ INSTANTIATE_TEST_SUITE_P( MQT_NAMED_BUILDER(qc::nestedIfOpForLoop)})); /// @} +/// \name QCOToQC/Operations/IfOp.cpp +/// @{ +INSTANTIATE_TEST_SUITE_P( + QCOIndexSwitchOpTest, QCOToQCTest, + testing::Values(QCOToQCTestCase{"SimpleIndexSwitchOp", + MQT_NAMED_BUILDER(qco::simpleIndexSwitch), + MQT_NAMED_BUILDER(qc::simpleIndexSwitch)})); +/// @} + /// \name QCOToQC/Operations/WhileOp.cpp /// @{ INSTANTIATE_TEST_SUITE_P( diff --git a/mlir/unittests/programs/qc_programs.cpp b/mlir/unittests/programs/qc_programs.cpp index c54cc9c2d9..6eea31c71b 100644 --- a/mlir/unittests/programs/qc_programs.cpp +++ b/mlir/unittests/programs/qc_programs.cpp @@ -14,6 +14,8 @@ #include #include +#include +#include #include #include @@ -1872,6 +1874,18 @@ 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}; +} + Value simpleWhileReset(QCProgramBuilder& b) { auto q = b.allocQubit(); b.h(q); diff --git a/mlir/unittests/programs/qc_programs.h b/mlir/unittests/programs/qc_programs.h index 4849d861c7..8d47c1f1a4 100644 --- a/mlir/unittests/programs/qc_programs.h +++ b/mlir/unittests/programs/qc_programs.h @@ -910,6 +910,11 @@ 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); + // --- WhileOp -------------------------------------------------------------- // /// Creates a circuit with a while operation using a while loop. diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 69542d9963..c390638b06 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -3182,6 +3183,34 @@ 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(ValueRange)>>{ + [&](ValueRange args) { + auto innerQubit = b.x(args[0]); + return SmallVector{innerQubit}; + }}, + [&](ValueRange args) { + auto innerQubit = b.z(args[0]); + return SmallVector{innerQubit}; + })[0]; + + std::tie(q, bit1) = b.measure(q); + + return {bit0, bit1}; +} + SmallVector qtensorAlloc(QCOProgramBuilder& b) { (void)b.qtensorAlloc(3); return measureAndReturn(b, {}); diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index 61b4ac7eee..ad25b93fe8 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -1101,6 +1101,11 @@ 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); + // --- WhileOp -------------------------------------------------------------- // /// Creates a circuit with a while operation using a while loop. From 7e92ab7c5d9a95bdf857239d777e0bad0bcc0f65 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Mon, 20 Jul 2026 10:53:32 +0200 Subject: [PATCH 04/22] Add QCO to QC conversion --- mlir/lib/Conversion/QCOToQC/QCOToQC.cpp | 50 +++++++++++++++++-- .../Conversion/QCOToQC/test_qco_to_qc.cpp | 2 +- .../Conversion/QCToQCO/test_qc_to_qco.cpp | 9 ++++ 3 files changed, 55 insertions(+), 6 deletions(-) diff --git a/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp b/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp index 768e2d2379..5a596510e0 100644 --- a/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp +++ b/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp @@ -698,7 +698,7 @@ struct ConvertQCOInvOp 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 @@ -715,7 +715,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); @@ -868,6 +868,45 @@ struct ConvertQCOIfOp final : OpConversionPattern { } }; +/** + * @brief Converts qco.index_switch to scf.index_switch + * + * @par Example: + * ```mlir + * TODO + * ``` + * is converted to + * ```mlir + * TODO + * ``` + */ +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(), {}, op.getArg(), + op.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 @@ -996,9 +1035,10 @@ struct QCOToQC final : impl::QCOToQCBase { #undef MQT_ADD_QCO_TO_QC_GATE patterns.add(typeConverter, context); + ConvertQCOYieldOp, ConvertQCOIfOp, ConvertQCOIndexSwitchOp, + ConvertQCOSCFWhileOp, ConvertQCOSCFConditionOp, + ConvertQCOSCFYieldOp, ConvertQCOSCFForOp>(typeConverter, + context); // Register operation conversion patterns that need state tracking patterns.add Date: Tue, 21 Jul 2026 13:32:13 +0200 Subject: [PATCH 05/22] Add QC to QCO conversion --- mlir/lib/Conversion/QCToQCO/QCToQCO.cpp | 97 ++++++++++++++++++++++--- 1 file changed, 87 insertions(+), 10 deletions(-) diff --git a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp index f778521d06..3f4132707f 100644 --- a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp +++ b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -1405,7 +1406,7 @@ struct ConvertSCFIfOp final : StatefulOpConversionPattern { // 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, @@ -1457,9 +1458,84 @@ struct ConvertSCFIfOp final : StatefulOpConversionPattern { } }; +/** + * @brief Converts scf.index_switch to qco.index_switch + * + * @par Example: + * ```mlir + * TODO + * ``` + * is converted to + * ```mlir + * TODO + * ``` + */ +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 types = ValueRange(targets).getTypes(); + const SmallVector locs(targets.size(), op->getLoc()); + llvm::dbgs() << targets.size() << '\n'; + auto newOp = + IndexSwitchOp::create(rewriter, op.getLoc(), types, 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)); + + // Extract all the previously inserted qubits again + extractQubitsAfterOp(state, operation, rewriter); + + const auto newCaseRegions = newOp.getCaseRegions(); + const auto oldCaseRegions = op.getCaseRegions(); + const auto buildRegion = [&](Region& oldRegion, Region& newRegion) { + Block* oldBlock = &(*oldRegion.begin()); + Block* newBlock = + rewriter.createBlock(&newRegion, {}, newOp.getResultTypes(), locs); + + rewriter.inlineBlockBefore(&oldRegion.getBlocks().front(), newBlock, + newBlock->begin()); + + pushModifierFrameWithRegisters( + state, qubits, registers, newBlock->getArguments().take_back(nqubits), + newBlock->getArguments().take_front(nregisters)); + }; + + rewriter.setInsertionPointAfter(newOp); + for (size_t i = 0; i < op.getNumCases(); ++i) { + buildRegion(oldCaseRegions[i], newCaseRegions[i]); + } + + 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 @@ -1482,7 +1558,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); @@ -1603,13 +1679,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>( From 63a2e1a636ac98808e11ac5dbccc7ebf82b603db Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Tue, 21 Jul 2026 13:52:17 +0200 Subject: [PATCH 06/22] Fix todos --- .../Dialect/QC/Builder/QCProgramBuilder.h | 13 +++++++++-- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 23 +++++++++++++++++-- mlir/lib/Conversion/QCOToQC/QCOToQC.cpp | 18 +++++++++++++-- mlir/lib/Conversion/QCToQCO/QCToQCO.cpp | 18 +++++++++++++-- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 2 +- 5 files changed, 65 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h b/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h index 754d2e7464..04a2f053b9 100644 --- a/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h @@ -1139,10 +1139,19 @@ class QCProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * TODO + * builder.scfIndexSwitch(index, + * SmallVector{0}, + * SmallVector>{[&] { b.x(q0); }}, + * [&] { b.z(q0); }); * ``` * ```mlir - * TODO + * scf.index_switch %condition + * case 0 { + * qc.x %q0 : !qc.qubit + * } + * default { + * qc.z %q0 : !qc.qubit + * } * ``` */ QCProgramBuilder& scfIndexSwitch(const std::variant& arg, diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 09b5a39f98..a6a9e8dca6 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1372,10 +1372,29 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * * @par Example: * ```c++ - * TODO + * result = b.qcoIndexSwitch(condition, 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 - * TODO + * %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 + * } * ``` */ ValueRange qcoIndexSwitch( diff --git a/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp b/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp index 5a596510e0..0ccbcf4aa1 100644 --- a/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp +++ b/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp @@ -873,11 +873,25 @@ struct ConvertQCOIfOp final : OpConversionPattern { * * @par Example: * ```mlir - * TODO + * %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 - * TODO + * scf.index_switch %condition + * case 0 { + * qc.x %q0 : !qc.qubit + * } + * default { + * qc.z %q0 : !qc.qubit + * } * ``` */ struct ConvertQCOIndexSwitchOp final : OpConversionPattern { diff --git a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp index 3f4132707f..41abb90af8 100644 --- a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp +++ b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp @@ -1463,11 +1463,25 @@ struct ConvertSCFIfOp final : StatefulOpConversionPattern { * * @par Example: * ```mlir - * TODO + * scf.index_switch %condition + * case 0 { + * qc.x %q0 : !qc.qubit + * } + * default { + * qc.z %q0 : !qc.qubit + * } * ``` * is converted to * ```mlir - * TODO + * %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 diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index db511e7b48..9730806aec 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -1154,7 +1154,7 @@ ValueRange QCOProgramBuilder::qcoIndexSwitch( buildRegion(switchOp.getDefaultRegion(), prev, defaultBody); updateQubitValueTracking(prev, switchOp.getResults()); - + return switchOp.getResults(); } From f5d6b2d7868c8bec4cb455bc5e0a302f680fb68b Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Wed, 22 Jul 2026 09:34:25 +0200 Subject: [PATCH 07/22] Add multi-case test --- mlir/lib/Conversion/QCToQCO/QCToQCO.cpp | 42 ++-- mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 4 +- mlir/lib/Support/IRVerification.cpp | 188 +++++++++--------- .../Conversion/QCOToQC/test_qco_to_qc.cpp | 6 +- .../Conversion/QCToQCO/test_qc_to_qco.cpp | 5 +- mlir/unittests/programs/qc_programs.cpp | 35 ++++ mlir/unittests/programs/qc_programs.h | 3 + mlir/unittests/programs/qco_programs.cpp | 49 +++++ mlir/unittests/programs/qco_programs.h | 3 + 9 files changed, 221 insertions(+), 114 deletions(-) diff --git a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp index 41abb90af8..87e4b7d966 100644 --- a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp +++ b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp @@ -1410,7 +1410,7 @@ struct ConvertSCFIfOp final : StatefulOpConversionPattern { 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); @@ -1504,11 +1504,11 @@ struct ConvertSCFIndexSwitchOp final insertQubitsBeforeOp(state, operation, rewriter); const auto targets = resolveAllValues(state, operation); - const auto types = ValueRange(targets).getTypes(); - const SmallVector locs(targets.size(), op->getLoc()); - llvm::dbgs() << targets.size() << '\n'; + const auto results = ValueRange(targets).getTypes(); + const SmallVector locs(targets.size(), op.getLoc()); + auto newOp = - IndexSwitchOp::create(rewriter, op.getLoc(), types, op.getArg(), + IndexSwitchOp::create(rewriter, op.getLoc(), results, op.getArg(), op.getCases(), targets, op.getNumCases()); assignMappedTensors(state, op.getOperation(), registerMap, @@ -1516,32 +1516,36 @@ struct ConvertSCFIndexSwitchOp final assignMappedQubits(state, op.getOperation(), qubitMap, newOp.getResults().take_back(nqubits)); - // Extract all the previously inserted qubits again + 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* oldBlock = &(*oldRegion.begin()); - Block* newBlock = + Block* const oldBlock = &(*oldRegion.begin()); + Block* const newBlock = rewriter.createBlock(&newRegion, {}, newOp.getResultTypes(), locs); - - rewriter.inlineBlockBefore(&oldRegion.getBlocks().front(), newBlock, - newBlock->begin()); - - pushModifierFrameWithRegisters( - state, qubits, registers, newBlock->getArguments().take_back(nqubits), - newBlock->getArguments().take_front(nregisters)); + const auto args = newBlock->getArguments(); + + // rewriter.inlineBlockBefore(oldBlock, newBlock, newBlock->begin()); + newBlock->getOperations().splice(newBlock->end(), + oldBlock->getOperations()); + pushModifierFrameWithRegisters(state, qubits, registers, + args.take_back(nqubits), + args.take_front(nregisters)); }; - rewriter.setInsertionPointAfter(newOp); - for (size_t i = 0; i < op.getNumCases(); ++i) { - buildRegion(oldCaseRegions[i], newCaseRegions[i]); - } + // 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(); } }; diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index e73212d3fe..7b1d9ea292 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -58,7 +58,7 @@ void IndexSwitchOp::getRegionInvocationBounds( const auto* it = llvm::find(getCases(), arg.getInt()); const auto liveIndex = it != getCases().end() ? std::distance(getCases().begin(), it) - : nregions - 1; // Default region. + : 0; // Default region. for (size_t i = 0; i < nregions; ++i) { bounds.emplace_back(/*lb=*/0, /*ub=*/i == liveIndex ? 1 : 0); @@ -83,7 +83,7 @@ void IndexSwitchOp::getEntrySuccessorRegions( const auto* it = llvm::find(getCases(), arg.getInt()); const auto liveIndex = it != getCases().end() ? std::distance(getCases().begin(), it) - : nregions - 1; // Default region. + : 0; // Default region. regions.emplace_back(&getRegion(liveIndex)); } diff --git a/mlir/lib/Support/IRVerification.cpp b/mlir/lib/Support/IRVerification.cpp index 79ec785ba7..c05165ffc9 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -19,6 +19,8 @@ #include #include #include +#include +#include #include #include #include @@ -154,101 +156,44 @@ 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. +/// 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 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); +getPermutation(const LhsRange& lhs, const RhsRange& rhs, const IRMapping& m, + const TensorMapping& tm) { + SmallVector permutation(lhs.size()); + for (const auto& [i, trgt] : llvm::enumerate(lhs)) { + const auto it = + hasTypeQubitTensor(trgt) + ? llvm::find_if(rhs, + [&](const auto v) { return tm.equals(trgt, v); }) + : llvm::find(rhs, m.lookup(trgt)); + 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)) { + workset.insert_range(rhs); + for (const auto lhsValue : lhs) { + const auto mapped = + hasTypeQubitTensor(lhsValue) + ? *llvm::find_if(rhs, + [&](const auto rhsValue) { + return tm.equals(lhsValue, rhsValue); + }) + : m.lookup(lhsValue); + if (!workset.contains(mapped)) { return false; } - workset.erase(v); + workset.erase(mapped); } - - 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)) { - return false; - } - workset.erase(v); - } - return workset.empty(); } @@ -413,12 +358,38 @@ 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)) { + } + } 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)); + if (!compareValueLists(cast(lhs).getTargets(), + cast(rhs).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 +521,30 @@ 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); + mapArguments( + lhs, rhs, + getPermutation(lhsCtrl.getTargets(), rhsCtrl.getTargets(), m, tm), 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); + 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); + 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); + mapArguments(lhs, rhs, permutation, m); } else { SmallVector permutation(lhs.getNumArguments()); std::iota(permutation.begin(), permutation.end(), 0); @@ -608,11 +597,28 @@ static bool compareBlocks(Block& lhs, Block& rhs, SmallVector permutation; permutation.reserve(lhsCtrl.getNumQubits()); - permutation.append(getControlPermutation(lhsCtrl, rhsCtrl, m)); - for (const auto i : getTargetPermutation(lhsCtrl, rhsCtrl, m)) { + permutation.append(getPermutation( + lhsCtrl.getInputControls(), rhsCtrl.getInputControls(), m, tm)); + for (const auto i : + getPermutation(lhsCtrl.getInputTargets(), + rhsCtrl.getInputTargets(), m, tm)) { 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); + 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); + 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 34b663411b..2fb8ede181 100644 --- a/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp +++ b/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp @@ -675,7 +675,11 @@ INSTANTIATE_TEST_SUITE_P( QCOIndexSwitchOpTest, QCOToQCTest, testing::Values(QCOToQCTestCase{"SimpleIndexSwitchOp", MQT_NAMED_BUILDER(qco::simpleIndexSwitch), - MQT_NAMED_BUILDER(qc::simpleIndexSwitch)})); + MQT_NAMED_BUILDER(qc::simpleIndexSwitch)}, + QCOToQCTestCase{ + "IndexSwitchMultiCase", + MQT_NAMED_BUILDER(qco::indexSwitchMultiCase), + MQT_NAMED_BUILDER(qc::indexSwitchMultiCase)})); /// @} /// \name QCOToQC/Operations/WhileOp.cpp diff --git a/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp b/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp index 5f9b6f9c5f..cb7fcc3d48 100644 --- a/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp +++ b/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp @@ -678,7 +678,10 @@ INSTANTIATE_TEST_SUITE_P(QCOIndexSwitchOpTest, QCToQCOTest, testing::Values(QCToQCOTestCase{ "SimpleIndexSwitchOp", MQT_NAMED_BUILDER(qc::simpleIndexSwitch), - MQT_NAMED_BUILDER(qco::simpleIndexSwitch)})); + MQT_NAMED_BUILDER(qco::simpleIndexSwitch)},QCToQCOTestCase{ + "IndexSwitchMultiCase", + MQT_NAMED_BUILDER(qc::indexSwitchMultiCase), + MQT_NAMED_BUILDER(qco::indexSwitchMultiCase)})); /// @} /// \name QCToQCO/Operations/WhileOp.cpp diff --git a/mlir/unittests/programs/qc_programs.cpp b/mlir/unittests/programs/qc_programs.cpp index 6eea31c71b..bad1937281 100644 --- a/mlir/unittests/programs/qc_programs.cpp +++ b/mlir/unittests/programs/qc_programs.cpp @@ -1886,6 +1886,41 @@ SmallVector simpleIndexSwitch(QCProgramBuilder& b) { 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 */ }); + + SmallVector bits(size); + for (int64_t i = 0; i < size; ++i) { + bits[i] = b.measure(reg[i]); + } + + return bits; +} + Value simpleWhileReset(QCProgramBuilder& b) { auto q = b.allocQubit(); b.h(q); diff --git a/mlir/unittests/programs/qc_programs.h b/mlir/unittests/programs/qc_programs.h index 8d47c1f1a4..1668d8f95e 100644 --- a/mlir/unittests/programs/qc_programs.h +++ b/mlir/unittests/programs/qc_programs.h @@ -915,6 +915,9 @@ Value nestedIfOpForLoop(QCProgramBuilder& b); /// Creates a circuit with an index switch operation with one qubit. SmallVector simpleIndexSwitch(QCProgramBuilder& b); +/// Creates a circuit with an index switch operation using three qubits. +SmallVector indexSwitchMultiCase(QCProgramBuilder& b); + // --- WhileOp -------------------------------------------------------------- // /// Creates a circuit with a while operation using a while loop. diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index c390638b06..9d72f0f856 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -3211,6 +3211,55 @@ SmallVector simpleIndexSwitch(QCOProgramBuilder& b) { 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; }); + + SmallVector bits(size); + for (int64_t i = 0; i < size; ++i) { + std::tie(reg[i], bits[i]) = b.measure(reg[i]); + } + + return bits; +} + SmallVector qtensorAlloc(QCOProgramBuilder& b) { (void)b.qtensorAlloc(3); return measureAndReturn(b, {}); diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index ad25b93fe8..939e8acdff 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -1106,6 +1106,9 @@ Value nestedIfOpForLoop(QCOProgramBuilder& b); /// Creates a circuit with an index switch operation with one qubit. SmallVector simpleIndexSwitch(QCOProgramBuilder& b); +/// Creates a circuit with an index switch operation using three qubits. +SmallVector indexSwitchMultiCase(QCOProgramBuilder& b); + // --- WhileOp -------------------------------------------------------------- // /// Creates a circuit with a while operation using a while loop. From 3a6e22ef9ad6a7515f2ad3a4849f4250dace58cc Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Wed, 22 Jul 2026 10:15:56 +0200 Subject: [PATCH 08/22] Add nested unit test --- .../Conversion/QCOToQC/test_qco_to_qc.cpp | 3 ++ .../Conversion/QCToQCO/test_qc_to_qco.cpp | 20 +++++---- mlir/unittests/programs/qc_programs.cpp | 26 ++++++++--- mlir/unittests/programs/qc_programs.h | 4 ++ mlir/unittests/programs/qco_programs.cpp | 43 ++++++++++++++++--- mlir/unittests/programs/qco_programs.h | 4 ++ 6 files changed, 80 insertions(+), 20 deletions(-) diff --git a/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp b/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp index 2fb8ede181..aef908ec33 100644 --- a/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp +++ b/mlir/unittests/Conversion/QCOToQC/test_qco_to_qc.cpp @@ -707,6 +707,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 cb7fcc3d48..bbf56582fe 100644 --- a/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp +++ b/mlir/unittests/Conversion/QCToQCO/test_qc_to_qco.cpp @@ -674,14 +674,15 @@ INSTANTIATE_TEST_SUITE_P( /// \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)})); +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 @@ -709,6 +710,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/programs/qc_programs.cpp b/mlir/unittests/programs/qc_programs.cpp index bad1937281..eabfb6d409 100644 --- a/mlir/unittests/programs/qc_programs.cpp +++ b/mlir/unittests/programs/qc_programs.cpp @@ -1913,12 +1913,7 @@ SmallVector indexSwitchMultiCase(QCProgramBuilder& b) { }}, [&] { /* no-op */ }); - SmallVector bits(size); - for (int64_t i = 0; i < size; ++i) { - bits[i] = b.measure(reg[i]); - } - - return bits; + return measureAndReturn(b, reg.qubits); } Value simpleWhileReset(QCProgramBuilder& b) { @@ -1986,6 +1981,25 @@ SmallVector nestedForLoopWhileOp(QCProgramBuilder& b) { return measureAndReturn(b, reg.qubits); } +SmallVector nestedForLoopSwitchOp(QCProgramBuilder& b) { + constexpr int64_t n = 32; + auto reg = b.allocQubitRegister(1); + 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 1668d8f95e..32da26c1a8 100644 --- a/mlir/unittests/programs/qc_programs.h +++ b/mlir/unittests/programs/qc_programs.h @@ -939,6 +939,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 9d72f0f856..7c5ec2cda1 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -3252,12 +3252,7 @@ SmallVector indexSwitchMultiCase(QCOProgramBuilder& b) { }}, [&](ValueRange args) { return args; }); - SmallVector bits(size); - for (int64_t i = 0; i < size; ++i) { - std::tie(reg[i], bits[i]) = b.measure(reg[i]); - } - - return bits; + return measureAndReturn(b, reg.qubits); } SmallVector qtensorAlloc(QCOProgramBuilder& b) { @@ -3457,6 +3452,42 @@ SmallVector nestedForLoopWhileOp(QCOProgramBuilder& b) { return measureAndReturnQTensor(b, scfFor[0], 2); } +SmallVector nestedForLoopSwitchOp(QCOProgramBuilder& b) { + constexpr int64_t n = 32; + auto reg = b.allocQubitRegister(1); + 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 [t0, q0] = b.qtensorExtract(iterArgs[0], iv); + q0 = b.qcoIndexSwitch( + rem, {q0}, SmallVector{1, 2, 3}, + 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(q0, t0, iv); + return SmallVector{insert}; + })[0]; + + return measureAndReturnQTensor(b, reg.value, 1); +} + 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 939e8acdff..dd643fb173 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -1130,6 +1130,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. From 59e6ec7be31b99fd95740187ed71e5579c18285a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 08:16:47 +0000 Subject: [PATCH 09/22] =?UTF-8?q?=F0=9F=8E=A8=20pre-commit=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mlir/include/mlir/Dialect/QCO/IR/QCOOps.td | 21 ++++---- .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 2 +- mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 2 +- mlir/unittests/programs/qco_programs.cpp | 53 +++++++++---------- 4 files changed, 38 insertions(+), 40 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td index 0480dc5039..f61a524343 100644 --- a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td +++ b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td @@ -1363,8 +1363,7 @@ def IndexSwitchOp traits = [DeclareOpInterfaceMethods< RegionBranchOpInterface, ["getRegionInvocationBounds", "getEntrySuccessorRegions"]>, - SingleBlockImplicitTerminator<"YieldOp">, - Pure, + SingleBlockImplicitTerminator<"YieldOp">, Pure, RecursiveMemoryEffects]> { let summary = "Index-based switch operation for linear (qubit) types"; @@ -1398,10 +1397,10 @@ def IndexSwitchOp }]; let arguments = (ins Index:$arg, DenseI64ArrayAttr:$cases, - Variadic:$targets); + Variadic:$targets); let results = (outs Variadic:$results); let regions = (region SizedRegion<1>:$defaultRegion, - VariadicRegion>:$caseRegions); + VariadicRegion>:$caseRegions); let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ @@ -1423,7 +1422,7 @@ def IndexSwitchOp /// 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); @@ -1431,7 +1430,7 @@ def IndexSwitchOp /// 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); @@ -1439,7 +1438,7 @@ def IndexSwitchOp /// 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); @@ -1456,19 +1455,19 @@ def IndexSwitchOp 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()); } diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 9730806aec..db511e7b48 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -1154,7 +1154,7 @@ ValueRange QCOProgramBuilder::qcoIndexSwitch( buildRegion(switchOp.getDefaultRegion(), prev, defaultBody); updateQubitValueTracking(prev, switchOp.getResults()); - + return switchOp.getResults(); } diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index 7b1d9ea292..445d5cc019 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -220,4 +220,4 @@ IndexSwitchOp::replaceWithAdditionalTargets(RewriterBase& rewriter, newSwitchOp.getResults().take_front(getNumResults())); return newSwitchOp; -} \ No newline at end of file +} diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 7c5ec2cda1..1ff963e10b 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -3457,33 +3457,32 @@ SmallVector nestedForLoopSwitchOp(QCOProgramBuilder& b) { auto reg = b.allocQubitRegister(1); 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 [t0, q0] = b.qtensorExtract(iterArgs[0], iv); - q0 = b.qcoIndexSwitch( - rem, {q0}, SmallVector{1, 2, 3}, - 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(q0, t0, iv); - return SmallVector{insert}; - })[0]; + reg.value = b.scfFor(0, n, 1, reg.value, [&](Value iv, ValueRange iterArgs) { + auto rem = arith::RemUIOp::create(b, {iv, c3}).getResult(); + auto [t0, q0] = b.qtensorExtract(iterArgs[0], iv); + q0 = b.qcoIndexSwitch( + rem, {q0}, SmallVector{1, 2, 3}, + 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(q0, t0, iv); + return SmallVector{insert}; + })[0]; return measureAndReturnQTensor(b, reg.value, 1); } From cdef797e7abec36612b989753dcf5bd5a7d9777b Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Wed, 22 Jul 2026 10:18:37 +0200 Subject: [PATCH 10/22] Update CHANGELOG.md --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5ee30434b3..9cb3938f06 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,7 +60,7 @@ releases may include breaking changes. [#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], [#1914]) ([**@burgholzer**], + [#1824], [#1830], [#1869], [#1872], [#1914], [#1925]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**]) @@ -639,6 +639,7 @@ changelogs._ +[#1925]: https://github.com/munich-quantum-toolkit/core/pull/1925 [#1914]: https://github.com/munich-quantum-toolkit/core/pull/1914 [#1911]: https://github.com/munich-quantum-toolkit/core/pull/1911 [#1904]: https://github.com/munich-quantum-toolkit/core/pull/1904 From b11d37f325c1bcfd03299d33a8163dc300288863 Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Wed, 22 Jul 2026 10:40:43 +0200 Subject: [PATCH 11/22] Fix linting --- mlir/lib/Conversion/QCToQCO/QCToQCO.cpp | 1 - mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp | 1 + mlir/lib/Dialect/QCO/IR/QCOOps.cpp | 4 ++++ mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 9 ++++++--- mlir/lib/Support/IRVerification.cpp | 3 --- mlir/unittests/programs/qc_programs.cpp | 2 +- 6 files changed, 12 insertions(+), 8 deletions(-) diff --git a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp index 87e4b7d966..1cdde4261f 100644 --- a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp +++ b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp @@ -21,7 +21,6 @@ #include #include -#include #include #include #include diff --git a/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp b/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp index 86c80bb4b2..1ce7763443 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 diff --git a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp index 4723f73fd0..cf58a15143 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 diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index 445d5cc019..162f29e2f2 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -11,6 +11,7 @@ #include "mlir/Dialect/QCO/IR/QCOOps.h" #include +#include #include #include #include @@ -21,6 +22,9 @@ #include #include +#include +#include + using namespace mlir; using namespace mlir::qco; @@ -79,7 +83,6 @@ void IndexSwitchOp::getEntrySuccessorRegions( // Otherwise, try to find a case with a matching value. If not, the // default region is the only successor. - const auto nregions = getNumRegions(); const auto* it = llvm::find(getCases(), arg.getInt()); const auto liveIndex = it != getCases().end() ? std::distance(getCases().begin(), it) @@ -185,8 +188,8 @@ IndexSwitchOp::replaceWithAdditionalTargets(RewriterBase& rewriter, const auto nregions = getNumRegions(); SmallVector newTargets; - newTargets.reserve(getTargets().size() + addons.size()); - newTargets.append(getTargets().begin(), getTargets().end()); + newTargets.reserve(targets.size() + addons.size()); + newTargets.append(targets.begin(), targets.end()); newTargets.append(addons.begin(), addons.end()); const auto newTargetTypes = ValueRange(newTargets).getTypes(); diff --git a/mlir/lib/Support/IRVerification.cpp b/mlir/lib/Support/IRVerification.cpp index c05165ffc9..73a6008ce6 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -16,11 +16,8 @@ #include "mlir/Dialect/QTensor/Utils/TensorIterator.h" #include -#include #include #include -#include -#include #include #include #include diff --git a/mlir/unittests/programs/qc_programs.cpp b/mlir/unittests/programs/qc_programs.cpp index eabfb6d409..f713c31379 100644 --- a/mlir/unittests/programs/qc_programs.cpp +++ b/mlir/unittests/programs/qc_programs.cpp @@ -15,10 +15,10 @@ #include #include #include -#include #include #include +#include #include namespace mlir::qc { From fdd169f6f26df4bdba600c2b4d93e5d355de07ec Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Wed, 22 Jul 2026 10:58:57 +0200 Subject: [PATCH 12/22] Fix issue in verifier --- mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 3 ++- mlir/lib/Support/IRVerification.cpp | 26 +++++++++++++------ 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index 162f29e2f2..b97ca28aa8 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -24,6 +24,7 @@ #include #include +#include using namespace mlir; using namespace mlir::qco; @@ -65,7 +66,7 @@ void IndexSwitchOp::getRegionInvocationBounds( : 0; // Default region. for (size_t i = 0; i < nregions; ++i) { - bounds.emplace_back(/*lb=*/0, /*ub=*/i == liveIndex ? 1 : 0); + bounds.emplace_back(/*lb=*/0, /*ub=*/std::cmp_equal(i, liveIndex) ? 1 : 0); } } diff --git a/mlir/lib/Support/IRVerification.cpp b/mlir/lib/Support/IRVerification.cpp index 73a6008ce6..e2980d8205 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -94,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(); @@ -160,12 +161,16 @@ static SmallVector getPermutation(const LhsRange& lhs, const RhsRange& rhs, const IRMapping& m, const TensorMapping& tm) { SmallVector permutation(lhs.size()); - for (const auto& [i, trgt] : llvm::enumerate(lhs)) { - const auto it = - hasTypeQubitTensor(trgt) - ? llvm::find_if(rhs, - [&](const auto v) { return tm.equals(trgt, v); }) - : llvm::find(rhs, m.lookup(trgt)); + 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)); const auto j = std::distance(rhs.begin(), it); permutation[i] = j; } @@ -178,11 +183,15 @@ static bool compareValueLists(const LhsRange& lhs, const RhsRange& rhs, const IRMapping& m, const TensorMapping& tm) { DenseSet workset; workset.insert_range(rhs); + for (const auto lhsValue : lhs) { const auto mapped = hasTypeQubitTensor(lhsValue) ? *llvm::find_if(rhs, [&](const auto rhsValue) { + if (!hasTypeQubitTensor(rhsValue)) { + return false; + } return tm.equals(lhsValue, rhsValue); }) : m.lookup(lhsValue); @@ -191,6 +200,7 @@ static bool compareValueLists(const LhsRange& lhs, const RhsRange& rhs, } workset.erase(mapped); } + return workset.empty(); } From 995afda34285ddbfaa4810c78432b64b3f5d32ef Mon Sep 17 00:00:00 2001 From: Matthias Reumann Date: Wed, 22 Jul 2026 13:37:16 +0200 Subject: [PATCH 13/22] Apply bunny fixes --- .../Dialect/QC/Builder/QCProgramBuilder.h | 7 +- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 9 +- mlir/lib/Conversion/QCOToQC/QCOToQC.cpp | 4 +- mlir/lib/Conversion/QCToQCO/QCToQCO.cpp | 6 +- .../Dialect/QC/Builder/QCProgramBuilder.cpp | 2 + mlir/lib/Dialect/QCO/IR/QCOOps.cpp | 100 ++++++++++++------ mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 4 +- mlir/lib/Support/IRVerification.cpp | 1 + mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 68 +++++++++++- mlir/unittests/programs/qc_programs.cpp | 4 +- mlir/unittests/programs/qc_programs.h | 2 +- mlir/unittests/programs/qco_programs.cpp | 14 +-- mlir/unittests/programs/qco_programs.h | 2 +- 13 files changed, 164 insertions(+), 59 deletions(-) diff --git a/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h b/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h index 04a2f053b9..667c00413d 100644 --- a/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h @@ -1132,9 +1132,10 @@ class QCProgramBuilder final : public ImplicitLocOpBuilder { /** * @brief Construct an scf.index_switch operation * - * @param arg Index argument for the index switch operation - * @param thenBody Function that builds the then body of the if operation - * @param elseBody Function that builds the else body of the if 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: diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 67d754cc7b..656703cf38 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1379,15 +1379,16 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * this operation are automatically inserted before the operation is * constructed. * - * @param arg Bool condition. + * @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 + * @return ValueRange of the results. * * @par Example: * ```c++ - * result = b.qcoIndexSwitch(condition, initTargets, + * result = b.qcoIndexSwitch(arg, initTargets, * SmallVector{0}, * SmallVector(ValueRange)>>{ * [&](ValueRange args) { @@ -1401,7 +1402,7 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { * }); * ``` * ```mlir - * %result = qco.index_switch %condition -> !qco.qubit + * %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 diff --git a/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp b/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp index 0ccbcf4aa1..8b31afd1ec 100644 --- a/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp +++ b/mlir/lib/Conversion/QCOToQC/QCOToQC.cpp @@ -901,8 +901,8 @@ struct ConvertQCOIndexSwitchOp final : OpConversionPattern { matchAndRewrite(IndexSwitchOp op, OpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override { auto newOp = - scf::IndexSwitchOp::create(rewriter, op.getLoc(), {}, op.getArg(), - op.getCases(), op.getNumCases()); + scf::IndexSwitchOp::create(rewriter, op.getLoc(), {}, adaptor.getArg(), + adaptor.getCases(), op.getNumCases()); const auto oldRegions = op.getCaseRegions(); const auto newCaseRegions = newOp.getCaseRegions(); diff --git a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp index 1cdde4261f..23acce8cd7 100644 --- a/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp +++ b/mlir/lib/Conversion/QCToQCO/QCToQCO.cpp @@ -1524,11 +1524,9 @@ struct ConvertSCFIndexSwitchOp final Block* const oldBlock = &(*oldRegion.begin()); Block* const newBlock = rewriter.createBlock(&newRegion, {}, newOp.getResultTypes(), locs); - const auto args = newBlock->getArguments(); + rewriter.inlineBlockBefore(oldBlock, newBlock, newBlock->begin()); - // rewriter.inlineBlockBefore(oldBlock, newBlock, newBlock->begin()); - newBlock->getOperations().splice(newBlock->end(), - oldBlock->getOperations()); + const auto args = newBlock->getArguments(); pushModifierFrameWithRegisters(state, qubits, registers, args.take_back(nqubits), args.take_front(nregisters)); diff --git a/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp b/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp index 1ce7763443..fe6c5dcc17 100644 --- a/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp +++ b/mlir/lib/Dialect/QC/Builder/QCProgramBuilder.cpp @@ -620,8 +620,10 @@ QCProgramBuilder::scfIndexSwitch(const std::variant& arg, 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] : diff --git a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp index cf58a15143..a6cf5d3c3a 100644 --- a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp +++ b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp @@ -175,31 +175,29 @@ void IfOp::print(OpAsmPrinter& p) { ParseResult IndexSwitchOp::parse(::mlir::OpAsmParser& parser, ::mlir::OperationState& result) { - auto& builder = parser.getBuilder(); OpAsmParser::UnresolvedOperand index; - - // Parse the index operand if (parser.parseOperand(index) || - parser.resolveOperand(index, builder.getIndexType(), result.operands)) { + parser.resolveOperand(index, parser.getBuilder().getIndexType(), + result.operands)) { return failure(); } - // Parse optional result type list if (parser.parseOptionalArrowTypeList(result.types)) { return failure(); } - // Parse optional attributes if (parser.parseOptionalAttrDict(result.attributes)) { return failure(); } - // Parse the case regions and default region + // Create default region here to ensure regions(0) = default. + Region* defaultRegion = result.addRegion(); + + SmallVector operands; SmallVector caseValues; SmallVector regionArgs; SmallVector regionOperands; - // Parse case regions while (succeeded(parser.parseOptionalKeyword("case"))) { int64_t caseValue = 0; if (parser.parseInteger(caseValue)) { @@ -212,27 +210,59 @@ ParseResult IndexSwitchOp::parse(::mlir::OpAsmParser& parser, return failure(); } - regionArgs.clear(); - regionOperands.clear(); - - // Parse assignment list for this case if (parser.parseAssignmentList(regionArgs, regionOperands)) { return failure(); } - // Set argument types - for (auto [iterArg, type] : llvm::zip_equal(regionArgs, result.types)) { - iterArg.type = type; + 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; } - // Add case region - Region* caseRegion = result.addRegion(); - if (parser.parseRegion(*caseRegion, regionArgs)) { + if (parser.parseRegion(*result.addRegion(), regionArgs)) { return failure(); } + + operands.clear(); + regionArgs.clear(); + regionOperands.clear(); } - // Parse default region + 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(); } @@ -241,29 +271,35 @@ ParseResult IndexSwitchOp::parse(::mlir::OpAsmParser& parser, return failure(); } - regionArgs.clear(); - regionOperands.clear(); - - // Parse assignment list for default if (parser.parseAssignmentList(regionArgs, regionOperands)) { return failure(); } - // Set argument types - for (auto [iterArg, type] : llvm::zip_equal(regionArgs, result.types)) { - iterArg.type = type; + // 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(); + } + + 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; } - // Add default region - Region* defaultRegion = result.addRegion(); if (parser.parseRegion(*defaultRegion, regionArgs)) { return failure(); } - // Set the cases attribute - auto casesAttr = DenseI64ArrayAttr::get(parser.getContext(), caseValues); - result.addAttribute("cases", casesAttr); - return success(); } diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index b97ca28aa8..83ab3185e7 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -155,7 +155,7 @@ BlockArgument IndexSwitchOp::getTiedCaseBlockArgument(OpOperand* target, OpOperand* IndexSwitchOp::getTiedCaseYieldedValue(BlockArgument bbArg, size_t i) { - if (bbArg.getOwner()->getParentOp() != getOperation() || i >= getNumCases()) { + if (i >= getNumCases() || bbArg.getOwner() != getCaseBlock(i)) { return nullptr; } @@ -171,7 +171,7 @@ BlockArgument IndexSwitchOp::getTiedDefaultBlockArgument(OpOperand* target) { } OpOperand* IndexSwitchOp::getTiedDefaultYieldedValue(BlockArgument bbArg) { - if (bbArg.getOwner()->getParentOp() != getOperation()) { + if (bbArg.getOwner() != getDefaultBlock()) { return nullptr; } diff --git a/mlir/lib/Support/IRVerification.cpp b/mlir/lib/Support/IRVerification.cpp index e2980d8205..52a1b1bcff 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -380,6 +380,7 @@ static bool compareOperations(Operation* lhs, Operation* rhs, rhsCtrl.getInputControls(), m, tm) || !compareValueLists(lhsCtrl.getInputTargets(), rhsCtrl.getInputTargets(), m, tm)) { + return false; } } else if (isa(lhs)) { assert(isa(rhs)); diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 5d829b6284..3d9bdf23a1 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -62,7 +63,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(); @@ -199,6 +200,71 @@ 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())); +} + /// \name QCO/SCF/IfOp.cpp /// @{ INSTANTIATE_TEST_SUITE_P( diff --git a/mlir/unittests/programs/qc_programs.cpp b/mlir/unittests/programs/qc_programs.cpp index f713c31379..7a9a8f3405 100644 --- a/mlir/unittests/programs/qc_programs.cpp +++ b/mlir/unittests/programs/qc_programs.cpp @@ -1982,8 +1982,8 @@ SmallVector nestedForLoopWhileOp(QCProgramBuilder& b) { } SmallVector nestedForLoopSwitchOp(QCProgramBuilder& b) { - constexpr int64_t n = 32; - auto reg = b.allocQubitRegister(1); + 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(); diff --git a/mlir/unittests/programs/qc_programs.h b/mlir/unittests/programs/qc_programs.h index 32da26c1a8..93e71cc6a3 100644 --- a/mlir/unittests/programs/qc_programs.h +++ b/mlir/unittests/programs/qc_programs.h @@ -915,7 +915,7 @@ Value nestedIfOpForLoop(QCProgramBuilder& b); /// Creates a circuit with an index switch operation with one qubit. SmallVector simpleIndexSwitch(QCProgramBuilder& b); -/// Creates a circuit with an index switch operation using three qubits. +/// Creates a circuit with an index switch operation with multiple cases. SmallVector indexSwitchMultiCase(QCProgramBuilder& b); // --- WhileOp -------------------------------------------------------------- // diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 1ff963e10b..8d39d97d0b 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -3453,15 +3453,15 @@ SmallVector nestedForLoopWhileOp(QCOProgramBuilder& b) { } SmallVector nestedForLoopSwitchOp(QCOProgramBuilder& b) { - constexpr int64_t n = 32; - auto reg = b.allocQubitRegister(1); + 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 [t0, q0] = b.qtensorExtract(iterArgs[0], iv); - q0 = b.qcoIndexSwitch( - rem, {q0}, SmallVector{1, 2, 3}, + auto [t, q] = b.qtensorExtract(iterArgs[0], iv); + q = b.qcoIndexSwitch( + rem, {q}, SmallVector{0, 1, 2}, SmallVector(ValueRange)>>{ [&](ValueRange args) { SmallVector qs(args); @@ -3480,11 +3480,11 @@ SmallVector nestedForLoopSwitchOp(QCOProgramBuilder& b) { return qs; }}, [&](ValueRange args) { return args; })[0]; - auto insert = b.qtensorInsert(q0, t0, iv); + auto insert = b.qtensorInsert(q, t, iv); return SmallVector{insert}; })[0]; - return measureAndReturnQTensor(b, reg.value, 1); + return measureAndReturnQTensor(b, reg.value, n); } Value nestedForLoopCtrlOpWithSeparateQubit(QCOProgramBuilder& b) { diff --git a/mlir/unittests/programs/qco_programs.h b/mlir/unittests/programs/qco_programs.h index dd643fb173..d49b5b054d 100644 --- a/mlir/unittests/programs/qco_programs.h +++ b/mlir/unittests/programs/qco_programs.h @@ -1106,7 +1106,7 @@ Value nestedIfOpForLoop(QCOProgramBuilder& b); /// Creates a circuit with an index switch operation with one qubit. SmallVector simpleIndexSwitch(QCOProgramBuilder& b); -/// Creates a circuit with an index switch operation using three qubits. +/// Creates a circuit with an index switch operation with multiple cases. SmallVector indexSwitchMultiCase(QCOProgramBuilder& b); // --- WhileOp -------------------------------------------------------------- // From 55987ce497f29f9181f1e0705042aec5ceb62d4f Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 02:46:02 +0200 Subject: [PATCH 14/22] =?UTF-8?q?=F0=9F=90=9B=20Fix=20index=20switch=20reg?= =?UTF-8?q?ion=20selection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GPT-5.6 via Codex --- mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 17 +++++-- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 51 +++++++++++++++++++ 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index 83ab3185e7..b73e993883 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -62,7 +62,7 @@ void IndexSwitchOp::getRegionInvocationBounds( const auto nregions = getNumRegions(); const auto* it = llvm::find(getCases(), arg.getInt()); const auto liveIndex = it != getCases().end() - ? std::distance(getCases().begin(), it) + ? std::distance(getCases().begin(), it) + 1 : 0; // Default region. for (size_t i = 0; i < nregions; ++i) { @@ -85,14 +85,21 @@ void IndexSwitchOp::getEntrySuccessorRegions( // default region is the only successor. const auto* it = llvm::find(getCases(), arg.getInt()); - const auto liveIndex = it != getCases().end() - ? std::distance(getCases().begin(), it) - : 0; // Default region. + if (it == getCases().end()) { + regions.emplace_back(&getDefaultRegion()); + return; + } - regions.emplace_back(&getRegion(liveIndex)); + 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"); + } + const auto targets = getTargets(); const auto ntargets = targets.size(); const auto results = getResults(); diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index fdfbf57264..97529de909 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -262,6 +263,56 @@ TEST_F(QCOTest, IndexSwitchParser) { refBuilder.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); + + 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()[0], 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()); +} + /// \name QCO/SCF/IfOp.cpp /// @{ INSTANTIATE_TEST_SUITE_P( From 68f06e7988c47459723eed56eb4f05c1f6812617 Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 02:46:38 +0200 Subject: [PATCH 15/22] =?UTF-8?q?=F0=9F=90=9B=20Parse=20default-only=20QCO?= =?UTF-8?q?=20index=20switches?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GPT-5.6 via Codex --- mlir/lib/Dialect/QCO/IR/QCOOps.cpp | 16 ++++++---- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 30 +++++++++++++++++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp index a6cf5d3c3a..0fb150ee4e 100644 --- a/mlir/lib/Dialect/QCO/IR/QCOOps.cpp +++ b/mlir/lib/Dialect/QCO/IR/QCOOps.cpp @@ -283,12 +283,16 @@ ParseResult IndexSwitchOp::parse(::mlir::OpAsmParser& parser, 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"); + 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"); + } } } diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 97529de909..91612dc6e5 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -19,6 +19,7 @@ #include "qco_programs.h" #include +#include #include #include #include @@ -263,6 +264,35 @@ TEST_F(QCOTest, IndexSwitchParser) { 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, IndexSwitchConstantSuccessor) { QCOProgramBuilder builder(context.get()); builder.initialize(); From 96a3d4f7f9bc5a75f39b8f89f2d030a813476e01 Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 02:48:18 +0200 Subject: [PATCH 16/22] =?UTF-8?q?=E2=9C=A8=20Support=20index=20switches=20?= =?UTF-8?q?in=20linear=20traversal?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GPT-5.6 via Codex --- mlir/include/mlir/Dialect/QCO/Utils/Drivers.h | 5 ++++ mlir/lib/Dialect/QCO/Utils/WireIterator.cpp | 10 +++++++ .../Dialect/QTensor/Utils/TensorIterator.cpp | 14 ++++++++++ mlir/lib/Support/IRVerification.cpp | 12 +++++++++ mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 27 +++++++++++++++++++ .../Dialect/QCO/Utils/test_drivers.cpp | 10 ++++++- .../Dialect/QCO/Utils/test_wireiterator.cpp | 25 ++++++++++++++--- .../QTensor/Utils/test_tensoriterator.cpp | 21 ++++++++++++--- 8 files changed, 116 insertions(+), 8 deletions(-) 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/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 52a1b1bcff..3a0f2d4de4 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -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())); diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 91612dc6e5..8b9a552e36 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -293,6 +293,33 @@ TEST_F(QCOTest, DefaultOnlyIndexSwitchParser) { 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, IndexSwitchConstantSuccessor) { QCOProgramBuilder builder(context.get()); builder.initialize(); diff --git a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp index 7bbc2437dd..00c896d8f3 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp @@ -109,7 +109,15 @@ 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..449e87ae72 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp @@ -80,8 +80,17 @@ 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..80c3c68d6d 100644 --- a/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp +++ b/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp @@ -101,7 +101,14 @@ 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 +149,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 +163,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); From f5740553ddf789b21388806104fd339463f6bf58 Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 02:50:09 +0200 Subject: [PATCH 17/22] =?UTF-8?q?=F0=9F=90=9B=20Handle=20missing=20tensor?= =?UTF-8?q?=20equivalence=20matches?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GPT-5.6 via Codex --- mlir/lib/Support/IRVerification.cpp | 83 ++++++++++++------- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 43 ++++++++-- 2 files changed, 89 insertions(+), 37 deletions(-) diff --git a/mlir/lib/Support/IRVerification.cpp b/mlir/lib/Support/IRVerification.cpp index 3a0f2d4de4..2382e4f83d 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -115,9 +115,9 @@ static void initEquivGroup(TypedValue v, size_t id, const auto idx = std::distance(targets.begin(), targetIt); for (Region* region : op.getRegions()) { - initEquivGroup(cast>( - region->getArgument(idx)), - id, group); + initEquivGroup( + cast>(region->getArgument(idx)), id, + group); } } else if (auto forOp = dyn_cast(it.operation())) { const auto& arg = @@ -169,7 +169,7 @@ static void mapArguments(Block& lhs, Block& rhs, ArrayRef permutation, /// 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 SmallVector +static FailureOr> getPermutation(const LhsRange& lhs, const RhsRange& rhs, const IRMapping& m, const TensorMapping& tm) { SmallVector permutation(lhs.size()); @@ -183,6 +183,9 @@ getPermutation(const LhsRange& lhs, const RhsRange& rhs, const IRMapping& m, 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; } @@ -197,16 +200,18 @@ static bool compareValueLists(const LhsRange& lhs, const RhsRange& rhs, workset.insert_range(rhs); for (const auto lhsValue : lhs) { - const auto mapped = - hasTypeQubitTensor(lhsValue) - ? *llvm::find_if(rhs, - [&](const auto rhsValue) { - if (!hasTypeQubitTensor(rhsValue)) { - return false; - } - return tm.equals(lhsValue, rhsValue); - }) - : m.lookup(lhsValue); + 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); + } if (!workset.contains(mapped)) { return false; } @@ -541,30 +546,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, - getPermutation(lhsCtrl.getTargets(), rhsCtrl.getTargets(), m, tm), 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()); const auto permutation = getPermutation(lhsCtrl.getInputTargets(), rhsCtrl.getInputTargets(), m, tm); - mapArguments(lhs, rhs, permutation, m); + 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); - mapArguments(lhs, rhs, permutation, m); + 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); - mapArguments(lhs, rhs, permutation, m); + if (failed(permutation)) { + return false; + } + mapArguments(lhs, rhs, *permutation, m); } else { SmallVector permutation(lhs.getNumArguments()); std::iota(permutation.begin(), permutation.end(), 0); @@ -615,13 +632,17 @@ 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(getPermutation( - lhsCtrl.getInputControls(), rhsCtrl.getInputControls(), m, tm)); - for (const auto i : - getPermutation(lhsCtrl.getInputTargets(), - rhsCtrl.getInputTargets(), m, tm)) { + for (const auto i : *targetPermutation) { permutation.emplace_back(lhsCtrl.getNumControls() + i); } mapResults(lhsCtrl, rhsCtrl, permutation, m); @@ -631,14 +652,20 @@ static bool compareBlocks(Block& lhs, Block& rhs, auto rhsIf = cast(rhsOp); const auto permutation = getPermutation(lhsIf.getQubits(), rhsIf.getQubits(), m, tm); - mapResults(lhsIf, rhsIf, permutation, m); + 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); - mapResults(lhsSwitch, rhsSwitch, permutation, m); + 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/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 8b9a552e36..4b141307d7 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -23,8 +23,8 @@ #include #include #include -#include #include +#include #include #include #include @@ -298,9 +298,7 @@ TEST_F(QCOTest, EquivalentTensorIndexSwitches) { QCOProgramBuilder builder(context.get()); builder.initialize(); - const auto identity = [](ValueRange args) { - return llvm::to_vector(args); - }; + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; const SmallVector(ValueRange)>> caseBodies{ identity}; @@ -320,19 +318,46 @@ TEST_F(QCOTest, EquivalentTensorIndexSwitches) { EXPECT_TRUE(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 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); + auto result = builder.qcoIndexSwitch(1, q0, SmallVector{0, 1}, + caseBodies, identity); builder.sink(result.front()); [[maybe_unused]] auto module = builder.finalize(); From 9d8739326090d17039fe49cd69d9fc8fcc9644b1 Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 02:50:52 +0200 Subject: [PATCH 18/22] =?UTF-8?q?=F0=9F=8E=A8=20Apply=20lint=20formatting?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GPT-5.6 via Codex --- CHANGELOG.md | 6 +++--- mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp | 8 +++----- mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp | 4 +--- .../Dialect/QTensor/Utils/test_tensoriterator.cpp | 4 +--- 4 files changed, 8 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f758d76d3..fc826d26e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -66,9 +66,9 @@ releases may include breaking changes. [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765], [#1780], [#1781], [#1782], [#1787], [#1806], [#1807], [#1815], [#1808], [#1823], [#1824], [#1830], [#1869], [#1872], [#1886], [#1914], [#1925]) - ([**@burgholzer**], - [**@denialhaag**], [**@taminob**], [**@DRovara**], [**@li-mingbao**], - [**@Ectras**], [**@MatthiasReumann**], [**@simon1hofmann**]) + ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**], + [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**], + [**@simon1hofmann**]) ### Changed diff --git a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp index 00c896d8f3..0715687185 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp @@ -109,13 +109,11 @@ TEST_F(DriversTest, ProgramGraphWalk) { return SmallVector{builder.id(args[0])}; })[0]; - const auto identity = [](ValueRange args) { - return llvm::to_vector(args); - }; + 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]; + const auto q06 = builder.qcoIndexSwitch(0, q05, SmallVector{0}, + caseBodies, identity)[0]; builder.measure(q06); builder.measure(forResults[1]); diff --git a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp index 449e87ae72..16df2d8307 100644 --- a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp +++ b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp @@ -80,9 +80,7 @@ TEST_P(WireIteratorTest, Traversal) { [&](ValueRange args) { return SmallVector{args[0], args[1]}; }); const auto q06 = ifOut[0]; const auto q13 = ifOut[1]; - const auto identity = [](ValueRange args) { - return llvm::to_vector(args); - }; + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; const SmallVector(ValueRange)>> caseBodies{ identity}; const auto switchOut = builder.qcoIndexSwitch( diff --git a/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp b/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp index 80c3c68d6d..a105e6bfa8 100644 --- a/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp +++ b/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp @@ -101,9 +101,7 @@ TEST_F(TensorIteratorTest, Traversal) { tensorElse2 = builder.qtensorInsert(q, tensorElse1, 0); return SmallVector{tensorElse2}; })[0]; - const auto identity = [](ValueRange args) { - return llvm::to_vector(args); - }; + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; const SmallVector(ValueRange)>> caseBodies{ identity}; const auto tensor9 = builder.qcoIndexSwitch( From 6379b9ae088a899eebc3b6199ab2c9a9a08ce2b6 Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 08:37:06 +0200 Subject: [PATCH 19/22] =?UTF-8?q?=F0=9F=90=9B=20Reject=20duplicate=20index?= =?UTF-8?q?=20switch=20cases?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Also correct the operation syntax example. Assisted-by: GPT-5.6 via Codex --- mlir/include/mlir/Dialect/QCO/IR/QCOOps.td | 2 +- mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 8 ++++++++ mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 4 ++++ 3 files changed, 13 insertions(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td index 4d21db142b..1e9a9b455c 100644 --- a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td +++ b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td @@ -1417,7 +1417,7 @@ def IndexSwitchOp Example: ```mlir - %result = qco.index_switch %index : index -> !qco.qubit + %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 diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index b73e993883..546113603a 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -10,6 +10,7 @@ #include "mlir/Dialect/QCO/IR/QCOOps.h" +#include #include #include #include @@ -100,6 +101,13 @@ LogicalResult IndexSwitchOp::verify() { "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(); diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 4b141307d7..62a7983591 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -393,6 +393,10 @@ TEST_F(QCOTest, IndexSwitchConstantSuccessor) { 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 From 4e8564997ee266369638a7e1f76477ea7a413d15 Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 08:37:39 +0200 Subject: [PATCH 20/22] =?UTF-8?q?=F0=9F=90=9B=20Compare=20index=20switch?= =?UTF-8?q?=20case=20values?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GPT-5.6 via Codex --- mlir/lib/Support/IRVerification.cpp | 7 +++-- mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 27 +++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Support/IRVerification.cpp b/mlir/lib/Support/IRVerification.cpp index 2382e4f83d..edf44be05e 100644 --- a/mlir/lib/Support/IRVerification.cpp +++ b/mlir/lib/Support/IRVerification.cpp @@ -407,8 +407,11 @@ static bool compareOperations(Operation* lhs, Operation* rhs, } } else if (isa(lhs)) { assert(isa(rhs)); - if (!compareValueLists(cast(lhs).getTargets(), - cast(rhs).getTargets(), m, tm)) { + 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)) { diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 62a7983591..d5388dcbed 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -318,6 +318,33 @@ TEST_F(QCOTest, EquivalentTensorIndexSwitches) { 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()); From 5ae46fae6c1b0feece739afe7bee9de1f2476c97 Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 08:42:13 +0200 Subject: [PATCH 21/22] =?UTF-8?q?=E2=9C=A8=20Add=20single-value=20index=20?= =?UTF-8?q?switch=20builders?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GPT-5.6 via Codex --- .../Dialect/QCO/Builder/QCOProgramBuilder.h | 29 ++++++++++++++ mlir/include/mlir/Dialect/QCO/IR/QCOOps.td | 5 +++ .../Dialect/QCO/Builder/QCOProgramBuilder.cpp | 38 +++++++++++++++++++ mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 29 ++++++++++++++ mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 27 +++++++++++-- mlir/unittests/programs/qco_programs.cpp | 15 ++------ 6 files changed, 129 insertions(+), 14 deletions(-) diff --git a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h index 37796a4b36..c292fa7098 100644 --- a/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h +++ b/mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h @@ -1561,6 +1561,35 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder { 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 1e9a9b455c..43674aeb61 100644 --- a/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td +++ b/mlir/include/mlir/Dialect/QCO/IR/QCOOps.td @@ -1439,6 +1439,11 @@ def IndexSwitchOp 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(); diff --git a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp index 1edb3d4308..22f6aa6af8 100644 --- a/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp +++ b/mlir/lib/Dialect/QCO/Builder/QCOProgramBuilder.cpp @@ -1213,6 +1213,44 @@ ValueRange QCOProgramBuilder::qcoIndexSwitch( 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, Value initArg, function_ref thenBody, diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index 546113603a..3ea7b7ad2b 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -30,6 +31,34 @@ 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 diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index d5388dcbed..9760ac59a0 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -154,6 +154,29 @@ 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, IfOpParser) { // Test IfOp parser const char* mlirCode = R"( @@ -323,9 +346,7 @@ TEST_F(QCOTest, IndexSwitchCaseValuesAffectEquivalence) { QCOProgramBuilder builder(context.get()); builder.initialize(); - const auto identity = [](ValueRange args) { - return llvm::to_vector(args); - }; + const auto identity = [](ValueRange args) { return llvm::to_vector(args); }; const SmallVector(ValueRange)>> caseBodies{ identity}; diff --git a/mlir/unittests/programs/qco_programs.cpp b/mlir/unittests/programs/qco_programs.cpp index 67bbbfab01..4b8f263f61 100644 --- a/mlir/unittests/programs/qco_programs.cpp +++ b/mlir/unittests/programs/qco_programs.cpp @@ -3284,17 +3284,10 @@ SmallVector simpleIndexSwitch(QCOProgramBuilder& b) { 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(ValueRange)>>{ - [&](ValueRange args) { - auto innerQubit = b.x(args[0]); - return SmallVector{innerQubit}; - }}, - [&](ValueRange args) { - auto innerQubit = b.z(args[0]); - return SmallVector{innerQubit}; - })[0]; + 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); From d1f4c6f387b8cbff5a909f7aa219b5853b9c6593 Mon Sep 17 00:00:00 2001 From: Lukas Burgholzer Date: Thu, 23 Jul 2026 09:31:48 +0200 Subject: [PATCH 22/22] =?UTF-8?q?=E2=9C=85=20Fix=20index=20switch=20lint?= =?UTF-8?q?=20and=20coverage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assisted-by: GPT-5.6 via Codex --- mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp | 1 + mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp | 106 +++++++++++++++++- .../Dialect/QCO/Utils/test_drivers.cpp | 1 + .../Dialect/QCO/Utils/test_wireiterator.cpp | 2 + .../QTensor/Utils/test_tensoriterator.cpp | 1 + 5 files changed, 110 insertions(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp index 3ea7b7ad2b..f57c1e7cb4 100644 --- a/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp +++ b/mlir/lib/Dialect/QCO/IR/SCF/IndexSwitchOp.cpp @@ -25,6 +25,7 @@ #include #include +#include #include #include diff --git a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp index 68f7d96f70..ac551542f9 100644 --- a/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp +++ b/mlir/unittests/Dialect/QCO/IR/test_qco_ir.cpp @@ -19,18 +19,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 @@ -177,6 +185,92 @@ TEST_F(QCOTest, DirectSingleTargetIndexSwitchBuilder) { 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"( @@ -412,6 +506,16 @@ TEST_F(QCOTest, IndexSwitchConstantSuccessor) { 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()); @@ -432,7 +536,7 @@ TEST_F(QCOTest, IndexSwitchConstantSuccessor) { } }; - checkConstant(0, &switchOp.getCaseRegions()[0], 1); + checkConstant(0, switchOp.getCaseRegions().data(), 1); checkConstant(1, &switchOp.getCaseRegions()[1], 2); checkConstant(2, &switchOp.getDefaultRegion(), 0); diff --git a/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp b/mlir/unittests/Dialect/QCO/Utils/test_drivers.cpp index 0715687185..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 diff --git a/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp b/mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp index 16df2d8307..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 diff --git a/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp b/mlir/unittests/Dialect/QTensor/Utils/test_tensoriterator.cpp index a105e6bfa8..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