[Mlir-commits] [mlir] [mlir][cf] Fix APInt comparison crash in `foldSwitch` (PR #206407)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-cf

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

Fix crash when switch value and case values have different APInt bit width. Fixes #<!-- -->205234.

---
Full diff: https://github.com/llvm/llvm-project/pull/206407.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp (+2-1) 
- (modified) mlir/test/Dialect/ControlFlow/canonicalize.mlir (+26) 


``````````diff
diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index 435c37bc95aac..53cf35f0599fe 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -808,9 +808,10 @@ dropSwitchCasesThatMatchDefault(SwitchOp op, PatternRewriter &rewriter) {
 /// -> br ^bb2
 static void foldSwitch(SwitchOp op, PatternRewriter &rewriter,
                        const APInt &caseValue) {
+  uint64_t caseLimitedValue = caseValue.getLimitedValue();
   auto caseValues = op.getCaseValues();
   for (const auto &it : llvm::enumerate(caseValues->getValues<APInt>())) {
-    if (it.value() == caseValue) {
+    if (it.value().getLimitedValue() == caseLimitedValue) {
       rewriter.replaceOpWithNewOp<BranchOp>(
           op, op.getCaseDestinations()[it.index()],
           op.getCaseOperands(it.index()));
diff --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
index c8cf5931a1297..b5a0e98b9bdc8 100644
--- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir
+++ b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
@@ -263,6 +263,32 @@ func.func @switch_on_const_with_match(%caseOperand0 : f32, %caseOperand1 : f32,
     "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
 }
 
+// CHECK-LABEL: func @switch_on_const_with_match_diff_width(
+// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
+func.func @switch_on_const_with_match_diff_width(%caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
+  // add predecessors for all blocks to avoid other canonicalizations.
+  "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4] : () -> ()
+  ^bb1:
+    // CHECK-NOT: cf.switch
+    // CHECK: cf.br ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
+    %c0_i32 = llvm.mlir.constant(1 : index) : i32
+    cf.switch %c0_i32 : i32, [
+      default: ^bb2(%caseOperand0 : f32),
+      -1: ^bb3(%caseOperand1 : f32),
+      1: ^bb4(%caseOperand2 : f32)
+    ]
+  ^bb2(%bb2Arg : f32):
+    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
+  ^bb3(%bb3Arg : f32):
+    "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
+  // CHECK: ^[[BB4]]({{.*}}):
+  // CHECK-NEXT: "foo.bb4Terminator"
+  ^bb4(%bb4Arg : f32):
+    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
+}
+
 // CHECK-LABEL: func @switch_passthrough(
 // CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]

``````````

</details>


https://github.com/llvm/llvm-project/pull/206407