[mlir][scf] Add an IndexSwitchOp

[Mogball]

Summary

The scf.index_switch is a control-flow operation that branches to one of the
given regions based on the values of the argument and the cases. The
argument is always of type index.

Example:

mlir
%0 = scf.index_switch %arg0 -> i32
case 2 {
  %1 = arith.constant 10 : i32
  scf.yield %1 : i32
}
case 5 {
  %2 = arith.constant 20 : i32
  scf.yield %2 : i32
}
default {
  %3 = arith.constant 30 : i32
  scf.yield %3 : i32
}

[Mogball]

fix typo

[Mogball]

put the trait on the wrong op

[Mogball]

actually allow any integer arg

[Mogball]

rebase

[Mogball]

fix mis-merge

[jpienaar]

Nice, I could have sworn we discussed an op like this ... It's a bit far out of my cache. I think this as is is equivalent than a nested sequence of scf conditionals wrt type constraints, is that accurate?

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

1022

Why is this allowing signless ints too?

1060–1061

Nit: keep sorted

mlir/lib/Dialect/SCF/IR/SCF.cpp

3435

So in the attribute we don't have fixed type but require all to be same? Wouldn't a dense array attr represent that Bette?

mlir/test/Dialect/SCF/invalid.mlir

431

Ensure this is in order too once switched

[jpienaar]

https://discourse.llvm.org/t/rfc-add-std-switch-and-scf-switch/3090/15 that was it

[Mogball]

review comments: revert back to index + dense array.

Also implement RegionBranchOpInterface and just use that for verification

[Mogball]

I think this as is is equivalent than a nested sequence of scf conditionals wrt type constraints, is that accurate?

Yes. Now that you mention it, I actually remember that RFC too. I can add cf.switch as per that RFC as well since that's what I was planning to do next anyways.

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

1022

I had it fixed to index originally. I'll change it back.

[rriddle]

cf.switch already exists. I would also expect an update posted to that RFC to make sure that agreement has been reached as per how scf.switch should be structured before moving forward.

[Mogball]

cf.switch already exists.

Oh right.

I would also expect an update posted to that RFC to make sure that agreement has been reached as per how scf.switch should be structured before moving forward.

I noticed there was a hangup on what exactly scf.switch should encompass. I posted this patch to that thread. IMO more "involved" switch operations can always be added and called something different.

[Mogball]

implement getRegionInvocationBounds while I'm at it

[Mogball]

use a regular verifier to get better error messages

[jpienaar]

Nice, I think this makes sense. Keeping it directed & inline with the others here. I would love pattern matching, but that indeed feels like a different level. Perhaps give a workday or two for folks to chime in.

mlir/lib/Dialect/SCF/IR/SCF.cpp

3510

Ha nice.

[mehdi_amini]

In D136003#3861388, @Mogball wrote:

I noticed there was a hangup on what exactly scf.switch should encompass. I posted this patch to that thread. IMO more "involved" switch operations can always be added and called something different.

Shall we name this one scf.indexed_switch or something unambiguous if we anticipate more a possible general switch?

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

1003

Description seems terse, for example it does not say anything about the default case.

[Mogball]

rename to index_switch and expand description

[rriddle]

In D136003#3864009, @mehdi_amini wrote:

In D136003#3861388, @Mogball wrote:

I noticed there was a hangup on what exactly scf.switch should encompass. I posted this patch to that thread. IMO more "involved" switch operations can always be added and called something different.

Shall we name this one scf.indexed_switch or something unambiguous if we anticipate more a possible general switch?

I don't think index_switch is a great name here. It kind of awkwardly diverges from the naming of the other ops, e.g. we don't have index_for even though it's not a truly general for loop. Can we consider renaming this to just switch?

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

