[mlir][ReshapeOpsUtils] ComposeExpandOfCollapseOp is not detecting complicated dynamic shapes

[hockyy]

module {
  func.func @broadcast_test(%arg0: tensor<4x8x?x3x?xf32>) -> tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32> {
    %c2 = arith.constant 2 : index
    %c4 = arith.constant 4 : index
    %dim = tensor.dim %arg0, %c4 : tensor<4x8x?x3x?xf32>
    %dim_0 = tensor.dim %arg0, %c2 : tensor<4x8x?x3x?xf32>
    %collapsed = tensor.collapse_shape %arg0 [[0], [1, 2, 3, 4]] : tensor<4x8x?x3x?xf32> into tensor<4x?xf32>
    %0 = tensor.empty(%dim_0, %dim) : tensor<4x3x8x?x3x?x14xf32>
    %collapsed_1 = tensor.collapse_shape %0 [[0], [1], [2, 3, 4, 5], [6]] : tensor<4x3x8x?x3x?x14xf32> into tensor<4x3x?x14xf32>
    %broadcasted = linalg.broadcast ins(%collapsed : tensor<4x?xf32>) outs(%collapsed_1 : tensor<4x3x?x14xf32>) dimensions = [1, 3]
    %collapsed_2 = tensor.collapse_shape %broadcasted [[0, 1, 2, 3]] : tensor<4x3x?x14xf32> into tensor<?xf32>
    %expanded = tensor.expand_shape %collapsed_2 [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]] output_shape [1, 4, 3, 8, 1, 1, %dim_0, 3, 1, 1, %dim, 1, 14] : tensor<?xf32> into tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>
    return %expanded : tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>
  }
}

    %collapsed_2 = tensor.collapse_shape %broadcasted [[0, 1, 2, 3]] : tensor<4x3x?x14xf32> into tensor<?xf32>
    %expanded = tensor.expand_shape %collapsed_2 [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]] output_shape [1, 4, 3, 8, 1, 1, %dim_0, 3, 1, 1, %dim, 1, 14] : tensor<?xf32> into tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>

This part is not canonicalized

[hockyy]

@jpienaar @pifon2a

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Hocky Yudhiono (hockyy)

``` module { func.func @broadcast_test(%arg0: tensor<4x8x?x3x?xf32>) -> tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32> { %c2 = arith.constant 2 : index %c4 = arith.constant 4 : index %dim = tensor.dim %arg0, %c4 : tensor<4x8x?x3x?xf32> %dim_0 = tensor.dim %arg0, %c2 : tensor<4x8x?x3x?xf32> %collapsed = tensor.collapse_shape %arg0 [[0], [1, 2, 3, 4]] : tensor<4x8x?x3x?xf32> into tensor<4x?xf32> %0 = tensor.empty(%dim_0, %dim) : tensor<4x3x8x?x3x?x14xf32> %collapsed_1 = tensor.collapse_shape %0 [[0], [1], [2, 3, 4, 5], [6]] : tensor<4x3x8x?x3x?x14xf32> into tensor<4x3x?x14xf32> %broadcasted = linalg.broadcast ins(%collapsed : tensor<4x?xf32>) outs(%collapsed_1 : tensor<4x3x?x14xf32>) dimensions = [1, 3] %collapsed_2 = tensor.collapse_shape %broadcasted [[0, 1, 2, 3]] : tensor<4x3x?x14xf32> into tensor<?xf32> %expanded = tensor.expand_shape %collapsed_2 [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]] output_shape [1, 4, 3, 8, 1, 1, %dim_0, 3, 1, 1, %dim, 1, 14] : tensor<?xf32> into tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32> return %expanded : tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32> } } ```

    %collapsed_2 = tensor.collapse_shape %broadcasted [[0, 1, 2, 3]] : tensor<4x3x?x14xf32> into tensor<?xf32>
    %expanded = tensor.expand_shape %collapsed_2 [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]] output_shape [1, 4, 3, 8, 1, 1, %dim_0, 3, 1, 1, %dim, 1, 14] : tensor<?xf32> into tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>

This part is not canonicalized

[CoTinker]

Maybe it's same to #109650.