✨ Add IndexSwitch operation

CHANGELOG.md

@@ -66,7 +66,7 @@ releases may include breaking changes.
   [#1624], [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700],
   [#1710], [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765],
   [#1780], [#1781], [#1782], [#1787], [#1806], [#1807], [#1815], [#1808],
-  [#1823], [#1824], [#1830], [#1869], [#1872], [#1886], [#1914])
+  [#1823], [#1824], [#1830], [#1869], [#1872], [#1886], [#1914], [#1925])
   ([**@burgholzer**], [**@denialhaag**], [**@taminob**], [**@DRovara**],
   [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**],
   [**@simon1hofmann**], [**@J4MMlE**])
@@ -646,6 +646,7 @@ changelogs._
 
 <!-- PR links -->
 
+[#1925]: https://github.com/munich-quantum-toolkit/core/pull/1925
 [#1924]: https://github.com/munich-quantum-toolkit/core/pull/1924
 [#1915]: https://github.com/munich-quantum-toolkit/core/pull/1915
 [#1914]: https://github.com/munich-quantum-toolkit/core/pull/1914

mlir/include/mlir/Dialect/QC/Builder/QCProgramBuilder.h

@@ -1239,6 +1239,37 @@ class QCProgramBuilder final : public ImplicitLocOpBuilder {
                           const function_ref<void()>& thenBody,
                           const function_ref<void()>& elseBody = nullptr);
 
+  /**
+   * @brief Construct an scf.index_switch operation
+   *
+   * @param arg Index argument.
+   * @param cases The individual switch cases.
+   * @param caseBodies An array of functions that build the case bodies.
+   * @param defaultBody Function that builds the default body.
+   * @return Reference to this builder for method chaining
+   *
+   * @par Example:
+   * ```c++
+   * builder.scfIndexSwitch(index,
+   *   SmallVector<int64_t>{0},
+   *   SmallVector<function_ref<void()>>{[&] { b.x(q0); }},
+   *   [&] { b.z(q0); });
+   * ```
+   * ```mlir
+   * scf.index_switch %condition
+   * case 0 {
+   *   qc.x %q0 : !qc.qubit
+   * }
+   * default {
+   *   qc.z %q0 : !qc.qubit
+   * }
+   * ```
+   */
+  QCProgramBuilder& scfIndexSwitch(const std::variant<int64_t, Value>& arg,
+                                   ArrayRef<int64_t> cases,
+                                   ArrayRef<function_ref<void()>> caseBodies,
+                                   const function_ref<void()>& defaultBody);
+
   /**
    * @brief Construct an scf.condition operation
    *

mlir/include/mlir/Dialect/QCO/Builder/QCOProgramBuilder.h

@@ -1556,6 +1556,87 @@ class QCOProgramBuilder final : public ImplicitLocOpBuilder {
               function_ref<Value(Value)> thenBody,
               function_ref<Value(Value)> elseBody = nullptr);
 
+  /**
+   * @brief Construct an index switch operation for qubits or tensors of qubits
+   * with linear typing.
+   *
+   * @details
+   * Constructs an index switch operation that takes an index Value and a range
+   * of qubit and qtensor values that are used in the case regions of this
+   * operation. The values are passed down as block arguments to each region.
+   * Qubits that were extracted from a tensor that is used as an argument for
+   * this operation are automatically inserted before the operation is
+   * constructed.
+   *
+   * @param arg Index argument.
+   * @param targets Initial arguments for the index switch branches.
+   * @param cases The individual switch cases.
+   * @param caseBodies An array of functions that build the case bodies.
+   * @param defaultBody Function that builds the default body.
+   * @return ValueRange of the results.
+   *
+   * @par Example:
+   * ```c++
+   * result = b.qcoIndexSwitch(arg, initTargets,
+   *   SmallVector<int64_t>{0},
+   *   SmallVector<function_ref<SmallVector<Value>(ValueRange)>>{
+   *     [&](ValueRange args) {
+   *       auto q1 = builder.x(args[0]);
+   *       return {q1};
+   *     }
+   *   },
+   *   [&](ValueRange args) {
+   *     auto q2 = builder.x(args[0]);
+   *     return {q2};
+   *   });
+   * ```
+   * ```mlir
+   * %result = qco.index_switch %arg -> !qco.qubit
+   * case 0 args(%arg0 = %q0) {
+   *   %q1 = qco.x %arg0 : !qco.qubit -> !qco.qubit
+   *   qco.yield %q1 : !qco.qubit
+   * }
+   * default args(%arg0 = %q0) {
+   *   %q2 = qco.z %arg0 : !qco.qubit -> !qco.qubit
+   *   qco.yield %q2 : !qco.qubit
+   * }
+   * ```
+   */
+  ValueRange qcoIndexSwitch(
+      const std::variant<int64_t, Value>& arg, ValueRange targets,
+      ArrayRef<int64_t> cases,
+      ArrayRef<function_ref<SmallVector<Value>(ValueRange)>> caseBodies,
+      function_ref<SmallVector<Value>(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<int64_t>{0},
+   *     SmallVector<function_ref<Value(Value)>>{
+   *         [&](Value value) { return builder.x(value); }},
+   *     [&](Value value) { return builder.z(value); });
+   * ```
+   */
+  Value qcoIndexSwitch(const std::variant<int64_t, Value>& arg, Value target,
+                       ArrayRef<int64_t> cases,
+                       ArrayRef<function_ref<Value(Value)>> caseBodies,
+                       function_ref<Value(Value)> defaultBody);
+
   /**
    * @brief Construct an scf.for operation
    *

mlir/include/mlir/Dialect/QCO/IR/QCOOps.td

@@ -1480,4 +1480,132 @@ def IfOp
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Index Switch Operation
+//===----------------------------------------------------------------------===//
+
+def IndexSwitchOp
+    : QCOOp<
+          "index_switch",
+          traits = [DeclareOpInterfaceMethods<
+                        RegionBranchOpInterface, ["getRegionInvocationBounds",
+                                                  "getEntrySuccessorRegions"]>,
+                    SingleBlockImplicitTerminator<"YieldOp">, Pure,
+                    RecursiveMemoryEffects]> {
+
+  let summary = "Index-based switch operation for linear (qubit) types";
+  let description = [{
+        The `qco.index_switch` operation provides multi-way branching based on an index
+        value, analogous to `scf.index_switch`. In addition to the index, the operation
+        takes a variadic number of qubits and qtensors as inputs that are required in all
+        case branches. These values are passed down to each case region as block arguments.
+        The number of results and the type of the results must be equivalent to the number
+        and types of the input qubit and qtensor values.
+
+        The operation has one region for each case value plus a default region. Each region
+        must terminate with a `qco.yield` operation that yields values matching the operation's
+        result types.
+
+        Example:
+        ```mlir
+        %result = qco.index_switch %index -> !qco.qubit
+        case 0 args(%arg0 = %q0) {
+          %q1 = qco.h %arg0 : !qco.qubit -> !qco.qubit
+          qco.yield %q1 : !qco.qubit
+        }
+        case 1 args(%arg0 = %q0) {
+          %q1 = qco.x %arg0 : !qco.qubit -> !qco.qubit
+          qco.yield %q1 : !qco.qubit
+        }
+        default args(%arg0 = %q0) {
+          qco.yield %arg0 : !qco.qubit
+        }
+        ```
+    }];
+
+  let arguments = (ins Index:$arg, DenseI64ArrayAttr:$cases,
+      Variadic<LinearType>:$targets);
+  let results = (outs Variadic<LinearType>:$results);
+  let regions = (region SizedRegion<1>:$defaultRegion,
+      VariadicRegion<SizedRegion<1>>:$caseRegions);
+  let hasCustomAssemblyFormat = 1;
+
+  let builders = [OpBuilder<(ins "Value":$arg, "Value":$target,
+      "ArrayRef<int64_t>":$cases,
+      "ArrayRef<function_ref<Value(Value)>>":$caseBuilders,
+      "function_ref<Value(Value)>":$defaultBuilder)>];
+
+  let extraClassDeclaration = [{
+        /// Return the number of case regions
+        size_t getNumCases();
+
+        /// Return the block for a specific case
+        Block* getCaseBlock(size_t index);
+
+        /// Return the yield operation for a specific case
+        YieldOp getCaseYield(size_t index);
+
+        /// Return the default block
+        Block* getDefaultBlock();
+
+        /// Return the default yield operation
+        YieldOp getDefaultYield();
+
+        /// Return the result that corresponds to the given target operand,
+        /// or "empty" OpResult on failure.
+        OpResult getTiedResult(OpOperand* target);
+
+        /// Return the target operand that corresponds to the given result,
+        /// or `nullptr` on failure.
+        OpOperand* getTiedTarget(OpResult result);
+
+        /// Return the argument that corresponds to the given target operand
+        /// for the i-th case, or "empty" BlockArgument on failure.
+        BlockArgument getTiedCaseBlockArgument(OpOperand* target, size_t i);
+
+        /// Return the yielded value that corresponds to the given argument
+        /// for the i-th case, or `nullptr` on failure.
+        OpOperand* getTiedCaseYieldedValue(BlockArgument bbArg, size_t i);
+
+        /// Return the argument that corresponds to the given target operand
+        /// for the default case, or "empty" BlockArgument on failure.
+        BlockArgument getTiedDefaultBlockArgument(OpOperand* target);
+
+        /// Return the yielded value that corresponds to the given argument
+        /// for the default case, or `nullptr` on failure.
+        OpOperand* getTiedDefaultYieldedValue(BlockArgument bbArg);
+
+        /// Append the specified additional linear operands: replace this
+        /// index_switch with a new index_switch that has the additional linear operands.
+        /// The operands can be of qubit or qtensor type.
+        /// The branch bodies of this index_switch are moved over to the new index_switch.
+        /// The newly added qubits are yielded from each branch.
+        IndexSwitchOp replaceWithAdditionalTargets(RewriterBase& rewriter, ValueRange addons);
+    }];
+
+  let extraClassDefinition = [{
+        size_t $cppClass::getNumCases() {
+          return getCaseRegions().size();
+        }
+
+        Block* $cppClass::getCaseBlock(size_t index) {
+          return &getCaseRegions()[index].front();
+        }
+
+        YieldOp $cppClass::getCaseYield(size_t index) {
+          return cast<YieldOp>(getCaseBlock(index)->getTerminator());
+        }
+
+        Block* $cppClass::getDefaultBlock() {
+          return &getDefaultRegion().front();
+        }
+
+        YieldOp $cppClass::getDefaultYield() {
+          return cast<YieldOp>(getDefaultBlock()->getTerminator());
+        }
+    }];
+
+  let hasVerifier = 1;
+}
+
 #endif // MLIR_DIALECT_QCO_IR_QCOOPS_TD

mlir/include/mlir/Dialect/QCO/Utils/Drivers.h

@@ -53,6 +53,11 @@ inline size_t getNumQubitArgs(Operation* op) {
           return isa<QubitType>(v.getType());
         });
       })
+      .Case<qco::IndexSwitchOp>([&](qco::IndexSwitchOp op) {
+        return llvm::count_if(op.getTargets(), [](Value v) {
+          return isa<QubitType>(v.getType());
+        });
+      })
       .Default([&](Operation* op) {
         const auto name = op->getName().getStringRef();
         reportFatalInternalError("unknown op: " + name);

mlir/lib/Conversion/QCOToQC/QCOToQC.cpp

@@ -737,7 +737,7 @@ struct ConvertQCOPowOp final : OpConversionPattern<qco::PowOp> {
 
 /**
  * @brief Converts qco.yield to qc.yield or to scf.yield if the parent is a
- * scf::IfOp
+ * scf::IfOp or scf::IndexSwitchOp.
  *
  * @par Example:
  * ```mlir
@@ -754,7 +754,7 @@ struct ConvertQCOYieldOp final : OpConversionPattern<qco::YieldOp> {
   LogicalResult
   matchAndRewrite(qco::YieldOp op, OpAdaptor /*adaptor*/,
                   ConversionPatternRewriter& rewriter) const override {
-    if (isa<scf::IfOp>(op->getParentOp())) {
+    if (isa<scf::IfOp, scf::IndexSwitchOp>(op->getParentOp())) {
       rewriter.replaceOpWithNewOp<scf::YieldOp>(op);
     } else {
       rewriter.replaceOpWithNewOp<qc::YieldOp>(op);
@@ -907,6 +907,59 @@ struct ConvertQCOIfOp final : OpConversionPattern<IfOp> {
   }
 };
 
+/**
+ * @brief Converts qco.index_switch to scf.index_switch
+ *
+ * @par Example:
+ * ```mlir
+ * %result = qco.index_switch %condition -> !qco.qubit
+ * case 0 args(%arg0 = %q0) {
+ *   %q1 = qco.x %arg0 : !qco.qubit -> !qco.qubit
+ *   qco.yield %q1 : !qco.qubit
+ * }
+ * default args(%arg0 = %q0) {
+ *   %q2 = qco.z %arg0 : !qco.qubit -> !qco.qubit
+ *   qco.yield %q2 : !qco.qubit
+ * }
+ * ```
+ * is converted to
+ * ```mlir
+ * scf.index_switch %condition
+ * case 0 {
+ *   qc.x %q0 : !qc.qubit
+ * }
+ * default {
+ *   qc.z %q0 : !qc.qubit
+ * }
+ * ```
+ */
+struct ConvertQCOIndexSwitchOp final : OpConversionPattern<IndexSwitchOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(IndexSwitchOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter& rewriter) const override {
+    auto newOp =
+        scf::IndexSwitchOp::create(rewriter, op.getLoc(), {}, adaptor.getArg(),
+                                   adaptor.getCases(), op.getNumCases());
+
+    const auto oldRegions = op.getCaseRegions();
+    const auto newCaseRegions = newOp.getCaseRegions();
+    for (size_t i = 0; i < op.getNumCases(); ++i) {
+      inlineRegion(oldRegions[i], newCaseRegions[i], 0,
+                   adaptor.getOperands().drop_front(1), rewriter);
+    }
+
+    inlineRegion(op.getDefaultRegion(), newOp.getDefaultRegion(), 0,
+                 adaptor.getOperands().drop_front(1), rewriter);
+
+    // Replace the qco results with the input qc values except the condition
+    rewriter.replaceOp(op, adaptor.getOperands().drop_front(1));
+
+    return success();
+  }
+};
+
 /**
  * @brief Converts scf.yield with value semantics to scf.yield with memory
  * semantics for qubit values. This currently assumes no mixed types as yielded
@@ -1036,9 +1089,9 @@ struct QCOToQC final : impl::QCOToQCBase<QCOToQC> {
 
     patterns.add<ConvertQCOBarrierOp, ConvertQCOCtrlOp, ConvertQCOInvOp,
                  ConvertQCOPowOp, ConvertQCOYieldOp, ConvertQCOIfOp,
-                 ConvertQCOSCFWhileOp, ConvertQCOSCFConditionOp,
-                 ConvertQCOSCFYieldOp, ConvertQCOSCFForOp>(typeConverter,
-                                                           context);
+                 ConvertQCOIndexSwitchOp, ConvertQCOSCFWhileOp,
+                 ConvertQCOSCFConditionOp, ConvertQCOSCFYieldOp,
+                 ConvertQCOSCFForOp>(typeConverter, context);
 
     // Register operation conversion patterns that need state tracking
     patterns.add<ConvertQTensorExtractOp, ConvertQTensorAllocOp,