[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

865–869

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

865–869

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/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) {
-        if (++currLhsDim == lhsShape.size()) {
-          break;
-        }
+        currLhsDim++;
+        if (currLhsDim < lhsShape.size()) {
           lhsSize *= lhsShape[currLhsDim];
-      } else {
-        if (++currRhsDim == rhsShape.size()) {
-          break;
         }
+      } 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