[mlir][LLVMIR] Add 'llvm.switch' op

[modocache]

Summary

The LLVM IR 'switch' instruction allows control flow to be transferred
to one of any number of branches depending on an integer control value,
or a default value if the control does not match any branch values. This patch
adds llvm.switch to the MLIR LLVMIR dialect, as well as translation routines
for lowering it to LLVM IR.

To store a variable number of operands for a variable number of branch
destinations, the new op makes use of the AttrSizedOperandSegments
trait. It stores its default branch operands as one segment, and all
remaining case branches' operands as another. It also stores pairs of
begin and end offset values to delineate the sub-range of each case branch's
operands. There's probably a better way to implement this, since the
offset computation complicates several parts of the op definition. This is the
approach I settled on because in doing so I was able to delegate to the default
op builder member functions. However, it may be preferable to instead specify
skipDefaultBuilders in the op's ODS, or use a completely separate
approach; feedback is welcome!

Another contentious part of this patch may be the custom printer and
parser functions for the op. Ideally I would have liked the MLIR to be
printed in this way:

llvm.switch %0, ^bb1(%1 : !llvm.i32) [
  1: ^bb2,
  2: ^bb3(%2, %3 : !llvm.i32, !llvm.i32)
]

The above would resemble how LLVM IR is formatted for the 'switch'
instruction. But I found it difficult to print and parse something like
this, whether I used the declarative assembly format or custom functions.
I also was not sure a multi-line format would be welcome -- it seems
like most MLIR ops do not use newlines. Again, I'd be happy to hear any
feedback here as well, or on any other aspect of the patch.

[ftynse]

Thanks, this look good already, but I have a set of comments.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

603

Let's mention caseValues here as well.

604–605

Let's have a consistent naming scheme for attributes, please! Check what the rest of this file uses predominantly.

614

Nit: caseDestinations.empty()

616

Do we also want to check that the number of branch weights matches the number of successors?

630–660

This should live in a C++ file

631

No need to prefix SmallVector with llvm::, it is re-exported into the mlir:: namespace

650–652

Can't we rather accept an ArrayRef<int32_t> and avoid casting here?

Also, llvm::to_vector.

673–677

This attribute looks redundant with operand_segment_sizes. Furthermore, it is only used in the parser, but ignored in the builder.

682

Nit: you can return pair<0,0> or any other empty range and drop Optional.

686

Same as above.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

281–282

Nit: I'd rather assert that index is in bounds

289

This looks like the pair is "begin, length" and not "begin, end" as the documentation claims.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

733

Nit: please only use auto when it improves readability, i.e. when the type is obvious from RHS (a cast) or is very long/impossible to spell out (iterators, lambdas).

739–740

I'm pretty sure DenseIntElementsAttr has a way of obtaining the value without going through (the creation under lock of) IntegerAttr, maybe through iterators.

750

If it has to be i32 in LLVM IR, we should add a corresponding check to the op verifier.

[rriddle]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

196

I think you could represent this with a mix of declarative assembly format and C++:

let assemblyFormat = [{
  $value `,` $defaultDestination (`(` $defaultOperands^ `)`)?
  `[` custom<SwitchOpCases>($caseDestinations, $caseOperands, type($caseOperands), $caseOperandOffsets) `]`
  attr-dict
}];

static ParseResult parseSwitchOpCases(OpAsmParser &parser,  SmallVectorImpl<Block *> &caseDestinations,
                                                                    SmallVectorImpl<OpAsmParser::OperandType> &caseOperands, 
                                                                    SmallVectorImpl<Type> &caseOperandTypes, ElementsAttr &caseOperandOffsets) {
  // Do the parsing of the following here in C++:
  //   1: ^bb2,
  //   2: ^bb3(%2, %3 : !llvm.i32, !llvm.i32)
}
static void printSwitchOpCases(OpAsmPrinter &printer, SwitchOp op, SuccessorRange caseDestinations, OperandRange caseOperands, TypeRange caseOperandTypes, 
                                                       ElementsAttr caseOperandOffsets) {
  // Do the printing of the following here in C++:
  //   1: ^bb2,
  //   2: ^bb3(%2, %3 : !llvm.i32, !llvm.i32)
}

The above might need slight tweaking given that I haven't tried to compile it.

For the newlines, I've had the thought for a while to allow operation parser to add newlines that get indented to the beginning of the operation line. I can add this in parallel to your change, and then you could pick it up when it lands.

[modocache]

Thanks for the review, @ftynse!

I noticed I didn't explicitly call this out in my commit message, so I'll do so now: with this patch as-is, the llvm.switch MLIR op has stricter constraints that the LLVM IR switch instruction: the LLVM IR switch can be used with no case values, whereas MLIR llvm.switch currently cannot. The following LLVM IR switch behaves the same way as an unconditional br %default instruction would:

switch i32 %0, label %default [
]

However, with this patch as-is, we reject the equivalent MLIR, an llvm.switch op with no case values and no case destinations:

llvm.switch %0, ^bb1, [], []

