[MLIR] Add a switch operation to the standard dialect

[GMNGeoffrey]

Summary

This is similar to the definition of llvm.switch, providing
unstructured branch-based control flow. It differs from the LLVM
operation in that it accepts any signless integer (not only an i32),
takes no branch weights (the same as the Branch and CondBranch ops),
and has a slightly different syntax for the default case that includes
it in the list of cases with an explicit default keyword.

Also included are several canonicalizers.

See https://llvm.discourse.group/t/rfc-add-std-switch-and-scf-switch/3090

[GMNGeoffrey]

Remove leftover debugging stuff

[GMNGeoffrey]

• Remove extraneous include

[rriddle]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2044

typo siwtch

2060

Not sure I understand the justification for choosing i32 here.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2316

Avoid computing the size during the loop.

2317

Drop trivial braces here and on the for.

[bondhugula]

Looks great - some minor comments.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2053

Do you want to instead just:

switch %flag [
  default: ^bb1(%a : i32)
  42: ^bb1(%b : i32),
  43: ^bb3(%c : i32)
]

2060

Use i64 instead? You could have a case for a large value 2^64-1 and then default for everything.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2175

unsigned?

2256

Nit: ' '

[ftynse]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2060

Can't it be just any integer? Same for case values, which can become IntElementsAttr. Just have a verifier check that the types match.

[GMNGeoffrey]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2060

Per the RFC, I modeled this directly off of llvm.switch. This could indeed be any Integer, or really any type at all. How much do we want to expand it?

[ftynse]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2060

Sorry, I didn't see the RFC before. I would have asked why do we need a carbon copy of the LLVM dialect operation in the Standard dialect if we can mix dialects instead (LLVM no longer needs special types...) I am fine with any integer because I don't see how to make the verification work for an arbitrary type.

[GMNGeoffrey]

Are these clang-tidy naming errors legit? I think the names match the ones used

Also terrifying that my test for passthrough isn't failing given that I apparently forgot to add the passthrough pattern.... It would be really nice to be able to limit canonicalization to only patterns on certain ops or something

[GMNGeoffrey]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2060

Seems like arbitrary type should be verifiable in just the same way? Check that the type of the ElementsAttr elements are the same as the type of value. The fact that ElementsAttr is limited to tensor or vector places some restrictions on the type, although we could even expand this to be an Array instead (there was discussion of having something that general in scf). That "feels" too high level for standard to me (switch on my custom proto type!), but I'm not sure I have a good argument for why.

[ftynse]

In D99925#2671612, @GMNGeoffrey wrote:

Are these clang-tidy naming errors legit? I think the names match the ones used

MLIR uses (lower) camelBack for function names.

Also terrifying that my test for passthrough isn't failing given that I apparently forgot to add the passthrough pattern.... It would be really nice to be able to limit canonicalization to only patterns on certain ops or something

I think br has a canonicalizer itself that removes useless branches.

[ftynse]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2060

Standard ops (mostly) work on built-in types only, so I'd stick to that for consistency.

[GMNGeoffrey]

In D99925#2671661, @ftynse wrote:

In D99925#2671612, @GMNGeoffrey wrote:

Are these clang-tidy naming errors legit? I think the names match the ones used

MLIR uses (lower) camelBack for function names.

Oh right because the other canonicalizers are structs and these are just free functions.

[rriddle]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2060

I'd stick to just builtin integers for now.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2214

Can you use succeeded/failed for these optional parse functions?

2217–2218

This would mean that no error message is emitted when the integer isn't there, but should be. Can you split these two cases and give a proper error message?

2225

Can you use succeeded/failed for these optional parse functions?

2269

Can you cache case_values() and caseDestinations()?

2353

Same comment on recomputing size here and below.

2361

Please avoid going through IntegerAttr, use APInt instead (assuming the cases can be any integer).

2550

Drop the llvm:: here.

[GMNGeoffrey]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2053

Do others have thoughts on the assembly format here? I was again using the same format as llvm.switch, but I do think that perhaps this way with a keyword is clearer.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2214

Will clean up the print/parse functions. I just copied them from llvm :-D

2269

Are these accessors expensive? I had assumed they're just giving references to existing vectors.

2316

Is size computation expensive for these? Again I assumed this would be a constant access to some offset.

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2269

The below applies to operations in general:

They can be, because the locations of the operands have to be computed. It isn't just a index into a field of Operation, it is also an address calculation for the trailing array. That also isn't factoring in when multiple operand groups of the operation are variadic, it has to compute the range of this operand group (which may involve accessing attributes from the operation).

2316

In this case, size() is likely "free" (case_values is not). In the broader LLVM style, we take the approach of never recomputing the size unless necessary for the algorithm being performed.

https://llvm.org/docs/CodingStandards.html#don-t-evaluate-end-every-time-through-a-loop

[GMNGeoffrey]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2269

Got it. Thanks for explaining :-)

2316

Got it, thanks for the pointer!

[GMNGeoffrey]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2175

The function being called here is getI32VectorAttr(ArrayRef<int32_t> values);, so I think int32_t is the correct type

2214

How do we feel about optional trailing commas when parsing? In general, I'm pro, but it means there are non-whitespace differences between different correct ways to write this, which I'm not sure yet exists in MLIR

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2214

I'm pro in programming language and other things like JSON (which does not have them and I hate it for this!).
I'm not sure it makes sense for an IR format though?

