wehub-resource-sync
79 lines
2.8 KiB
C++
79 lines
2.8 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/constraints_manager.h"
|
|
#include "paddle/pir/include/dialect/shape/utils/dim_expr_builder.h"
|
|
|
|
namespace symbol::test {
|
|
|
|
TEST(ConstraintsManager, EqualCstr) {
|
|
ConstraintsManager cstr_mgr;
|
|
DimExprBuilder builder;
|
|
|
|
// Eq(Mul(S0,S1),Mul(S2,S3))
|
|
DimExpr sym_expr_0 = builder.Symbol("S0");
|
|
DimExpr sym_expr_1 = builder.Symbol("S1");
|
|
DimExpr sym_expr_2 = builder.Symbol("S2");
|
|
DimExpr sym_expr_3 = builder.Symbol("S3");
|
|
DimExpr mul_expr_0 = builder.Mul(sym_expr_0, sym_expr_1);
|
|
DimExpr mul_expr_1 = builder.Mul(sym_expr_2, sym_expr_3);
|
|
cstr_mgr.AddEqCstr(mul_expr_0, mul_expr_1);
|
|
ASSERT_TRUE(cstr_mgr.IsEqual(mul_expr_0, mul_expr_1));
|
|
|
|
// Eq(Add(S0,S1),Add(S0,1))
|
|
DimExpr int_expr_1 = builder.ConstSize(1);
|
|
DimExpr add_expr_0 = builder.Add(sym_expr_0, sym_expr_1);
|
|
DimExpr add_expr_1 = builder.Add(sym_expr_0, int_expr_1);
|
|
cstr_mgr.AddEqCstr(add_expr_0, add_expr_1);
|
|
ASSERT_FALSE(cstr_mgr.IsEqual(add_expr_0, add_expr_1));
|
|
DimExpr mul_expr_2 = builder.Mul(sym_expr_0, int_expr_1);
|
|
ASSERT_TRUE(cstr_mgr.IsEqual(mul_expr_2, mul_expr_1));
|
|
}
|
|
|
|
TEST(ConstraintsManager, GreatThanOneCstr) {
|
|
ConstraintsManager cstr_mgr;
|
|
DimExpr sym_expr_0 = DimExpr("S0");
|
|
DimExpr int_expr = DimExpr(5);
|
|
ASSERT_TRUE(cstr_mgr.IsGTOne(int_expr + sym_expr_0));
|
|
ASSERT_TRUE(cstr_mgr.IsGTOne(int_expr * sym_expr_0));
|
|
cstr_mgr.AddGTOneCstr(sym_expr_0);
|
|
ASSERT_TRUE(cstr_mgr.IsGTOne(sym_expr_0));
|
|
}
|
|
|
|
TEST(ConstraintsManager, BroadcastableCstr) {
|
|
ConstraintsManager cstr_mgr;
|
|
DimExpr sym_expr_0 = DimExpr("S0");
|
|
DimExpr int_expr = DimExpr(5);
|
|
cstr_mgr.AddBroadcastableCstr(sym_expr_0, int_expr);
|
|
ASSERT_TRUE(cstr_mgr.IsBroadcastable(sym_expr_0, int_expr));
|
|
}
|
|
|
|
TEST(ConstraintsManager, Case1) {
|
|
// BC(S0, S1) == S2 => BC(S0, S2) == S2
|
|
ConstraintsManager cstr_mgr;
|
|
DimExpr s0 = DimExpr("S0");
|
|
DimExpr s1 = DimExpr("S1");
|
|
DimExpr s2 = DimExpr("S2");
|
|
DimExpr bc_s0_s1 = Broadcast<DimExpr>{{s0, s1}};
|
|
cstr_mgr.AddEqCstr(bc_s0_s1, s2);
|
|
cstr_mgr.AddBroadcastableCstr(s0, s1);
|
|
DimExpr bc_s0_s2 = Broadcast<DimExpr>{{s0, s2}};
|
|
cstr_mgr.AddBroadcastableCstr(s0, s2);
|
|
ASSERT_TRUE(cstr_mgr.IsEqual(bc_s0_s2, s2));
|
|
}
|
|
|
|
} // namespace symbol::test
|