[mlir][llvm] Fix import of SwitchOp

[gysit]

Summary

This revision ensures SwitchOps with case and condition
bitwidths other than 32-bit are imported properly. It adds an
APInt based builder to the SwitchOp and implements
a verifier that checks that the condition and the case
value types match.

[Dinistro]

Thanks for fixing this, LGTM!

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

 
 def LLVM_SwitchOp : LLVM_TerminatorOp<"switch",
     [AttrSizedOperandSegments, DeclareOpInterfaceMethods<BranchOpInterface>,
      Pure]> {
   let arguments = (ins
     AnyInteger:$value,
     Variadic<AnyType>:$defaultOperands,
     VariadicOfVariadic<AnyType, "case_operand_segments">:$caseOperands,
-    OptionalAttr<ElementsAttr>:$case_values,
+    OptionalAttr<AnyIntElementsAttr>:$case_values,
     DenseI32ArrayAttr:$case_operand_segments,
     OptionalAttr<ElementsAttr>:$branch_weights
   );
   let successors = (successor
     AnySuccessor:$defaultDestination,
     VariadicSuccessor<AnySuccessor>:$caseDestinations
   );
 
   let assemblyFormat = [{
     $value `:` type($value) `,`
     $defaultDestination (`(` $defaultOperands^ `:` type($defaultOperands) `)`)?
     `[` `\n` custom<SwitchOpCases>(ref(type($value)), $case_values, $caseDestinations,
                                    $caseOperands, type($caseOperands)) `]`
     attr-dict
   }];
   let hasVerifier = 1;
 
   let builders = [
     OpBuilder<(ins "Value":$value,
       "Block *":$defaultDestination,
       "ValueRange":$defaultOperands,
+      CArg<"ArrayRef<APInt>", "{}">:$caseValues,
+      CArg<"BlockRange", "{}">:$caseDestinations,
+      CArg<"ArrayRef<ValueRange>", "{}">:$caseOperands,
+      CArg<"ArrayRef<int32_t>", "{}">:$branchWeights)>,
+    OpBuilder<(ins "Value":$value,
+      "Block *":$defaultDestination,
+      "ValueRange":$defaultOperands,
       CArg<"ArrayRef<int32_t>", "{}">:$caseValues,
       CArg<"BlockRange", "{}">:$caseDestinations,
       CArg<"ArrayRef<ValueRange>", "{}">:$caseOperands,
       CArg<"ArrayRef<int32_t>", "{}">:$branchWeights)>,
+    OpBuilder<(ins "Value":$value,
+      "Block *":$defaultDestination,
+      "ValueRange":$defaultOperands,
+      CArg<"DenseIntElementsAttr", "{}">:$caseValues,
+      CArg<"BlockRange", "{}">:$caseDestinations,
+      CArg<"ArrayRef<ValueRange>", "{}">:$caseOperands,
+      CArg<"ArrayRef<int32_t>", "{}">:$branchWeights)>,
     LLVM_TerminatorPassthroughOpBuilder
   ];
 
   let extraClassDeclaration = [{
     /// Return the operands for the case destination block at the given index.
     OperandRange getCaseOperands(unsigned index) {
       return getCaseOperands()[index];
     }

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

 }
 
 //===----------------------------------------------------------------------===//
 // LLVM::SwitchOp
 //===----------------------------------------------------------------------===//
 
 void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
                      Block *defaultDestination, ValueRange defaultOperands,
-                     ArrayRef<int32_t> caseValues, BlockRange caseDestinations,
+                     DenseIntElementsAttr caseValues,
+                     BlockRange caseDestinations,
                      ArrayRef<ValueRange> caseOperands,
                      ArrayRef<int32_t> branchWeights) {
-  ElementsAttr caseValuesAttr;
-  if (!caseValues.empty())
-    caseValuesAttr = builder.getI32VectorAttr(caseValues);
-
   ElementsAttr weightsAttr;
   if (!branchWeights.empty())
     weightsAttr = builder.getI32VectorAttr(llvm::to_vector<4>(branchWeights));
 
-  build(builder, result, value, defaultOperands, caseOperands, caseValuesAttr,
+  build(builder, result, value, defaultOperands, caseOperands, caseValues,
         weightsAttr, defaultDestination, caseDestinations);
 }
 
+void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
+                     Block *defaultDestination, ValueRange defaultOperands,
+                     ArrayRef<APInt> caseValues, BlockRange caseDestinations,
+                     ArrayRef<ValueRange> caseOperands,
+                     ArrayRef<int32_t> branchWeights) {
+  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, branchWeights);
+}
+
+void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
+                     Block *defaultDestination, ValueRange defaultOperands,
+                     ArrayRef<int32_t> caseValues, BlockRange caseDestinations,
+                     ArrayRef<ValueRange> caseOperands,
+                     ArrayRef<int32_t> branchWeights) {
+  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, branchWeights);
+}
+
 /// <cases> ::= integer `:` bb-id (`(` ssa-use-and-type-list `)`)?
 ///             ( `,` integer `:` bb-id (`(` ssa-use-and-type-list `)`)? )?
 static ParseResult parseSwitchOpCases(
-    OpAsmParser &parser, Type flagType, ElementsAttr &caseValues,
+    OpAsmParser &parser, Type flagType, DenseIntElementsAttr &caseValues,
     SmallVectorImpl<Block *> &caseDestinations,
     SmallVectorImpl<SmallVector<OpAsmParser::UnresolvedOperand>> &caseOperands,
     SmallVectorImpl<SmallVector<Type>> &caseOperandTypes) {
   SmallVector<APInt> values;
   unsigned bitWidth = flagType.getIntOrFloatBitWidth();
   do {
     int64_t value = 0;
     OptionalParseResult integerParseResult = parser.parseOptionalInteger(value);
static ParseResult parseSwitchOpCases(
 
   ShapedType caseValueType =
       VectorType::get(static_cast<int64_t>(values.size()), flagType);
   caseValues = DenseIntElementsAttr::get(caseValueType, values);
   return success();
 }
 
 static void printSwitchOpCases(OpAsmPrinter &p, SwitchOp op, Type flagType,
-                               ElementsAttr caseValues,
+                               DenseIntElementsAttr caseValues,
                                SuccessorRange caseDestinations,
                                OperandRangeRange caseOperands,
                                const TypeRangeRange &caseOperandTypes) {
   if (!caseValues)
     return;
 
   size_t index = 0;
   llvm::interleave(
-      llvm::zip(llvm::cast<DenseIntElementsAttr>(caseValues), caseDestinations),
+      llvm::zip(caseValues, caseDestinations),
       [&](auto i) {
         p << "  ";
         p << std::get<0>(i).getLimitedValue();
         p << ": ";
         p.printSuccessorAndUseList(std::get<1>(i), caseOperands[index++]);
       },
       [&] {
         p << ',';
         p.printNewline();
       });
   p.printNewline();
 }
 
 LogicalResult SwitchOp::verify() {
   if ((!getCaseValues() && !getCaseDestinations().empty()) ||
       (getCaseValues() &&
        getCaseValues()->size() !=
            static_cast<int64_t>(getCaseDestinations().size())))
     return emitOpError("expects number of case values to match number of "
                        "case destinations");
   if (getBranchWeights() && getBranchWeights()->size() != getNumSuccessors())
     return emitError("expects number of branch weights to match number of "
                      "successors: ")
            << getBranchWeights()->size() << " vs " << getNumSuccessors();
+  if (getCaseValues() &&
+      getValue().getType() != getCaseValues()->getElementType())
+    return emitError("expects case value type to match condition value type");
   return success();
 }
 
 SuccessorOperands SwitchOp::getSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
   return SuccessorOperands(index == 0 ? getDefaultOperandsMutable()
                                       : getCaseOperandsMutable(index - 1));
 }

mlir/lib/Target/LLVMIR/ModuleImport.cpp

if (inst->getOpcode() == llvm::Instruction::Switch) {
     llvm::BasicBlock *defaultBB = swInst->getDefaultDest();
     if (failed(convertBranchArgs(swInst, defaultBB, defaultBlockArgs)))
       return failure();
 
     // Process the cases.
     unsigned numCases = swInst->getNumCases();
     SmallVector<SmallVector<Value>> caseOperands(numCases);
     SmallVector<ValueRange> caseOperandRefs(numCases);
-    SmallVector<int32_t> caseValues(numCases);
+    SmallVector<APInt> caseValues(numCases);
     SmallVector<Block *> caseBlocks(numCases);
     for (const auto &it : llvm::enumerate(swInst->cases())) {
       const llvm::SwitchInst::CaseHandle &caseHandle = it.value();
       llvm::BasicBlock *succBB = caseHandle.getCaseSuccessor();
       if (failed(convertBranchArgs(swInst, succBB, caseOperands[it.index()])))
         return failure();
       caseOperandRefs[it.index()] = caseOperands[it.index()];
-      caseValues[it.index()] = caseHandle.getCaseValue()->getSExtValue();
+      caseValues[it.index()] = caseHandle.getCaseValue()->getValue();
       caseBlocks[it.index()] = lookupBlock(succBB);
     }
 
     auto switchOp = builder.create<SwitchOp>(
         loc, *condition, lookupBlock(defaultBB), defaultBlockArgs, caseValues,
         caseBlocks, caseOperandRefs);
     mapNoResultOp(inst, switchOp);
     return success();

mlir/test/Dialect/LLVMIR/invalid.mlir

^bb1: // pred: ^bb0
   llvm.return
 
 ^bb2(%1: i32, %2: i32): // pred: ^bb0
   llvm.return
 }
 
 // -----
 
+func.func @switch_case_type_mismatch(%arg0 : i64) {
+  // expected-error@below {{expects case value type to match condition value type}}
+  "llvm.switch"(%arg0)[^bb1, ^bb2] <{case_operand_segments = array<i32: 0>, case_values = dense<42> : vector<1xi32>, operand_segment_sizes = array<i32: 1, 0, 0>}> : (i64) -> ()
+^bb1: // pred: ^bb0
+  llvm.return
+^bb2: // pred: ^bb0
+  llvm.return
+}
+
+// -----
+
 // expected-error@below {{expected zero value for 'common' linkage}}
 llvm.mlir.global common @non_zero_global_common_linkage(42 : i32) : i32
 
 // -----
 
 // expected-error@below {{expected zero value for 'common' linkage}}
 llvm.mlir.global common @non_zero_compound_global_common_linkage(dense<[0, 0, 0, 1, 0]> : vector<5xi32>) : !llvm.array<5 x i32>
 

mlir/test/Target/LLVMIR/Import/control-flow.ll

bb2:
   %1 = add i64 %arg2, 3
   br label %bb1
 }
 
 ; // -----
 
 ; CHECK-LABEL: @simple_switch(
 ; CHECK-SAME:  %[[ARG1:[a-zA-Z0-9]+]]
-define i32 @simple_switch(i32 %arg1) {
+define i64 @simple_switch(i64 %arg1) {
   ; CHECK: %[[VAL1:.+]] = llvm.add
   ; CHECK: %[[VAL2:.+]] = llvm.sub
   ; CHECK: %[[VAL3:.+]] = llvm.mul
-  %1 = add i32 %arg1, 42
-  %2 = sub i32 %arg1, 42
-  %3 = mul i32 %arg1, 42
-  ; CHECK: llvm.switch %[[ARG1]] : i32, ^[[BBD:.+]] [
+  %1 = add i64 %arg1, 42
+  %2 = sub i64 %arg1, 42
+  %3 = mul i64 %arg1, 42
+  ; CHECK: llvm.switch %[[ARG1]] : i64, ^[[BBD:.+]] [
   ; CHECK:   0: ^[[BB1:.+]],
   ; CHECK:   9: ^[[BB2:.+]]
   ; CHECK: ]
-  switch i32 %arg1, label %bbd [
-    i32 0, label %bb1
-    i32 9, label %bb2
+  switch i64 %arg1, label %bbd [
+    i64 0, label %bb1
+    i64 9, label %bb2
   ]
 bb1:
   ; CHECK: ^[[BB1]]:
   ; CHECK: llvm.return %[[VAL1]]
-  ret i32 %1
+  ret i64 %1
 bb2:
   ; CHECK: ^[[BB2]]:
   ; CHECK: llvm.return %[[VAL2]]
-  ret i32 %2
+  ret i64 %2
 bbd:
   ; CHECK: ^[[BBD]]:
   ; CHECK: llvm.return %[[VAL3]]
-  ret i32 %3
+  ret i64 %3
 }
 
 ; // -----
 
 ; CHECK-LABEL: @switch_args
 ; CHECK-SAME:  %[[ARG1:[a-zA-Z0-9]+]]
 define i32 @switch_args(i32 %arg1) {
 entry: