[CINN] Add SimplifyWithObviousGreaterThan (#71341)

* Add SimplifyWithObviousGreaterThan * refine Min(S0, Add(S0,S1), Mul(Add(S1,S2),S2)) => S0 * fix typo * Add more tests * fix * fix bugs * log * apply review * clean * fix bugs * fix * make sure accuracy * fix typo * may be * Add a test * finally test * deal zero, one * fix bug for if condition

Author: ooooo-create

paddle/pir/include/dialect/shape/utils/dim_expr_util.h

@@ -39,8 +39,10 @@ IR_API PriorityComparisonStatus CompareDimExprPriority(const DimExpr& lhs,
 
 enum class DimExprCompareResult {
  GT, // lhs is greater than rhs
+ GE, // lhs is greater than or equal to rhs
  EQ, // lhs and rhs is equal
  LT, // lhs is less than rhs
+ LE, // lhs is less than or equal to rhs
  UNKNOWN, // lhs and rhs is not comparable
 };
 IR_API DimExprCompareResult Compare(const DimExpr& lhs, const DimExpr& rhs);

paddle/pir/src/dialect/shape/utils/dim_expr_util.cc

@@ -1005,11 +1005,12 @@ struct FoldRedundantSymbolicBroadcast {
  * Simplify Example:
  * Broadcast(S0,S0,S1) => Broadcast(S0,S1)
  */
-struct FoldRedundantBroadcast {
- using dim_expr_type = Broadcast<DimExpr>;
+template <template <typename> class Op>
+struct FoldRepetitiveSymbol {
+ using dim_expr_type = Op<DimExpr>;
 
  DimExpr Rewrite(const DimExpr& expr) {
- const auto& [operands] = expr.Get<Broadcast<DimExpr>>();
+ const auto& [operands] = expr.Get<Op<DimExpr>>();
  while (operands->size() > 1) {
  int pos_index = SearchSameIndex(operands);
  if (pos_index < 0) {
@@ -1020,7 +1021,7 @@ struct FoldRedundantBroadcast {
  if (operands->size() == 1) {
  return operands->at(0);
  } else {
- return Broadcast<DimExpr>{operands};
+ return Op<DimExpr>{operands};
  }
  PADDLE_THROW(common::errors::Fatal("Dead code."));
  }
@@ -1040,6 +1041,197 @@ struct FoldRedundantBroadcast {
  }
 };
 
+DimExprCompareResult EasyCompareAddWithZero(const Add<DimExpr>& add) {
+ // Only return GT, GE, UNKNOWN.
+ List<DimExpr> operands = add.operands;
+ for (const auto& operand : *operands) {
+ if (operand.isa<std::string>()) {
+ continue;
+ }
+ if (operand.isa<std::int64_t>() && operand.dyn_cast<int64_t>() > 0) {
+ continue;
+ }
+ return DimExprCompareResult::UNKNOWN;
+ }
+ return DimExprCompareResult::GT;
+}
+
+DimExprCompareResult EasyCompareMulWithOne(const Mul<DimExpr>& mul) {
+ // Only return GT, GE, UNKNOWN.
+ List<DimExpr> operands = mul.operands;
+ int64_t const_result = 1;
+ for (const auto& operand : *operands) {
+ if (operand.isa<std::string>()) {
+ continue;
+ }
+ if (operand.isa<std::int64_t>() && operand.dyn_cast<int64_t>() > 1) {
+ const_result = operand.dyn_cast<int64_t>();
+ continue;
+ }
+
+ return DimExprCompareResult::UNKNOWN;
+ }
+ if (const_result == 1) {
+ return DimExprCompareResult::GE;
+ } else {
+ return DimExprCompareResult::GT;
+ }
+}
+
+bool EasyIsGtWithZero(const DimExpr& expr) {
+ auto ExprVisit = common::Overloaded{
+ [](const std::int64_t& expr) { return expr > 0; },
+ [](const std::string& expr) { return true; },
+ [](const Mul<DimExpr>& expr) {
+ return EasyCompareMulWithOne(expr) != DimExprCompareResult::UNKNOWN;
+ },
+ [](const Add<DimExpr>& expr) {
+ return EasyCompareAddWithZero(expr) == DimExprCompareResult::GT;
+ },
+ [](const Broadcast<DimExpr>& expr) { return true; },
+ [](const auto& expr) { return false; }};
+ return std::visit(ExprVisit, expr.variant());
+}
+
+DimExprCompareResult EasyCompareGtOrGe(const DimExpr& lhs,
+ const DimExpr& rhs,
+ bool is_broadcast = false) {
+ // TODO(ooooo): not perfect but ensures accuracy now.Such as:
+ // S0 < Add(S0, Mul(S1, S2)), S2 also can be Add(S4, S5, -1)
+ // range info may be used.
+ auto CompareDivResult = common::Overloaded{
+ [](const std::int64_t& expr) {
+ // trick for Min(Mul(5, S0), Mul(3, S0))
+ return expr >= 1 ? DimExprCompareResult::GT
+ : DimExprCompareResult::UNKNOWN;
+ },
+ [](const std::string& expr) { return DimExprCompareResult::GE; },
+ [&](const Mul<DimExpr>& expr) { return EasyCompareMulWithOne(expr); },
+ [](const auto& expr) { return DimExprCompareResult::UNKNOWN; }};
+
+ auto CompareSubResult = common::Overloaded{
+ [](const std::int64_t& expr) {
+ return expr > 0 ? DimExprCompareResult::GT
+ : DimExprCompareResult::UNKNOWN;
+ },
+ [](const std::string& expr) { return DimExprCompareResult::GT; },
+ [&](const Add<DimExpr>& expr) { return EasyCompareAddWithZero(expr); },
+ [](const auto& expr) { return DimExprCompareResult::UNKNOWN; }};
+
+ if (lhs.isa<std::string>() && rhs.isa<std::string>()) {
+ return DimExprCompareResult::UNKNOWN;
+ }
+
+ auto IsAddOrMul = [](const DimExpr& expr) {
+ return expr.isa<Add<DimExpr>>() || expr.isa<Mul<DimExpr>>();
+ };
+ auto IsOneOrZero = [](const DimExpr& expr) {
+ return expr == DimExpr{1} || expr == DimExpr{0};
+ };
+
+ if (!IsAddOrMul(lhs) && !IsAddOrMul(rhs) && !IsOneOrZero(lhs) &&
+ !IsOneOrZero(rhs)) {
+ return DimExprCompareResult::UNKNOWN;
+ }
+
+ // check with Sub
+ DimExpr simplified_result_sub = SimplifyDimExpr(DimExpr{lhs} - DimExpr{rhs});
+ auto sub_compare =
+ std::visit(CompareSubResult, simplified_result_sub.variant());
+ if (sub_compare != DimExprCompareResult::UNKNOWN) {
+ return sub_compare;
+ }
+ if (rhs != symbol::DimExpr{0}) {
+ if (!is_broadcast && !EasyIsGtWithZero(rhs)) {
+ // assume operands in broadcast is always positive.
+ return DimExprCompareResult::UNKNOWN;
+ }
+ DimExpr simplified_result_div =
+ SimplifyDimExpr(DimExpr{lhs} / DimExpr{rhs});
+ auto div_compare =
+ std::visit(CompareDivResult, simplified_result_div.variant());
+ return div_compare;
+ } else {
+ return DimExprCompareResult::UNKNOWN;
+ }
+}
+
+struct SimplifyMaxWithGE {
+ using dim_expr_type = Max<DimExpr>;
+ static List<DimExpr> SearchErasable(const List<DimExpr>& operands) {
+ List<DimExpr> simplified_operands{};
+ for (std::size_t i = 0; i < operands->size(); ++i) {
+ bool is_redundant = false;
+ for (std::size_t j = 0; j < operands->size(); ++j) {
+ if (i == j) {
+ continue;
+ }
+ auto compare_j_i = EasyCompareGtOrGe(operands->at(j), operands->at(i));
+ if (compare_j_i == DimExprCompareResult::GT ||
+ compare_j_i == DimExprCompareResult::GE) {
+ is_redundant = true;
+ break;
+ }
+ }
+ if (!is_redundant) {
+ simplified_operands->push_back(operands->at(i));
+ }
+ }
+ return simplified_operands;
+ }
+
+ DimExpr Rewrite(const DimExpr& expr) {
+ const auto [operands] = expr.Get<Max<DimExpr>>();
+ List<DimExpr> simplified_operands = SearchErasable(operands);
+
+ if (simplified_operands->size() == 1) {
+ return simplified_operands->at(0);
+ } else {
+ return Max<DimExpr>{simplified_operands};
+ }
+ }
+};
+
+/*
+ * Simplify Example:
+ * Min(S0, Mul(S0, S1)) => S0
+ */
+struct SimplifyMinWithGE {
+ using dim_expr_type = Min<DimExpr>;
+ static List<DimExpr> SearchErasable(const List<DimExpr>& operands) {
+ List<DimExpr> simplified_operands{};
+ for (std::size_t i = 0; i < operands->size(); ++i) {
+ bool is_redundant = false;
+ for (std::size_t j = 0; j < operands->size(); ++j) {
+ if (i == j) {
+ continue;
+ }
+ auto compare_i_j = EasyCompareGtOrGe(operands->at(i), operands->at(j));
+ if (compare_i_j == DimExprCompareResult::GT ||
+ compare_i_j == DimExprCompareResult::GE) {
+ is_redundant = true;
+ break;
+ }
+ }
+ if (!is_redundant) {
+ simplified_operands->push_back(operands->at(i));
+ }
+ }
+ return simplified_operands;
+ }
+
+ DimExpr Rewrite(const DimExpr& expr) {
+ const auto [operands] = expr.Get<Min<DimExpr>>();
+ List<DimExpr> simplified_operands = SearchErasable(operands);
+
+ if (simplified_operands->size() == 1) {
+ return simplified_operands->at(0);
+ } else {
+ return Min<DimExpr>{simplified_operands};
+ }
+ }
+};
+
 /*
  * Simplify Example:
  * Broadcast(S0, Mul(S0, S1)) => Mul(S0, S1)
@@ -1266,6 +1458,8 @@ DimExpr Simplify(const DimExpr& expr) {
  DoPass<SimplifyOperands<Mul>>(&keep_rewrite, &ret);
  DoPass<SimplifyOperands<Div>>(&keep_rewrite, &ret);
  DoPass<SimplifyOperands<Broadcast>>(&keep_rewrite, &ret);
+ DoPass<SimplifyOperands<Min>>(&keep_rewrite, &ret);
+ DoPass<SimplifyOperands<Max>>(&keep_rewrite, &ret);
  DoPass<SortOperands<Add>>(&keep_rewrite, &ret);
  DoPass<SortOperands<Mul>>(&keep_rewrite, &ret);
  DoPass<SortOperands<Broadcast>>(&keep_rewrite, &ret);
@@ -1283,9 +1477,13 @@ DimExpr Simplify(const DimExpr& expr) {
  DoPass<FoldConstants<Min>>(&keep_rewrite, &ret);
  DoPass<FoldConstants<Broadcast>>(&keep_rewrite, &ret);
  DoPass<FoldInversedPairToUnit<Add>>(&keep_rewrite, &ret);
- DoPass<FoldRedundantBroadcast>(&keep_rewrite, &ret);
+ DoPass<FoldRepetitiveSymbol<Broadcast>>(&keep_rewrite, &ret);
+ DoPass<FoldRepetitiveSymbol<Min>>(&keep_rewrite, &ret);
+ DoPass<FoldRepetitiveSymbol<Max>>(&keep_rewrite, &ret);
  DoPass<FoldRedundantSymbolicBroadcast>(&keep_rewrite, &ret);
  DoPass<SimplifyBroadcast>(&keep_rewrite, &ret);
+ DoPass<SimplifyMinWithGE>(&keep_rewrite, &ret);
+ DoPass<SimplifyMaxWithGE>(&keep_rewrite, &ret);
  DoPass<SimplifyDiv>(&keep_rewrite, &ret);
  if (expr_before_run_pipeline == ret) break;
  }

test/cpp/pir/shape_dialect/simplify_dim_expr_test.cc

@@ -213,12 +213,60 @@ TEST(Simplify, FoldBroadcast) {
  ASSERT_TRUE(simplify_broadcast3 == add);
 }
 
-TEST(Simplify, FoldRedundantBroadcast) {
+TEST(Simplify, FoldRepetitiveSymbol) {
  DimExpr S0{"S0"};
  DimExpr S1{"S1"};
  DimExpr bc{Broadcast<DimExpr>{{S0, S0, S1, S1}}};
+ DimExpr max{Max<DimExpr>{{S0, S0, S1, S1}}};
+ DimExpr min{Min<DimExpr>{{S0, S0, S1, S1}}};
  DimExpr simplify_bc = SimplifyDimExpr(bc);
+ DimExpr simplify_max = SimplifyDimExpr(max);
+ DimExpr simplify_min = SimplifyDimExpr(min);
  ASSERT_TRUE((simplify_bc == Broadcast<DimExpr>{{S0, S1}}));
+ ASSERT_TRUE((simplify_max == Max<DimExpr>{{S0, S1}}));
+ ASSERT_TRUE((simplify_min == Min<DimExpr>{{S0, S1}}));
+}
+
+TEST(Simplify, SimplifyMaxAndMinWithGE) {
+ DimExpr S0{"S0"};
+ DimExpr S1{"S1"};
+ DimExpr S2{"S2"};
+ DimExpr add1{Add<DimExpr>{{S0, S1}}};
+ DimExpr max1{Max<DimExpr>{{S0, add1, S2}}};
+ DimExpr min1{Min<DimExpr>{{S0, add1, S2}}};
+ // Min(S0, Add(S0, S1), S2) => Min(S0, S2)
+ ASSERT_TRUE((SimplifyDimExpr(max1) == Max<DimExpr>{{add1, S2}}));
+ ASSERT_TRUE((SimplifyDimExpr(min1) == Min<DimExpr>{{S0, S2}}));
+
+ // Min(S0, Add(S0,S1), Mul(Add(S1, S2), S2)) => S0
+ DimExpr mul{Mul<DimExpr>{{add1, S2}}};
+ DimExpr max2{Max<DimExpr>{{S0, add1, mul}}};
+ DimExpr min2{Min<DimExpr>{{S0, add1, mul}}};
+ ASSERT_TRUE((SimplifyDimExpr(max2) == mul));
+ ASSERT_TRUE((SimplifyDimExpr(min2) == S0));
+
+ // Min(S0, Add(S0, -1)) => Add(S0, -1)
+ DimExpr add2{Add<DimExpr>{{S0, Negative<DimExpr>{DimExpr(1)}}}};
+ ASSERT_TRUE(
+ (SimplifyDimExpr(Min<DimExpr>{{S0, add2}}) == Add<DimExpr>{{S0, -1}}));
+
+ // Min(S0, Add(S0, -S1)) => Add(S0, -S1)
+ DimExpr add3{Add<DimExpr>{{S0, Negative<DimExpr>{S1}}}};
+ ASSERT_TRUE((SimplifyDimExpr(Min<DimExpr>{{S0, add3}}) == add3));
+
+ // Min(S0, 0) => 0, Max(S0, 0) => S0
+ ASSERT_TRUE((SimplifyDimExpr(Min<DimExpr>{{S0, DimExpr(0)}}) == DimExpr(0)));
+ ASSERT_TRUE((SimplifyDimExpr(Max<DimExpr>{{S0, DimExpr(0)}}) == S0));
+
+ // Min(S0, 1) => 1, Max(S0, 1) => S0
+ ASSERT_TRUE((SimplifyDimExpr(Min<DimExpr>{{S0, DimExpr(1)}}) == DimExpr(1)));
+ ASSERT_TRUE((SimplifyDimExpr(Max<DimExpr>{{S0, DimExpr(1)}}) == S0));
+
+ // Min(Mul(S0, S1), 0) => Min(Mul(S0, S1), 0)
+ // Now simplify ability is limited.
+ DimExpr mul2{Mul<DimExpr>{{S0, S1}}};
+ ASSERT_TRUE((SimplifyDimExpr(Min<DimExpr>{{mul2, DimExpr(0)}}) ==
+ Min<DimExpr>{{mul2, DimExpr(0)}}));
 }
 
 TEST(Simplify, SimplifyDoubleNegForMulAndDiv) {