Files
paddlepaddle--paddle/test/cpp/cinn/adt/tree_test.cc
T
2026-07-13 12:40:42 +08:00

200 lines
7.6 KiB
C++

// 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