[Mlir-commits] [mlir] [MLIR][Parser] Fix result numbers not being allowed in switch statements (PR #224581)

[Yue Huang]

https://github.com/AdUhTkJm created https://github.com/llvm/llvm-project/pull/224581

PR #87658 fixed this for the `case` branch for `cf.switch`, but not the `default` case. Moreover, `llvm.switch` has the same problem. This PR fixes them.

>From 26668ec8ca4b303712b038336cec4b142d90b10a Mon Sep 17 00:00:00 2001
From: Yue Huang <yh548 at cam.ac.uk>
Date: Fri, 18 Sep 2026 10:42:48 +0100
Subject: [PATCH] [MLIR][Parser] Fix result numbers not being allowed in switch
 statements

---
 mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp |  4 ++--
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp         |  3 +--
 mlir/test/Dialect/ControlFlow/ops.mlir             | 13 +++++++++++++
 mlir/test/Dialect/LLVMIR/roundtrip.mlir            | 13 +++++++++++++
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index 435c37bc95aac..9016d1929f859 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -638,8 +638,8 @@ static ParseResult parseSwitchOpCases(
       parser.parseSuccessor(defaultDestination))
     return failure();
   if (succeeded(parser.parseOptionalLParen())) {
-    if (parser.parseOperandList(defaultOperands, OpAsmParser::Delimiter::None,
-                                /*allowResultNumber=*/false) ||
+    if (parser.parseOperandList(defaultOperands,
+                                OpAsmParser::Delimiter::None) ||
         parser.parseColonTypeList(defaultOperandTypes) || parser.parseRParen())
       return failure();
   }
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 114800058f7e6..e70d09bcdb0cc 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -565,8 +565,7 @@ ParseResult mlir::LLVM::parseSwitchOpCases(
     if (parser.parseColon() || parser.parseSuccessor(destination))
       return failure();
     if (!parser.parseOptionalLParen()) {
-      if (parser.parseOperandList(operands, OpAsmParser::Delimiter::None,
-                                  /*allowResultNumber=*/false) ||
+      if (parser.parseOperandList(operands, OpAsmParser::Delimiter::None) ||
           parser.parseColonTypeList(operandTypes) || parser.parseRParen())
         return failure();
     }
diff --git a/mlir/test/Dialect/ControlFlow/ops.mlir b/mlir/test/Dialect/ControlFlow/ops.mlir
index 160534240e0fa..03091a95b456e 100644
--- a/mlir/test/Dialect/ControlFlow/ops.mlir
+++ b/mlir/test/Dialect/ControlFlow/ops.mlir
@@ -52,6 +52,19 @@ func.func @switch_result_number(%arg0: i32) {
     return
 }
 
+// CHECK-LABEL: func @switch_result_number_default
+func.func @switch_result_number_default(%arg0: i32) {
+  %0:2 = "test.op_with_two_results"() : () -> (i32, i32)
+  cf.switch %arg0 : i32, [
+    default: ^bb1(%0#0 : i32),
+    0: ^bb2(%0#1 : i32)
+  ]
+  ^bb1(%1: i32):
+    return
+  ^bb2(%2: i32):
+    return
+}
+
 // CHECK-LABEL: func @cond_weights
 func.func @cond_weights(%cond: i1) {
 // CHECK: cf.cond_br %{{.*}} weights([60, 40]), ^{{.*}}, ^{{.*}}
diff --git a/mlir/test/Dialect/LLVMIR/roundtrip.mlir b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
index f33cce1b8928e..02e0bdaf3b694 100644
--- a/mlir/test/Dialect/LLVMIR/roundtrip.mlir
+++ b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
@@ -1334,3 +1334,16 @@ llvm.func @masked_nontemporal_no_alignment(%ptr: !llvm.ptr, %mask: vector<7xi1>)
   llvm.intr.masked.store(%0, %ptr, %mask) <nontemporal> : vector<7xf32>, vector<7xi1> into !llvm.ptr
   llvm.return
 }
+
+// CHECK-LABEL: @switch_result_number
+llvm.func @switch_result_number(%arg0: i32) -> i32 {
+  %0:2 = "test.op_with_two_results"() : () -> (i32, i32)
+  // CHECK: llvm.switch
+  llvm.switch %arg0 : i32, ^bb1(%0#0 : i32) [
+    0: ^bb2(%0#1 : i32)
+  ]
+^bb1(%1: i32):
+  llvm.return %1 : i32
+^bb2(%2: i32):
+  llvm.return %2 : i32
+}