PaddlePaddle
diff --git a/‎paddle/cinn/common/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/common/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+112-9Lines changed: 112 additions & 9 deletions b/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+112-9Lines changed: 112 additions & 9 deletions
diff --git a/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+32-2Lines changed: 32 additions & 2 deletions b/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+32-2Lines changed: 32 additions & 2 deletions
diff --git a/‎paddle/cinn/common/simplify_corner_case.cc renamed to ‎paddle/cinn/common/simplify_special_pattern.cc
Copy file name to clipboard
+58-12Lines changed: 58 additions & 12 deletions b/‎paddle/cinn/common/simplify_corner_case.cc renamed to ‎paddle/cinn/common/simplify_special_pattern.cc
Copy file name to clipboard
+58-12Lines changed: 58 additions & 12 deletions
diff --git a/‎paddle/cinn/common/simplify_corner_case.h renamed to ‎paddle/cinn/common/simplify_special_pattern.h
Copy file name to clipboard b/‎paddle/cinn/common/simplify_corner_case.h renamed to ‎paddle/cinn/common/simplify_special_pattern.h
Copy file name to clipboard
@@ -26,7 +26,7 @@ gather_srcs(
   dim_expr_converter.cc
   broadcast_tree.cc
   iter_simplify.cc
-  simplify_corner_case.cc)
+  simplify_special_pattern.cc)
 
 cinn_cc_test(test_equation_graph_topo_walker SRCS
              equation_graph_topo_walker_test.cc DEPS gtest glog)
 
@@ -19,7 +19,7 @@
 #include <unordered_set>
 
 #include "paddle/cinn/common/cas.h"
-#include "paddle/cinn/common/simplify_corner_case.h"
+#include "paddle/cinn/common/simplify_special_pattern.h"
 #include "paddle/cinn/ir/ir_mutator.h"
 #include "paddle/cinn/ir/ir_printer.h"
 #include "paddle/cinn/ir/op/ir_operators.h"
