PaddlePaddle
diff --git a/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+136Lines changed: 136 additions & 0 deletions b/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+136Lines changed: 136 additions & 0 deletions
diff --git a/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+50Lines changed: 50 additions & 0 deletions b/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+50Lines changed: 50 additions & 0 deletions
diff --git a/‎paddle/cinn/common/iter_simplify.cc
Copy file name to clipboard
+25-25Lines changed: 25 additions & 25 deletions b/‎paddle/cinn/common/iter_simplify.cc
Copy file name to clipboard
+25-25Lines changed: 25 additions & 25 deletions
diff --git a/‎paddle/cinn/common/iter_simplify.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/common/iter_simplify.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
@@ -21,6 +21,7 @@
 #include "paddle/cinn/ir/ir_mutator.h"
 #include "paddle/cinn/ir/ir_printer.h"
 #include "paddle/cinn/ir/op/ir_operators.h"
+#include "paddle/cinn/ir/utils/ir_compare.h"
 #include "paddle/common/enforce.h"
 namespace cinn {
 namespace common {
@@ -501,5 +502,140 @@ Expr min(Expr a, Expr b) {
   return ir::Min::Make(a, b);
 }
 
+int32_t CalculateExprComplexity(const Expr &expr, int count) {
+  switch (expr.node_type()) {
+    case ir::IrNodeTy::_Var_:
+    case ir::IrNodeTy::IntImm:
+      return count + 1;
+    case ir::IrNodeTy::Add:
+    case ir::IrNodeTy::Mul:
+    case ir::IrNodeTy::Div:
+    case ir::IrNodeTy::Mod: {
+      int lhs_count = CalculateExprComplexity(expr->operand(0), count);
+      int rhs_count = CalculateExprComplexity(expr->operand(1), count);
+      return lhs_count + rhs_count + 1;
+    }
+  }
+}
+
+bool IsCorrectPriority(const Expr &lhs, const Expr &rhs) {
+  if (lhs.node_type() == ir::IrNodeTy::IntImm &&
+      rhs.node_type() != ir::IrNodeTy::IntImm)
+    return false;
+  if (rhs.node_type() == ir::IrNodeTy::IntImm &&
+      lhs.node_type() != ir::IrNodeTy::IntImm)
+    return true;
+  if (auto lhsVar = lhs.As<ir::_Var_>())
+    if (auto rhsVar = rhs.As<ir::_Var_>())
+      return std::make_tuple(lhsVar->name.length(), lhsVar->name) <=
+             std::make_tuple(rhsVar->name.length(), rhsVar->name);
+  auto lhsComplexity = CalculateExprComplexity(lhs);
+  auto rhsComplexity = CalculateExprComplexity(rhs);
+  if (lhsComplexity < rhsComplexity) return false;
+  // Mul < Div < Mod.
+  else if (lhsComplexity == rhsComplexity)
+    return lhs.node_type() <= rhs.node_type();
+  else
+    return true;
+}
+
+ir::IndexExpr MulAndNormalize(const ir::IndexExpr &lhs,
+                              const ir::IndexExpr &rhs) {
+  int64_t cscale = 1;
+  ir::IndexExpr res = ir::One(lhs.type());
+  auto fcollect = [&](ir::IndexExpr val) {
+    if (const auto *intimm = val.As<ir::IntImm>()) {
+      cscale *= intimm->value;
+    } else {
+      res = res * val;
+    }
+  };
+  UnpackReduction<ir::Mul>(lhs, fcollect);
+  UnpackReduction<ir::Mul>(rhs, fcollect);
+  if (cscale != 1) {
+    res = res * ir::IndexExpr(make_shared<ir::IntImm>(res.type(), cscale));
+  }
+  return res;
+}
+
+bool IsSumPartialBySymbol(const ir::IndexExpr &expr,
+                          const ir::IndexExpr &symbol) {
+  // TODO(liujinnan): Check Ty
+  switch (expr.node_type()) {
+    case ir::IrNodeTy::IntImm: {
+      return false;
+    }
+    case ir::IrNodeTy::_Var_:
+      return expr == symbol;
+    case ir::IrNodeTy::Add:
+      [[fallthrough]];
+    case ir::IrNodeTy::Sub:
+      return IsSumPartialBySymbol(expr->operand(0).as_index(), symbol) ||
+             IsSumPartialBySymbol(expr->operand(1).as_index(), symbol);
+    case ir::IrNodeTy::Mul:
+      return expr->operand(0).as_index() == symbol ||
+             expr->operand(0).as_index() == symbol;
+    case ir::IrNodeTy::Div: {
+      return IsSumPartialBySymbol(expr->operand(0).as_index(), symbol);
+    }
+    case ir::IrNodeTy::Mod:
+      return false;
+  }
+}
+
+bool IsDivisiblieBySymbol(const ir::IndexExpr &expr,
+                          const ir::IndexExpr &symbol,
+                          const ir::IrNodeTy &ty) {
+  // TODO(liujinnan): Check Ty
+  switch (expr.node_type()) {
+    case ir::IrNodeTy::IntImm: {
+      auto imm = expr.As<ir::IntImm>();
+      return imm->value == 0;
+    }
+    case ir::IrNodeTy::_Var_:
+      return expr == symbol;
+    case ir::IrNodeTy::Add:
+      return IsDivisiblieBySymbol(expr->operand(0).as_index(), symbol, ty) &&
+             IsDivisiblieBySymbol(expr->operand(1).as_index(), symbol, ty);
+    case ir::IrNodeTy::Mul:
+      return IsDivisiblieBySymbol(expr->operand(0).as_index(), symbol, ty) ||
+             IsDivisiblieBySymbol(expr->operand(1).as_index(), symbol, ty);
+    case ir::IrNodeTy::Mod:
+      return IsDivisiblieBySymbol(
+                 expr->operand(0).as_index(), symbol, expr.node_type()) &&
+             IsDivisiblieBySymbol(
+                 expr->operand(1).as_index(), symbol, expr.node_type());
+    case ir::IrNodeTy::Div: {
+      if (ty != expr.node_type()) return false;
+      return IsDivisiblieBySymbol(
+          expr->operand(0).as_index(), symbol, expr.node_type());
+    }
+  }
+}
+
+bool ProveDivisible(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs) {
+  // TODO(liujinnan): corner case, it will upgrade to a more general solution
+  // `expr.Match(xxxx)` later
+  if (auto lhs_sub = lhs.As<ir::Sub>()) {
+    auto llhs = lhs_sub->a();
+    auto lrhs = lhs_sub->b();
+    if (auto llhs_mod = llhs.As<ir::Mod>()) {
+      return (ProveDivisible(llhs_mod->b(), rhs) && llhs_mod->a() == lrhs);
+    }
+    if (auto lrhs_mod = lrhs.As<ir::Mod>()) {
+      return (ProveDivisible(lrhs_mod->b(), rhs) && lrhs_mod->a() == llhs);
+    }
+  }
+
+  if (auto rhs_imm = rhs.As<ir::IntImm>()) {
+    if (lhs.as_index().GetLargestMutiplyPart() % rhs_imm->value == 0)
+      return true;
+    return IsZero(AutoSimplify(lhs % rhs));
+  } else if (rhs.is_var()) {
+    return IsDivisiblieBySymbol(lhs, rhs, ir::IrNodeTy::Div);
+  } else {
+    return false;
+  }
+}
 }  // namespace common
 }  // namespace cinn
