From 0b120424f43aa30d7bd6e75b6898e171ea295ac3 Mon Sep 17 00:00:00 2001 From: mencotton 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(op->getResult(0).getType())) + return failure(); auto bcastOp = op->getOperand(0).getDefiningOp(); 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, %arg1: index, %arg vector.store %0, %arg0[%arg1] : memref, 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 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, %arg1: index, %arg vector.store %0, %arg0[%arg1] : memref, 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>> -}