Fix scf-to-cfg conversion for IndexSwitchOp

mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp

IndexSwitchLowering::matchAndRewrite(IndexSwitchOp op,
   // Convert the default region.
   FailureOr<Block *> defaultBlock = convertRegion(op.getDefaultRegion());
   if (failed(defaultBlock))
     return failure();
 
   // Create the switch.
   rewriter.setInsertionPointToEnd(condBlock);
   SmallVector<ValueRange> caseOperands(caseSuccessors.size(), {});
+
+  // Cast switch index to integer case value.
+  Value caseValue = rewriter.create<arith::IndexCastOp>(
+      op.getLoc(), rewriter.getI32Type(), op.getArg());
+
   rewriter.create<cf::SwitchOp>(
-      op.getLoc(), op.getArg(), *defaultBlock, ValueRange(),
+      op.getLoc(), caseValue, *defaultBlock, ValueRange(),
       rewriter.getDenseI32ArrayAttr(caseValues), caseSuccessors, caseOperands);
   rewriter.replaceOp(op, continueBlock->getArguments());
   return success();
 }
 
 void mlir::populateSCFToControlFlowConversionPatterns(
     RewritePatternSet &patterns) {
   patterns.add<ForLowering, IfLowering, ParallelLowering, WhileLowering,
                ExecuteRegionLowering, IndexSwitchLowering>(
       patterns.getContext());
   patterns.add<DoWhileLowering>(patterns.getContext(), /*benefit=*/2);
 }
 
 void SCFToControlFlowPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   populateSCFToControlFlowConversionPatterns(patterns);
 
   // Configure conversion to lower out SCF operations.
   ConversionTarget target(getContext());
-  target.addIllegalOp<scf::ForOp, scf::IfOp, scf::ParallelOp, scf::WhileOp,
-                      scf::ExecuteRegionOp>();
+  target.addIllegalOp<scf::ForOp, scf::IfOp, scf::IndexSwitchOp,
+                      scf::ParallelOp, scf::WhileOp, scf::ExecuteRegionOp>();
   target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
   if (failed(
           applyPartialConversion(getOperation(), target, std::move(patterns))))
     signalPassFailure();
 }
 
 std::unique_ptr<Pass> mlir::createConvertSCFToCFPass() {
   return std::make_unique<SCFToControlFlowPass>();
 }

mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir

 // CHECK:     cf.br ^[[bb3:.+]](%[[x]] : i64)
 // CHECK:   ^[[bb2]]:
 // CHECK:     %[[y:.+]] = "test.val2"
 // CHECK:     cf.br ^[[bb3]](%[[y:.+]] : i64)
 // CHECK:   ^[[bb3]](%[[z:.+]]: i64):
 // CHECK:     "test.bar"(%[[z]])
 // CHECK:     return
 
-// SWITCH-LABEL: @index_switch
+// CHECK-LABEL: @index_switch
 func.func @index_switch(%i: index, %a: i32, %b: i32, %c: i32) -> i32 {
-  // SWITCH: cf.switch %arg0 : index
-  // SWITCH-NEXT: default: ^bb3
-  // SWITCH-NEXT: 0: ^bb1
-  // SWITCH-NEXT: 1: ^bb2
+  // CHECK: %[[CASE:.*]] = arith.index_cast %arg0 : index to i32
+  // CHECK: cf.switch %[[CASE]] : i32
+  // CHECK-NEXT: default: ^[[DEFAULT:.+]],
+  // CHECK-NEXT: 0: ^[[bb1:.+]],
+  // CHECK-NEXT: 1: ^[[bb2:.+]]
   %0 = scf.index_switch %i -> i32
-  // SWITCH: ^bb1:
+  // CHECK: ^[[bb1]]:
   case 0 {
-    // SWITCH-NEXT: llvm.br ^bb4(%arg1
+    // CHECK-NEXT: br ^[[bb4:.+]](%arg1 : i32)
     scf.yield %a : i32
   }
-  // SWITCH: ^bb2:
+  // CHECK: ^[[bb2]]:
   case 1 {
-    // SWITCH-NEXT: llvm.br ^bb4(%arg2
+    // CHECK-NEXT: br ^[[bb4]](%arg2 : i32)
     scf.yield %b : i32
   }
-  // SWITCH: ^bb3:
+  // CHECK: ^[[DEFAULT]]:
   default {
-    // SWITCH-NEXT: llvm.br ^bb4(%arg3
+    // CHECK-NEXT: br ^[[bb4]](%arg3 : i32)
     scf.yield %c : i32
   }
-  // SWITCH: ^bb4(%[[V:.*]]: i32
-  // SWITCH-NEXT: return %[[V]]
+  // CHECK: ^[[bb4]](%[[V:.*]]: i32
+  // CHECK-NEXT: return %[[V]]
   return %0 : i32
 }