@@ -160,5 +160,55 @@ Expr FoldExpr(FuncOp func_op, const std::vector<Expr> &values) {
   return init_value;
 }
 
+inline bool IsIterExpr(const Expr &a, const Expr &b) {
+  return a.As<ir::IterSplit>() || a.As<ir::IterSum>() ||
+         b.As<ir::IterSplit>() || b.As<ir::IterSum>();
+}
+
+inline bool IsOne(const Expr &expr) {
+  if (expr.is_constant() && expr.get_constant() == 1) {
+    return true;
+  }
+  return false;
+}
+inline bool IsZero(const Expr &expr) {
+  if (expr.is_constant() && expr.get_constant() == 0) {
+    return true;
+  }
+  return false;
+}
+
+template <typename TNode, typename FLeaf>
+inline void UnpackReduction(const ir::IndexExpr &value, FLeaf fleaf) {
+  if (const TNode *node = value.As<TNode>()) {
+    UnpackReduction<TNode, FLeaf>(node->a(), fleaf);
+    UnpackReduction<TNode, FLeaf>(node->b(), fleaf);
+  } else {
+    fleaf(value);
+  }
+}
+
+// TODO(liuruyan): canby simplify into IndexExpr multiply.
+ir::IndexExpr MulAndNormalize(const ir::IndexExpr &lhs,
+                              const ir::IndexExpr &rhs);
+
+int32_t CalculateExprComplexity(const Expr &expr, int count = 0);
+
+// True means don't change sequence
+bool IsCorrectPriority(const Expr &lhs, const Expr &rhs);
+
+bool IsSumPartialBySymbol(const ir::IndexExpr &expr,
+                          const ir::IndexExpr &symbol);
+
+// If true is returned, the operation will be attempted on each subpart in
+// outter `simplify` function. Note: this func dont deal the corner case, e.g.
+// `IsDivisiblieBySymbol(f % S0  - f, S0)` is `false`. please use
+// `ProveDivisible` for exact result.
+bool IsDivisiblieBySymbol(const ir::IndexExpr &expr,
+                          const ir::IndexExpr &symbol,
+                          const ir::IrNodeTy &ty);
+
+bool ProveDivisible(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
+
 }  // namespace common
 }  // namespace cinn