(I'm sure many custom parser are already accepting non-whitespace difference though, I think I wrote such code myself on occasions)

[GMNGeoffrey]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2214

Yes that's how I feel as well. I write mlir enough that it feels like a programming language sometimes :-D One of the reasons I ask is that it can make the parser potentially simpler if I'm looking for the right square bracket to end the loop rather than it being dependent on commas:

c++
  while(failed(parser.parseOptionalRSquare())) {
    if (failed(parser.parseInteger(value)))
      return failure();
    values.push_back(value);

    Block *destination;
    SmallVector<OpAsmParser::OperandType> operands;
    if (failed(parser.parseColon()) ||
        failed(parser.parseSuccessor(destination)))
      return failure();
    if (succeeded(parser.parseOptionalLParen())) {
      if (failed(parser.parseRegionArgumentList(operands)) ||
          failed(parser.parseColonTypeList(caseOperandTypes)) ||
          failed(parser.parseRParen()))
        return failure();
    }
    caseDestinations.push_back(destination);
    caseOperands.append(operands.begin(), operands.end());
    offsets.push_back(offset);
    offset += operands.size();
    if (failed(parser.parseOptionalComma())) {
      if(failed(parser.parseRSquare()))
        return failure();
      break;
    }
  }

Maybe makes sense to see how people feel about Uday's suggestion to move the default case within the brackets as well. One other wrinkle with the format is that if we allow any integer I think we have to include the type. Something like this is now what I'm thinking:

mlir
switch %flag, [
  default: ^bb1(%a : f32),
  42: ^bb2(%b : f32)
] : i32

Or maybe it's better to leave the type by the flag:

mlir
switch %flag : i32, [
  default: ^bb1(%a : f32),
  42: ^bb2(%b : f32)
]

[GMNGeoffrey]

Respond to review comments

[GMNGeoffrey]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2053

I switched to this format. I like it better.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2256

Sorry I don't understand. This is printing two spaces, so can't be a char, if that's what you mean. Let me know if this comment still applies to the updated printer.

2317

Done. Though I still hate this :-D

[GMNGeoffrey]

Fix function naming

[bondhugula]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2053

Looks great (the syntax I found in the test case). This snippet needs to be updated.

mlir/test/Dialect/Standard/parser.mlir

41

The syntax in the op doc hasn't been updated to reflect this I think.

[GMNGeoffrey]

• Add unknown op predecessor to tests to block other canonicalizations
• Fix passthrough canonicalization
• Clean up test cases

mlir/test/Dialect/Standard/parser.mlir

41

Thanks for catching :-)

[rriddle]

LGTM in general.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2044

The description here is out-of-date.

2092–2094

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2345–2346

nit: Can you spell out auto here?

2346

Hoist this to the top of the function to avoid calling case_values twice?

2350

This assumes that the values are int64_t, I think will assert if not.

2550

nit: Spell out auto here.

2609

nit: Drop trivial braces here.

[GMNGeoffrey]

Address review comments

[GMNGeoffrey]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2044

Thanks for catching. Fixed.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2350

