[Mlir-commits] [mlir] b27d97b - [mlir][tensor] Add test for invalid tensor.unpack + update error msg (#118275)

[llvmlistbot at llvm.org]

Author: Andrzej Warzyński
Date: 2024-12-02T15:33:05Z
New Revision: b27d97bef782a8405e80e7b9f60b44d0537f43b6

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

LOG: [mlir][tensor] Add test for invalid tensor.unpack + update error msg (#118275)

Adds a new test for invalid `tensor.unpack` operations where the output
rank does not match the expected rank (input rank + num inner tile
sizes). For example:

```mlir
tensor.unpack %output
  inner_dims_pos = [0, 1]
  inner_tiles = [32, 16]
  into %input : tensor<64x32x16xf32> -> tensor<256x128xf32>
```

In addition, updates the corresponding error message to make it more
informative:

BEFORE:
```mlir
error: packed rank must equal unpacked rank + tiling factors}
```

AFTER:
```mlir
error: packed rank != (unpacked rank + num tiling factors), got 3 != 4
```

Added: 
    

Modified: 
    mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
    mlir/test/Dialect/Tensor/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 616d4a7d0a0ab5..9bb628781342ca 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3983,9 +3983,11 @@ static LogicalResult commonVerifierPackAndUnPackOp(OpTy packOrUnPack) {
                               : packOrUnPack.getSourceType();
   size_t packedRank = packedType.getRank();
   // Require output rank to match input rank + number of blocking factors.
-  if (unpackedRank + mixedTiles.size() != packedRank) {
+  size_t expectedPackedRank = unpackedRank + mixedTiles.size();
+  if (expectedPackedRank != packedRank) {
     return op->emitError(
-        "packed rank must equal unpacked rank + tiling factors");
+               "packed rank != (unpacked rank + num tiling factors), got ")
+           << packedRank << " != " << expectedPackedRank;
   }
 
   // Verify result shape is greater than the minimum expected

diff  --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index be470ce2af9b31..77cae1cc5f242d 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt <%s -split-input-file -verify-diagnostics
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 // Asking the dimension of a 0-D shape doesn't make sense.
 func.func @dim_0_ranked(%arg : tensor<f32>, %arg1 : index) {
@@ -692,13 +692,21 @@ func.func @pack_invalid_duplicate_element_in_outer_perm(%input: tensor<256x128xf
 // -----
 
 func.func @pack_invalid_output_rank(%input: tensor<256x128xf32>, %output: tensor<64x32x16xf32>) -> tensor<64x32x16xf32> {
-  // expected-error at +1 {{packed rank must equal unpacked rank + tiling factors}}
+  // expected-error at +1 {{packed rank != (unpacked rank + num tiling factors), got 3 != 4}}
   %0 = tensor.pack %input inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %output : tensor<256x128xf32> -> tensor<64x32x16xf32>
   return %0 : tensor<64x32x16xf32>
 }
 
 // -----
 
+func.func @pack_invalid_output_rank(%input: tensor<256x128xf32>, %output: tensor<64x32x16xf32>) -> tensor<256x128xf32> {
+  // expected-error at +1 {{packed rank != (unpacked rank + num tiling factors), got 3 != 4}}
+  %0 = tensor.unpack %output inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %input : tensor<64x32x16xf32> -> tensor<256x128xf32>
+  return %0 : tensor<256x128xf32>
+}
+
+// -----
+
 func.func @unpack_invalid_out_of_bound_outer_perm(%input: tensor<256x128xf32>, %output: tensor<8x8x32x16xf32>) -> tensor<8x8x32x16xf32> {
   // expected-error at +1 {{invalid outer_dims_perm vector}}
   %0 = tensor.unpack %output outer_dims_perm = [2, 1] inner_dims_pos = [0, 1] inner_tiles = [2, 2] into %input : tensor<8x8x32x16xf32> -> tensor<256x128xf32>