@@ -176,15 +176,15 @@ Expr RampRelatedMul(Expr a, Expr b) {
 }  // namespace
 
 static void MergeMulModInsertElements(
-    const std::vector<ir::IndexExpr> &eles,
+    const std::vector<ir::IndexExpr> &elems,
     std::list<ir::IndexExpr> *mult_exprs,
     std::list<std::pair<ir::IndexExpr, ir::IndexExpr>> *mod_exprs,
     ir::IndexExpr *no_opt_sum,
     bool *has_mult,
     bool *has_mod) {
   *has_mult = false;
   *has_mod = false;
-  for (const ir::IndexExpr ele : eles) {
+  for (const ir::IndexExpr ele : elems) {
     auto mod_ptr = ele.As<ir::Mod>();
     auto mult_ptr = ele.As<ir::Mul>();
     if (mod_ptr) {
@@ -234,6 +234,14 @@ static std::optional<ir::IndexExpr> MergeMulModInner(
     } else if (inner_div_ptr) {
       ir::IndexExpr overall_mult =
           mult_inner.get() ? mult_inner * mult_outer : mult_outer;
+      VLOG(5) << "inner_div_ptr_b: " << inner_div_ptr->b().as_index();
+      VLOG(5) << "overall_mult: " << overall_mult;
+      VLOG(5) << "mod_r_expr: " << mod_r_expr;
+      VLOG(5) << "inner_div_ptr_a - mod_l_expr: "
+              << inner_div_ptr->a().as_index() - mod_l_expr;
+      VLOG(5) << "ProveDivisible: "
+              << ProveDivisible(inner_div_ptr->a().as_index() - mod_l_expr,
+                                mod_r_expr);
       if (overall_mult == inner_div_ptr->b().as_index() &&
           overall_mult == mod_r_expr &&
           ProveDivisible(inner_div_ptr->a().as_index() - mod_l_expr,
@@ -273,14 +281,14 @@ static std::optional<ir::IndexExpr> MergeMulModInner(
 ir::IndexExpr MergeMulMod(SymbolicExprAnalyzer *analyzer,
                           const ir::IndexExpr &base) {
   ir::IndexExpr simplified_base = base.as_index().Normalize();
-  std::vector<ir::IndexExpr> eles = GetFlattenExprs<ir::Add>(simplified_base);
+  std::vector<ir::IndexExpr> elems = GetFlattenExprs<ir::Add>(simplified_base);
   std::list<ir::IndexExpr> mult_exprs;
   std::list<std::pair<ir::IndexExpr, ir::IndexExpr>> mod_exprs;
   ir::IndexExpr no_opt_sum;
   bool has_mult;
   bool has_mod;
   MergeMulModInsertElements(
-      eles, &mult_exprs, &mod_exprs, &no_opt_sum, &has_mult, &has_mod);
+      elems, &mult_exprs, &mod_exprs, &no_opt_sum, &has_mult, &has_mod);
   bool find_opt = false;
   std::list<std::pair<ir::IndexExpr, ir::IndexExpr>>::iterator search_mod_it =
       mod_exprs.begin();
@@ -297,9 +305,9 @@ ir::IndexExpr MergeMulMod(SymbolicExprAnalyzer *analyzer,
         ++search_mod_it;
         mod_exprs.erase(temp_mod_it);
         mult_exprs.erase(mult_it);
-        std::vector<ir::IndexExpr> ret_eles =
+        std::vector<ir::IndexExpr> ret_elems =
             GetFlattenExprs<ir::Add>(ret.value());
-        MergeMulModInsertElements(ret_eles,
+        MergeMulModInsertElements(ret_elems,
                                   &mult_exprs,
                                   &mod_exprs,
                                   &no_opt_sum,
@@ -348,7 +356,7 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
                     ::common::errors::InvalidArgument(
                         "The size of shape should be less than or "
                         "equal to the size of indices."));
-  Expr res;
+  Expr res(0);
   ir::TryElevateInt32ToInt64(shape);
   common::cas_intervals_t var_intervals =
       common::CollectVarIntervalsOfExprs(indices);
@@ -363,7 +371,8 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
             "the current data type of shape[{}] is {}",
             i,
             shape[i].type()));
-    Expr indice_cast = indices[i];
+
+    ir::IndexExpr indice_cast = indices[i];
     optim::SimplifyCast(&indice_cast);
     if (res.defined()) {
       res = RampRelatedAdd(RampRelatedMul(res, shape[i]), indice_cast);
@@ -373,6 +382,7 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
     } else {
       res = indice_cast;
     }
+
     if (i > 0) {
       if (res.is_index()) {
         res = MergeMulMod(&analyzer, res.as_index()).as_index().Normalize();
@@ -660,6 +670,7 @@ bool ComparePriority(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs) {
 
 bool IsSumPartialBySymbol(const ir::IndexExpr &expr,
                           const ir::IndexExpr &symbol) {
+  if (expr == symbol) return true;
   // TODO(liujinnan): Check Ty
   switch (expr.node_type()) {
     case ir::IrNodeTy::IntImm: {
@@ -690,10 +701,102 @@ bool IsSumPartialBySymbol(const ir::IndexExpr &expr,
           expr));
   }
 }
+ir::IndexExpr SimplifySymbolicAdd(const ir::IndexExpr &lhs,
+                                  const ir::IndexExpr &sym,
+                                  const ir::IndexExpr &outter_mul_factor) {
+  if (lhs == sym) return sym * (outter_mul_factor + ir::IndexExpr(1));
+  switch (lhs.node_type()) {
+    case ir::IrNodeTy::IntImm: {
+      auto imm = lhs.As<ir::IntImm>();
+      if (imm->value != 0)
+        PADDLE_THROW(::common::errors::Fatal("Error in SimplifySymbolicAdd!"));
+      return ir::IndexExpr(0);
+    }
+    case ir::IrNodeTy::_Var_: {
+      return sym * (outter_mul_factor + ir::IndexExpr(1));
+    }
+    case ir::IrNodeTy::Add: {
+      if (!common::IsSumPartialBySymbol(lhs->operand(0).as_index(), sym))
+        return lhs->operand(0).as_index() +
+               SimplifySymbolicAdd(
+                   lhs->operand(1).as_index(), sym, outter_mul_factor);
+      return SimplifySymbolicAdd(
+                 lhs->operand(0).as_index(), sym, outter_mul_factor) +
+             lhs->operand(1).as_index();
+    }
+    case ir::IrNodeTy::Mul: {
+      if (lhs->operand(1).is_constant() &&
+          lhs->operand(1).get_constant() == -1) {
+        return SimplifySymbolicAdd(
+                   lhs->operand(0).as_index(), sym, -outter_mul_factor) *
+               lhs->operand(1).as_index();
+      }
+      if (lhs->operand(0).as_index() == sym)
+        return lhs->operand(0).as_index() *
+               (lhs->operand(1).as_index() + outter_mul_factor);
+      return (lhs->operand(0).as_index() + outter_mul_factor) *
+             lhs->operand(1).as_index();
+    }
+    case ir::IrNodeTy::Mod:
+      PADDLE_THROW(::common::errors::Fatal("Error in SimplifySymbolicAdd!"));
+    case ir::IrNodeTy::Div: {
+      return SimplifySymbolicAdd(
+                 lhs->operand(0).as_index(),
+                 sym,
+                 lhs->operand(1).as_index() * outter_mul_factor) /
+             lhs->operand(1).as_index();
+    }
+    default:
+      PADDLE_THROW(::common::errors::InvalidArgument(
+          "Unsupported type of lhs in SimplifySymbolicAdd which is: %s", lhs));
+  }
+}
+
+ir::IndexExpr SimplifySymbolicDivide(const ir::IndexExpr &lhs,
+                                     const ir::IndexExpr &sym,
+                                     const ir::IrNodeTy &ty) {
+  if (lhs == sym) return ir::IndexExpr(1);
+  switch (lhs.node_type()) {
+    case ir::IrNodeTy::IntImm: {
+      auto imm = lhs.As<ir::IntImm>();
+      if (imm->value != 0)
+        PADDLE_THROW(
+            ::common::errors::Fatal("Error in SimplifySymbolicDivide!"));
+      return ir::IndexExpr(0);
+    }
+    case ir::IrNodeTy::_Var_:
+      return ir::IndexExpr(1);
+    case ir::IrNodeTy::Add:
+      return SimplifySymbolicDivide(lhs->operand(0).as_index(), sym, ty) +
+             SimplifySymbolicDivide(lhs->operand(1).as_index(), sym, ty);
+    case ir::IrNodeTy::Mul: {
+      if (!common::IsDivisiblieBySymbol(lhs->operand(0).as_index(), sym, ty))
+        return lhs->operand(0).as_index() *
+               SimplifySymbolicDivide(lhs->operand(1).as_index(), sym, ty);
+      return SimplifySymbolicDivide(lhs->operand(0).as_index(), sym, ty) *
+             lhs->operand(1).as_index();
+    }
+    case ir::IrNodeTy::Mod:
+      return SimplifySymbolicDivide(
+                 lhs->operand(0).as_index(), sym, lhs.node_type()) %
+             SimplifySymbolicDivide(
+                 lhs->operand(1).as_index(), sym, lhs.node_type());
+    case ir::IrNodeTy::Div: {
+      return SimplifySymbolicDivide(
+                 lhs->operand(0).as_index(), sym, lhs.node_type()) /
+             lhs->operand(1).as_index();
+    }
+    default:
+      PADDLE_THROW(::common::errors::InvalidArgument(
+          "Unsupported type of lhs in SimplifySymbolicDivide which is: %s",
+          lhs));
+  }
+}
 
 bool IsDivisiblieBySymbol(const ir::IndexExpr &expr,
                           const ir::IndexExpr &symbol,
                           const ir::IrNodeTy &ty) {
+  if (expr == symbol) return true;
   // TODO(liujinnan): Check Ty
   switch (expr.node_type()) {
     case ir::IrNodeTy::IntImm: {
 
@@ -262,6 +262,22 @@ bool ComparePriority(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
 bool IsSumPartialBySymbol(const ir::IndexExpr &expr,
                           const ir::IndexExpr &symbol);
 
+/*!
+ * \brief Simplify the `lhs` by symbol `sym`. Usually run after
+ * `IsSumPartialBySymbol`
+ *
+ * \param lhs The expression to be simplified.
+ * \param sym  The symbol to be checked.
+ *    it may be `i, j ..` or  `S0, S1 ..` or other symbolic expr.
+ * \param outter_mul_factor The scale of symbolic expr.
+ *    e.g. `S0 * 4` ===> sym == S0, outter_mul_factor == 4
+ * \return The expr after simplification.
+ */
+ir::IndexExpr SimplifySymbolicAdd(
+    const ir::IndexExpr &lhs,
+    const ir::IndexExpr &sym,
+    const ir::IndexExpr &outter_mul_factor = ir::IndexExpr(1));
+
 /*!
  * \brief Determines whether there are sub-parts in the `expr` that can be
  * simplified by `Div` operation with the input `symbol`. If true is returned,
@@ -289,6 +305,20 @@ bool IsDivisiblieBySymbol(const ir::IndexExpr &expr,
                           const ir::IndexExpr &symbol,
                           const ir::IrNodeTy &ty);
 
+/*!
+ * \brief Simplify the `lhs` by symbol `sym`. Usually run after
+ * `IsDivisiblieBySymbol`
+ *
+ * \param lhs The expression to be simplified.
+ * \param sym  The symbol to be checked.
+ *    it may be `i, j ..` or  `S0, S1 ..` or other symbolic expr.
+ * \param ty ty is `Mod` or `Div`.
+ * \return The expr after simplification.
+ */
+ir::IndexExpr SimplifySymbolicDivide(const ir::IndexExpr &lhs,
+                                     const ir::IndexExpr &sym,
+                                     const ir::IrNodeTy &ty);
+
 /*!
  * \brief Determine whether `lhs` is divisible by `rhs`, regardless of whether
  * `rhs` is a constant or a symbol.
@@ -300,8 +330,8 @@ bool ProveDivisible(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
 
 /*!
  * \brief Judge whether `candidate` is a negated index expression.
- * \param lhs The expression to be checked.
- * \param rhs The positive part
+ * \param candidate The expression to be checked.
+ * \param expr The positive part
  * \return A boolean value indicating whether `candidate` is negative.
  */
 bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
 
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 #pragma once
-#include "paddle/cinn/common/simplify_corner_case.h"
+#include "paddle/cinn/common/simplify_special_pattern.h"
 #include <optional>
 #include <stack>
 #include <unordered_map>
@@ -58,6 +58,43 @@ std::optional<ir::IndexExpr> DivMulAddModCornerCase(const ir::IndexExpr& lhs,
   return std::nullopt;
 }
 
+// (S0 * 8 + S1 * 2 + S2) + (S1 * 2 + S2) * (-1) ===> 0
+std::optional<ir::IndexExpr> AddMulCornerCase(
+    const ir::IndexExpr& lhs,
+    const ir::IndexExpr& rhs,
+    const ir::IndexExpr& scale = ir::IndexExpr(1)) {
+  auto rhsMul = rhs.As<ir::Mul>();
+  if (!rhsMul) return std::nullopt;
+  if (!rhsMul->b().is_constant()) return std::nullopt;
+
+  auto scale_ = scale * rhsMul->b().as_index();
+  auto flatten = GetFlattenExprs<ir::Add>(rhsMul->a());
+  std::optional<ir::IndexExpr> resOpt;
+  ir::IndexExpr res = lhs;
+  for (const auto& expr : flatten) {
+    if (auto innerMul = expr.As<ir::Mul>()) {
+      if (!innerMul->b().is_constant()) return std::nullopt;
+      auto resOpt = AddMulCornerCase(res, expr, scale_);
+      if (!resOpt.has_value())
+        return std::nullopt;
+      else
+        res = resOpt.value();
+    } else {
+      if (!IsSumPartialBySymbol(res, expr)) return std::nullopt;
+    }
+  }
+
+  for (const auto& expr : flatten) {
+    if (expr.As<ir::Mul>()) continue;
+    if (expr.is_constant()) {
+      res = res + expr * scale_;
+      continue;
+    }
+    res = SimplifySymbolicAdd(res, expr, scale_);
+  }
+  return res;
+}
+
 // (S0 + S1 - (S0 + S1) % S2) % S2 == 0
 // (S0 + S1 - (S0 + S1) % S2) / S2 == (S0 + S1) / S2
 std::optional<ir::IndexExpr> SubModCornerCase(const ir::IndexExpr& lhs,
@@ -86,6 +123,20 @@ std::optional<ir::IndexExpr> SubModCornerCase(const ir::IndexExpr& lhs,
       return isNeg ? innerMod->a().as_index() / rhs
                    : -(innerMod->a().as_index() / rhs);
     }
+
+    // For simplify mod case: ((S0 * 256 + S1) % 512 - S1) % 32 == 0
+    if (!isDiv) {
+      auto diffBeforeNegation = diff;
+      auto isDiffNeg = IsNegatedIndexExpr(diff, diffBeforeNegation);
+      if (isDiffNeg) diff = diffBeforeNegation;
+      auto flatten_diff = GetFlattenExprs<ir::Add>(diff);
+      bool isDivisible = true;
+      for (const auto& expr : flatten_diff) {
+        if (!isDivisible) break;
+        if (!ProveDivisible(expr, rhs)) isDivisible = false;
+      }
+      if (isDivisible) return ir::IndexExpr(0);
+    }
   }
   return std::nullopt;
 }
@@ -125,9 +176,8 @@ std::optional<ir::IndexExpr> SimplifyCornerCase(const ir::IndexExpr& expr) {
 
 std::optional<ir::IndexExpr> SimplifyAddCornerCase(const ir::IndexExpr& lhs,
                                                    const ir::IndexExpr& rhs) {
-  if (DivMulAddModCornerCase(lhs, rhs).has_value())
-    return DivMulAddModCornerCase(lhs, rhs).value();
-
+  if (auto res = DivMulAddModCornerCase(lhs, rhs)) return res.value();
+  if (auto res = AddMulCornerCase(lhs, rhs)) return res.value();
   // Add other corner cases
   return std::nullopt;
 }
@@ -140,21 +190,17 @@ std::optional<ir::IndexExpr> SimplifyMulCornerCase(const ir::IndexExpr& lhs,
 
 std::optional<ir::IndexExpr> SimplifyDivCornerCase(const ir::IndexExpr& lhs,
                                                    const ir::IndexExpr& rhs) {
-  if (SubModCornerCase(lhs, rhs, true).has_value())
-    return SubModCornerCase(lhs, rhs, true).value();
-  if (MultiArgsDivAndMod(lhs, rhs, true).has_value())
-    return MultiArgsDivAndMod(lhs, rhs, true).value();
+  if (auto res = SubModCornerCase(lhs, rhs, true)) return res.value();
+  if (auto res = MultiArgsDivAndMod(lhs, rhs, true)) return res.value();
   // Add other corner cases
   return std::nullopt;
 }
 
 std::optional<ir::IndexExpr> SimplifyModCornerCase(const ir::IndexExpr& lhs,
                                                    const ir::IndexExpr& rhs) {
-  if (SubModCornerCase(lhs, rhs, false).has_value())
-    return SubModCornerCase(lhs, rhs, false).value();
+  if (auto res = SubModCornerCase(lhs, rhs, false)) return res.value();
   // Add other corner cases
-  if (MultiArgsDivAndMod(lhs, rhs, false).has_value())
-    return MultiArgsDivAndMod(lhs, rhs, false).value();
+  if (auto res = MultiArgsDivAndMod(lhs, rhs, false)) return res.value();
   return std::nullopt;
 }