wehub-resource-sync
112 lines
4.0 KiB
C++
112 lines
4.0 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 "gtest/gtest.h"
|
|
|
|
#include "paddle/pir/include/dialect/shape/utils/dim_expr_builder.h"
|
|
#include "paddle/pir/include/dialect/shape/utils/dim_expr_util.h"
|
|
|
|
namespace symbol::test {
|
|
|
|
namespace {
|
|
|
|
// (S0 - S1) * 2 / S0
|
|
DimExpr CreateExampleDimExpr() {
|
|
DimExpr sym0 = DimExpr("S0");
|
|
DimExpr sym1 = DimExpr("S1");
|
|
DimExpr constant = DimExpr(2);
|
|
return (sym0 - sym1) * constant / sym0;
|
|
}
|
|
} // namespace
|
|
|
|
TEST(DimExprUtil, SimplifyNeg) {
|
|
DimExpr dim_expr = Negative<DimExpr>{-1};
|
|
DimExpr ret = SimplifyDimExpr(dim_expr);
|
|
ASSERT_TRUE(ret.Has<std::int64_t>());
|
|
ASSERT_EQ(ret.Get<std::int64_t>(), 1);
|
|
DimExpr double_neg_expr = Negative<DimExpr>{dim_expr};
|
|
ret = SimplifyDimExpr(double_neg_expr);
|
|
ASSERT_TRUE(ret.Has<std::int64_t>());
|
|
ASSERT_EQ(ret.Get<std::int64_t>(), -1);
|
|
}
|
|
|
|
TEST(DimExprUtil, Substitute) {
|
|
DimExpr dim_expr = CreateExampleDimExpr();
|
|
std::unordered_map<symbol::DimExpr, symbol::DimExpr> naive_to_full_name{
|
|
{DimExpr("S0"), DimExpr("symbol0")}, {DimExpr("S1"), DimExpr("symbol1")}};
|
|
std::unordered_map<symbol::DimExpr, symbol::DimExpr> full_name_to_naive{
|
|
{DimExpr("symbol0"), DimExpr("S0")}, {DimExpr("symbol1"), DimExpr("S1")}};
|
|
|
|
const auto& mid_expr = SubstituteDimExpr(dim_expr, naive_to_full_name);
|
|
const auto& ret_expr = SubstituteDimExpr(mid_expr, full_name_to_naive);
|
|
ASSERT_EQ(ret_expr, dim_expr);
|
|
}
|
|
|
|
TEST(DimExprUtil, Calculate) {
|
|
// (S0 - S1) * 2 / S0
|
|
DimExpr dim_expr = CreateExampleDimExpr();
|
|
// (4 - 2) * 2 / 4 => 1
|
|
DimExpr substitute_expr = SubstituteDimExpr(dim_expr, {{"S0", 4}, {"S1", 2}});
|
|
DimExpr ret = SimplifyDimExpr(substitute_expr);
|
|
ASSERT_TRUE(ret.Has<std::int64_t>());
|
|
ASSERT_EQ(ret.Get<std::int64_t>(), 1);
|
|
}
|
|
|
|
TEST(DimExprUtil, GetDimExprPriority) {
|
|
DimExprBuilder builder;
|
|
int priority_int = GetDimExprPriority(builder.ConstSize(1));
|
|
ASSERT_EQ(priority_int, 0);
|
|
int priority_sym = GetDimExprPriority(builder.Symbol("S0"));
|
|
ASSERT_EQ(priority_sym, 1);
|
|
int priority_add =
|
|
GetDimExprPriority(builder.Add(DimExpr("S1"), DimExpr("S2")));
|
|
ASSERT_EQ(priority_add, 2);
|
|
int priority_bc =
|
|
GetDimExprPriority(builder.Broadcast(DimExpr("S3"), DimExpr("S4")));
|
|
ASSERT_EQ(priority_bc, 2);
|
|
}
|
|
|
|
TEST(DimExprUtil, CompareDimExprPriority) {
|
|
DimExprBuilder builder;
|
|
DimExpr sym_expr_0 = builder.Symbol("S0");
|
|
DimExpr sym_expr_1 = builder.Symbol("S1");
|
|
DimExpr add_expr = builder.Add(DimExpr("S2"), DimExpr("S3"));
|
|
DimExpr bc_expr = builder.Broadcast(DimExpr("S4"), DimExpr("S5"));
|
|
ASSERT_EQ(CompareDimExprPriority(sym_expr_0, sym_expr_1),
|
|
PriorityComparisonStatus::HIGHER);
|
|
ASSERT_EQ(CompareDimExprPriority(add_expr, sym_expr_0),
|
|
PriorityComparisonStatus::LOWER);
|
|
ASSERT_EQ(CompareDimExprPriority(add_expr, bc_expr),
|
|
PriorityComparisonStatus::EQUAL);
|
|
}
|
|
|
|
TEST(DimExpr, CollectDimExprSymbol) {
|
|
DimExpr dim_expr = [&]() -> DimExpr {
|
|
DimExprBuilder builder;
|
|
DimExpr max_expr = builder.Max(DimExpr("S2"), DimExpr("S3"));
|
|
DimExpr min_expr = builder.Min(max_expr, DimExpr("S4"));
|
|
DimExpr broadcast_expr = builder.Broadcast(min_expr, DimExpr("S5"));
|
|
return CreateExampleDimExpr() + broadcast_expr;
|
|
}();
|
|
std::unordered_set<std::string> symbols = CollectDimExprSymbols(dim_expr);
|
|
std::unordered_set<std::string> expected = {
|
|
"S0", "S1", "S2", "S3", "S4", "S5"};
|
|
ASSERT_EQ(symbols.size(), 6UL);
|
|
for (const auto& symbol : symbols) {
|
|
ASSERT_TRUE(expected.find(symbol) != expected.end());
|
|
}
|
|
}
|
|
|
|
} // namespace symbol::test
|