PaddlePaddle
diff --git a/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+183-3Lines changed: 183 additions & 3 deletions b/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+183-3Lines changed: 183 additions & 3 deletions
diff --git a/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎paddle/cinn/common/iter_simplify.cc
Copy file name to clipboard
+8-4Lines changed: 8 additions & 4 deletions b/‎paddle/cinn/common/iter_simplify.cc
Copy file name to clipboard
+8-4Lines changed: 8 additions & 4 deletions
diff --git a/‎paddle/cinn/common/iter_simplify.h
Copy file name to clipboard
+5-3Lines changed: 5 additions & 3 deletions b/‎paddle/cinn/common/iter_simplify.h
Copy file name to clipboard
+5-3Lines changed: 5 additions & 3 deletions
@@ -15,6 +15,7 @@
 #include "paddle/cinn/common/ir_util.h"
 
 #include <algorithm>
+#include <stack>
 #include <unordered_set>
 
 #include "paddle/cinn/common/cas.h"
@@ -174,6 +175,169 @@ Expr RampRelatedMul(Expr a, Expr b) {
 
 }  // namespace
 
+static void MergeMulModInsertElements(
+    const std::vector<ir::IndexExpr> &eles,
+    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) {
+    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(
+    SymbolicExprAnalyzer *analyzer,
+    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();
+  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.as_index();
+    } 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;
+      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(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::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);
+  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(
+          analyzer, *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_eles =
+            GetFlattenExprs<ir::Add>(ret.value());
+        MergeMulModInsertElements(ret_eles,
+                                  &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";
@@ -186,7 +350,11 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
                         "equal to the size of indices."));
   Expr res;
   ir::TryElevateInt32ToInt64(shape);
-  for (int i = 0; i < shape.size(); i++) {
+  common::cas_intervals_t var_intervals =
+      common::CollectVarIntervalsOfExprs(indices);
+  common::SymbolicExprAnalyzer analyzer{var_intervals};
+
+  for (int32_t i = 0; i < shape.size(); i++) {
     PADDLE_ENFORCE_EQ(
         shape[i].type() == Int(64) || shape[i].type() == Int(32),
         true,
@@ -199,12 +367,16 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
     optim::SimplifyCast(&indice_cast);
     if (res.defined()) {
       res = RampRelatedAdd(RampRelatedMul(res, shape[i]), indice_cast);
+      if (res.is_index()) {
+        res = res.as_index().Normalize();
+      }
     } else {
       res = indice_cast;
     }
-
     if (i > 0) {
-      res = cinn::common::AutoSimplify(res);
+      if (res.is_index()) {
+        res = MergeMulMod(&analyzer, res.as_index()).as_index().Normalize();
+      }
     }
   }
 
@@ -512,6 +684,10 @@ bool IsSumPartialBySymbol(const ir::IndexExpr &expr,
     }
     case ir::IrNodeTy::Mod:
       return false;
+    default:
+      PADDLE_THROW(::common::errors::InvalidArgument(
+          "Unsupported type of expr in IsSumPartialBySymbol which is: %s",
+          expr));
   }
 }
 
@@ -544,6 +720,10 @@ bool IsDivisiblieBySymbol(const ir::IndexExpr &expr,
       return IsDivisiblieBySymbol(
           expr->operand(0).as_index(), symbol, expr.node_type());
     }
+    default:
+      PADDLE_THROW(::common::errors::InvalidArgument(
+          "Unsupported type of expr in IsDivisiblieBySymbol which is: %s",
+          expr));
   }
 }
 
 
@@ -22,6 +22,7 @@
 
 #include "paddle/cinn/common/bfloat16.h"
 #include "paddle/cinn/common/float16.h"
+#include "paddle/cinn/common/integer_set.h"
 #include "paddle/cinn/ir/ir.h"
 
 namespace cinn {
@@ -37,6 +38,9 @@ Expr PrecedingAxisToAbsOffset(const std::vector<Expr> &shape,
 
 Expr CastIfNeeded(Expr body, Type type);
 
+ir::IndexExpr MergeMulMod(SymbolicExprAnalyzer *analyzer,
+                          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.
 
@@ -49,11 +49,11 @@ ir::IndexExpr IterMapToExprNormalizer::ConvertIterSum(ir::IterSum* expr) {
 }
 
 ir::IndexExpr IterMapToExprNormalizer::ConvertIterSplit(ir::IterSplit* expr) {
-  // quick branch
-  if (IsZero(expr->scale) || IsOne(expr->extent)) return ir::IndexExpr(0);
   ir::IndexExpr source;
   ir::IterMark* mark = expr->source.As<ir::IterMark>();
   if (auto opt = mark->source.As<ir::_Var_>()) {
+    // For unit loop, we can't simplify loop_var to `0`.
+    if (IsOne(mark->extent)) return opt;
     source = opt;
   } else if (auto opt = mark->source.As<ir::IterSum>()) {
     source = ConvertIterSum(opt);
@@ -62,6 +62,11 @@ ir::IndexExpr IterMapToExprNormalizer::ConvertIterSplit(ir::IterSplit* expr) {
     Visit(&(mark->source), &(mark->source));
     source = mark->source;
   }
+
+  // quick branch
+  if (IsZero(expr->scale) || IsOne(expr->extent))
+    return ir::Zero(expr->extent.type());
+
   if (analyzer_.ProveEQ(expr->extent, mark->extent).value_or(false) &&
       IsOne(expr->lower_factor)) {
     return source * expr->scale;
@@ -642,8 +647,7 @@ void IterMapRewriter::AddToLhs(ir::IterSum* lhs,
   if (sign > 0) {
     lhs->args.push_back(rhs_expr);
   } else {
-    rhs_expr.As<ir::IterSplit>()->scale =
-        ir::Zero(rhs.scale.type()).as_index() - rhs.scale;
+    rhs_expr.As<ir::IterSplit>()->scale = -rhs.scale;
     lhs->args.push_back(rhs_expr);
   }
 }
 
@@ -49,9 +49,11 @@ class IterMapRewriter : public ir::IRMutator<> {
                            const SymbolicExprAnalyzer& analyzer)
       : analyzer_(analyzer) {
     for (const auto& iter : input_iters) {
-      if (IsOne(iter->upper_bound)) {
-        var_map_[iter->name] = ir::IterSum::Make({}, iter->lower_bound);
-      } else if (IsZero(iter->lower_bound)) {
+      // Remove unit extent `for` ?
+      // if (IsOne(iter->upper_bound)) {
+      //   var_map_[iter->name] = ir::IterSum::Make({}, iter->lower_bound);
+      // }
+      if (IsZero(iter->lower_bound)) {
         auto tmp =
             ir::IterMark::Make(ir::IndexExpr(iter.ptr()), iter->upper_bound);
         auto mark = tmp.As<ir::IterMark>();