chore: import upstream snapshot with attribution

test/cpp/pir/shape_dialect/constraints_manager_test.cc

@@ -0,0 +1,78 @@
+// 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