[mlir][LLVMIR] Add support for translating Switch instruction

[myhsu]

Summary

Add support for translating llvm::SwitchInst.

[Mogball]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

866

add braces

872

spell out the auto

873

add braces

884

llvm::enumerate(swInst->cases()) or llvm::zip?

885

spell out the auto

886

add braces

[myhsu]

Addressed coding style issues.

[Mogball]

Thanks for all the fixes/additions to the LLVM importer

[ftynse]

Please add tests for user-visible error messages.

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

867–868

The (unwritten) MLIR convention is to have a test for every user-visible error message. Here, I suspect we will see two error messages in a row because whatever failed in processValue should have reported the error as well. Callers usually just propagate the error state.

mlir/test/Target/LLVMIR/Import/switch.ll

6 ↗

(On Diff #425904)

Nit: can't this just live with the rest of ops in import.ll?

[myhsu]

• Integrate test cases into test/Target/LLVMIR/Import/basic.ll
• Add a new test case to test llvm.switch with arguments
• Remove error message printing while translating predicate values and block arguments to avoid showing (similar) error messages twice.

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

867–868

Here, I suspect we will see two error messages in a row because whatever failed in processValue should have reported the error as well. Callers usually just propagate the error state.

I feel like a better solution will be showing "failed to process the condition value of...", in addition to the error message from processValue, as a note (similar to what clang does) so that user will have a quicker grasp on the problem. But that's a different story I guess.

I will remove the emitError here, the one at line 821 and 836 to avoid confusion (i.e. two error messages). But I'll still add a test case that contains block arguments

[ftynse]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

867–868

Yes, it would be better to have the general capability of spreading the construction of a diagnostic across multiple calls, but we don't have it now.

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

 // Maps from LLVM opcode to MLIR OperationName. This is deliberately ordered
 // as in llvm/IR/Instructions.def to aid comprehension and spot missing
 // instructions.
 #define INST(llvm_n, mlir_n)                                                   \
   { llvm::Instruction::llvm_n, LLVM::mlir_n##Op::getOperationName() }
   static const DenseMap<unsigned, StringRef> opcMap = {
       // Ret is handled specially.
       // Br is handled specially.
-      // FIXME: switch
+      // Switch is handled specially.
       // FIXME: indirectbr
       // FIXME: invoke
       INST(Resume, Resume),
       // FIXME: unreachable
       // FIXME: cleanupret
       // FIXME: catchret
       // FIXME: catchswitch
       // FIXME: callbr
case llvm::Instruction::Br: {
     if (brInst->isConditional()) {
       state.addAttribute(LLVM::CondBrOp::getOperandSegmentSizeAttr(),
                          b.getI32VectorAttr(operandSegmentSizes));
     }
 
     b.create(state);
     return success();
   }
+  case llvm::Instruction::Switch: {
+    auto *swInst = cast<llvm::SwitchInst>(inst);
+    // Process the condition value.
+    Value condition = processValue(swInst->getCondition());
+    if (!condition)
+      return failure();
+
+    SmallVector<Value> defaultBlockArgs;
+    // Process the default case.
+    llvm::BasicBlock *defaultBB = swInst->getDefaultDest();
+    if (failed(processBranchArgs(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<Block *> caseBlocks(numCases);
+    for (const auto &en : llvm::enumerate(swInst->cases())) {
+      const llvm::SwitchInst::CaseHandle &caseHandle = en.value();
+      unsigned i = en.index();
+      llvm::BasicBlock *succBB = caseHandle.getCaseSuccessor();
+      if (failed(processBranchArgs(swInst, succBB, caseOperands[i])))
+        return failure();
+      caseOperandRefs[i] = caseOperands[i];
+      caseValues[i] = caseHandle.getCaseValue()->getSExtValue();
+      caseBlocks[i] = blocks[succBB];
+    }
+
+    b.create<SwitchOp>(loc, condition, blocks[defaultBB], defaultBlockArgs,
+                       caseValues, caseBlocks, caseOperandRefs);
+    return success();
+  }
   case llvm::Instruction::PHI: {
     Type type = processType(inst->getType());
     if (!type)
       return failure();
     instMap[inst] = b.getInsertionBlock()->addArgument(
         type, processDebugLoc(inst->getDebugLoc(), inst));
     return success();
   }

mlir/test/Target/LLVMIR/Import/basic.ll

define i32 @useFenceInst() {
   ;CHECK: llvm.fence syncscope("agent") seq_cst
   fence syncscope("agent") seq_cst
   ;CHECK: llvm.fence release
   fence release
   ;CHECK: llvm.fence seq_cst
   fence syncscope("") seq_cst
   ret i32 0
 }
+
+; Switch instruction
+declare void @g(i32)
+
+; CHECK-LABEL: llvm.func @simple_switch(%arg0: i32) {
+define void @simple_switch(i32 %val) {
+; CHECK: %[[C0:.+]] = llvm.mlir.constant(11 : i32) : i32
+; CHECK: %[[C1:.+]] = llvm.mlir.constant(87 : i32) : i32
+; CHECK: %[[C2:.+]] = llvm.mlir.constant(78 : i32) : i32
+; CHECK: %[[C3:.+]] = llvm.mlir.constant(94 : i32) : i32
+; CHECK: %[[C4:.+]] = llvm.mlir.constant(1 : i32) : i32
+; CHECK: llvm.switch %arg0 : i32, ^[[BB5:.+]] [
+; CHECK:   0: ^[[BB1:.+]],
+; CHECK:   9: ^[[BB2:.+]],
+; CHECK:   994: ^[[BB3:.+]],
+; CHECK:   1154: ^[[BB4:.+]]
+; CHECK: ]
+  switch i32 %val, label %def [
+    i32 0, label %one
+    i32 9, label %two
+    i32 994, label %three
+    i32 1154, label %four
+  ]
+
+; CHECK: ^[[BB1]]:
+; CHECK: llvm.call @g(%[[C4]]) : (i32) -> ()
+; CHECK: llvm.return
+one:
+  call void @g(i32 1)
+  ret void
+; CHECK: ^[[BB2]]:
+; CHECK: llvm.call @g(%[[C3]]) : (i32) -> ()
+; CHECK: llvm.return
+two:
+  call void @g(i32 94)
+  ret void
+; CHECK: ^[[BB3]]:
+; CHECK: llvm.call @g(%[[C2]]) : (i32) -> ()
+; CHECK: llvm.return
+three:
+  call void @g(i32 78)
+  ret void
+; CHECK: ^[[BB4]]:
+; CHECK: llvm.call @g(%[[C1]]) : (i32) -> ()
+; CHECK: llvm.return
+four:
+  call void @g(i32 87)
+  ret void
+; CHECK: ^[[BB5]]:
+; CHECK: llvm.call @g(%[[C0]]) : (i32) -> ()
+; CHECK: llvm.return
+def:
+  call void @g(i32 11)
+  ret void
+}
+
+; CHECK-LABEL: llvm.func @switch_args(%arg0: i32) {
+define void @switch_args(i32 %val) {
+  ; CHECK: %[[C0:.+]] = llvm.mlir.constant(44 : i32) : i32
+  ; CHECK: %[[C1:.+]] = llvm.mlir.constant(34 : i32) : i32
+  ; CHECK: %[[C2:.+]] = llvm.mlir.constant(33 : i32) : i32
+  %pred = icmp ult i32 %val, 87
+  br i1 %pred, label %bbs, label %bb1
+
+bb1:
+  %vx = add i32 %val, 22
+  %pred2 = icmp ult i32 %val, 94
+  br i1 %pred2, label %bb2, label %bb3
+
+bb2:
+  %vx0 = add i32 %val, 23
+  br label %one
+
+bb3:
+  br label %def
+
+; CHECK: %[[V1:.+]] = llvm.add %arg0, %[[C2]] : i32
+; CHECK: %[[V2:.+]] = llvm.add %arg0, %[[C1]] : i32
+; CHECK: %[[V3:.+]] = llvm.add %arg0, %[[C0]] : i32
+; CHECK: llvm.switch %arg0 : i32, ^[[BBD:.+]](%[[V3]] : i32) [
+; CHECK:   0: ^[[BB1:.+]](%[[V1]], %[[V2]] : i32, i32)
+; CHECK: ]
+bbs:
+  %vy = add i32 %val, 33
+  %vy0 = add i32 %val, 34
+  %vz = add i32 %val, 44
+  switch i32 %val, label %def [
+    i32 0, label %one
+  ]
+
+; CHECK: ^[[BB1]](%[[BA0:.+]]: i32, %[[BA1:.+]]: i32):
+one: ; pred: bb2, bbs
+  %v0 = phi i32 [%vx, %bb2], [%vy, %bbs]
+  %v1 = phi i32 [%vx0, %bb2], [%vy0, %bbs]
+  ; CHECK: llvm.add %[[BA0]], %[[BA1]]  : i32
+  %vf = add i32 %v0, %v1
+  call void @g(i32 %vf)
+  ret void
+
+; CHECK: ^[[BBD]](%[[BA2:.+]]: i32):
+def: ; pred: bb3, bbs
+  %v2 = phi i32 [%vx, %bb3], [%vz, %bbs]
+  ; CHECK: llvm.call @g(%[[BA2]])
+  call void @g(i32 %v2)
+  ret void
+}