[Mlir-commits] [mlir] 0377562 - [mlir]Add a check to ensure bailing out when reducing to a scalar (#129694)

[llvmlistbot at llvm.org]

Author: Prakhar Dixit
Date: 2025-03-08T16:34:01Z
New Revision: 037756242f6eb3cdd6e73ab6597651f87aa7b933

URL: https://github.com/llvm/llvm-project/commit/037756242f6eb3cdd6e73ab6597651f87aa7b933
DIFF: https://github.com/llvm/llvm-project/commit/037756242f6eb3cdd6e73ab6597651f87aa7b933.diff

LOG: [mlir]Add a check to ensure bailing out when reducing to a scalar (#129694)

Fixes issue #64075
Referencing this comment for more detailed view ->
https://github.com/llvm/llvm-project/issues/64075#issuecomment-2694112594

**Minimal example crashing :** 
```
func.func @multi_reduction(%0: vector<4x2xf32>, %acc1: f32) -> f32 {
  %2 = vector.multi_reduction <add>, %0, %acc1 [0, 1] : vector<4x2xf32> to f32
  return %2 : f32
}
```

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
    mlir/test/Dialect/Vector/vector-unroll-options.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 08ba972b12ce6..04c38f9f7b2e3 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -355,6 +355,11 @@ struct UnrollMultiReductionPattern
 
   LogicalResult matchAndRewrite(vector::MultiDimReductionOp reductionOp,
                                 PatternRewriter &rewriter) const override {
+    auto resultType = reductionOp->getResult(0).getType();
+    if (resultType.isIntOrFloat()) {
+      return rewriter.notifyMatchFailure(reductionOp,
+                                         "Unrolling scalars is not supported");
+    }
     std::optional<SmallVector<int64_t>> targetShape =
         getTargetShape(options, reductionOp);
     if (!targetShape)

diff  --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 16d30aec7c041..9c158d05b723c 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -222,6 +222,15 @@ func.func @vector_multi_reduction(%v : vector<4x6xf32>, %acc: vector<4xf32>) ->
 //       CHECK:   %[[V2:.*]] = vector.insert_strided_slice %[[R5]], %[[V1]] {offsets = [2], strides = [1]} : vector<2xf32> into vector<4xf32>
 //       CHECK:   return %[[V2]] : vector<4xf32>
 
+// This is a negative test case to ensure that further unrolling is not performed. Since the vector.multi_reduction
+// operation has already been unrolled, attempting additional unrolling should not be allowed.
+func.func @negative_vector_multi_reduction(%v: vector<4x2xf32>, %acc: f32) -> f32 {
+  %0 = vector.multi_reduction #vector.kind<add>, %v, %acc [0, 1] : vector<4x2xf32> to f32
+  return %0 : f32
+}
+// CHECK-LABEL: func @negative_vector_multi_reduction
+//  CHECK-NEXT:   %[[R0:.*]] = vector.multi_reduction <add>, %{{.*}}, %{{.*}} [0, 1] : vector<4x2xf32> to f32
+//  CHECK-NEXT:   return %[[R0]] : f32
 
 func.func @vector_reduction(%v : vector<8xf32>) -> f32 {
   %0 = vector.reduction <add>, %v : vector<8xf32> into f32