Fix scf-to-cfg conversion for IndexSwitchOp

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D158212

Author: ezhulenev

mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp

@@ -665,8 +665,13 @@ IndexSwitchLowering::matchAndRewrite(IndexSwitchOp op,
   // 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();
@@ -686,8 +691,8 @@ void SCFToControlFlowPass::runOnOperation() {
 
   // 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))))

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

@@ -621,29 +621,30 @@ func.func @func_execute_region_elim_multi_yield() {
 // 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
 }