diff --git a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h index af575e10acc8e..61c2a50e514ca 100644 --- a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h +++ b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h @@ -387,11 +387,14 @@ struct ComposeExpandOfCollapseOp : public OpRewritePattern { auto resultSubShape = resultShape.slice(resultIndices.front(), resultIndices.size()); + if (llvm::count_if(srcSubShape, ShapedType::isDynamic) >= 2 && + llvm::count_if(resultSubShape, ShapedType::isDynamic) >= 2) + return std::nullopt; + if (srcSubShape.size() == resultSubShape.size()) { - if (srcSubShape != resultSubShape || - llvm::count_if(srcSubShape, ShapedType::isDynamic) >= 2) { + if (srcSubShape != resultSubShape) return std::nullopt; - } + for (auto index : llvm::seq(0, srcSubShape.size())) { composedReassociation.emplace_back(1, srcIndices.front() + index); } diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir index 646b2197d9aa6..7e5e57423f44d 100644 --- a/mlir/test/Dialect/Tensor/canonicalize.mlir +++ b/mlir/test/Dialect/Tensor/canonicalize.mlir @@ -1319,6 +1319,20 @@ func.func @compose_expand_of_collapse_dynamic(%arg0 : tensor<4x?x10x64x2xf16>, % // ----- +func.func @no_compose_collapse_of_expand_dynamic(%arg0 : tensor, %arg1: index) -> tensor { + %collapse = tensor.collapse_shape %arg0 [[0, 1, 2, 3]] : tensor into tensor + %expanded_19 = tensor.expand_shape %collapse [[0, 1, 2]] output_shape [%arg1, 8, %arg1] : tensor into tensor + return %expanded_19 : tensor +} +// CHECK-LABEL: func @no_compose_collapse_of_expand_dynamic +// CHECK-SAME: %[[ARG0:.+]]: tensor +// CHECK-SAME: %[[ARG1:.+]]: index +// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]] +// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]] +// CHECK: return %[[EXPAND]] + +// ----- + // CHECK-LABEL: func @zero_rank_reshape_multi func.func @zero_rank_reshape_multi(%arg0: tensor) -> tensor { // CHECK: return %arg0