@@ -14,15 +14,15 @@
 
 #include "paddle/cinn/common/iter_simplify.h"
 #include "paddle/cinn/common/const_fold.h"
-#include "paddle/cinn/common/integer_set.h"
+#include "paddle/cinn/common/ir_util.h"
 #include "paddle/cinn/ir/ir_mutator.h"
-#include "paddle/cinn/ir/ir_utils.h"
 #include "paddle/cinn/ir/ir_visitor.h"
 #include "paddle/cinn/ir/op/ir_operators.h"
 #include "paddle/cinn/ir/utils/ir_compare.h"
 #include "paddle/cinn/ir/utils/ir_copy.h"
 #include "paddle/common/enforce.h"
 #include "paddle/common/errors.h"
+
 namespace cinn {
 namespace common {
 
@@ -62,11 +62,11 @@ ir::IndexExpr IterMapToExprNormalizer::ConvertIterSplit(ir::IterSplit* expr) {
     Visit(&(mark->source), &(mark->source));
     source = mark->source;
   }
-  if (ProveEQ(expr->extent, mark->extent, analyzer_) &&
+  if (analyzer_.ProveEQ(expr->extent, mark->extent).value_or(false) &&
       IsOne(expr->lower_factor)) {
     return source * expr->scale;
-  } else if (ProveLE(
-                 mark->extent, expr->lower_factor * expr->extent, analyzer_)) {
+  } else if (analyzer_.ProveLE(mark->extent, expr->lower_factor * expr->extent)
+                 .value_or(false)) {
     if (IsOne(expr->extent)) {
       return ir::Zero(expr->extent.type());
     }
@@ -267,18 +267,17 @@ ir::IndexExpr IterMapRewriter::SplitDivConst(ir::IndexExpr lhs_expr,
 
   auto lhs = lhs_expr.As<ir::IterSplit>();
   if (!IsOne(lhs->scale)) {
-    if (ProveDivisible(lhs->scale, rhs, analyzer_) && IsZero(base)) {
+    if (ProveDivisible(lhs->scale, rhs) && IsZero(base)) {
       lhs->scale = lhs->scale / rhs;
       return lhs;
-    } else if (ProveDivisible(lhs->scale, rhs, analyzer_) &&
-               ProveDivisible(base, rhs, analyzer_)) {
+    } else if (ProveDivisible(lhs->scale, rhs) && ProveDivisible(base, rhs)) {
       lhs->scale = lhs->scale / rhs;
       return ir::IterSum::Make({lhs}, base / rhs);
-    } else if (ProveDivisible(rhs, lhs->scale, analyzer_) && IsZero(base)) {
+    } else if (ProveDivisible(rhs, lhs->scale) && IsZero(base)) {
       rhs = rhs / lhs->scale;
       lhs->scale = ir::One(rhs.type());
-    } else if (ProveDivisible(rhs, lhs->scale, analyzer_) &&
-               ProveDivisible(base, lhs->scale, analyzer_)) {
+    } else if (ProveDivisible(rhs, lhs->scale) &&
+               ProveDivisible(base, lhs->scale)) {
       base = base / lhs->scale;
       rhs = rhs / lhs->scale;
       lhs->scale = ir::One(rhs.type());
@@ -293,7 +292,7 @@ ir::IndexExpr IterMapRewriter::SplitDivConst(ir::IndexExpr lhs_expr,
   // divided by divisor now.
 
   ir::IndexExpr new_split;
-  if (!ProveDivisible(base, rhs, analyzer_)) {
+  if (!ProveDivisible(base, rhs)) {
     // padding base to divisor later. Treat the whole expr as IterMark now.
     return ir::IterSum::Make(
         {ir::IterSplit::Make(
@@ -305,13 +304,13 @@ ir::IndexExpr IterMapRewriter::SplitDivConst(ir::IndexExpr lhs_expr,
         ir::Zero(rhs.type()));
   }
 
-  if (ProveDivisible(lhs->extent, rhs, analyzer_)) {
+  if (ProveDivisible(lhs->extent, rhs)) {
     new_split = ir::IterSplit::Make(
         lhs->source, lhs->lower_factor * rhs, lhs->extent / rhs, lhs->scale);
   } else if (IsOne(lhs->lower_factor) &&
-             ProveEQ(lhs->extent,
-                     lhs->source.As<ir::IterMark>()->extent,
-                     analyzer_)) {
+             analyzer_
+                 .ProveEQ(lhs->extent, lhs->source.As<ir::IterMark>()->extent)
+                 .value_or(false)) {
     new_split = ir::IterSplit::Make(
         lhs->source, rhs, (lhs->extent + rhs - 1) / rhs, lhs->scale);
   } else {
@@ -334,15 +333,14 @@ ir::IndexExpr IterMapRewriter::SplitModConst(ir::IndexExpr lhs_expr,
 
   auto lhs = lhs_expr.As<ir::IterSplit>();
   if (!IsOne(lhs->scale)) {
-    if (ProveDivisible(lhs->scale, rhs, analyzer_) && IsZero(base)) {
+    if (ProveDivisible(lhs->scale, rhs) && IsZero(base)) {
       return ir::Zero(lhs_expr.type());
-    } else if (ProveDivisible(lhs->scale, rhs, analyzer_) &&
-               ProveDivisible(base, rhs, analyzer_)) {
+    } else if (ProveDivisible(lhs->scale, rhs) && ProveDivisible(base, rhs)) {
       return ir::Zero(lhs_expr.type());
-    } else if (ProveDivisible(rhs, lhs->scale, analyzer_) && IsZero(base)) {
+    } else if (ProveDivisible(rhs, lhs->scale) && IsZero(base)) {
       rhs = rhs / lhs->scale;
-    } else if (ProveDivisible(rhs, lhs->scale, analyzer_) &&
-               ProveDivisible(base, lhs->scale, analyzer_)) {
+    } else if (ProveDivisible(rhs, lhs->scale) &&
+               ProveDivisible(base, lhs->scale)) {
       base = base / lhs->scale;
       rhs = rhs / lhs->scale;
     } else {
@@ -352,7 +350,7 @@ ir::IndexExpr IterMapRewriter::SplitModConst(ir::IndexExpr lhs_expr,
     }
   }
 
-  if (!ProveDivisible(base, rhs, analyzer_)) {
+  if (!ProveDivisible(base, rhs)) {
     auto lhs_s1 = ir::IterSplit::Make(
         lhs->source, lhs->lower_factor, lhs->extent, ir::One(lhs_expr.type()));
     // padding base to divisor later. Treat the whole expr as IterMark now.
@@ -394,7 +392,7 @@ int32_t IterMapRewriter::FindSplitWithExactScale(
     auto split = expr.args[j].As<ir::IterSplit>();
     if (match_source.defined() && match_source != split->source) continue;
     const ir::IndexExpr& cur_scale = split->scale;
-    if (ProveEQ(cur_scale, expected_scale, analyzer_)) {
+    if (analyzer_.ProveEQ(cur_scale, expected_scale).value_or(false)) {
       if (IsOne(split->extent)) return j;
       // We prefer the unit extent Iter. just search when extent != 1.
       if (matched_pos == -1) {
@@ -599,7 +597,9 @@ std::optional<ir::IndexExpr> IterMapRewriter::TryFuseSameSource(
       auto lhs_iter = iter_sum->args[matched_index].As<ir::IterSplit>();
       ir::IndexExpr lhs_lower_factor =
           MulAndNormalize(rhs_iter->lower_factor, rhs_iter->extent);
-      if (!ProveEQ(lhs_iter->lower_factor, lhs_lower_factor, analyzer_)) break;
+      if (!analyzer_.ProveEQ(lhs_iter->lower_factor, lhs_lower_factor)
+               .value_or(false))
+        break;
       visited[matched_index] = true;
 
       rhs_iter->extent = MulAndNormalize(lhs_iter->extent, rhs_iter->extent);
 
@@ -17,7 +17,7 @@
 #include <unordered_map>
 #include <vector>
 #include "paddle/cinn/common/integer_set.h"
-#include "paddle/cinn/common/iter_util.h"
+#include "paddle/cinn/common/ir_util.h"
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/ir_base.h"
 #include "paddle/cinn/ir/ir_mutator.h"