Legalization of 'tosa.reduce' results in known expand shape op

[mariecwhite]

What happened?

When compiling ResNet-V1-101 with dynamic shapes, getting the error:

resnet_v1_101_dynamic.mlir:223:12: error: failed to legalize operation 'tosa.max_pool2d'
    %220 = "tosa.max_pool2d"(%219) {kernel = [3, 3], pad = [1, 1, 1, 1], stride = [2, 2]} : (tensor<?x?x?x64xf32>) -> tensor<?x?x?x64xf32>
           ^
resnet_v1_101_dynamic.mlir:223:12: note: see current operation: %272 = "tosa.max_pool2d"(%271) {kernel = [3, 3], pad = [1, 1, 1, 1], stride = [2, 2]} : (tensor<?x?x?x64xf32>) -> tensor<?x?x?x64xf32>
compilation failed

TFLite and MLIR files: https://storage.googleapis.com/iree-model-artifacts/resnet_v1_101_dynamic.zip

[NatashaKnk]

Just to make sure, this and 10291 are associated with newly-added models, correct? We didn't support dynamic non-batch dims for those ops (based on our conservative approach for dynamic dim support), but since those are now needed I'm working on a fix.

[mariecwhite]

Hi Natasha, these are newly added models where ResNet-V1-101 is dynamic in batch, width and height. And SSD-MobileNetV2 has dynamic width and height.

[NatashaKnk]

Should be fixed in D133389.

[rsuderman]

Ran the model through and found a new loading failure. We reuse the tosa.reshape that performs a multiple collapse down / back up. We should directly lower to Expand/Collapse so that we don't need to go via tosa.reshape (due to minimal shape support). Simple sample case below:

func.func public @main(%arg0: tensor<?x?x?x2048xf32>) -> tensor<?x1x?x2048xf32> {
  %0 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<?x?x?x2048xf32>) -> tensor<?x1x?x2048xf32>
  return %0 : tensor<?x1x?x2048xf32>
}