Yikes! Thanks for catching. What's the right way to test this? It doesn't seem right that every op (especially those outside of core dialects) should be tested as part of the sccp pass test.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let extraClassDeclaration = extraBaseClassDeclaration # [{
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
+
+//===----------------------------------------------------------------------===//
+// SwitchOp
+//===----------------------------------------------------------------------===//
+
+def SwitchOp : Std_Op<"switch",
+    [AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>,
+     NoSideEffect, Terminator]> {
+  let summary = "switch operation";
+  let description = [{
+    The `switch` terminator operation represents a switch on a signless integer
+    value. If the flag matches one of the specified cases, then the
+    corresponding destination is jumped to. If the flag does not match any of
+    the cases, the default destination is jumped to. The count and types of
+    operands must align with the arguments in the corresponding target blocks.
+
+    Example:
+
+    ```mlir
+    switch %flag : i32, [
+      default: ^bb1(%a : i32),
+      42: ^bb1(%b : i32),
+      43: ^bb3(%c : i32)
+    ]
+    ```
+  }];
+
+  let arguments = (ins AnyInteger:$flag,
+                       Variadic<AnyType>:$defaultOperands,
+                       Variadic<AnyType>:$caseOperands,
+                       OptionalAttr<AnyIntElementsAttr>:$case_values,
+                       OptionalAttr<I32ElementsAttr>:$case_operand_offsets);
+  let successors = (successor
+                        AnySuccessor:$defaultDestination,
+                        VariadicSuccessor<AnySuccessor>:$caseDestinations);
+  let builders = [
+    OpBuilder<(ins "Value":$flag,
+      "Block *":$defaultDestination,
+      "ValueRange":$defaultOperands,
+      CArg<"ArrayRef<APInt>", "{}">:$caseValues,
+      CArg<"BlockRange", "{}">:$caseDestinations,
+      CArg<"ArrayRef<ValueRange>", "{}">:$caseOperands)>,
+    OpBuilder<(ins "Value":$flag,
+      "Block *":$defaultDestination,
+      "ValueRange":$defaultOperands,
+      CArg<"ArrayRef<int32_t>", "{}">:$caseValues,
+      CArg<"BlockRange", "{}">:$caseDestinations,
+      CArg<"ArrayRef<ValueRange>", "{}">:$caseOperands)>,
+    OpBuilder<(ins "Value":$flag,
+      "Block *":$defaultDestination,
+      "ValueRange":$defaultOperands,
+      CArg<"DenseIntElementsAttr", "{}">:$caseValues,
+      CArg<"BlockRange", "{}">:$caseDestinations,
+      CArg<"ArrayRef<ValueRange>", "{}">:$caseOperands)>
+  ];
+
+  let assemblyFormat = [{
+    $flag `:` type($flag) `,` `[` `\n`
+      custom<SwitchOpCases>(ref(type($flag)),$defaultDestination,
+                            $defaultOperands,
+                            type($defaultOperands),
+                            $case_values,
+                            $caseDestinations,
+                            $caseOperands,
+                            type($caseOperands),
+                            $case_operand_offsets)
+   `]`
+    attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    /// Return the operands for the case destination block at the given index.
+    OperandRange getCaseOperands(unsigned index);
+
+    /// Return a mutable range of operands for the case destination block at the
+    /// given index.
+    MutableOperandRange getCaseOperandsMutable(unsigned index);
+  }];
+
+  let hasCanonicalizer = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // TruncateIOp
 //===----------------------------------------------------------------------===//
 
 def TruncateIOp : Std_Op<"trunci", [NoSideEffect,
     DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
     ElementwiseMappable.traits> {
   let summary = "integer truncation operation";

mlir/include/mlir/IR/OpBase.td

 
 def IndexElementsAttr
     : IntElementsAttrBase<CPred<[{$_self.cast<::mlir::DenseIntElementsAttr>()
                                       .getType()
                                       .getElementType()
                                       .isIndex()}]>,
                           "index elements attribute">;
 
-class AnyIntElementsAttr<int width> : IntElementsAttrBase<
+def AnyIntElementsAttr : IntElementsAttrBase<CPred<"true">, "integer elements attribute">;
+
+class IntElementsAttrOf<int width> : IntElementsAttrBase<
   CPred<"$_self.cast<::mlir::DenseIntElementsAttr>().getType()."
         "getElementType().isInteger(" # width # ")">,
   width # "-bit integer elements attribute">;
 
-def AnyI32ElementsAttr : AnyIntElementsAttr<32>;
-def AnyI64ElementsAttr : AnyIntElementsAttr<64>;
+def AnyI32ElementsAttr : IntElementsAttrOf<32>;
+def AnyI64ElementsAttr : IntElementsAttrOf<64>;
 
 class SignlessIntElementsAttr<int width> : IntElementsAttrBase<
   CPred<"$_self.cast<::mlir::DenseIntElementsAttr>().getType()."
         "getElementType().isSignlessInteger(" # width # ")">,
   width # "-bit signless integer elements attribute"> {
 
   // Note that this is only constructing scalar elements attribute.
   let constBuilderCall = "::mlir::DenseElementsAttr::get("

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 
 //===----------------------------------------------------------------------===//
 // BranchOp
 //===----------------------------------------------------------------------===//
 
 /// Given a successor, try to collapse it to a new destination if it only
 /// contains a passthrough unconditional branch. If the successor is
 /// collapsable, `successor` and `successorOperands` are updated to reference
-/// the new destination and values. `argStorage` is an optional storage to use
-/// if operands to the collapsed successor need to be remapped.
+/// the new destination and values. `argStorage` is used as storage if operands
+/// to the collapsed successor need to be remapped. It must outlive uses of
+/// successorOperands.
 static LogicalResult collapseBranch(Block *&successor,
                                     ValueRange &successorOperands,
                                     SmallVectorImpl<Value> &argStorage) {
   // Check that the successor only contains a unconditional branch.
   if (std::next(successor->begin()) != successor->end())
     return failure();
   // Check that the terminator is an unconditional branch.
   BranchOp successorBranch = dyn_cast<BranchOp>(successor->getTerminator());

 
 void SubTensorInsertOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                     MLIRContext *context) {
   results.add<SubTensorInsertOpConstantArgumentFolder,
               SubTensorInsertOpCastFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// SwitchOp
+//===----------------------------------------------------------------------===//
+
+void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
+                     Block *defaultDestination, ValueRange defaultOperands,
+                     DenseIntElementsAttr caseValues,
+                     BlockRange caseDestinations,
+                     ArrayRef<ValueRange> caseOperands) {
+  SmallVector<Value> flattenedCaseOperands;
+  SmallVector<int32_t> caseOperandOffsets;
+  int32_t offset = 0;
+  for (ValueRange operands : caseOperands) {
+    flattenedCaseOperands.append(operands.begin(), operands.end());
+    caseOperandOffsets.push_back(offset);
+    offset += operands.size();
+  }
+  DenseIntElementsAttr caseOperandOffsetsAttr;
+  if (!caseOperandOffsets.empty())
+    caseOperandOffsetsAttr = builder.getI32VectorAttr(caseOperandOffsets);
+
+  build(builder, result, value, defaultOperands, flattenedCaseOperands,
+        caseValues, caseOperandOffsetsAttr, defaultDestination,
+        caseDestinations);
+}
+
+void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
+                     Block *defaultDestination, ValueRange defaultOperands,
+                     ArrayRef<APInt> caseValues, BlockRange caseDestinations,
+                     ArrayRef<ValueRange> caseOperands) {
+  DenseIntElementsAttr caseValuesAttr;
+  if (!caseValues.empty()) {
+    ShapedType caseValueType = VectorType::get(
+        static_cast<int64_t>(caseValues.size()), value.getType());
+    caseValuesAttr = DenseIntElementsAttr::get(caseValueType, caseValues);
+  }
+  build(builder, result, value, defaultDestination, defaultOperands,
+        caseValuesAttr, caseDestinations, caseOperands);
+}
+
+/// <cases> ::= `default` `:` bb-id (`(` ssa-use-and-type-list `)`)?
+///             ( `,` integer `:` bb-id (`(` ssa-use-and-type-list `)`)? )*
+static ParseResult
+parseSwitchOpCases(OpAsmParser &parser, Type &flagType,
+                   Block *&defaultDestination,
+                   SmallVectorImpl<OpAsmParser::OperandType> &defaultOperands,
+                   SmallVectorImpl<Type> &defaultOperandTypes,
+                   DenseIntElementsAttr &caseValues,
+                   SmallVectorImpl<Block *> &caseDestinations,
+                   SmallVectorImpl<OpAsmParser::OperandType> &caseOperands,
+                   SmallVectorImpl<Type> &caseOperandTypes,
+                   DenseIntElementsAttr &caseOperandOffsets) {
+
+  if (failed(parser.parseKeyword("default")) || failed(parser.parseColon()) ||
+      failed(parser.parseSuccessor(defaultDestination)))
+    return failure();
+  if (succeeded(parser.parseOptionalLParen())) {
+    if (failed(parser.parseRegionArgumentList(defaultOperands)) ||
+        failed(parser.parseColonTypeList(defaultOperandTypes)) ||
+        failed(parser.parseRParen()))
+      return failure();
+  }
+
+  SmallVector<APInt> values;
+  SmallVector<int32_t> offsets;
+  unsigned bitWidth = flagType.getIntOrFloatBitWidth();
+  int64_t offset = 0;
+  while (succeeded(parser.parseOptionalComma())) {
+    int64_t value = 0;
+    if (failed(parser.parseInteger(value)))
+      return failure();
+    values.push_back(APInt(bitWidth, value));
+
+    Block *destination;
+    SmallVector<OpAsmParser::OperandType> operands;
+    if (failed(parser.parseColon()) ||
+        failed(parser.parseSuccessor(destination)))
+      return failure();
+    if (succeeded(parser.parseOptionalLParen())) {
+      if (failed(parser.parseRegionArgumentList(operands)) ||
+          failed(parser.parseColonTypeList(caseOperandTypes)) ||
+          failed(parser.parseRParen()))
+        return failure();
+    }
+    caseDestinations.push_back(destination);
+    caseOperands.append(operands.begin(), operands.end());
+    offsets.push_back(offset);
+    offset += operands.size();
+  }
+
+  if (values.empty())
+    return success();
+
+  Builder &builder = parser.getBuilder();
+  ShapedType caseValueType =
+      VectorType::get(static_cast<int64_t>(values.size()), flagType);
+  caseValues = DenseIntElementsAttr::get(caseValueType, values);
+  caseOperandOffsets = builder.getI32VectorAttr(offsets);
+
+  return success();
+}
+
+static void printSwitchOpCases(
+    OpAsmPrinter &p, SwitchOp op, Type flagType, Block *defaultDestination,
+    OperandRange defaultOperands, TypeRange defaultOperandTypes,
+    DenseIntElementsAttr caseValues, SuccessorRange caseDestinations,
+    OperandRange caseOperands, TypeRange caseOperandTypes,
+    ElementsAttr caseOperandOffsets) {
+  p << "  default: ";
+  p.printSuccessorAndUseList(defaultDestination, defaultOperands);
+
+  if (!caseValues)
+    return;
+
+  for (int64_t i = 0, size = caseValues.size(); i < size; ++i) {
+    p << ',';
+    p.printNewline();
+    p << "  ";
+    p << caseValues.getValue<APInt>(i).getLimitedValue();
+    p << ": ";
+    p.printSuccessorAndUseList(caseDestinations[i], op.getCaseOperands(i));
+  }
+  p.printNewline();
+}
+
+static LogicalResult verify(SwitchOp op) {
+  auto caseValues = op.case_values();
+  auto caseDestinations = op.caseDestinations();
+
+  if (!caseValues && caseDestinations.empty())
+    return success();
+
+  Type flagType = op.flag().getType();
+  Type caseValueType = caseValues->getType().getElementType();
+  if (caseValueType != flagType)
+    return op.emitOpError()
+           << "'flag' type (" << flagType << ") should match case value type ("
+           << caseValueType << ")";
+
+  if (caseValues &&
+      caseValues->size() != static_cast<int64_t>(caseDestinations.size()))
+    return op.emitOpError() << "number of case values (" << caseValues->size()
+                            << ") should match number of "
+                               "case destinations ("
+                            << caseDestinations.size() << ")";
+  return success();
+}
+
+OperandRange SwitchOp::getCaseOperands(unsigned index) {
+  return getCaseOperandsMutable(index);
+}
+
+MutableOperandRange SwitchOp::getCaseOperandsMutable(unsigned index) {
+  MutableOperandRange caseOperands = caseOperandsMutable();
+  if (!case_operand_offsets()) {
+    assert(caseOperands.size() == 0 &&
+           "non-empty case operands must have offsets");
+    return caseOperands;
+  }
+
+  ElementsAttr offsets = case_operand_offsets().getValue();
+  assert(index < offsets.size() && "invalid case operand offset index");
+
+  int64_t begin = offsets.getValue(index).cast<IntegerAttr>().getInt();
+  int64_t end = index + 1 == offsets.size()
+                    ? caseOperands.size()
+                    : offsets.getValue(index + 1).cast<IntegerAttr>().getInt();
+  return caseOperandsMutable().slice(begin, end - begin);
+}
+
+Optional<MutableOperandRange>
+SwitchOp::getMutableSuccessorOperands(unsigned index) {
+  assert(index < getNumSuccessors() && "invalid successor index");
+  return index == 0 ? defaultOperandsMutable()
+                    : getCaseOperandsMutable(index - 1);
+}
+
+Block *SwitchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {
+  Optional<DenseIntElementsAttr> caseValues = case_values();
+
+  if (!caseValues)
+    return defaultDestination();
+
+  SuccessorRange caseDests = caseDestinations();
+  if (auto value = operands.front().dyn_cast_or_null<IntegerAttr>()) {
+    for (int64_t i = 0, size = case_values()->size(); i < size; ++i)
+      if (value == caseValues->getValue<IntegerAttr>(i))
+        return caseDests[i];
+    return defaultDestination();
+  }
+  return nullptr;
+}
+
+/// switch %flag : i32, [
+///   default:  ^bb1
+/// ]
+///  -> br ^bb1
+static LogicalResult simplifySwitchWithOnlyDefault(SwitchOp op,
+                                                   PatternRewriter &rewriter) {
+  if (!op.caseDestinations().empty())
+    return failure();
+
+  rewriter.replaceOpWithNewOp<BranchOp>(op, op.defaultDestination(),
+                                        op.defaultOperands());
+  return success();
+}
+
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb1,
+///   43: ^bb2
+/// ]
+/// ->
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   43: ^bb2
+/// ]
+static LogicalResult
+dropSwitchCasesThatMatchDefault(SwitchOp op, PatternRewriter &rewriter) {
+  SmallVector<Block *> newCaseDestinations;
+  SmallVector<ValueRange> newCaseOperands;
+  SmallVector<APInt> newCaseValues;
+  bool requiresChange = false;
+  auto caseValues = op.case_values();
+  auto caseDests = op.caseDestinations();
+
+  for (int64_t i = 0, size = caseValues->size(); i < size; ++i) {
+    if (caseDests[i] == op.defaultDestination() &&
+        op.getCaseOperands(i) == op.defaultOperands()) {
+      requiresChange = true;
+      continue;
+    }
+    newCaseDestinations.push_back(caseDests[i]);
+    newCaseOperands.push_back(op.getCaseOperands(i));
+    newCaseValues.push_back(caseValues->getValue<APInt>(i));
+  }
+
+  if (!requiresChange)
+    return failure();
+
+  rewriter.replaceOpWithNewOp<SwitchOp>(op, op.flag(), op.defaultDestination(),
+                                        op.defaultOperands(), newCaseValues,
+                                        newCaseDestinations, newCaseOperands);
+  return success();
+}
+
+/// Helper for folding a switch with a constant value.
+/// switch %c_42 : i32, [
+///   default: ^bb1 ,
+///   42: ^bb2,
+///   43: ^bb3
+/// ]
+/// -> br ^bb2
+static void foldSwitch(SwitchOp op, PatternRewriter &rewriter,
+                       APInt caseValue) {
+  auto caseValues = op.case_values();
+  for (int64_t i = 0, size = caseValues->size(); i < size; ++i) {
+    if (caseValues->getValue<APInt>(i) == caseValue) {
+      rewriter.replaceOpWithNewOp<BranchOp>(op, op.caseDestinations()[i],
+                                            op.getCaseOperands(i));
+      return;
+    }
+  }
+  rewriter.replaceOpWithNewOp<BranchOp>(op, op.defaultDestination(),
+                                        op.defaultOperands());
+}
+
+/// switch %c_42 : i32, [
+///   default: ^bb1,
+///   42: ^bb2,
+///   43: ^bb3
+/// ]
+/// -> br ^bb2
+static LogicalResult simplifyConstSwitchValue(SwitchOp op,
+                                              PatternRewriter &rewriter) {
+  APInt caseValue;
+  if (!matchPattern(op.flag(), m_ConstantInt(&caseValue)))
+    return failure();
+
+  foldSwitch(op, rewriter, caseValue);
+  return success();
+}
+
+/// switch %c_42 : i32, [
+///   default: ^bb1,
+///   42: ^bb2,
+/// ]
+/// ^bb2:
+///   br ^bb3
+/// ->
+/// switch %c_42 : i32, [
+///   default: ^bb1,
+///   42: ^bb3,
+/// ]
+static LogicalResult simplifyPassThroughSwitch(SwitchOp op,
+                                               PatternRewriter &rewriter) {
+
+  SmallVector<Block *> newCaseDests;
+  SmallVector<ValueRange> newCaseOperands;
+  SmallVector<SmallVector<Value>> argStorage;
+  auto caseValues = op.case_values();
+  auto caseDests = op.caseDestinations();
+  bool requiresChange = false;
+  for (int64_t i = 0, size = caseValues->size(); i < size; ++i) {
+    Block *caseDest = caseDests[i];
+    ValueRange caseOperands = op.getCaseOperands(i);
+    argStorage.emplace_back();
+    if (succeeded(collapseBranch(caseDest, caseOperands, argStorage.back())))
+      requiresChange = true;
+
+    newCaseDests.push_back(caseDest);
+    newCaseOperands.push_back(caseOperands);
+  }
+
+  Block *defaultDest = op.defaultDestination();
+  ValueRange defaultOperands = op.defaultOperands();
+  argStorage.emplace_back();
+
+  if (succeeded(
+          collapseBranch(defaultDest, defaultOperands, argStorage.back())))
+    requiresChange = true;
+
+  if (!requiresChange)
+    return failure();
+
+  rewriter.replaceOpWithNewOp<SwitchOp>(op, op.flag(), defaultDest,
+                                        defaultOperands, caseValues.getValue(),
+                                        newCaseDests, newCaseOperands);
+  return success();
+}
+
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb2,
+/// ]
+/// ^bb2:
+///   switch %flag : i32, [
+///     default: ^bb3,
+///     42: ^bb4
+///   ]
+/// ->
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb2,
+/// ]
+/// ^bb2:
+///   br ^bb4
+///
+///  and
+///
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb2,
+/// ]
+/// ^bb2:
+///   switch %flag : i32, [
+///     default: ^bb3,
+///     43: ^bb4
+///   ]
+/// ->
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb2,
+/// ]
+/// ^bb2:
+///   br ^bb3
+static LogicalResult
+simplifySwitchFromSwitchOnSameCondition(SwitchOp op,
+                                        PatternRewriter &rewriter) {
+  // Check that we have a single distinct predecessor.
+  Block *currentBlock = op->getBlock();
+  Block *predecessor = currentBlock->getSinglePredecessor();
+  if (!predecessor)
+    return failure();
+
+  // Check that the predecessor terminates with a switch branch to this block
+  // and that it branches on the same condition and that this branch isn't the
+  // default destination.
+  auto predSwitch = dyn_cast<SwitchOp>(predecessor->getTerminator());
+  if (!predSwitch || op.flag() != predSwitch.flag() ||
+      predSwitch.defaultDestination() == currentBlock)
+    return failure();
+
+  // Fold this switch to an unconditional branch.
+  APInt caseValue;
+  bool isDefault = true;
+  SuccessorRange predDests = predSwitch.caseDestinations();
+  Optional<DenseIntElementsAttr> predCaseValues = predSwitch.case_values();
+  for (int64_t i = 0, size = predCaseValues->size(); i < size; ++i) {
+    if (currentBlock == predDests[i]) {
+      caseValue = predCaseValues->getValue<APInt>(i);
+      isDefault = false;
+      break;
+    }
+  }
+  if (isDefault)
+    rewriter.replaceOpWithNewOp<BranchOp>(op, op.defaultDestination(),
+                                          op.defaultOperands());
+  else
+    foldSwitch(op, rewriter, caseValue);
+  return success();
+}
+
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb2
+/// ]
+/// ^bb1:
+///   switch %flag : i32, [
+///     default: ^bb3,
+///     42: ^bb4,
+///     43: ^bb5
+///   ]
+/// ->
+/// switch %flag : i32, [
+///   default: ^bb1,
+///   42: ^bb2,
+/// ]
+/// ^bb1:
+///   switch %flag : i32, [
+///     default: ^bb3,
+///     43: ^bb5
+///   ]
+static LogicalResult
+simplifySwitchFromDefaultSwitchOnSameCondition(SwitchOp op,
+                                               PatternRewriter &rewriter) {
+  // Check that we have a single distinct predecessor.
+  Block *currentBlock = op->getBlock();
+  Block *predecessor = currentBlock->getSinglePredecessor();
+  if (!predecessor)
+    return failure();
+
+  // Check that the predecessor terminates with a switch branch to this block
+  // and that it branches on the same condition and that this branch is the
+  // default destination.
+  auto predSwitch = dyn_cast<SwitchOp>(predecessor->getTerminator());
+  if (!predSwitch || op.flag() != predSwitch.flag() ||
+      predSwitch.defaultDestination() != currentBlock)
+    return failure();
+
+  // Delete case values that are not possible here.
+  DenseSet<APInt> caseValuesToRemove;
+  auto predDests = predSwitch.caseDestinations();
+  auto predCaseValues = predSwitch.case_values();
+  for (int64_t i = 0, size = predCaseValues->size(); i < size; ++i)
+    if (currentBlock != predDests[i])
+      caseValuesToRemove.insert(predCaseValues->getValue<APInt>(i));
+
+  SmallVector<Block *> newCaseDestinations;
+  SmallVector<ValueRange> newCaseOperands;
+  SmallVector<APInt> newCaseValues;
+  bool requiresChange = false;
+
+  auto caseValues = op.case_values();
+  auto caseDests = op.caseDestinations();
+  for (int64_t i = 0, size = caseValues->size(); i < size; ++i) {
+    if (caseValuesToRemove.contains(caseValues->getValue<APInt>(i))) {
+      requiresChange = true;
+      continue;
+    }
+    newCaseDestinations.push_back(caseDests[i]);
+    newCaseOperands.push_back(op.getCaseOperands(i));
+    newCaseValues.push_back(caseValues->getValue<APInt>(i));
+  }
+
+  if (!requiresChange)
+    return failure();
+
+  rewriter.replaceOpWithNewOp<SwitchOp>(op, op.flag(), op.defaultDestination(),
+                                        op.defaultOperands(), newCaseValues,
+                                        newCaseDestinations, newCaseOperands);
+  return success();
+}
+
+void SwitchOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                           MLIRContext *context) {
+  results.add(&simplifySwitchWithOnlyDefault)
+      .add(&dropSwitchCasesThatMatchDefault)
+      .add(&simplifyConstSwitchValue)
+      .add(&simplifyPassThroughSwitch)
+      .add(&simplifySwitchFromSwitchOnSameCondition)
+      .add(&simplifySwitchFromDefaultSwitchOnSameCondition);
+}
+
+//===----------------------------------------------------------------------===//
 // TruncateIOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(TruncateIOp op) {
   auto srcType = getElementTypeOrSelf(op.getOperand().getType());
   auto dstType = getElementTypeOrSelf(op.getType());
 
   if (srcType.isa<IndexType>())

mlir/test/Dialect/Standard/canonicalize-cf.mlir

-// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline='func(canonicalize)' -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline='func(canonicalize)' -split-input-file | FileCheck --dump-input-context 20 %s
 
 /// Test the folding of BranchOp.
 
 // CHECK-LABEL: func @br_folding(
 func @br_folding() -> i32 {
   // CHECK-NEXT: %[[CST:.*]] = constant 0 : i32
   // CHECK-NEXT: return %[[CST]] : i32
   %c0_i32 = constant 0 : i32
^bb1:
   // Successors can't be collapsed if they contain other operations.
   "foo.op"() : () -> ()
   br ^bb2
 
 ^bb2:
   return
 }
 
+
+/// Test the folding of SwitchOp
+
+// CHECK-LABEL: func @switch_only_default(
+// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+func @switch_only_default(%flag : i32, %caseOperand0 : f32) {
+  // add predecessors for all blocks to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2] : () -> ()
+  ^bb1:
+    // CHECK-NOT: switch
+    // CHECK: br ^[[BB2:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
+    switch %flag : i32, [
+      default: ^bb2(%caseOperand0 : f32)
+    ]
+  // CHECK: ^[[BB2]]({{.*}}):
+  ^bb2(%bb2Arg : f32):
+    // CHECK-NEXT: "foo.bb2Terminator"
+    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
+}
+
+
+// CHECK-LABEL: func @switch_case_matching_default(
+// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
+func @switch_case_matching_default(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32) {
+  // add predecessors for all blocks to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2, ^bb3] : () -> ()
+  ^bb1:
+    // CHECK: switch %[[FLAG]]
+    // CHECK-NEXT:   default: ^[[BB1:.+]](%[[CASE_OPERAND_0]] : f32)
+    // CHECK-NEXT:   10: ^[[BB2:.+]](%[[CASE_OPERAND_1]] : f32)
+    // CHECK-NEXT: ]
+    switch %flag : i32, [
+      default: ^bb2(%caseOperand0 : f32),
+      42: ^bb2(%caseOperand0 : f32),
+      10: ^bb3(%caseOperand1 : f32),
+      17: ^bb2(%caseOperand0 : f32)
+    ]
+  ^bb2(%bb2Arg : f32):
+    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
+  ^bb3(%bb3Arg : f32):
+    "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
+}
+
+
+// CHECK-LABEL: func @switch_on_const_no_match(
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
+func @switch_on_const_no_match(%caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
+  // add predecessors for all blocks to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4] : () -> ()
+  ^bb1:
+    // CHECK-NOT: switch
+    // CHECK: br ^[[BB2:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
+    %c0_i32 = constant 0 : i32
+    switch %c0_i32 : i32, [
+      default: ^bb2(%caseOperand0 : f32),
+      -1: ^bb3(%caseOperand1 : f32),
+      1: ^bb4(%caseOperand2 : f32)
+    ]
+  // CHECK: ^[[BB2]]({{.*}}):
+  // CHECK-NEXT: "foo.bb2Terminator"
+  ^bb2(%bb2Arg : f32):
+    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
+  ^bb3(%bb3Arg : f32):
+    "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
+  ^bb4(%bb4Arg : f32):
+    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
+}
+
+// CHECK-LABEL: func @switch_on_const_with_match(
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
+func @switch_on_const_with_match(%caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
+  // add predecessors for all blocks to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4] : () -> ()
+  ^bb1:
+    // CHECK-NOT: switch
+    // CHECK: br ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
+    %c0_i32 = constant 1 : i32
+    switch %c0_i32 : i32, [
+      default: ^bb2(%caseOperand0 : f32),
+      -1: ^bb3(%caseOperand1 : f32),
+      1: ^bb4(%caseOperand2 : f32)
+    ]
+  ^bb2(%bb2Arg : f32):
+    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
+  ^bb3(%bb3Arg : f32):
+    "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
+  // CHECK: ^[[BB4]]({{.*}}):
+  // CHECK-NEXT: "foo.bb4Terminator"
+  ^bb4(%bb4Arg : f32):
+    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
+}
+
+// CHECK-LABEL: func @switch_passthrough(
+// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_3:[a-zA-Z0-9_]+]]
+func @switch_passthrough(%flag : i32,
+                         %caseOperand0 : f32,
+                         %caseOperand1 : f32,
+                         %caseOperand2 : f32,
+                         %caseOperand3 : f32) {
+  // add predecessors for all blocks to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4, ^bb5, ^bb6] : () -> ()
+
+  ^bb1:
+  //      CHECK: switch %[[FLAG]]
+  // CHECK-NEXT:   default: ^[[BB5:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
+  // CHECK-NEXT:   43: ^[[BB6:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_1]]
+  // CHECK-NEXT:   44: ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
+  // CHECK-NEXT: ]
+    switch %flag : i32, [
+      default: ^bb2(%caseOperand0 : f32),
+      43: ^bb3(%caseOperand1 : f32),
+      44: ^bb4(%caseOperand2 : f32)
+    ]
+  ^bb2(%bb2Arg : f32):
+    br ^bb5(%bb2Arg : f32)
+  ^bb3(%bb3Arg : f32):
+    br ^bb6(%bb3Arg : f32)
+  ^bb4(%bb4Arg : f32):
+    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
+
+  // CHECK: ^[[BB5]]({{.*}}):
+  // CHECK-NEXT: "foo.bb5Terminator"
+  ^bb5(%bb5Arg : f32):
+    "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
+
+  // CHECK: ^[[BB6]]({{.*}}):
+  // CHECK-NEXT: "foo.bb6Terminator"
+  ^bb6(%bb6Arg : f32):
+    "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
+}
+
+// CHECK-LABEL: func @switch_from_switch_with_same_value_with_match(
+// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
+func @switch_from_switch_with_same_value_with_match(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32) {
+  // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5] : () -> ()
+
+  ^bb1:
+    // CHECK: switch %[[FLAG]]
+    switch %flag : i32, [
+      default: ^bb2,
+      42: ^bb3
+    ]
+
+  ^bb2:
+    "foo.bb2Terminator"() : () -> ()
+  ^bb3:
+    // prevent this block from being simplified away
+    "foo.op"() : () -> ()
+    // CHECK-NOT: switch %[[FLAG]]
+    // CHECK: br ^[[BB5:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_1]]
+    switch %flag : i32, [
+      default: ^bb4(%caseOperand0 : f32),
+      42: ^bb5(%caseOperand1 : f32)
+    ]
+
+  ^bb4(%bb4Arg : f32):
+    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
+
+  // CHECK: ^[[BB5]]({{.*}}):
+  // CHECK-NEXT: "foo.bb5Terminator"
+  ^bb5(%bb5Arg : f32):
+    "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
+}
+
+// CHECK-LABEL: func @switch_from_switch_with_same_value_no_match(
+// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
+func @switch_from_switch_with_same_value_no_match(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
+  // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5, ^bb6] : () -> ()
+
+  ^bb1:
+    // CHECK: switch %[[FLAG]]
+    switch %flag : i32, [
+      default: ^bb2,
+      42: ^bb3
+    ]
+
+  ^bb2:
+    "foo.bb2Terminator"() : () -> ()
+  ^bb3:
+    "foo.op"() : () -> ()
+    // CHECK-NOT: switch %[[FLAG]]
+    // CHECK: br ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
+    switch %flag : i32, [
+      default: ^bb4(%caseOperand0 : f32),
+      0: ^bb5(%caseOperand1 : f32),
+      43: ^bb6(%caseOperand2 : f32)
+    ]
+
+  // CHECK: ^[[BB4]]({{.*}})
+  // CHECK-NEXT: "foo.bb4Terminator"
+  ^bb4(%bb4Arg : f32):
+    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
+
+  ^bb5(%bb5Arg : f32):
+    "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
+
+  ^bb6(%bb6Arg : f32):
+    "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
+}
+
+// CHECK-LABEL: func @switch_from_switch_default_with_same_value(
+// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
+func @switch_from_switch_default_with_same_value(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
+  // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5, ^bb6] : () -> ()
+
+  ^bb1:
+    // CHECK: switch %[[FLAG]]
+    switch %flag : i32, [
+      default: ^bb3,
+      42: ^bb2
+    ]
+
+  ^bb2:
+    "foo.bb2Terminator"() : () -> ()
+  ^bb3:
+    "foo.op"() : () -> ()
+    // CHECK: switch %[[FLAG]]
+    // CHECK-NEXT: default: ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
+    // CHECK-NEXT: 43: ^[[BB6:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
+    // CHECK-NOT: 42
+    switch %flag : i32, [
+      default: ^bb4(%caseOperand0 : f32),
+      42: ^bb5(%caseOperand1 : f32),
+      43: ^bb6(%caseOperand2 : f32)
+    ]
+
+  // CHECK: ^[[BB4]]({{.*}}):
+  // CHECK-NEXT: "foo.bb4Terminator"
+  ^bb4(%bb4Arg : f32):
+    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
+
+  ^bb5(%bb5Arg : f32):
+    "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
+
+  // CHECK: ^[[BB6]]({{.*}}):
+  // CHECK-NEXT: "foo.bb6Terminator"
+  ^bb6(%bb6Arg : f32):
+    "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
+}
+
 /// Test folding conditional branches that are successors of conditional
 /// branches with the same condition.
 
 // CHECK-LABEL: func @cond_br_from_cond_br_with_same_condition
 func @cond_br_from_cond_br_with_same_condition(%cond : i1) {
   // CHECK:   cond_br %{{.*}}, ^bb1, ^bb2
   // CHECK: ^bb1:
   // CHECK:   return

mlir/test/Dialect/Standard/ops.mlir

 }
 
 // CHECK-LABEL: func @read_global_memref
 func @read_global_memref() {
   %0 = memref.get_global @memref0 : memref<2xf32>
   %1 = memref.tensor_load %0 : memref<2xf32>
   return
 }
