wehub-resource-sync
119 lines
4.4 KiB
C++
119 lines
4.4 KiB
C++
// Copyright (c) 2024 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/common/broadcast_tree.h"
|
|
|
|
#include "gtest/gtest.h"
|
|
|
|
namespace cinn::common {
|
|
using namespace symbol; // NOLINT
|
|
|
|
namespace {
|
|
|
|
DimExpr MakeBroadcastDimExpr(const DimExpr& expr1, const DimExpr& expr2) {
|
|
List<DimExpr> operands{expr1, expr2};
|
|
return Broadcast<DimExpr>{operands};
|
|
}
|
|
|
|
bool DimExprNonBroadcast(const DimExpr& dim_expr) {
|
|
if (dim_expr.Has<Broadcast<DimExpr>>()) {
|
|
return false;
|
|
} else {
|
|
return true;
|
|
}
|
|
}
|
|
|
|
void CheckLeafNonBroadcast(const BroadcastLeaf& leaf) {
|
|
for (const auto& operands : *leaf) {
|
|
for (const auto& operand : operands) {
|
|
ASSERT_TRUE(DimExprNonBroadcast(operand));
|
|
}
|
|
}
|
|
}
|
|
|
|
void CheckInnerBranchNonBroadcast(
|
|
const BroadcastBranch<BroadcastTree>& branch) {
|
|
const auto& [_, lhs_eq_rhs_tree, lhs_eq_one_tree, rhs_eq_one_tree] =
|
|
branch.tuple();
|
|
ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastLeaf>());
|
|
ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastLeaf>());
|
|
ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastLeaf>());
|
|
CheckLeafNonBroadcast(lhs_eq_rhs_tree.Get<BroadcastLeaf>());
|
|
CheckLeafNonBroadcast(lhs_eq_one_tree.Get<BroadcastLeaf>());
|
|
CheckLeafNonBroadcast(rhs_eq_one_tree.Get<BroadcastLeaf>());
|
|
}
|
|
|
|
} // namespace
|
|
|
|
TEST(BroadcastTree, Naive) {
|
|
DimExpr expr1("S1");
|
|
DimExpr expr2("S2");
|
|
DimExpr expr3("S3");
|
|
DimExpr expr4("S4");
|
|
std::vector<DimExpr> tensor_shape{expr1,
|
|
expr2,
|
|
MakeBroadcastDimExpr(expr1, expr2),
|
|
MakeBroadcastDimExpr(expr3, expr4)};
|
|
BroadcastLeaf leaf = adt::List<std::vector<DimExpr>>{tensor_shape};
|
|
int num_of_leaves = 0;
|
|
BroadcastTree tree = ConstructBroadcastTree(leaf, &num_of_leaves);
|
|
ASSERT_TRUE(tree.Has<BroadcastBranch<BroadcastTree>>());
|
|
const auto& branch = tree.Get<BroadcastBranch<BroadcastTree>>();
|
|
const auto& [cstr_broadcastable,
|
|
lhs_eq_rhs_tree,
|
|
lhs_eq_one_tree,
|
|
rhs_eq_one_tree] = branch.tuple();
|
|
ASSERT_EQ(cstr_broadcastable->lhs, DimExpr("S1"));
|
|
ASSERT_EQ(cstr_broadcastable->rhs, DimExpr("S2"));
|
|
ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastBranch<BroadcastTree>>());
|
|
ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastBranch<BroadcastTree>>());
|
|
ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastBranch<BroadcastTree>>());
|
|
CheckInnerBranchNonBroadcast(
|
|
lhs_eq_rhs_tree.Get<BroadcastBranch<BroadcastTree>>());
|
|
CheckInnerBranchNonBroadcast(
|
|
lhs_eq_one_tree.Get<BroadcastBranch<BroadcastTree>>());
|
|
CheckInnerBranchNonBroadcast(
|
|
rhs_eq_one_tree.Get<BroadcastBranch<BroadcastTree>>());
|
|
}
|
|
|
|
TEST(BroadcastTree, SimplifyConstantBroadcast) {
|
|
DimExpr expr1("S1");
|
|
DimExpr expr2("S2");
|
|
DimExpr expr3("S3");
|
|
DimExpr expr4(4);
|
|
std::vector<DimExpr> tensor_shape{expr1,
|
|
expr2,
|
|
MakeBroadcastDimExpr(expr1, expr2),
|
|
MakeBroadcastDimExpr(expr3, expr4)};
|
|
BroadcastLeaf leaf = adt::List<std::vector<DimExpr>>{tensor_shape};
|
|
int num_of_leaves = 0;
|
|
BroadcastTree tree = ConstructBroadcastTree(leaf, &num_of_leaves);
|
|
ASSERT_TRUE(tree.Has<BroadcastBranch<BroadcastTree>>());
|
|
const auto& branch = tree.Get<BroadcastBranch<BroadcastTree>>();
|
|
const auto& [cstr_broadcastable,
|
|
lhs_eq_rhs_tree,
|
|
lhs_eq_one_tree,
|
|
rhs_eq_one_tree] = branch.tuple();
|
|
ASSERT_EQ(cstr_broadcastable->lhs, DimExpr("S1"));
|
|
ASSERT_EQ(cstr_broadcastable->rhs, DimExpr("S2"));
|
|
ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastLeaf>());
|
|
ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastLeaf>());
|
|
ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastLeaf>());
|
|
CheckLeafNonBroadcast(lhs_eq_rhs_tree.Get<BroadcastLeaf>());
|
|
CheckLeafNonBroadcast(lhs_eq_one_tree.Get<BroadcastLeaf>());
|
|
CheckLeafNonBroadcast(rhs_eq_one_tree.Get<BroadcastLeaf>());
|
|
}
|
|
|
|
} // namespace cinn::common
|