It's not difficult to modify this patch to allow for a case-less llvm.switch MLIR op, but I wasn't sure it was a good idea to do so. It seems to me that any reasonable frontend could avoid emitting llvm.switch %0, ^bb1, [], [] and instead emit llvm.br ^bb1. So, for now, case-less llvm.switch MLIR ops are rejected, but if anyone disagrees, please let me know.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

673–677

Am I understanding you right that you're saying I ought to remove case_operand_offsets, and instead use operand_segment_sizes? I ask because case_operand_offsets *is* used in the builder, I think: it's passed along to a default build member function provided by MLIR TableGen. So I wanted to make sure.

I touched on this in my commit message:

To store a variable number of operands for a variable number of branch destinations, the new op makes use of the AttrSizedOperandSegments trait. It stores its default branch operands as one segment, and all remaining case branches' operands as another. It also stores pairs of begin and end offset values to delineate the sub-range of each case branch's operands. There's probably a better way to implement this, since the offset computation complicates several parts of the op definition. This is the approach I settled on because in doing so I was able to delegate to the default op builder member functions. However, it may be preferable to instead specify skipDefaultBuilders in the op's ODS, or use a completely separate approach; feedback is welcome!

I think it should be possible to rewrite this patch to use operand_segment_sizes instead of case_operand_offsets, but I want to make sure that's what you're requesting before I do, because you also commented on some of the offset calculation logic. All of that code will probably go away if I rewrite this to use operand_segment_sizes, I think -- I appreciate the tips all the same, but just thought I'd mention that in case you were envisioning a patch that removes the case_operand_offsets attribute but still calculates offsets.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

733

Will do, thanks! Based on your advice I left this alone: the type is a little long, llvm::Optional<mlir::ElementsAttr>, and the CondBrOp branch immediately above this also uses auto weights = ... for the same type, so auto here seemed consistent. That being said I'd be happy to change just this auto, or both of them, just let me know which you'd prefer.

I did change many other uses of auto in this patch to use the specific type name instead, and I agree it improves readability, so I appreciate the suggestion! (For example, for (int32_t caseValue : caseValues) ... instead of auto, in SwitchOp::build.)

750

I had thought that was taken care of by marking the argument as let arguments = (ins LLVM_i32:$value, ...); in the ODS for SwitchOp, but I may be mistaken. When I attempt to create an llvm.switch with an i1 value operand, I get:

error: 'llvm.switch' op operand #0 must be LLVM 32-bit integer type, but got '!llvm.i1'

I hadn't noticed that mlir/test/Dialect/LLVMIR/invalid.mlir tests these invalid cases. I added a test there that demonstrates the existing verification, please let me know if I should add anything to the verifier in addition.

[rriddle]

In D93005#2446898, @modocache wrote:

Thanks for the review, @ftynse!

I noticed I didn't explicitly call this out in my commit message, so I'll do so now: with this patch as-is, the llvm.switch MLIR op has stricter constraints that the LLVM IR switch instruction: the LLVM IR switch can be used with no case values, whereas MLIR llvm.switch currently cannot. The following LLVM IR switch behaves the same way as an unconditional br %default instruction would:

switch i32 %0, label %default [
]

However, with this patch as-is, we reject the equivalent MLIR, an llvm.switch op with no case values and no case destinations:

llvm.switch %0, ^bb1, [], []

It's not difficult to modify this patch to allow for a case-less llvm.switch MLIR op, but I wasn't sure it was a good idea to do so. It seems to me that any reasonable frontend could avoid emitting llvm.switch %0, ^bb1, [], [] and instead emit llvm.br ^bb1. So, for now, case-less llvm.switch MLIR ops are rejected, but if anyone disagrees, please let me know.

If we do decide to go this direction, any variations from the LLVM operation should be documented in the description of the op. I'll defer to @ftynse for their preference here.

[modocache]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

196

Oh, cool! I hadn't realized that custom C++ could be mixed into the declarative assembly format. I'll give that a try, thanks!

For the newlines, I've had the thought for a while to allow operation parser to add newlines that get indented to the beginning of the operation line. I can add this in parallel to your change, and then you could pick it up when it lands.

That sounds great, thanks! With your change I'd probably change the assembly format to include newlines.

[modocache]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

673–677

Oooh, I see what you're saying now. I had thought that using operand_segment_sizes would obviate the need for this kind of offset calculation, but writing the patch now I can see that's not the case. Sorry about that, and thanks for the suggestion! I'll upload a better version of this patch in a few hours.

[modocache]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

673–677

Actually, sorry: I'm no longer sure operand_segment_sizes will work here. In OpOperandAdaptorEmitter::addVerification, a verify function is built for ops with the AttrSizedOperandSegments trait, which verifies that the number of operands matches the number of segment sizes. However, in the ODS for this patch, I specified the LLVM_SwitchOp with the following operands:

let arguments = (ins LLVM_i32:$value, // operand #1
                 Variadic<AnyType>:$defaultOperands, // #2
                 Variadic<AnyType>:$caseOperands, // #3
                 ArrayAttr:$case_values,
                 OptionalAttr<ElementsAttr>:$branch_weights);