using BodyBuilderFn =
         function_ref<void(OpBuilder &, Location, Value, ValueRange)>;
 
     Value getInductionVar() { return getBody()->getArgument(0); }
     Block::BlockArgListType getRegionIterArgs() {
       return getBody()->getArguments().drop_front(getNumInductionVars());
     }
     /// Return the `index`-th region iteration argument.
     BlockArgument getRegionIterArg(unsigned index) {
-      assert(index < getNumRegionIterArgs() && 
+      assert(index < getNumRegionIterArgs() &&
         "expected an index less than the number of region iter args");
       return getBody()->getArguments().drop_front(getNumInductionVars())[index];
     }
     Operation::operand_range getIterOperands() {
       return getOperands().drop_front(getNumControlOperands());
     }
     MutableArrayRef<OpOperand> getIterOpOperands() {
       return
%matmul_and_pointwise:2 = scf.foreach_thread (%thread_id_1, %thread_id_2) in
       }
     }
     // Implicit synchronization point.
     // Sequential context.
     //
     ```
 
     Example with thread_dim_mapping attribute:
-    
     ```mlir
     //
     // Sequential context.
     //
     %matmul_and_pointwise:2 = scf.foreach_thread (%thread_id_1, %thread_id_2) in
         (%num_threads_1, %numthread_id_2) shared_outs(...)
       -> (tensor<?x?xT>, tensor<?xT>) {
       //
       // Parallel context, each thread with id = **(%thread_id_2, %thread_id_1)**
       // runs its version of the code.
       //
        scf.foreach_thread.perform_concurrently {
          ...
       }
     } { thread_dim_mapping = [1, 0] }
     // Implicit synchronization point.
     // Sequential context.
     //
     ```
 
     Example with privatized tensors:
-    
     ```mlir
     %t0 = ...
     %t1 = ...
     %r = scf.foreach_thread ... shared_outs(%o = t0) -> tensor<?xf32> {
       // %t0 and %t1 are privatized. %t0 is definitely copied for each thread
       // because the scf.foreach_thread op's %t0 use bufferizes to a memory
       // write. In the absence of other conflicts, %t1 is copied only if there
       // are uses of %t1 in the body that bufferize to a memory read and to a
let extraClassDeclaration = [{
     ::mlir::Value getThreadIndex(int64_t idx) {
       return getThreadIndices()[idx];
     }
 
     ::mlir::Block::BlockArgListType getRegionOutArgs() {
       return getBody()->getArguments().drop_front(getRank());
     }
 
-    /// Return the thread indices in the order specified by the 
-    /// thread_dim_mapping attribute. Return failure is 
+    /// Return the thread indices in the order specified by the
+    /// thread_dim_mapping attribute. Return failure is
     /// thread_dim_mapping is not a valid permutation.
     FailureOr<SmallVector<Value>> getPermutedThreadIndices();
 
     /// Return the number of threads in the order specified by the
     /// thread_dim_mapping attribute.
     /// Return failure is thread_dim_mapping is not a valid permutation.
     FailureOr<SmallVector<OpFoldResult>> getPermutedNumThreads(OpBuilder &b);
 
def WhileOp : SCF_Op<"while",
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// IndexSwitchOp
+//===----------------------------------------------------------------------===//
+
+def IndexSwitchOp : SCF_Op<"index_switch", [RecursiveMemoryEffects,
+    SingleBlockImplicitTerminator<"scf::YieldOp">,
+    DeclareOpInterfaceMethods<RegionBranchOpInterface,
+                              ["getRegionInvocationBounds"]>]> {
+  let summary = "switch-case operation on an index argument";
+  let description = [{
+    The `scf.index_switch` is a control-flow operation that branches to one of
+    the given regions based on the values of the argument and the cases. The
+    argument is always of type `index`.
+
+    The operation always has a "default" region and any number of case regions
+    denoted by integer constants. Control-flow transfers to the case region
+    whose constant value equals the value of the argument. If the argument does
+    not equal any of the case values, control-flow transfer to the "default"
+    region.
+
+    Example:
+
+    ```mlir
+    %0 = scf.index_switch %arg0 : index -> i32
+    case 2 {
+      %1 = arith.constant 10 : i32
+      scf.yield %1 : i32
+    }
+    case 5 {
+      %2 = arith.constant 20 : i32
+      scf.yield %2 : i32
+    }
+    default {
+      %3 = arith.constant 30 : i32
+      scf.yield %3 : i32
+    }
+    ```
+  }];
+
+  let arguments = (ins Index:$arg, DenseI64ArrayAttr:$cases);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region SizedRegion<1>:$defaultRegion,
+                        VariadicRegion<SizedRegion<1>>:$caseRegions);
+
+  let assemblyFormat = [{
+    $arg attr-dict (`->` type($results)^)?
+    custom<SwitchCases>($cases, $caseRegions) `\n`
+    `` `default` $defaultRegion
+  }];
+
+  let extraClassDeclaration = [{
+    /// Get the number of cases.
+    unsigned getNumCases();
+
+    /// Get the default region body.
+    Block &getDefaultBlock();
+
+    /// Get the body of a case region.
+    Block &getCaseBlock(unsigned idx);
+  }];
+
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 def YieldOp : SCF_Op<"yield", [Pure, ReturnLike, Terminator,
-                               ParentOneOf<["ExecuteRegionOp, ForOp",
-                                            "IfOp, ParallelOp, WhileOp"]>]> {
+    ParentOneOf<["ExecuteRegionOp, ForOp", "IfOp", "IndexSwitchOp",
+                 "ParallelOp", "WhileOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "scf.yield" yields an SSA value from the SCF dialect op region and
     terminates the regions. The semantics of how the values are yielded is
     defined by the parent operation.
     If "scf.yield" has any operands, the operands must match the parent
     operation's results.
     If the parent operation defines no values, then the "scf.yield" may be

mlir/lib/Dialect/SCF/IR/SCF.cpp

 void WhileOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
   results.add<RemoveLoopInvariantArgsFromBeforeBlock,
               RemoveLoopInvariantValueYielded, WhileConditionTruth,
               WhileCmpCond, WhileUnusedResult>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// IndexSwitchOp
+//===----------------------------------------------------------------------===//
+
+/// Parse the case regions and values.
+static ParseResult
+parseSwitchCases(OpAsmParser &p, DenseI64ArrayAttr &cases,
+                 SmallVectorImpl<std::unique_ptr<Region>> &caseRegions) {
+  SmallVector<int64_t> caseValues;
+  while (succeeded(p.parseOptionalKeyword("case"))) {
+    int64_t value;
+    Region &region =
+        *caseRegions.emplace_back(std::make_unique<Region>()).get();
+    if (p.parseInteger(value) || p.parseRegion(region, /*arguments=*/{}))
+      return failure();
+    caseValues.push_back(value);
+  }
+  cases = p.getBuilder().getDenseI64ArrayAttr(caseValues);
+  return success();
+}
+
+/// Print the case regions and values.
+static void printSwitchCases(OpAsmPrinter &p, Operation *op,
+                             DenseI64ArrayAttr cases, RegionRange caseRegions) {
+  for (auto [value, region] : llvm::zip(cases.asArrayRef(), caseRegions)) {
+    p.printNewline();
+    p << "case " << value << ' ';
+    p.printRegion(*region, /*printEntryBlockArgs=*/false);
+  }
+}
+
+LogicalResult scf::IndexSwitchOp::verify() {
+  if (getCases().size() != getCaseRegions().size()) {
+    return emitOpError("has ")
+           << getCaseRegions().size() << " case regions but "
+           << getCases().size() << " case values";
+  }
+
+  DenseSet<int64_t> valueSet;
+  for (int64_t value : getCases())
+    if (!valueSet.insert(value).second)
+      return emitOpError("has duplicate case value: ") << value;
+
+  auto verifyRegion = [&](Region &region, const Twine &name) -> LogicalResult {
+    auto yield = cast<YieldOp>(region.front().getTerminator());
+    if (yield.getNumOperands() != getNumResults()) {
+      return (emitOpError("expected each region to return ")
+              << getNumResults() << " values, but " << name << " returns "
+              << yield.getNumOperands())
+                 .attachNote(yield.getLoc())
+             << "see yield operation here";
+    }
+    for (auto [idx, result, operand] :
+         llvm::zip(llvm::seq<unsigned>(0, getNumResults()), getResultTypes(),
+                   yield.getOperandTypes())) {
+      if (result == operand)
+        continue;
+      return (emitOpError("expected result #")
+              << idx << " of each region to be " << result)
+                 .attachNote(yield.getLoc())
+             << name << " returns " << operand << " here";
+    }
+    return success();
+  };
+
+  if (failed(verifyRegion(getDefaultRegion(), "default region")))
+    return failure();
+  for (auto &[idx, caseRegion] : llvm::enumerate(getCaseRegions()))
+    if (failed(verifyRegion(caseRegion, "case region #" + Twine(idx))))
+      return failure();
+
+  return success();
+}
+
+unsigned scf::IndexSwitchOp::getNumCases() { return getCases().size(); }
+
+Block &scf::IndexSwitchOp::getDefaultBlock() {
+  return getDefaultRegion().front();
+}
+
+Block &scf::IndexSwitchOp::getCaseBlock(unsigned idx) {
+  assert(idx < getNumCases() && "case index out-of-bounds");
+  return getCaseRegions()[idx].front();
+}
+
+void IndexSwitchOp::getSuccessorRegions(
+    Optional<unsigned> index, ArrayRef<Attribute> operands,
+    SmallVectorImpl<RegionSuccessor> &successors) {
+  // All regions branch back to the parent op.
+  if (index) {
+    successors.emplace_back(getResults());
+    return;
+  }
+
+  // If a constant was not provided, all regions are possible successors.
+  auto operandValue = operands.front().dyn_cast_or_null<IntegerAttr>();
+  if (!operandValue) {
+    for (Region &caseRegion : getCaseRegions())
+      successors.emplace_back(&caseRegion);
+    successors.emplace_back(&getDefaultRegion());
+    return;
+  }
+
+  // Otherwise, try to find a case with a matching value. If not, the default
+  // region is the only successor.
+  for (auto [caseValue, caseRegion] : llvm::zip(getCases(), getCaseRegions())) {
+    if (caseValue == operandValue.getInt()) {
+      successors.emplace_back(&caseRegion);
+      return;
+    }
+  }
+  successors.emplace_back(&getDefaultRegion());
+}
+
+void IndexSwitchOp::getRegionInvocationBounds(
+    ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
+  auto operandValue = operands.front().dyn_cast_or_null<IntegerAttr>();
+  if (!operandValue) {
+    // All regions are invoked at most once.
+    bounds.append(getNumRegions(), InvocationBounds(/*lb=*/0, /*ub=*/1));
+    return;
+  }
+
+  unsigned liveIndex = getNumRegions() - 1;
+  auto it = llvm::find(getCases(), operandValue.getInt());
+  if (it != getCases().end())
+    liveIndex = std::distance(getCases().begin(), it);
+  for (unsigned i = 0, e = getNumRegions(); i < e; ++i)
+    bounds.emplace_back(/*lb=*/0, /*ub=*/i == liveIndex);
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SCF/IR/SCFOps.cpp.inc"

mlir/test/Dialect/SCF/invalid.mlir

func.func @parallel_invalid_yield(
   }
   return
 }
 
 // -----
 
 func.func @yield_invalid_parent_op() {
   "my.op"() ({
-   // expected-error@+1 {{'scf.yield' op expects parent op to be one of 'scf.execute_region, scf.for, scf.if, scf.parallel, scf.while'}}
+   // expected-error@+1 {{'scf.yield' op expects parent op to be one of 'scf.execute_region, scf.for, scf.if, scf.index_switch, scf.parallel, scf.while'}}
    scf.yield
   }) : () -> ()
   return
 }
 
 // -----
 
 func.func @while_parser_type_mismatch() {
%result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> (tensor<100xf32>) {
       scf.foreach_thread.perform_concurrently {
         tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] :
           tensor<1xf32> into tensor<100xf32>
         %0 = arith.constant 1: index
       }
   }
   return
 }
+
+// -----
+
+func.func @switch_wrong_case_count(%arg0: index) {
+  // expected-error @below {{'scf.index_switch' op has 0 case regions but 1 case values}}
+  "scf.index_switch"(%arg0) ({
+    scf.yield
+  }) {cases = array<i64: 1>} : (index) -> ()
+  return
+}
+
+// -----
+
+func.func @switch_duplicate_case(%arg0: index) {
+  // expected-error @below {{'scf.index_switch' op has duplicate case value: 0}}
+  scf.index_switch %arg0
+  case 0 {
+    scf.yield
+  }
+  case 0 {
+    scf.yield
+  }
+  default {
+    scf.yield
+  }
+  return
+}
+
+// -----
+
+func.func @switch_wrong_types(%arg0: index) {
+  // expected-error @below {{'scf.index_switch' op expected each region to return 0 values, but default region returns 1}}
+  scf.index_switch %arg0
+  default {
+    // expected-note @below {{see yield operation here}}
+    scf.yield %arg0 : index
+  }
+  return
+}
+
+// -----
+
+func.func @switch_wrong_types(%arg0: index, %arg1: i32) {
+  // expected-error @below {{'scf.index_switch' op expected result #0 of each region to be 'index'}}
+  scf.index_switch %arg0 -> index
+  case 0 {
+    // expected-note @below {{case region #0 returns 'i32' here}}
+    scf.yield %arg1 : i32
+  }
+  default {
+    scf.yield %arg0 : index
+  }
+  return
+}

mlir/test/Dialect/SCF/ops.mlir

func.func @elide_terminator() -> () {
   // CHECK-NEXT:  } {thread_dim_mapping = [42]}
   // CHECK-NEXT:  return
   scf.foreach_thread (%thread_idx) in (%num_threads) {
     scf.foreach_thread.perform_concurrently {
     }
   } {thread_dim_mapping = [42]}
   return
 }
+
+// CHECK-LABEL: @switch
+func.func @switch(%arg0: index) -> i32 {
+  // CHECK: %{{.*}} = scf.index_switch %arg0 -> i32
+  %0 = scf.index_switch %arg0 -> i32
+  // CHECK-NEXT: case 2 {
+  case 2 {
+    // CHECK-NEXT: arith.constant
+    %c10_i32 = arith.constant 10 : i32
+    // CHECK-NEXT: scf.yield %{{.*}} : i32
+    scf.yield %c10_i32 : i32
+    // CHECK-NEXT: }
+  }
+  // CHECK-NEXT: case 5 {
+  case 5 {
+    %c20_i32 = arith.constant 20 : i32
+    scf.yield %c20_i32 : i32
+  }
+  // CHECK: default {
+  default {
+    %c30_i32 = arith.constant 30 : i32
+    scf.yield %c30_i32 : i32
+  }
+
+  // CHECK: scf.index_switch %arg0
+  scf.index_switch %arg0
+  // CHECK-NEXT: default {
+  default {
+    scf.yield
+  }
+
+  return %0 : i32
+}