PaddlePaddle
diff --git a/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+5-187Lines changed: 5 additions & 187 deletions b/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+5-187Lines changed: 5 additions & 187 deletions
diff --git a/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions b/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions
@@ -175,175 +175,6 @@ Expr RampRelatedMul(Expr a, Expr b) {
 
 }  // namespace
 
-static void MergeMulModInsertElements(
-    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 : elems) {
-    auto mod_ptr = ele.As<ir::Mod>();
-    auto mult_ptr = ele.As<ir::Mul>();
-    if (mod_ptr) {
-      *has_mod = true;
-      mod_exprs->emplace_back(
-          std::make_pair(std::move(mod_ptr->a().as_index()),
-                         std::move(mod_ptr->b().as_index())));
-    } else if (mult_ptr) {
-      *has_mult = true;
-      mult_exprs->emplace_back(ele);
-    } else {
-      *no_opt_sum = no_opt_sum->get() ? *no_opt_sum + ele : ele;
-    }
-  }
-}
-
-static std::optional<ir::IndexExpr> MergeMulModInner(
-    const ir::IndexExpr &mult_expr,
-    const ir::IndexExpr &mod_l_expr,
-    const ir::IndexExpr &mod_r_expr) {
-  const ir::Mul *mult_ptr = mult_expr.As<ir::Mul>();
-  if (!mult_ptr) return std::nullopt;
-  ir::IndexExpr mult_outer = mult_ptr->b().as_index();
-  ir::IndexExpr inner = mult_ptr->a().as_index();
-
-  while (true) {
-    mult_ptr = inner.As<ir::Mul>();
-    if (mult_ptr) {
-      inner = mult_ptr->a().as_index();
-      mult_outer = mult_ptr->b().as_index() * mult_outer;
-    } else {
-      break;
-    }
-  }
-
-  ir::IndexExpr search_ptr = inner;
-  ir::IndexExpr mult_inner;  // The inner multiplication factor
-  ir::IndexExpr no_opt_sum;  // Sum of the exprs that cannot be optimized
-
-  while (true) {
-    auto inner_div_ptr = search_ptr.As<ir::Div>();
-    auto inner_mult_ptr = search_ptr.As<ir::Mul>();
-    auto inner_add_ptr = search_ptr.As<ir::Add>();
-    if (!inner_div_ptr && !inner_mult_ptr && !inner_add_ptr) {
-      return std::nullopt;
-    } 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,
-                         mod_r_expr)) {
-        // Found!
-        return no_opt_sum.get()
-                   ? no_opt_sum * mult_outer + inner_div_ptr->a().as_index()
-                   : inner_div_ptr->a().as_index();
-      } else {
-        return std::nullopt;
-      }
-    } else if (inner_mult_ptr) {
-      mult_inner = mult_inner.get()
-                       ? inner_mult_ptr->b().as_index() * mult_inner
-                       : inner_mult_ptr->b().as_index();
-      search_ptr = inner_mult_ptr->a().as_index();
-    } else if (inner_add_ptr) {
-      if (mult_inner.get()) {
-        return std::nullopt;
-      }
-      auto lhs = inner_add_ptr->a().as_index();
-      auto rhs = inner_add_ptr->b().as_index();
-      if (inner_add_ptr->b().as_index().is_constant()) {
-        std::swap(lhs, rhs);
-      } else if (inner_add_ptr->b().as_index().length() < mod_r_expr.length()) {
-        std::swap(lhs, rhs);
-      }
-      no_opt_sum = no_opt_sum.get() ? no_opt_sum + lhs : lhs;
-      search_ptr = rhs;
-    } else {
-      break;
-    }
-  }
-  return std::nullopt;
-}
-
-ir::IndexExpr MergeMulMod(const ir::IndexExpr &base) {
-  ir::IndexExpr simplified_base = base.Normalize();
-  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(
-      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();
-
-  while (search_mod_it != mod_exprs.end()) {
-    std::list<ir::IndexExpr>::iterator mult_it = mult_exprs.begin();
-    bool inner_find_opt = false;
-    while (mult_it != mult_exprs.end()) {
-      auto ret = MergeMulModInner(
-          *mult_it, search_mod_it->first, search_mod_it->second);
-      if (ret.has_value()) {
-        inner_find_opt = true;
-        auto temp_mod_it = search_mod_it;
-        ++search_mod_it;
-        mod_exprs.erase(temp_mod_it);
-        mult_exprs.erase(mult_it);
-        std::vector<ir::IndexExpr> ret_elems =
-            GetFlattenExprs<ir::Add>(ret.value());
-        MergeMulModInsertElements(ret_elems,
-                                  &mult_exprs,
-                                  &mod_exprs,
-                                  &no_opt_sum,
-                                  &has_mult,
-                                  &has_mod);
-        if (has_mult) {
-          search_mod_it = mod_exprs.begin();
-        } else if (has_mod && search_mod_it == mod_exprs.end()) {
-          search_mod_it--;
-        }
-        break;
-      } else {
-        ++mult_it;
-      }
-    }
-    find_opt = find_opt || inner_find_opt;
-    if (!inner_find_opt) {
-      ++search_mod_it;
-    }
-  }
-  if (!find_opt) {
-    return simplified_base;
-  }
-  for (std::list<ir::IndexExpr>::iterator it = mult_exprs.begin();
-       it != mult_exprs.end();
-       ++it) {
-    no_opt_sum = no_opt_sum.get() ? no_opt_sum + *it : *it;
-  }
-  for (std::list<std::pair<ir::IndexExpr, ir::IndexExpr>>::iterator it =
-           mod_exprs.begin();
-       it != mod_exprs.end();
-       ++it) {
-    no_opt_sum = no_opt_sum.get() ? no_opt_sum + it->first % it->second
-                                  : it->first % it->second;
-  }
-  return no_opt_sum;
-}
-
 Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
                        const std::vector<Expr> &indices) {
   VLOG(3) << "Begin IndiceToAbsOffset";
@@ -373,15 +204,9 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
     optim::SimplifyCast(&indice_cast);
     res = RampRelatedAdd(RampRelatedMul(res, shape[i]), indice_cast);
     if (res.is_index()) {
-      res = res.as_index().Normalize();
-    }
-
-    if (i > 0) {
-      if (res.is_index()) {
-        res = MergeMulMod(res).Normalize();
-      } else {
-        VLOG(8) << "**** expr is not index ****: " << res;
-      }
+      res = res.as_index().Normalize(ir::IndexExpr::OptLevel::Level2);
+    } else {
+      VLOG(8) << "**** expr is not index ****: " << res;
     }
   }
 
