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

[Andrzej Warzyński]

https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/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
```


>From 8f3aa9bc2c4a98133431bc581bdf87b27b97fb5d Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Mon, 2 Dec 2024 09:28:58 +0000
Subject: [PATCH] [mlir][tensor] Add test for invalid tensor.unpack + update
 error msg

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
```
---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp |  6 ++++--
 mlir/test/Dialect/Tensor/invalid.mlir    | 12 ++++++++++--
 2 files changed, 14 insertions(+), 4 deletions(-)

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>