[MLIR] Add test cases for illegal output shape argument

Signed-Off-by: Gaurav Shukla<[email protected]>

Author: Shukla-Gaurav

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

@@ -762,14 +762,14 @@ validateDynamicDimExpansion(LinalgOp linalgOp,
       continue;
     bool foundDynamic = false;
     for (int64_t shape : expandedShape) {
-      if (ShapedType::isDynamic(shape)) {
-        if (foundDynamic) {
-          return rewriter.notifyMatchFailure(
-              linalgOp, "cannot infer expanded shape with multiple dynamic "
-                        "dims in the same reassociation group");
-        }
-        foundDynamic = true;
+      if (!ShapedType::isDynamic(shape))
+        continue;
+      if (foundDynamic) {
+        return rewriter.notifyMatchFailure(
+            linalgOp, "cannot infer expanded shape with multiple dynamic "
+                      "dims in the same reassociation group");
       }
+      foundDynamic = true;
     }
   }
   return success();

mlir/test/Dialect/MemRef/invalid.mlir

@@ -408,6 +408,14 @@ func.func @expand_shape(%arg0: memref<f32>) {
 
 // -----
 
+func.func @expand_shape_illegal_output_shape(%arg0: memref<2xf32>) {
+  // expected-error @+1 {{expected number of static shape bounds to be equal to the output rank (3) but found 2 inputs instead}}
+  %0 = memref.expand_shape %arg0 [[0, 1, 2]] output_shape [1, 2] : memref<2xf32> into memref<1x1x2xf32>
+  return
+}
+
+// -----
+
 func.func @collapse_shape_out_of_bounds(%arg0: memref<?x?xf32>) {
   // expected-error @+1 {{op reassociation index 2 is out of bounds}}
   %0 = memref.collapse_shape %arg0 [[0, 1, 2]] : memref<?x?xf32> into memref<?xf32>

mlir/test/Dialect/Tensor/invalid.mlir

@@ -311,6 +311,15 @@ func.func @illegal_expanding_reshape_mixed_tensor_2(%arg0 : tensor<?x?xf32>, %sz
   return %0 : tensor<?x4x5xf32>
 }
 
+// -----
+
+func.func @expand_shape_illegal_output_shape(%arg0: tensor<2xf32>) {
+  // expected-error @+1 {{expected number of static shape dims to be equal to the output rank (3) but found 2 inputs instead}}
+  %0 = tensor.expand_shape %arg0 [[0, 1, 2]] output_shape [1, 2] : tensor<2xf32> into tensor<1x1x2xf32>
+  return
+}
+
+
 // -----
 
 func.func @illegal_collapsing_reshape_mixed_tensor(%arg0 : tensor<?x4x5xf32>) -> tensor<?x?xf32> {