+
+// CHECK-LABEL: func @switch(
+func @switch(%flag : i32, %caseOperand : i32) {
+  switch %flag : i32, [
+    default: ^bb1(%caseOperand : i32),
+    42: ^bb2(%caseOperand : i32),
+    43: ^bb3(%caseOperand : i32)
+  ]
+
+  ^bb1(%bb1arg : i32):
+    return
+  ^bb2(%bb2arg : i32):
+    return
+  ^bb3(%bb3arg : i32):
+    return
+}
+
+// CHECK-LABEL: func @switch_i64(
+func @switch_i64(%flag : i64, %caseOperand : i32) {
+  switch %flag : i64, [
+    default: ^bb1(%caseOperand : i32),
+    42: ^bb2(%caseOperand : i32),
+    43: ^bb3(%caseOperand : i32)
+  ]
+
+  ^bb1(%bb1arg : i32):
+    return
+  ^bb2(%bb2arg : i32):
+    return
+  ^bb3(%bb3arg : i32):
+    return
+}

mlir/test/Dialect/Standard/parser.mlir

+// RUN: mlir-opt -verify-diagnostics -split-input-file %s
+
+func @switch_missing_case_value(%flag : i32, %caseOperand : i32) {
+  switch %flag : i32, [
+    default: ^bb1(%caseOperand : i32),
+    45: ^bb2(%caseOperand : i32),
+    // expected-error@+1 {{expected integer value}}
+    : ^bb3(%caseOperand : i32)
+  ]
+
+  ^bb1(%bb1arg : i32):
+    return
+  ^bb2(%bb2arg : i32):
+    return
+  ^bb3(%bb3arg : i32):
+    return
+}
+
+// -----
+
+func @switch_wrong_type_case_value(%flag : i32, %caseOperand : i32) {
+  switch %flag : i32, [
+    default: ^bb1(%caseOperand : i32),
+    // expected-error@+1 {{expected integer value}}
+    "hello": ^bb2(%caseOperand : i32)
+  ]
+
+  ^bb1(%bb1arg : i32):
+    return
+  ^bb2(%bb2arg : i32):
+    return
+  ^bb3(%bb3arg : i32):
+    return
+}
+
+// -----
+
+func @switch_missing_comma(%flag : i32, %caseOperand : i32) {
+  switch %flag : i32, [
+    default: ^bb1(%caseOperand : i32),
+    45: ^bb2(%caseOperand : i32)
+    // expected-error@+1 {{expected ']'}}
+    43: ^bb3(%caseOperand : i32)
+  ]
+
+  ^bb1(%bb1arg : i32):
+    return
+  ^bb2(%bb2arg : i32):
+    return
+  ^bb3(%bb3arg : i32):
+    return
+}
+
+// -----
+
+func @switch_missing_default(%flag : i32, %caseOperand : i32) {
+  switch %flag : i32, [
+    // expected-error@+1 {{expected 'default'}}
+    45: ^bb2(%caseOperand : i32)
+    43: ^bb3(%caseOperand : i32)
+  ]
+
+  ^bb1(%bb1arg : i32):
+    return
+  ^bb2(%bb2arg : i32):
+    return
+  ^bb3(%bb3arg : i32):
+    return
+}