Add a negative test and modify the return statement

Author: Prakhar-Dixit

mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp

@@ -356,9 +356,9 @@ struct UnrollMultiReductionPattern
   LogicalResult matchAndRewrite(vector::MultiDimReductionOp reductionOp,
                                 PatternRewriter &rewriter) const override {
     auto resultType = reductionOp->getResult(0).getType();
-    if (mlir::isa<mlir::FloatType>(resultType) ||
-        mlir::isa<mlir::IntegerType>(resultType)) {
-      return failure();
+    if (resultType.isIntOrFloat()) {
+      return rewriter.notifyMatchFailure(reductionOp,
+                                         "Unrolling scalars is not supported");
     }
     std::optional<SmallVector<int64_t>> targetShape =
         getTargetShape(options, reductionOp);

mlir/test/Dialect/Vector/vector-unroll-options.mlir

@@ -222,6 +222,13 @@ 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>
 
+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:   %[[R0:.*]] = vector.multi_reduction <add>, %{{.*}}, %{{.*}} [0, 1] : vector<4x2xf32> to f32
+//       CHECK:   return %[[R0]] : f32
 
 func.func @vector_reduction(%v : vector<8xf32>) -> f32 {
   %0 = vector.reduction <add>, %v : vector<8xf32> into f32