So the SwitchOpAdaptor::verify member function is TableGen'ed to verify that the operand_segment_sizes is composed of 3 values. This doesn't match up with how Variadic<AnyType>:$caseOperands is being actually being used for the switch op: as a range of operand ranges. Whether I choose to amend this patch to add each case branch operand as an operand of the switch op, or each range of case operands as an operand, if the total number of switch operands isn't 3, the op will fail to verify:

llvm.switch %0, ^bb1, [1 : i32, 2 : i32, 3 : i32] [
  ^bb2(%1, %2: !llvm.i32, !llvm.i32),
  ^bb3(%3, %4, %5 : !llvm.i32, !llvm.i32, !llvm.i32)
  ^bb3(%6 : !llvm.i32)
]
// If each case operand is added as an operand of the switch op:
// error: 'operand_segment_sizes' attribute for specifying operand segments must have 3 elements, not 6
// If each range of case operands is added as an operand of the switch op:
// error: 'operand_segment_sizes' attribute for specifying operand segments must have 3 elements, not 4

I think the underlying mismatch here is that AttrSizedOperandSegments is intended to define the segment size of a compile-time-fixed number of variadic operands, whereas this switch op models a variable number of variadic operands. I'm not super knowledgeable about ODS, but I'm guessing that if I stick to using ODS, it's not possible to define an op using AttrSizedOperandSegments that has a variable, dynamic number of segments...?

I could look into defining an op outside of ODS, but I wonder if this inelegant case_operand_offsets approach is a simpler solution for my use case. As far as I can tell, adding a new trait for Variadic<Variadic<AnyType>> could be a large addition, whereas I'm looking to simply lower a C switch statement in my C MLIR to LLVM.

[modocache]

I'm adding an unrelated change to the AttrSizedOperandSegments verifier that prints how many elements it found in operand_segment_sizes. It's not necessarily part of my patch, but it helped me when debugging. I'd be happy to split it out into its own if anyone would prefer that.

[modocache]

Oops, forgot to clang-format.

[modocache]

Add to the test for operand segment sizes.

[modocache]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

196

Is it possible, using this technique, to add an attribute to the SwitchOp from within the parseSwitchOpCases function? I can see that parsing attributes from within custom assembly format printer/parsers is possible thanks to the FormatCustomDirectiveAttributes test case in mlir/test/lib/Dialect/Test/TestOps.td. But that test uses OpAsmParser::parseAttribute to add an attribute to its op, whereas in this case, I need to collect all the case values, and then add an ArrayAttr to the switch op. From what I can tell, the custom parser functions are not passed an OperationState &, so they're not able to add attributes to it directly.

As a result, I'm not sure it's possible to parse case values in the way you described. 9Maybe the custom parser functions could be passed a reference to the state?)

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

619

This variable should be camelCase, same below.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

313

Can you just merge this function into getCaseOperandsMutable? getCaseOperands can just call the mutable version.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

342

Can you just use std::adjacent_find on all of the successors of terminator instead? That would cover all possible terminator types.

733

nit: Merge this variable into the if below:

if (auto weights = switchOp.branch_weights())

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

2214 ↗

