[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

2315

Avoid computing the size during the loop.

2316

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

2174

unsigned?

2255

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

2213

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

2216–2217

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?

2224

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

2268

Can you cache case_values() and caseDestinations()?

2352

Same comment on recomputing size here and below.

2360

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

2549

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

2213

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

2268

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

2315

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

2268

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).

2315

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

2268

Got it. Thanks for explaining :-)

2315

Got it, thanks for the pointer!

[GMNGeoffrey]

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

2174

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

2213

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

2213

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

2213

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

2255

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.

2316

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

2344–2345

nit: Can you spell out auto here?

2345

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

2349

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

2549

nit: Spell out auto here.

2608

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

2349

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 32-bit 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, ^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

 
 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) {
+  if (!case_values())
+    return defaultDestination();
+
+  auto caseDests = caseDestinations();
+  auto caseValues = case_values();
+  if (int64_t value =
+          operands.front().dyn_cast_or_null<IntegerAttr>().getSInt()) {
+    for (int64_t i = 0, size = case_values()->size(); i < size; ++i)
+      if (value == caseValues->getValue<int64_t>(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;
+  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);
+    SmallVector<Value> argStorage;
+    if (succeeded(collapseBranch(caseDest, caseOperands, argStorage)))
+      requiresChange = true;
+
+    newCaseDests.push_back(caseDest);
+    newCaseOperands.push_back(caseOperands);
+  }
+
+  Block *defaultDest = op.defaultDestination();
+  ValueRange defaultOperands = op.defaultOperands();
+  SmallVector<Value> argStorage;
+
+  if (succeeded(collapseBranch(defaultDest, defaultOperands, argStorage)))
+    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;
+  auto predDests = predSwitch.caseDestinations();
+  auto 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

^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_fold_only_default(
+func @switch_fold_only_default(%arg0 : i32) -> i32 {
+  // CHECK-NEXT: %[[CST:.*]] = constant 1 : i32
+  // CHECK-NEXT: return %[[CST]] : i32
+  %c1_i32 = constant 1 : i32
+  switch %arg0 : i32, [
+    default: ^bb1(%c1_i32 : i32)
+  ]
+  ^bb1(%x : i32):
+    return %x : i32
+}
+
+
+// CHECK-LABEL: func @switch_fold_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_fold_case_matching_default(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32) -> f32 {
+  // 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: ^bb1(%caseOperand0 : f32),
+    42: ^bb1(%caseOperand0 : f32),
+    10: ^bb2(%caseOperand1 : f32),
+    17: ^bb1(%caseOperand0 : f32)
+  ]
+  ^bb1(%x : f32):
+    return %x : f32
+  ^bb2(%y : f32):
+    return %y : f32
+}
+
+
+// CHECK-LABEL: func @switch_const_folding(
+func @switch_const_folding() -> i32 {
+  // CHECK-NEXT: %[[CST:.*]] = constant 1 : i32
+  // CHECK-NEXT: return %[[CST]] : i32
+  %c0_i32 = constant 0 : i32
+  %c1_i32 = constant 1 : i32
+  %c2_i32 = constant 2 : i32
+  %c3_i32 = constant 3 : i32
+  switch %c0_i32 : i32, [
+    default: ^bb1(%c1_i32 : i32),
+    -1: ^bb2(%c2_i32 : i32),
+    1: ^bb3(%c3_i32 : i32)
+  ]
+  ^bb1(%x : i32):
+    return %x : i32
+  ^bb2(%y : i32):
+    return %y : i32
+  ^bb3(%z : i32):
+    return %z : i32
+}
+
+// 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) -> f32{
+  //      CHECK: switch %[[FLAG]]
+  // CHECK-NEXT:   default: ^[[BB4:.+]](%[[CASE_OPERAND_0]] : f32)
+  // CHECK-NEXT:   43: ^[[BB5:.+]](%[[CASE_OPERAND_1]] : f32)
+  // CHECK-NEXT:   44: ^[[BB3:.+]](%[[CASE_OPERAND_2]] : f32)
+  // CHECK-NEXT: ]
+  switch %flag : i32, [
+    default: ^bb1(%caseOperand0 : f32),
+    43: ^bb2(%caseOperand1 : f32),
+    44: ^bb3(%caseOperand2 : f32)
+  ]
+  ^bb1(%bb1Arg : f32):
+    br ^bb4(%bb1Arg : f32)
+  ^bb2(%bb2Arg : f32):
+    br ^bb5(%bb2Arg : f32)
+  ^bb3(%bb3Arg : f32):
+    return %bb3Arg : f32
+
+  // CHECK: ^[[BB4]](
+  // CHECK-SAME: %[[BB4_ARG:[a-zA-Z0-9_]+]]
+  // CHECK-NEXT: %[[BB4_RES:.+]] = "foo.op1"(%[[BB4_ARG]])
+  // CHECK-NEXT: return %[[BB4_RES]]
+  ^bb4(%bb4Arg : f32):
+    %bb4Res = "foo.op1"(%bb4Arg) : (f32) -> f32
+    return %bb4Res : f32
+
+  // CHECK: ^[[BB5]](
+  // CHECK-SAME: %[[BB5_ARG:[a-zA-Z0-9_]+]]
+  // CHECK-NEXT: %[[BB5_RES:.+]] = "foo.op2"(%[[BB5_ARG]])
+  // CHECK-NEXT: return %[[BB5_RES]]
+  ^bb5(%bb5Arg : f32):
+    %bb5Res = "foo.op2"(%bb5Arg) : (f32) -> f32
+    return %bb5Res : f32
+}
+
+// CHECK-LABEL: func @switch_from_switch_with_same_value(
+//  CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+func @switch_from_switch_with_same_value(%flag : i32) -> f32 {
+  // CHECK: switch %[[FLAG]]
+  switch %flag : i32, [
+    default: ^bb1,
+    42: ^bb2
+  ]
+
+  ^bb1:
+    // Add predecessors to avoid these getting folded by other patterns.
+    "foo.terminator1"() [^bb3, ^bb4] : () -> ()
+  ^bb2:
+    "foo.op2"() : () -> ()
+    // CHECK-NOT: switch %[[FLAG]]
+    // CHECK: br ^[[BB4:.+]]
+    switch %flag : i32, [
+      default: ^bb3,
+      42: ^bb4
+    ]
+
+  ^bb3:
+    "foo.terminator3"() : () -> ()
+
+  // CHECK: ^[[BB4]]:
+  // CHECK-NEXT: "foo.terminator4"
+  ^bb4:
+    "foo.terminator4"() : () -> ()
+}
+
+// CHECK-LABEL: func @switch_from_switch_with_same_value_no_match(
+//  CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+func @switch_from_switch_with_same_value_no_match(%flag : i32) -> f32 {
+  // CHECK: switch %[[FLAG]]
+  switch %flag : i32, [
+    default: ^bb1,
+    42: ^bb2
+  ]
+
+  ^bb1:
+    // Add predecessors to avoid these getting folded by other patterns.
+    "foo.terminator1"() [^bb3, ^bb4] : () -> ()
+  ^bb2:
+    "foo.op2"() : () -> ()
+    // CHECK-NOT: switch %[[FLAG]]
+    // CHECK: br ^[[BB3:.+]]
+    switch %flag : i32, [
+      default: ^bb3,
+      0: ^bb5,
+      43: ^bb4
+    ]
+
+  // CHECK: ^[[BB3]]:
+  // CHECK-NEXT: "foo.terminator3"
+  ^bb3:
+    "foo.terminator3"() : () -> ()
+
+  ^bb4:
+    "foo.terminator4"() : () -> ()
+
+  ^bb5:
+    "foo.terminator5"() : () -> ()
+}
+
+// CHECK-LABEL: func @switch_from_switch_default_with_same_value(
+//  CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
+func @switch_from_switch_default_with_same_value(%flag : i32) -> f32 {
+  // CHECK: switch %[[FLAG]]
+  switch %flag : i32, [
+    default: ^bb2,
+    42: ^bb1
+  ]
+
+  ^bb1:
+    // Add predecessors to avoid these getting folded by other patterns.
+    "foo.terminator1"() [^bb3, ^bb4, ^bb5] : () -> ()
+  ^bb2:
+    "foo.op2"() : () -> ()
+    // CHECK: switch %[[FLAG]]
+    // CHECK-NOT: 42:
+    switch %flag : i32, [
+      default: ^bb3,
+      42: ^bb4,
+      43: ^bb5
+    ]
+
+  ^bb3:
+    "foo.terminator3"() : () -> ()
+
+  ^bb4:
+    "foo.terminator4"() : () -> ()
+
+  ^bb5:
+    "foo.terminator5"() : () -> ()
+}
+
 /// 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
+}