@@ -838,20 +663,13 @@ bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
 bool VerifyIndex(const ir::Expr &expr) {
   switch (expr.node_type()) {
     case ir::IrNodeTy::_Var_:
-    case ir::IrNodeTy::IntImm: {
+    case ir::IrNodeTy::IntImm:
+    case ir::IrNodeTy::Load: {
       if (expr.type().is_index_type()) return true;
       return false;
     }
     case ir::IrNodeTy::Cast:
       return VerifyIndex(expr->operand(0));
-    case ir::IrNodeTy::Load: {
-      bool is_index = true;
-      auto load_op = expr.As<ir::Load>();
-      for (const auto &indice : load_op->indices) {
-        if (!VerifyIndex(indice)) return false;
-      }
-      return true;
-    }
     case ir::IrNodeTy::Add:
     case ir::IrNodeTy::Sub:
     case ir::IrNodeTy::Mul:
 
@@ -38,8 +38,6 @@ Expr PrecedingAxisToAbsOffset(const std::vector<Expr> &shape,
 
 Expr CastIfNeeded(Expr body, Type type);
 
-ir::IndexExpr MergeMulMod(const ir::IndexExpr &base);
-
 //! Substitute vars to other expressions.
 //! @param expr The expression to do modification.
 //! @param var_map The map from variables to the target expressions.
@@ -254,6 +252,11 @@ bool ComparePriority(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
  * 6. `IsSumPartialBySymbol(S0 / 3 + S1, S0)` return true;
  * 7. `IsSumPartialBySymbol(S0 % 3, S0)` return false;
  *
+ * Note: For performance reasons, special patterns will not be matched here.
+ * This can be allowed for extreme optimization.
+ * For example:
+ * `IsSumPartialBySymbol((S0 + S1 / 5 * 25) / 5, S1 % 5)` return false;
+ *
  * \param expr The expression to be checked.
  * \param symbol  The symbol to be checked.
  * \return True means there are sub-parts in the `expr` that can be simplified.