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

[Longsheng Mou]

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

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

>From 37f933ad2c687f38cd7799731ae93d75509056b3 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 29 Jun 2026 14:41:02 +0800
Subject: [PATCH 1/2] [mlir][cf] Fix APInt comparison crash

Fix crash when switch value and case values have different APInt bit width.
---
 mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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()));

>From 57fb513fbb3549cebda9bc44c901fd8bd8560f1d Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 29 Jun 2026 14:42:49 +0800
Subject: [PATCH 2/2] add test

---
 .../Dialect/ControlFlow/canonicalize.mlir     | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

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_]+]]