Update return method in VectorUnroll.cpp and modify test in vector-unroll-options.mlir

Author: Prakhar-Dixit

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

@@ -441,7 +441,7 @@ struct UnrollElementwisePattern : public RewritePattern {
     // fact that `ExtractStridedSlice` requires the rank for the input and
     // output to match. If needed, we can relax this later.
     if (originalSize.size() != targetShape->size())
-      return failure();
+    return rewriter.notifyMatchFailure(op, "expected input vector rank to match target shape rank");
     Location loc = op->getLoc();
     // Prepare the result vector.
     Value result = rewriter.create<arith::ConstantOp>(

mlir/test/Dialect/Vector/invalid.mlir

@@ -769,9 +769,8 @@ func.func @extract_strided_slice(%arg0: vector<4x8x16xf32>) {
 // -----
 
 func.func @extract_strided_slice(%arg0: vector<3x2x2xf32>) {
-  // expected-error@+1 {{expected input vector rank to match target shape rank}}
-  %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]}:
-         vector<3x2x2xf32> to vector<2x2xf32>
+  // expected-error@+1 {{op expected input vector rank to match target shape rank}}
+  %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]}: vector<3x2x2xf32> to vector<2x2xf32>
   return
 }
 

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

@@ -189,13 +189,13 @@ func.func @vector_fma(%a: vector<4x4xf32>, %b: vector<4x4xf32>, %c: vector<4x4xf
 // CHECK-COUNT-4: vector.fma %{{.+}}, %{{.+}}, %{{.+}} : vector<2x2xf32>
 
 // TODO: We should be able to unroll this like the example above - this will require extending UnrollElementwisePattern.
-func.func @negative_vector_fma_3d(%a: vector<3x2x2xf32>) {
+func.func @negative_vector_fma_3d(%a: vector<3x2x2xf32>) -> vector<3x2x2xf32>{
   %0 = vector.fma %a, %a, %a : vector<3x2x2xf32>
-  return 
+  return %0 : vector<3x2x2xf32>
 }
 // CHECK-LABEL: func @negative_vector_fma_3d
 //   CHECK-NOT: vector.extract_strided_slice
-//       CHECK: vector.fma %{{.+}}, %{{.+}}, %{{.+}} : vector<3x2x2xf32>
+//       CHECK: %[[R0:.*]] = vector.fma %{{.+}} : vector<3x2x2xf32>
 //       CHECK: return 
 
 func.func @vector_multi_reduction(%v : vector<4x6xf32>, %acc: vector<4xf32>) -> vector<4xf32> {