diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp index 616d4a7d0a0ab..9bb628781342c 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 be470ce2af9b3..77cae1cc5f242 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, %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@+1 {{packed rank must equal unpacked rank + tiling factors}} + // expected-error@+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@+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@+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>