[Mlir-commits] [mlir] [mlir][tensor] Loosen restrictions on folding dynamic reshapes (PR #137963)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 2 18:46:08 PDT 2025
================
@@ -0,0 +1,200 @@
+//===- ReshapeOpsUtilsTest.cpp - ReshapeOpsUtils unit tests ---------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "llvm/ADT/STLExtras.h"
+#include "gtest/gtest.h"
+#include <optional>
+
+using namespace mlir;
+
+/// Helper to make constructing
+/// `std::optional<SmallVector<ReassociationIndices>>` more readable.
+static std::optional<SmallVector<ReassociationIndices>>
+makeOptionalIndices(std::initializer_list<ReassociationIndices> list) {
+ return std::optional<SmallVector<ReassociationIndices>>(list);
+}
+
+TEST(ReassociationIndicesForCollapse, ScalarTest) {
+ EXPECT_EQ(getReassociationIndicesForCollapse({1}, {}),
+ makeOptionalIndices({{0}}));
+ EXPECT_EQ(getReassociationIndicesForCollapse({1, 1}, {}),
+ makeOptionalIndices({{0, 1}}));
+ EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic}, {}),
+ makeOptionalIndices({{0}}));
+ EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic,
+ ShapedType::kDynamic, 1,
+ ShapedType::kDynamic},
+ {}),
+ makeOptionalIndices({{0, 1, 2, 3, 4}}));
+}
+
+TEST(ReassociationIndicesForCollapse, ScalarTestFailure) {
+ EXPECT_EQ(getReassociationIndicesForCollapse({}, {}), std::nullopt);
+ EXPECT_EQ(getReassociationIndicesForCollapse({}, {1}), std::nullopt);
+ EXPECT_EQ(getReassociationIndicesForCollapse({2}, {}), std::nullopt);
+ EXPECT_EQ(
+ getReassociationIndicesForCollapse({1, 2, ShapedType::kDynamic, 1}, {}),
+ std::nullopt);
+}
+
+TEST(ReassociationIndicesForCollapse, StaticTest) {
+ EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {200}),
+ makeOptionalIndices({{0, 1}}));
+ EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {10, 600}),
+ makeOptionalIndices({{0}, {1, 2}}));
+ EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {200, 30}),
+ makeOptionalIndices({{0, 1}, {2}}));
+}
+
+TEST(ReassociationIndicesForCollapse, StaticTestFailure) {
+ EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10}), std::nullopt);
+ EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10, 20}),
----------------
MaheshRavishankar wrote:
I think this should just return `{{0}, {1}}` ?
https://github.com/llvm/llvm-project/pull/137963
More information about the Mlir-commits
mailing list