(On Diff #311281)

nit: I'd change this to must have {1} elements, but got " as that better matches some of the other error messages we have.

Also, can you split this into a different change?

[rriddle]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

196

The parse* function for the custom directive will be provided whatever elements you specify, so you can just add $case_values to that list. See here for the proper docs: https://mlir.llvm.org/docs/OpDefinitions/#custom-directives

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

604

Can you change this to an ElementsAttr instead? It's significantly more efficient for lists of integers.

[modocache]

So, unfortunately I don't think it's feasible to use the declarative assembly format for this op, at least not without making some TableGen modifications to pass custom directives a reference to the OperationState. I'm also not certain it's possible to get rid of case_operand_offsets in favor of operand_segment_sizes, since this op seems unique in storing a variable list of variadic operands. But, aside from those two points, I think I've addressed all of your review comments so far. Thank you, @ftynse and @rriddle! Please take another look at this when you have some time.

[rriddle]

In D93005#2449630, @modocache wrote:

So, unfortunately I don't think it's feasible to use the declarative assembly format for this op, at least not without making some TableGen modifications to pass custom directives a reference to the OperationState. I'm also not certain it's possible to get rid of case_operand_offsets in favor of operand_segment_sizes, since this op seems unique in storing a variable list of variadic operands. But, aside from those two points, I think I've addressed all of your review comments so far. Thank you, @ftynse and @rriddle! Please take another look at this when you have some time.

See my earlier comment for handling case_values, you just need to provide that to the custom directive in the declarative form to get a parameter for it in the C++ parser/printer functions. For operand_segment_sizes, that is specific to how ODS handles variadics, and given there is no support for variadics of variadics I don't think you'll be able to reuse that right now.

[modocache]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

196

Oh, excellent! I think I have it working. It was PEBKAC, I had previously passed in $case_values as you described, but in my parse* function definition I marked it ArrayAttr, not ArrayAttr&, and so I must have been assigning values to a copy. Thanks!

[modocache]

Thanks for answering all my questions! I think I've responded to all of your feedback, let me know if I missed anything. I also have a question about parsing and printing integers that I'll post as an inline comment.

[modocache]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

236

Here I'm parsing IntegerAttr, but ideally I'd like to parse integer literals. However, I didn't see a parser function to do so. Aesthetically I don't mind the op including the integer attribute type information (so cases look like [1 : i32 = ^bb1, ... ]), but the uglier part comes when I translate these into an ElementsAttr below...

257

...here I translate between parsed IntegerAttr, their values, and an i32 vector attribute. I don't know if it's possible to parse integer literals right now, but if not, do you know of a better way to convert an array of IntegerAttr into an ElementsAttr?

[modocache]

Rebase past D92845, add instantiated SwitchInst to the branch mapping.

[rriddle]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

257

Not at the moment, IntegerAttr is the only way to currently get an integer. Sent out D93152 for this.

Also sent out D93151 for newline support.

[ftynse]

In D93005#2446920, @rriddle wrote:

In D93005#2446898, @modocache wrote:

Thanks for the review, @ftynse!

I noticed I didn't explicitly call this out in my commit message, so I'll do so now: with this patch as-is, the llvm.switch MLIR op has stricter constraints that the LLVM IR switch instruction: the LLVM IR switch can be used with no case values, whereas MLIR llvm.switch currently cannot. The following LLVM IR switch behaves the same way as an unconditional br %default instruction would:

switch i32 %0, label %default [
]

However, with this patch as-is, we reject the equivalent MLIR, an llvm.switch op with no case values and no case destinations:

llvm.switch %0, ^bb1, [], []

It's not difficult to modify this patch to allow for a case-less llvm.switch MLIR op, but I wasn't sure it was a good idea to do so. It seems to me that any reasonable frontend could avoid emitting llvm.switch %0, ^bb1, [], [] and instead emit llvm.br ^bb1. So, for now, case-less llvm.switch MLIR ops are rejected, but if anyone disagrees, please let me know.

If we do decide to go this direction, any variations from the LLVM operation should be documented in the description of the op. I'll defer to @ftynse for their preference here.

I would argue that it's not a good idea to call something "LLVM dialect" and diverge from LLVM IR. Yes, it has a lot of quirks. "Fixing" them in MLIR but not in LLVM IR only makes the situation worse. No one not only has to remember all the quirks, but also how they are different in MLIR. So I have a very strong preference for supporting everything LLVM IR can. We can add more "usability" features like not repeating the type excessively or conditional branching to the same block with different arguments, but we should not decrease expressivity.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

673–677

Okay, thanks for the explanation! I don't think adding infrastructural support because of one op is justified.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

750

ODS indeed adds a verifier. And we don't check autogenerated verifiers in all uses, so no need to put this in invalid.mlir.

I suppose I reacted to getting the type unconditionally, I'd rather query the operand and get its type.

[modocache]

Sounds good, thanks! I agree, I'll modify this patch later today such that the MLIR llvm.switch op accepts zero cases, just as LLVM IR switch instruction does.

[modocache]

Thanks for the integer and newline patches @rriddle! I've updated this patch to use those new features. And thanks for another round of reviews @ftynse, I think I've responded to all the feedback given so far, let me know if I've missed anything.

[modocache]

Fix comment typo in test.

[ftynse]

Thanks for quick iteration!

[modocache]

Two small tweaks: give the BlockRange caseDestinations parameter of the op builder a default empty value, and "indent" cases before printing them.

[modocache]

Not at all, thank you for all the great feedback!

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

280

One last thing before I commit this: @rriddle, would it be possible to add an indent member function to the printer? I noticed that without adding spaces here, the op was being printed like this:

llvm.switch %0, ^bb1 [
1: ^bb2 // Oops! Not indented.
]

So I decided to add two spaces manually here, but in your addition of printNewline in D93151, you indent based on indentation widths that are shared throughout the printer's internals. If MLIR began indenting by 4, this switch op case printer would begin to look funny.

Let me know if you'd prefer this manual indentation to be cleaned up before this lands, and if you'd like me to try to write a patch to expose indenting to the printer's public interface.

[modocache]

I think I'll commit this later today, but let me know if you have any thoughts on adding an interface for indenting text to the printer, @rriddle. I'd be happy to try my hand at sending a patch to do so, as well.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_ResumeOp : LLVM_TerminatorOp<"resume", []> {
   let assemblyFormat = "$value attr-dict `:` type($value)";
 }
 def LLVM_UnreachableOp : LLVM_TerminatorOp<"unreachable", []> {
   string llvmBuilder = [{ builder.CreateUnreachable(); }];
   let parser = [{ return success(); }];
   let printer = [{ p << getOperationName(); }];
 }
 
+def LLVM_SwitchOp : LLVM_TerminatorOp<"switch",
+    [AttrSizedOperandSegments, DeclareOpInterfaceMethods<BranchOpInterface>,
+     NoSideEffect]> {
+  let arguments = (ins LLVM_i32:$value,
+                   Variadic<AnyType>:$defaultOperands,
+                   Variadic<AnyType>:$caseOperands,
+                   OptionalAttr<ElementsAttr>:$caseOperandOffsets,
+                   OptionalAttr<ElementsAttr>:$branch_weights);
+  let successors = (successor
+                    AnySuccessor:$defaultDestination,
+                    VariadicSuccessor<AnySuccessor>:$caseDestinations);
+
+  let verifier = [{
+    if (getCaseValuesAttr().size() != caseDestinations().size())
+      return emitOpError("expects number of case values to match number of "
+                         "case destinations");
+    if (caseDestinations().size() == 0)
+      return emitOpError("expects at least one case");
+    return success();
+  }];
+
+  let printer = [{ printSwitchOp(p, *this); }];
+  let parser = [{ return parseSwitchOp(parser, result); }];
+
+  let builders = [
+    OpBuilderDAG<(ins "Value":$value,
+      "Block *":$defaultDestination,
+      "ValueRange":$defaultOperands,
+      "ArrayRef<int32_t>":$caseValues,
+      "BlockRange":$caseDestinations,
+      CArg<"ArrayRef<ValueRange>", "{}">:$caseOperands,
+      CArg<"ArrayRef<uint32_t>", "{}">:$branch_weights),
+    [{
+        llvm::SmallVector<Attribute> caseValueAttrs;
+        for (auto caseValue : caseValues)
+          caseValueAttrs.push_back($_builder.getI32IntegerAttr(caseValue));
+        $_state.addAttribute(getCaseValuesAttrName(),
+                             $_builder.getArrayAttr(caseValueAttrs));
+
+        llvm::SmallVector<Value> flattenedCaseOperands;
+        llvm::SmallVector<int32_t> caseOperandOffsets;
+        int32_t offset = 0;
+        for (auto operands : caseOperands) {
+          flattenedCaseOperands.append(operands.begin(), operands.end());
+          caseOperandOffsets.push_back(offset);
+          offset += operands.size();
+        }
+        ElementsAttr caseOperandOffsetsAttr;
+        if (!caseOperandOffsets.empty())
+          caseOperandOffsetsAttr =
+              $_builder.getI32VectorAttr(caseOperandOffsets);
+
+        llvm::SmallVector<int32_t> weights;
+        for (auto branchWeight : branch_weights)
+          weights.push_back(static_cast<int32_t>(branchWeight));
+        ElementsAttr weightsAttr;
+        if (!weights.empty())
+          weightsAttr = $_builder.getI32VectorAttr(weights);
+
+        build($_builder, $_state, value, defaultOperands, flattenedCaseOperands,
+              caseOperandOffsetsAttr, weightsAttr, defaultDestination,
+              caseDestinations);
+    }]>,
+    LLVM_TerminatorPassthroughOpBuilder
+  ];
+
+  let extraClassDeclaration = [{
+    /// Return the name of the attribute used for case values.
+    static llvm::StringRef getCaseValuesAttrName() { return "case_values"; }
+
+    /// Return the array of case values.
+    ArrayAttr getCaseValuesAttr() {
+      return (*this)->getAttrOfType<ArrayAttr>(getCaseValuesAttrName());
+    }
+
+    /// Return the name of the attribute used to delineate ranges of case
+    /// destination operands.
+    static llvm::StringRef getCaseOperandOffsetsAttrName() {
+      return "caseOperandOffsets";
+    }
+
+    /// Return the [begin, end) offsets of the case operands for the case
+    /// destination at the given index, or none if there are no operands for
+    /// that case.
+    Optional<std::pair<unsigned, unsigned>>
+    getCaseOperandOffsets(unsigned index);
+
+    /// Return the operands for the case destination block at the given index.
+    Optional<OperandRange> getCaseOperands(unsigned index);
+
+    /// Return a mutable range of operands for the case destination block at the
+    /// given index.
+    Optional<MutableOperandRange> getCaseOperandsMutable(unsigned index);
+  }];
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Auxiliary operations (do not appear in LLVM IR but necessary for the dialect
 // to work correctly).
 ////////////////////////////////////////////////////////////////////////////////
 
 // Linkage attribute is used on functions and globals. The order follows that of
 // https://llvm.org/docs/LangRef.html#linkage-types. The names are equivalent to
 // visible names in the IR rather than to enum values names in llvm::GlobalValue

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

 
 Optional<MutableOperandRange>
 CondBrOp::getMutableSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
   return index == 0 ? trueDestOperandsMutable() : falseDestOperandsMutable();
 }
 
 //===----------------------------------------------------------------------===//
+// LLVM::SwitchOp
+//===----------------------------------------------------------------------===//
+
+static void printSwitchOp(OpAsmPrinter &p, SwitchOp &op) {
+  p << op.getOperationName() << ' ' << op.value() << ", ";
+  p.printSuccessorAndUseList(op.defaultDestination(), op.defaultOperands());
+  p << ", ";
+
+  p.printAttribute(op.getCaseValuesAttr());
+  p << ", ";
+
+  p << '[';
+  size_t index = 0;
+  llvm::interleaveComma(op.caseDestinations(), p, [&](auto *destination) {
+    if (auto operands = op.getCaseOperands(index++))
+      p.printSuccessorAndUseList(destination, operands.getValue());
+    else
+      p.printSuccessor(destination);
+  });
+  p << ']';
+
+  p.printOptionalAttrDict(op.getAttrs(), {op.getOperandSegmentSizeAttr(),
+                                          op.getCaseValuesAttrName(),
+                                          op.getCaseOperandOffsetsAttrName()});
+}
+
+/// <operation> ::= `llvm.switch`
+///                   ssa-use `,` bb-id (`(` ssa-use-and-type-list `)`)? `,`
+///                   `[` integer-literal (`,` integer-literal)* `]` `,`
+///                   `[` bb-id (`(` ssa-use-and-type-list `)`)?
+///                       (`,` bb-id (`(` ssa-use-and-type-list `)`)? )? `]`
+///                    attribute-dict?
+static ParseResult parseSwitchOp(OpAsmParser &parser, OperationState &result) {
+  auto &builder = parser.getBuilder();
+
+  OpAsmParser::OperandType value;
+  Block *defaultDestination, *caseDestination;
+  llvm::SmallVector<Value> defaultOperands;
+  ArrayAttr caseValuesAttr;
+  if (parser.parseOperand(value) || parser.parseComma() ||
+      parser.parseSuccessorAndUseList(defaultDestination, defaultOperands) ||
+      parser.parseComma() ||
+      parser.parseAttribute(caseValuesAttr, SwitchOp::getCaseValuesAttrName(),
+                            result.attributes) ||
+      parser.parseComma() || parser.parseLSquare())
+    return failure();
+
+  llvm::SmallVector<Block *> caseDestinations;
+  llvm::SmallVector<Value> caseOperands;
+  llvm::SmallVector<int32_t> caseOperandOffsets;
+  int32_t offset = 0;
+  do {
+    if (parser.parseSuccessorAndUseList(caseDestination, caseOperands))
+      return failure();
+    caseDestinations.push_back(caseDestination);
+    caseOperandOffsets.push_back(offset);
+    offset += caseOperands.size();
+  } while (!parser.parseOptionalComma());
+
+  if (parser.parseRSquare() ||
+      parser.parseOptionalAttrDict(result.attributes) ||
+      parser.resolveOperand(value, LLVMType::getInt32Ty(builder.getContext()),
+                            result.operands))
+    return failure();
+
+  result.addSuccessors(defaultDestination);
+  result.addOperands(defaultOperands);
+
+  result.addSuccessors(caseDestinations);
+  result.addOperands(caseOperands);
+  result.addAttribute(SwitchOp::getCaseOperandOffsetsAttrName(),
+                      builder.getI32VectorAttr(caseOperandOffsets));
+
+  result.addAttribute(
+      SwitchOp::getOperandSegmentSizeAttr(),
+      builder.getI32VectorAttr({1, static_cast<int32_t>(defaultOperands.size()),
+                                static_cast<int32_t>(caseOperands.size())}));
+
+  return success();
+}
+
+Optional<std::pair<unsigned, unsigned>>
+SwitchOp::getCaseOperandOffsets(unsigned index) {
+  auto caseOperands = caseOperandsMutable();
+  if (!caseOperandOffsets())
+    return llvm::None;
+
+  auto offsets = caseOperandOffsets().getValue();
+  if (index >= offsets.size())
+    return llvm::None;
+
+  auto begin = offsets.getValue(index).cast<IntegerAttr>().getInt();
+  auto end = index + 1 == offsets.size()
+                 ? caseOperands.size()
+                 : offsets.getValue(index + 1).cast<IntegerAttr>().getInt();
+  return std::make_pair(static_cast<unsigned>(begin),
+                        static_cast<unsigned>(end - begin));
+}
+
+Optional<OperandRange> SwitchOp::getCaseOperands(unsigned index) {
+  if (auto offsets = getCaseOperandOffsets(index))
+    return caseOperands().slice(offsets->first, offsets->second);
+  return llvm::None;
+}
+
+Optional<MutableOperandRange> SwitchOp::getCaseOperandsMutable(unsigned index) {
+  if (auto offsets = getCaseOperandOffsets(index))
+    return caseOperandsMutable().slice(offsets->first, offsets->second);
+  return llvm::None;
+}
+
+Optional<MutableOperandRange>
+SwitchOp::getMutableSuccessorOperands(unsigned index) {
+  assert(index < getNumSuccessors() && "invalid successor index");
+  return index == 0 ? defaultOperandsMutable()
+                    : getCaseOperandsMutable(index - 1);
+}
+
+//===----------------------------------------------------------------------===//
 // Builder, printer and parser for for LLVM::LoadOp.
 //===----------------------------------------------------------------------===//
 
 void LoadOp::build(OpBuilder &builder, OperationState &result, Type t,
                    Value addr, unsigned alignment, bool isVolatile,
                    bool isNonTemporal) {
   result.addOperands(addr);
   result.addTypes(t);
   if (isVolatile)
     result.addAttribute(kVolatileAttrName, builder.getUnitAttr());

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 /// Get the SSA value passed to the current block from the terminator operation
 /// of its predecessor.
 static Value getPHISourceValue(Block *current, Block *pred,
                                unsigned numArguments, unsigned index) {
   Operation &terminator = *pred->getTerminator();
   if (isa<LLVM::BrOp>(terminator))
     return terminator.getOperand(index);
 
-  // For conditional branches, we need to check if the current block is reached
-  // through the "true" or the "false" branch and take the relevant operands.
-  auto condBranchOp = dyn_cast<LLVM::CondBrOp>(terminator);
-  assert(condBranchOp &&
-         "only branch operations can be terminators of a block that "
-         "has successors");
+  // For instructions that branch based on a condition value, we need to take
+  // the operands for the branch that was taken.
+  if (auto condBranchOp = dyn_cast<LLVM::CondBrOp>(terminator)) {
+    // For conditional branches, we take the operands from either the "true" or
+    // the "false" branch.
     assert((condBranchOp.getSuccessor(0) != condBranchOp.getSuccessor(1)) &&
            "successors with arguments in LLVM conditional branches must be "
            "different blocks");
 
     return condBranchOp.getSuccessor(0) == current
                ? condBranchOp.trueDestOperands()[index]
                : condBranchOp.falseDestOperands()[index];
+  } else if (auto switchOp = dyn_cast<LLVM::SwitchOp>(terminator)) {
+    // For switches, we take the operands from either the default case, or from
+    // the case branch that was taken.
+#ifndef NDEBUG
+    for (size_t i = 0; i < switchOp.caseDestinations().size(); ++i) {
+      assert(switchOp.defaultDestination() != switchOp.caseDestinations()[i] &&
+             "successors with arguments in LLVM switch instructions must be "
+             "different blocks");
+      for (size_t j = i + 1; j < switchOp.caseDestinations().size(); ++j) {
+        assert(switchOp.caseDestinations()[i] !=
+                   switchOp.caseDestinations()[j] &&
+               "successors with arguments in LLVM switch instructions must be "
+               "different blocks");
+      }
+    }
+#endif
+
+    if (switchOp.defaultDestination() == current)
+      return switchOp.defaultOperands()[index];
+    for (auto i : llvm::enumerate(switchOp.caseDestinations()))
+      if (i.value() == current)
+        return switchOp.getCaseOperands(i.index()).getValue()[index];
+  }
+
+  llvm_unreachable("only branch or switch operations can be terminators of a "
+                   "block that has successors");
 }
 
 /// Connect the PHI nodes to the results of preceding blocks.
 template <typename T>
 static void
 connectPHINodes(T &func, const DenseMap<Value, llvm::Value *> &valueMapping,
                 const DenseMap<Block *, llvm::BasicBlock *> &blockMapping) {
   // Skip the first block, it cannot be branched to and its arguments correspond
if (weights) {
               .createBranchWeights(static_cast<uint32_t>(trueWeight),
                                    static_cast<uint32_t>(falseWeight));
     }
     builder.CreateCondBr(valueMapping.lookup(condbrOp.getOperand(0)),
                          blockMapping[condbrOp.getSuccessor(0)],
                          blockMapping[condbrOp.getSuccessor(1)], branchWeights);
     return success();
   }
+  if (auto switchOp = dyn_cast<LLVM::SwitchOp>(opInst)) {
+    auto weights = switchOp.branch_weights();
+    llvm::MDNode *branchWeights = nullptr;
+    if (weights) {
+      llvm::SmallVector<uint32_t> weightValues;
+      weightValues.reserve(weights->size());
+      for (int64_t i = 0; i < weights->size(); ++i)
+        weightValues.push_back(static_cast<uint32_t>(
+            weights->getValue(i).cast<IntegerAttr>().getInt()));
+      branchWeights = llvm::MDBuilder(llvmModule->getContext())
+                          .createBranchWeights(weightValues);
+    }
+
+    auto *switchInst =
+        builder.CreateSwitch(valueMapping[switchOp.value()],
+                             blockMapping[switchOp.defaultDestination()],
+                             switchOp.caseDestinations().size(), branchWeights);
+
+    auto *int32Type = llvm::Type::getInt32Ty(llvmModule->getContext());
+    for (auto i :
+         llvm::zip(switchOp.getCaseValuesAttr(), switchOp.caseDestinations()))
+      switchInst->addCase(
+          llvm::ConstantInt::get(int32Type,
+                                 std::get<0>(i).cast<IntegerAttr>().getInt()),
+          blockMapping[std::get<1>(i)]);
+
+    return success();
+  }
 
   // Emit addressof.  We need to look up the global value referenced by the
   // operation and store it in the MLIR-to-LLVM value mapping.  This does not
   // emit any LLVM instruction.
   if (auto addressOfOp = dyn_cast<LLVM::AddressOfOp>(opInst)) {
     LLVM::GlobalOp global = addressOfOp.getGlobal();
     LLVM::LLVMFuncOp function = addressOfOp.getFunction();
 

mlir/test/Dialect/LLVMIR/roundtrip.mlir

 
 // Terminator operations and their successors.
 //
 // CHECK: llvm.br ^[[BB1:.*]]
   llvm.br ^bb1
 
 // CHECK: ^[[BB1]]
 ^bb1:
-// CHECK: llvm.cond_br %7, ^[[BB2:.*]], ^[[BB1]]
-  llvm.cond_br %7, ^bb2, ^bb1
+// CHECK: llvm.cond_br %7, ^[[BB2:.*]], ^[[BB3:.*]]
+  llvm.cond_br %7, ^bb2, ^bb3
 
 // CHECK: ^[[BB2]]
 ^bb2:
 // CHECK: %{{.*}} = llvm.mlir.undef : !llvm.struct<(i32, double, i32)>
 // CHECK: %{{.*}} = llvm.mlir.constant(42 : i64) : !llvm.i47
   %22 = llvm.mlir.undef : !llvm.struct<(i32, double, i32)>
   %23 = llvm.mlir.constant(42) : !llvm.i47
+  // CHECK: llvm.switch %0, ^[[BB3]], [1, 2], [^[[BB4:.*]], ^[[BB5:.*]]]
+  llvm.switch %0, ^bb3, [1, 2], [^bb4, ^bb5]
 
+// CHECK: ^[[BB3]]
+^bb3:
+// CHECK: llvm.br ^[[BB6:.*]]
+  llvm.br ^bb6
+
+// CHECK: ^[[BB4]]
+^bb4:
+  llvm.br ^bb6
+
+// CHECK: ^[[BB5]]
+^bb5:
+  llvm.br ^bb6
+
+// CHECK: ^[[BB6]]
+^bb6:
 // Misc operations.
 // CHECK: %{{.*}} = llvm.select %{{.*}}, %{{.*}}, %{{.*}} : !llvm.i1, !llvm.i32
   %24 = llvm.select %7, %0, %1 : !llvm.i1, !llvm.i32
 
 // Integer to pointer and pointer to integer conversions.
 //
 // CHECK: %[[PTR:.*]] = llvm.inttoptr %[[I32]] : !llvm.i32 to !llvm.ptr<i32>
 // CHECK: %{{.*}} = llvm.ptrtoint %[[PTR]] : !llvm.ptr<i32> to !llvm.i32

mlir/test/Target/llvmir.mlir

llvm.func @useInlineAsm(%arg0: !llvm.i32) {
   %4 = llvm.inline_asm asm_dialect = "intel" "foo", "=r,r,r" %arg0, %arg0 : (!llvm.i32, !llvm.i32) -> !llvm.i8
 
   // CHECK-NEXT:  call { i8, i8 } asm "foo", "=r,=r,r"(i32 {{.*}}), !dbg !14
   %5 = llvm.inline_asm "foo", "=r,=r,r" %arg0 : (!llvm.i32) -> !llvm.struct<(i8, i8)>
 
   llvm.return
 }
 
+// -----
+
+// CHECK-LABEL: @switch
+llvm.func @switch(%arg0: !llvm.i32) {
+  %0 = llvm.mlir.constant(5 : i32) : !llvm.i32
+  %1 = llvm.mlir.constant(7 : i32) : !llvm.i32
+  %2 = llvm.mlir.constant(11 : i32) : !llvm.i32
+  llvm.switch %arg0, ^bb1, [-1, 1], [^bb2(%0 : !llvm.i32), ^bb3(%1, %2 : !llvm.i32, !llvm.i32)]
+
+^bb1:  // pred: ^bb0
+  llvm.return %arg0 : !llvm.i32
+
+^bb2(%3: !llvm.i32): // pred: ^bb0
+  llvm.return %1 : !llvm.i32
+
+^bb3(%4: !llvm.i32, %5: !llvm.i32): // pred: ^bb0
+  llvm.return %4 : !llvm.i32
+}
+
+// CHECK-LABEL: @switchWeights
+llvm.func @switchWeights(%arg0: !llvm.i32) {
+  %0 = llvm.mlir.constant(19 : i32) : !llvm.i32
+  %1 = llvm.mlir.constant(23 : i32) : !llvm.i32
+  %2 = llvm.mlir.constant(29 : i32) : !llvm.i32
+  // CHECK: !prof ![[SWITCH_WEIGHT_NODE:[0-9]+]]
+  llvm.switch %arg0, ^bb1(%0 : !llvm.i32), [9, 99], [^bb2(%1, %2 : !llvm.i32, !llvm.i32), ^bb3] {branch_weights = dense<[13, 17, 19]> : vector<3xi32>}
+
+^bb1(%3: !llvm.i32):  // pred: ^bb0
+  llvm.return %3 : !llvm.i32
+
+^bb2(%4: !llvm.i32, %5: !llvm.i32): // pred: ^bb0
+  llvm.return %5 : !llvm.i32
+
+^bb3: // pred: ^bb0
+  llvm.return %arg0 : !llvm.i32
+}
+
+// CHECK: ![[SWITCH_WEIGHT_NODE]] = !{!"branch_weights", i32 13, i32 17, i32 19}