llvm / llvm-project

The LLVM Project is a collection of modular and reusable compiler and toolchain technologies.
http://llvm.org
Other
29.21k stars 12.05k forks source link

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

Open hockyy opened 4 hours ago

hockyy commented 4 hours ago
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 commented 4 hours ago

@jpienaar @pifon2a