[Mlir-commits] [mlir] [mlir][vector] Fix crash in ReorderCastOpsOnBroadcast with non-vector result (PR #170985)

[llvmlistbot at llvm.org]

https://github.com/Men-cotton updated https://github.com/llvm/llvm-project/pull/170985

>From 0b120424f43aa30d7bd6e75b6898e171ea295ac3 Mon Sep 17 00:00:00 2001
From: mencotton <mencotton0410 at gmail.com>
Date: Sat, 6 Dec 2025 22:21:24 +0900
Subject: [PATCH 1/2] [mlir][vector] Fix crash in ReorderCastOpsOnBroadcast
 with non-vector result

---
 .../Vector/Transforms/VectorTransforms.cpp      |  2 ++
 mlir/test/Dialect/Vector/vector-sink.mlir       | 17 +++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 726da1e9a3d14..ad16b80a732b3 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -453,6 +453,8 @@ struct ReorderCastOpsOnBroadcast
                                 PatternRewriter &rewriter) const override {
     if (op->getNumOperands() != 1)
       return failure();
+    if (!isa<VectorType>(op->getResult(0).getType()))
+      return failure();
     auto bcastOp = op->getOperand(0).getDefiningOp<vector::BroadcastOp>();
     if (!bcastOp)
       return failure();
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index beaba52af1841..50ff63b44901a 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -850,3 +850,20 @@ func.func @negative_store_no_single_use(%arg0: memref<?xf32>, %arg1: index, %arg
   vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<1xf32>
   return %0 : vector<1xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func.func @broadcast_cast_non_vector_result
+// CHECK-SAME: (%[[ARG:.*]]: i64)
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG]] : i64 to vector<26x7xi64>
+// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[BCAST]] : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>>
+// CHECK: return %[[CAST]] : !llvm.array<26 x vector<7xi64>>
+/// This test ensures that the `ReorderCastOpsOnBroadcast` pattern does not
+/// attempt to reorder a cast operation that produces a non-vector result type.
+/// Previously, this would crash because the pattern assumed the result was a
+/// vector type when creating the new inner broadcast.
+func.func @broadcast_cast_non_vector_result(%arg0: i64) -> !llvm.array<26 x vector<7xi64>> {
+  %0 = vector.broadcast %arg0 : i64 to vector<26x7xi64>
+  %1 = builtin.unrealized_conversion_cast %0 : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>>
+  return %1 : !llvm.array<26 x vector<7xi64>>
+}

>From 22dc42886066b9dc543f03e43e80be824059c8ed Mon Sep 17 00:00:00 2001
From: mencotton <mencotton0410 at gmail.com>
Date: Tue, 9 Dec 2025 09:49:47 +0900
Subject: [PATCH 2/2] fix: refine test docs

---
 mlir/test/Dialect/Vector/vector-sink.mlir | 32 +++++++++++------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index 50ff63b44901a..69fba88a14048 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -382,6 +382,21 @@ func.func @broadcast_scalar_extsi_scalable(%a : i8) -> vector<2x[4]xi32> {
   return %r : vector<2x[4]xi32>
 }
 
+// -----
+
+// CHECK-LABEL: func.func @negative_broadcast_cast_non_vector_result
+// CHECK-SAME: (%[[ARG:.*]]: i64)
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG]] : i64 to vector<26x7xi64>
+// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[BCAST]] : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>>
+// CHECK: return %[[CAST]] : !llvm.array<26 x vector<7xi64>>
+/// This test ensures that the `ReorderCastOpsOnBroadcast` pattern does not
+/// attempt to reorder a cast operation that produces a non-vector result type.
+func.func @negative_broadcast_cast_non_vector_result(%arg0: i64) -> !llvm.array<26 x vector<7xi64>> {
+  %0 = vector.broadcast %arg0 : i64 to vector<26x7xi64>
+  %1 = builtin.unrealized_conversion_cast %0 : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>>
+  return %1 : !llvm.array<26 x vector<7xi64>>
+}
+
 //===----------------------------------------------------------------------===//
 // [Pattern: ReorderElementwiseOpsOnTranspose]
 //===----------------------------------------------------------------------===//
@@ -850,20 +865,3 @@ func.func @negative_store_no_single_use(%arg0: memref<?xf32>, %arg1: index, %arg
   vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<1xf32>
   return %0 : vector<1xf32>
 }
-
-// -----
-
-// CHECK-LABEL: func.func @broadcast_cast_non_vector_result
-// CHECK-SAME: (%[[ARG:.*]]: i64)
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG]] : i64 to vector<26x7xi64>
-// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[BCAST]] : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>>
-// CHECK: return %[[CAST]] : !llvm.array<26 x vector<7xi64>>
-/// This test ensures that the `ReorderCastOpsOnBroadcast` pattern does not
-/// attempt to reorder a cast operation that produces a non-vector result type.
-/// Previously, this would crash because the pattern assumed the result was a
-/// vector type when creating the new inner broadcast.
-func.func @broadcast_cast_non_vector_result(%arg0: i64) -> !llvm.array<26 x vector<7xi64>> {
-  %0 = vector.broadcast %arg0 : i64 to vector<26x7xi64>
-  %1 = builtin.unrealized_conversion_cast %0 : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>>
-  return %1 : !llvm.array<26 x vector<7xi64>>
-}