[mlir][tosa] Fix out-of-boundaries iteration for tosa-to-linalg

[a-puschin]

Summary

When the number of elements of two shapes are not equal, a Reshape operation cannot be used to transfer one into another

Function findIntermediateShape(...) can cause out-of-boundaries operator[] call if the abovementioned condition strikes

The test-case I used now causes no error as its root-cause was an issue in Tosa dialect with padded Conv2D operations lowering which is already solved in commit 69c984b6

[jpienaar]

Could you add a test as well?

[a-puschin]

In D140013#3995044, @jpienaar wrote:

Could you add a test as well?

There is a very similar test already:
https://github.com/llvm/llvm-project/blob/15a6e3c636977dc962a415c067182e6d57242116/mlir/test/Dialect/Tosa/tosa-decompose-depthwise.mlir#L41

My test was with [0, 1, 0, 1] padding:

// ----// IR Dump Before TosaToLinalg (tosa-to-linalg) //---- //
func.func @openvino_mlir_infer(%arg0: tensor<1x8x30x30xf32>) -> tensor<1x8x31x31xf32> {
  %0 = "tosa.const"() {value = dense<0.000000e+00> : tensor<1x1x1x8xf32>} : () -> tensor<1x1x1x8xf32>
  %1 = "tosa.const"() {value = dense<[[[[[1.000000e+00], [9.97466373], [9.39301586], [2.153120e+00], [9.99136447], [3.12480068], [4.56922674], [1.000000e+01]]]]]> : tensor<1x1x1x8x1xf32>} : () -> tensor<1x1x1x8x1xf32>
  %2 = "tosa.const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64>
  %3 = "tosa.const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} : () -> tensor<4xi64>
  %4 = "tosa.transpose"(%arg0, %3) : (tensor<1x8x30x30xf32>, tensor<4xi64>) -> tensor<1x30x30x8xf32>
  %5 = "tosa.reshape"(%4) {new_shape = [1, 30, 30, 8, 1]} : (tensor<1x30x30x8xf32>) -> tensor<1x30x30x8x1xf32>
  %6 = "tosa.mul"(%5, %1) {shift = 0 : i32} : (tensor<1x30x30x8x1xf32>, tensor<1x1x1x8x1xf32>) -> tensor<1x30x30x8x1xf32>

// Here's the broken Reshape, as there was no handling for padding
  %7 = "tosa.reshape"(%6) {new_shape = [1, 31, 31, 8]} : (tensor<1x30x30x8x1xf32>) -> tensor<1x31x31x8xf32>
//

  %8 = "tosa.add"(%7, %0) : (tensor<1x31x31x8xf32>, tensor<1x1x1x8xf32>) -> tensor<1x31x31x8xf32>
  %9 = "tosa.transpose"(%8, %2) : (tensor<1x31x31x8xf32>, tensor<4xi64>) -> tensor<1x8x31x31xf32>
  return %9 : tensor<1x8x31x31xf32>
}

[rsuderman]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

866

Generally it is ill advised to mutate and compare in a single if condition, especially with pre/post incrementations. It should be good enough to do the following:

currLhsDim++;
if (currLhsDim < lhsShape.size())
  lhsSize *= lhsShape[curreLhsDim]

Overall this is simpler and the extra break is not that beneficial.

[rsuderman]

Overall this should not be handle in the tosa-to-linalg lowering and instead the tosa.reshape verifier should check the number of elements and error out in other cases.

[a-puschin]

In D140013#3996142, @rsuderman wrote:

Overall this should not be handle in the tosa-to-linalg lowering and instead the tosa.reshape verifier should check the number of elements and error out in other cases.

Can only partially agree:of course a check in the verifier would also be nice to have. Yet currently the iteration is unsafe so a check here is still necessary

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

866

Agree, the suggested if-based form is more elegant

[a-puschin]

• Fix after review

[a-puschin]

Split fixes into 2 commits

[a-puschin]

Retry

[a-puschin]

Rebased

[frgossen]

This failed our internal tests. I think the issue is *not* with your revision but a test case would be nice here!
Could you add a test for the reshape op explicitly?

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

def Tosa_ReshapeOp: Tosa_Op<"reshape", [
   let description = [{
     Returns a tensor with the same type/values as the input, with a new shape
     specified by the shape argument. Reshape may operate on tensors of any rank.
     No data conversion happens during a reshape operation.
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
+  let hasVerifier = 1;
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     I64ArrayAttr:$new_shape
   );
 
   let results = (outs
     Tosa_RankedTensor:$output

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

static bool findIntermediateShape(ArrayRef<int64_t> lhsShape,
 
   unsigned currLhsDim = 0, currRhsDim = 0;
   while (currLhsDim < lhsShape.size() && currRhsDim < rhsShape.size()) {
     int64_t rhsSize = rhsShape[currRhsDim];
     int64_t lhsSize = lhsShape[currLhsDim];
     while (lhsSize != rhsSize && currLhsDim < lhsShape.size() &&
            currRhsDim < rhsShape.size()) {
       if (lhsSize < rhsSize) {
         currLhsDim++;
+        if (currLhsDim < lhsShape.size()) {
           lhsSize *= lhsShape[currLhsDim];
+        }
       } else {
         currRhsDim++;
+        if (currRhsDim < rhsShape.size()) {
           rhsSize *= rhsShape[currRhsDim];
         }
       }
+    }
     if (lhsSize == rhsSize) {
       intermediateShape.push_back(lhsSize);
     }
     currRhsDim++;
     currLhsDim++;
   }
 
   // If the iterators didn't reach the end and their leftover dimensions are not

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

for (auto &val : newShapeValue) {
     if (ShapedType::isDynamic(val))
       val = numElements / staticMul;
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
   return success();
 }
 
+mlir::LogicalResult tosa::ReshapeOp::verify() {
+  ShapedType inputType = getInput1().getType().cast<ShapedType>();
+  ShapedType outputType = getType().cast<ShapedType>();
+
+  if (inputType.hasStaticShape() && outputType.hasStaticShape()) {
+    int64_t inputElementsNum = inputType.getNumElements();
+    int64_t outputElementsNum = outputType.getNumElements();
+    if (inputElementsNum != outputElementsNum) {
+      return emitOpError() << "Cannot reshape " << inputElementsNum
+                           << " elements into " << outputElementsNum;
+    }
+  }
+  return mlir::success();
+}
+
 LogicalResult tosa::TransposeOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   ShapeAdaptor inputShape = operands.getShape(0);
   ShapeAdaptor permsShape = operands.getShape(1);
 
   // If input rank and permutation length is unknown, the output rank is