TosaToTensor regression from #85798

[bjacob]

The following test starts failing as of #85798, with a failing assertion introduced in that PR:

  func.func @mlp_invocation(%lhs: tensor<2x4xf32>, %rhs : tensor<4x8xf32>) -> tensor<2x8xf32> {
    %lhs_3D = tosa.reshape %lhs {new_shape = array<i64 : 1, 2, 2>} : (tensor<2x4xf32>) -> tensor<1x2x4xf32>
    %rhs_3D = tosa.reshape %rhs {new_shape = array<i64 : 1, 2, 2>} : (tensor<4x8xf32>) -> tensor<1x4x8xf32>
    %0 = tosa.matmul %lhs_3D, %rhs_3D : (tensor<1x2x4xf32>, tensor<1x4x8xf32>) -> tensor<1x2x8xf32>
    %1 = tosa.clamp %0 {
        min_int = 0 : i64, max_int = 9223372036854775807 : i64,
        min_fp = 0.0 : f32, max_fp = 3.4028235e+38 : f32}
        : (tensor<1x2x8xf32>) -> tensor<1x2x8xf32>
    %2 = tosa.negate %1 : (tensor<1x2x8xf32>) -> tensor<1x2x8xf32>
    %3 = tosa.reshape %2 {new_shape = array<i64 : 2, 2>}  : (tensor<1x2x8xf32>) -> tensor<2x8xf32>
    return %3 : tensor<2x8xf32>
  }

Run:

mlir-opt --tosa-to-tensor /tmp/a.mlir

Result:

mlir-opt: /home/benoit/megabump/work/iree/third_party/llvm-project/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp:196: SmallVector<ReassociationExprs> (anonymous namespace)::createReassociationMapForCollapse(OpBuilder &, Type, Type): Assertion `currSrcDim == srcShape.size() && currDstDim == dstShape.size()' failed.

[rafaelubalmw]

Hi Benoit.

It looks like all 3 tosa.reshape ops in that code are malformed, in the sense that the new_shape attribute does not match the statically sized return tensor. The product of new_shape elements also does not equal the total size of the input tensor.

Even in that case, it is true that mlir-opt shouldn't be asserting. I propose a fix that involves extending the tosa.reshape op verifier as follows:

• Element i of new_shape (if not -1 placeholder) must match dimension size i of the return type (if not dynamic).

• An inferred return shape is reconstructed either from non-placeholder elements of new_shape or static dimensions of the return type. If the input and inferred return shape are both static, their total size must match.

Attempting to run your command with this fix would cause the op verifier to flag the input code as invalid before the --tosa-to-tensor pass executes. Please let me know your thoughts.

[bjacob]

Thanks for catching this and the explanation. This does unblock me either way but thanks for also making sure this this code doesn't assert on an incorrect input.