chore: import upstream snapshot with attribution

test/cpp/cinn/adt/CMakeLists.txt

@@ -0,0 +1,4 @@
+cinn_cc_test(equation_value_match_trait_test SRCS
+             equation_value_match_trait_test.cc DEPS gtest glog)
+
+cinn_cc_test(tree_test SRCS tree_test.cc DEPS gtest glog)

test/cpp/cinn/adt/equation_value_match_trait_test.cc

@@ -0,0 +1,143 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/cinn/adt/equation_value_match_trait.h"
+#include "gtest/gtest.h"
+#include "paddle/cinn/adt/equation_value.h"
+#include "paddle/cinn/adt/match.h"
+
+namespace cinn::adt::test {
+
+TEST(Match, Union) {
+  using Pattern = Union<IndexUnDotValue<Value, List<DimExpr>>,
+                        IndexDotValue<Value, List<DimExpr>>>;
+  {
+    Value expr = IndexUnDotValue<Value, List<DimExpr>>{
+        Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
+
+    bool ret = cinn::adt::Match<Pattern>(expr);
+    ASSERT_TRUE(ret);
+  }
+  {
+    Value expr = IndexDotValue<Value, List<DimExpr>>{
+        Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
+
+    bool ret = cinn::adt::Match<Pattern>(expr);
+    ASSERT_TRUE(ret);
+  }
+  {
+    Value expr = List<Value>{Value{Ok()}, Value{Ok()}, Value{Ok()}};
+
+    bool ret = cinn::adt::Match<Pattern>(expr);
+    ASSERT_FALSE(ret);
+  }
+}
+
+TEST(Match, index_undot) {
+  Value expr = IndexUnDotValue<Value, List<DimExpr>>{
+      Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
+
+  bool ret = cinn::adt::Match<IndexUnDotValue<Value, List<DimExpr>>>(expr);
+  ASSERT_TRUE(ret);
+}
+
+TEST(Match, index_dot) {
+  Value expr = IndexDotValue<Value, List<DimExpr>>{
+      Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
+
+  bool ret = cinn::adt::Match<IndexDotValue<Value, List<DimExpr>>>(expr);
+  ASSERT_TRUE(ret);
+}
+
+TEST(Match, list) {
+  Value expr = List<Value>{Value{Ok()}, Value{Ok()}, Value{Ok()}};
+
+  bool ret = cinn::adt::Match<List<Value>>(expr);
+  ASSERT_TRUE(ret);
+}
+
+TEST(Match, list_get_item) {
+  Value list = List<Value>{Value{Ok()}, Value{Ok()}, Value{Ok()}};
+  Value expr = ListGetItem<Value, DimExpr>{list, DimExpr{std::int64_t(1)}};
+
+  bool ret = cinn::adt::Match<ListGetItem<Value, std::int64_t>>(expr);
+  ASSERT_TRUE(ret);
+}
+
+TEST(Match, list_get_item_index_undot) {
+  Value undot1 = IndexUnDotValue<Value, List<DimExpr>>{
+      Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
+  ASSERT_TRUE(
+      (cinn::adt::Match<IndexUnDotValue<Value, List<DimExpr>>>(undot1)));
+
+  Value expr = ListGetItem<Value, DimExpr>{undot1, DimExpr{std::int64_t(1)}};
+  ASSERT_TRUE(
+      (cinn::adt::Match<
+          ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
+          expr)));
+}
+
+// List<ListGetItem<IndexUnDotValue<Value>, std::int64_t>>
+TEST(Match, list_list_get_item_index_undot) {
+  Value undot = IndexUnDotValue<Value, List<DimExpr>>{
+      Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
+  ASSERT_TRUE((cinn::adt::Match<IndexUnDotValue<Value, List<DimExpr>>>(undot)));
+  Value expr1 = ListGetItem<Value, DimExpr>{undot, DimExpr{std::int64_t(0)}};
+  ASSERT_TRUE(
+      (cinn::adt::Match<
+          ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
+          expr1)));
+  Value expr2 = ListGetItem<Value, DimExpr>{undot, DimExpr{std::int64_t(1)}};
+  ASSERT_TRUE(
+      (cinn::adt::Match<
+          ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
+          expr2)));
+  Value list = List<Value>{expr1, expr2};
+  ASSERT_TRUE(
+      (cinn::adt::Match<List<
+           ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>>(
+          list)));
+}
+
+// IndexDotValue<List<ListGetItem<IndexUnDotValue<Value>, std::int64_t>>>
+TEST(Match, index_dot_list_list_get_item_index_undot) {
+  Value undot1 = IndexUnDotValue<Value, List<DimExpr>>{
+      Value{Ok()}, List<DimExpr>{std::int64_t(1)}};
+  ASSERT_TRUE(
+      (cinn::adt::Match<IndexUnDotValue<Value, List<DimExpr>>>(undot1)));
+  Value expr1 = ListGetItem<Value, DimExpr>{undot1, DimExpr{std::int64_t(0)}};
+  ASSERT_TRUE(
+      (cinn::adt::Match<
+          ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
+          expr1)));
+  Value expr2 = ListGetItem<Value, DimExpr>{undot1, DimExpr{std::int64_t(1)}};
+  ASSERT_TRUE(
+      (cinn::adt::Match<
+          ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>(
+          expr2)));
+  Value list = List<Value>{expr1, expr2};
+  ASSERT_TRUE(
+      (cinn::adt::Match<List<
+           ListGetItem<IndexUnDotValue<Value, List<DimExpr>>, std::int64_t>>>(
+          list)));
+  Value dot =
+      IndexDotValue<Value, List<DimExpr>>{list, List<DimExpr>{std::int64_t(1)}};
+  ASSERT_TRUE(
+      (cinn::adt::Match<
+          IndexDotValue<List<ListGetItem<IndexUnDotValue<Value, List<DimExpr>>,
+                                         std::int64_t>>,
+                        List<DimExpr>>>(dot)));
+}
+
+}  // namespace cinn::adt::test

test/cpp/cinn/adt/tree_test.cc

@@ -0,0 +1,199 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/cinn/adt/tree.h"
+
+#include "gtest/gtest.h"
+
+namespace cinn::adt {
+
+namespace test {
+
+using IntTreeLeafT = std::vector<int>;
+using IntTreeInnerDataT = std::vector<int>;
+using IntVecTree = Tree<TreeInner<IntTreeInnerDataT>::Node, IntTreeLeafT>;
+using IntTreeInnerT = TreeInner<IntTreeInnerDataT>::template Node<IntVecTree>;
+
+}  // namespace test
+
+template <>
+struct TreeMerger<test::IntVecTree> {
+  using tree_type = test::IntVecTree;
+  using inner_type = typename TreeTrait<test::IntVecTree>::inner_type;
+  using leaf_type = typename TreeTrait<test::IntVecTree>::leaf_type;
+  using inner_data_type = typename inner_type::value_type;
+
+  inner_data_type GetInnerDataForLeaf(const leaf_type& leaf) const {
+    return leaf;
+  }
+
+  inner_type MakeInnerNode(const inner_data_type& inner_data,
+                           const List<test::IntVecTree>& children) const {
+    return inner_type{inner_data, children};
+  }
+
+  using MergeResult = std::tuple<tCommon<inner_data_type>,
+                                 tLhsRemainder<inner_data_type>,
+                                 tRhsRemainder<inner_data_type>>;
+
+  MergeResult MergeInnerValue(const inner_data_type& lhs,
+                              const inner_data_type& rhs) const {
+    inner_data_type common{};
+    inner_data_type lhs_remainder{};
+    inner_data_type rhs_remainder{};
+    int min_size = std::min(lhs.size(), rhs.size());
+    int idx = 0;
+    for (; idx < min_size; ++idx) {
+      if (lhs.at(idx) == rhs.at(idx)) {
+        common.emplace_back(lhs.at(idx));
+      } else {
+        break;
+      }
+    }
+    for (int lhs_idx = idx; lhs_idx < lhs.size(); ++lhs_idx) {
+      lhs_remainder.emplace_back(lhs.at(lhs_idx));
+    }
+    for (int rhs_idx = idx; rhs_idx < rhs.size(); ++rhs_idx) {
+      rhs_remainder.emplace_back(rhs.at(rhs_idx));
+    }
+    return MergeResult{common, lhs_remainder, rhs_remainder};
+  }
+};
+
+namespace test {
+
+TEST(IntVecTree, naive) {
+  List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3}, IntTreeLeafT{4, 5, 6}};
+  TreeMerger<test::IntVecTree> tree_merger{};
+  List<IntVecTree> ret = MakeMergedTrees(tree_merger, leaves);
+  ASSERT_EQ(ret->size(), 2);
+
+  ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
+  const auto& [inner_data0, children0] =
+      ret->at(0).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
+  ASSERT_TRUE((children0->size() == 1));
+  ASSERT_TRUE((children0->at(0).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children0->at(0).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
+
+  ASSERT_TRUE(ret->at(1).Has<IntTreeInnerT>());
+  const auto& [inner_data1, children1] =
+      ret->at(1).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data1 == IntTreeLeafT{4, 5, 6}));
+  ASSERT_TRUE((children1->size() == 1));
+  ASSERT_TRUE((children1->at(0).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children1->at(0).Get<IntTreeLeafT>() == IntTreeLeafT{4, 5, 6}));
+}
+
+TEST(IntVecTree, left_equal_right) {
+  List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3}, IntTreeLeafT{1, 2, 3}};
+  List<IntVecTree> ret =
+      MakeMergedTrees(TreeMerger<test::IntVecTree>{}, leaves);
+  ASSERT_EQ(ret->size(), 1);
+
+  ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
+  const auto& [inner_data0, children0] =
+      ret->at(0).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
+  ASSERT_TRUE((children0->size() == 2));
+  ASSERT_TRUE((children0->at(0).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children0->at(0).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
+  ASSERT_TRUE((children0->at(1).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children0->at(1).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
+}
+
+TEST(IntVecTree, left_gt_right) {
+  List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3, 4, 5}, IntTreeLeafT{1, 2, 3}};
+  List<IntVecTree> ret =
+      MakeMergedTrees(TreeMerger<test::IntVecTree>{}, leaves);
+  ASSERT_EQ(ret->size(), 1);
+
+  ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
+  const auto& [inner_data0, children0] =
+      ret->at(0).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
+  ASSERT_TRUE((children0->size() == 2));
+
+  ASSERT_TRUE((children0->at(0).Has<IntTreeInnerT>()));
+  const auto& [inner_data_left0, children_left0] =
+      children0->at(0).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data_left0 == IntTreeLeafT{4, 5}));
+  ASSERT_TRUE((children_left0->size() == 1));
+  ASSERT_TRUE((children_left0->at(0).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children_left0->at(0).Get<IntTreeLeafT>() ==
+               IntTreeLeafT{1, 2, 3, 4, 5}));
+
+  ASSERT_TRUE((children0->at(1).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children0->at(1).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
+}
+
+TEST(IntVecTree, left_lt_right) {
+  List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3}, IntTreeLeafT{1, 2, 3, 4, 5}};
+  List<IntVecTree> ret =
+      MakeMergedTrees(TreeMerger<test::IntVecTree>{}, leaves);
+  ASSERT_EQ(ret->size(), 1);
+
+  ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
+  const auto& [inner_data0, children0] =
+      ret->at(0).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
+  ASSERT_TRUE((children0->size() == 2));
+
+  ASSERT_TRUE((children0->at(0).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children0->at(0).Get<IntTreeLeafT>() == IntTreeLeafT{1, 2, 3}));
+
+  ASSERT_TRUE((children0->at(1).Has<IntTreeInnerT>()));
+  const auto& [inner_data_right0, children_right0] =
+      children0->at(1).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data_right0 == IntTreeLeafT{4, 5}));
+  ASSERT_TRUE((children_right0->size() == 1));
+  ASSERT_TRUE((children_right0->at(0).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children_right0->at(0).Get<IntTreeLeafT>() ==
+               IntTreeLeafT{1, 2, 3, 4, 5}));
+}
+
+TEST(IntVecTree, left_ne_right) {
+  List<IntTreeLeafT> leaves{IntTreeLeafT{1, 2, 3, 4, 5},
+                            IntTreeLeafT{1, 2, 3, 6, 7}};
+  List<IntVecTree> ret =
+      MakeMergedTrees(TreeMerger<test::IntVecTree>{}, leaves);
+  ASSERT_EQ(ret->size(), 1);
+
+  ASSERT_TRUE(ret->at(0).Has<IntTreeInnerT>());
+  const auto& [inner_data0, children0] =
+      ret->at(0).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data0 == IntTreeLeafT{1, 2, 3}));
+  ASSERT_TRUE((children0->size() == 2));
+
+  ASSERT_TRUE((children0->at(0).Has<IntTreeInnerT>()));
+  const auto& [inner_data_left0, children_left0] =
+      children0->at(0).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data_left0 == IntTreeLeafT{4, 5}));
+  ASSERT_TRUE((children_left0->size() == 1));
+  ASSERT_TRUE((children_left0->at(0).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children_left0->at(0).Get<IntTreeLeafT>() ==
+               IntTreeLeafT{1, 2, 3, 4, 5}));
+
+  ASSERT_TRUE((children0->at(1).Has<IntTreeInnerT>()));
+  const auto& [inner_data_right0, children_right0] =
+      children0->at(1).Get<IntTreeInnerT>().tuple();
+  ASSERT_TRUE((inner_data_right0 == IntTreeLeafT{6, 7}));
+  ASSERT_TRUE((children_right0->size() == 1));
+  ASSERT_TRUE((children_right0->at(0).Has<IntTreeLeafT>()));
+  ASSERT_TRUE((children_right0->at(0).Get<IntTreeLeafT>() ==
+               IntTreeLeafT{1, 2, 3, 6, 7}));
+}
+
+}  // namespace test
+}  // namespace cinn::adt