[Mlir-commits] [mlir] 03b43d3 - Fix scf-to-cfg conversion for IndexSwitchOp

[Eugene Zhulenev]

Author: Eugene Zhulenev
Date: 2023-08-18T16:18:43-07:00
New Revision: 03b43d3dfbb38819c210a7f996602d9c25268da3

URL: https://github.com/llvm/llvm-project/commit/03b43d3dfbb38819c210a7f996602d9c25268da3
DIFF: https://github.com/llvm/llvm-project/commit/03b43d3dfbb38819c210a7f996602d9c25268da3.diff

LOG: Fix scf-to-cfg conversion for IndexSwitchOp

Reviewed By: mehdi_amini

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

Added: 
    

Modified: 
    mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
    mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index 51849518be6649..91dbdb429f948e 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/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))))

diff  --git a/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir b/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
index 94bacd258470f6..36307a910a6cad 100644
--- a/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
+++ b/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
 }