PaddlePaddle
diff --git a/‎paddle/cinn/common/axis.cc
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions b/‎paddle/cinn/common/axis.cc
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions
diff --git a/‎paddle/cinn/common/cas.cc
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions b/‎paddle/cinn/common/cas.cc
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions
diff --git a/‎paddle/cinn/common/const_fold.h
Copy file name to clipboard
+24Lines changed: 24 additions & 0 deletions b/‎paddle/cinn/common/const_fold.h
Copy file name to clipboard
+24Lines changed: 24 additions & 0 deletions
diff --git a/‎paddle/cinn/common/dim_expr_converter.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/common/dim_expr_converter.cc
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
+52-20Lines changed: 52 additions & 20 deletions b/‎paddle/cinn/common/ir_util.cc
Copy file name to clipboard
+52-20Lines changed: 52 additions & 20 deletions
diff --git a/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+9-2Lines changed: 9 additions & 2 deletions b/‎paddle/cinn/common/ir_util.h
Copy file name to clipboard
+9-2Lines changed: 9 additions & 2 deletions
diff --git a/‎paddle/cinn/common/iter_simplify.cc
Copy file name to clipboard
+11-12Lines changed: 11 additions & 12 deletions b/‎paddle/cinn/common/iter_simplify.cc
Copy file name to clipboard
+11-12Lines changed: 11 additions & 12 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/trivial_op_util.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/cinn/hlir/framework/pir/trivial_op_util.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
@@ -68,7 +68,8 @@ std::string axis_name(int level) {
 std::vector<ir::Var> GenDefaultAxis(int naxis) {
   std::vector<ir::Var> axis;
   for (int i = 0; i < naxis; i++) {
-    axis.emplace_back(cinn::common::axis_name(i));
+    Var ax(cinn::common::axis_name(i));
+    axis.emplace_back(ax.set_index(true));
     PADDLE_ENFORCE_EQ(axis.back()->type().valid(),
                       true,
                       ::common::errors::InvalidArgument(
@@ -83,7 +84,7 @@ std::vector<ir::Expr> GenDefaultAxisAsExpr(int naxis) {
   auto vars = GenDefaultAxis(naxis);
   std::vector<Expr> res;
   for (auto& v : vars) {
-    res.push_back(Expr(v));
+    res.push_back(Expr(v.set_index(true)));
   }
   return res;
 }
 
@@ -36,6 +36,7 @@ using namespace ir;  // NOLINT
 Expr AutoSimplify(
     const Expr& u,
     const absl::flat_hash_map<std::string, CasInterval>& var_intervals) {
+  bool is_index = u.is_index();
   VLOG(7) << "Begin AutoSimplify: " << u;
   Expr copied = ir::ir_utils::IRCopy(u);
   if (copied.type().is_float()) {
@@ -56,7 +57,7 @@ Expr AutoSimplify(
   copied = CasSimplify(copied, s_var_intervals);
   copied = detail::ConvertCasToCinn(copied);
   VLOG(7) << "End AutoSimplify " << copied;
-  return copied;
+  return is_index ? copied.set_index(true) : copied;
 }
 
 int gcd(int a, int b) {
@@ -625,7 +626,6 @@ Expr CasSimplifyMutator::SimplifySum(Expr u) {
   if (!temp.As<Sum>()) return temp;
 
   operands = temp.As<Sum>()->operands();
-
   auto args = SimplifySumRec(operands);
   if (args.empty()) return make_const(u.type(), 0);
   if (args.size() == 1) return args[0];
@@ -1654,6 +1654,9 @@ Expr ConvertCinnToCAS(Expr expr) {
     void operator()(Expr* expr) { Visit(expr); }
     void Visit(Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
 
+    // Because indice of `Load` is simplify by IndexExpr, we just skip it.
+    void Visit(const Load* op, Expr* expr) override { return; }
+
    private:
     void Visit(const Add* op, Expr* expr) override {
       auto a = op->a();
 
@@ -107,5 +107,29 @@ inline std::optional<ir::Expr> TryConstFold<ir::Mod>(ir::Expr a, ir::Expr b) {
   return std::nullopt;
 }
 
+template <>
+inline std::optional<ir::Expr> TryConstFold<ir::Min>(ir::Expr a, ir::Expr b) {
+  const ir::IntImm* pa = a.As<ir::IntImm>();
+  const ir::IntImm* pb = b.As<ir::IntImm>();
+  const auto& rtype = a.type();
+  if (pa && pb) {
+    int64_t res = std::min(pa->value, pb->value);
+    return cinn::common::make_shared<ir::IntImm>(rtype, res);
+  }
+  return std::nullopt;
+}
+
+template <>
+inline std::optional<ir::Expr> TryConstFold<ir::Max>(ir::Expr a, ir::Expr b) {
+  const ir::IntImm* pa = a.As<ir::IntImm>();
+  const ir::IntImm* pb = b.As<ir::IntImm>();
+  const auto& rtype = a.type();
+  if (pa && pb) {
+    int64_t res = std::max(pa->value, pb->value);
+    return cinn::common::make_shared<ir::IntImm>(rtype, res);
+  }
+  return std::nullopt;
+}
+
 }  // namespace common
 }  // namespace cinn
@@ -37,7 +37,7 @@ struct DimExprToIrExprVisitor {
                             dim_expr,
                             /* is_reduce  = */ false,
                             /* is_symbolic_constant = */ true);
-    return x;
+    return x.set_index(true);
   }
 
   ir::Expr operator()(const Negative<DimExpr>& dim_expr) {
 
@@ -202,7 +202,6 @@ static void MergeMulModInsertElements(
 }
 
 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) {
@@ -278,8 +277,7 @@ static std::optional<ir::IndexExpr> MergeMulModInner(
   return std::nullopt;
 }
 
-ir::IndexExpr MergeMulMod(SymbolicExprAnalyzer *analyzer,
-                          const ir::IndexExpr &base) {
+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;
@@ -298,7 +296,7 @@ ir::IndexExpr MergeMulMod(SymbolicExprAnalyzer *analyzer,
     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);
+          *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;
@@ -357,10 +355,6 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
                         "The size of shape should be less than or "
                         "equal to the size of indices."));
   Expr res(0);
-  ir::TryElevateInt32ToInt64(shape);
-  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(
@@ -372,22 +366,21 @@ Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
             i,
             shape[i].type()));
 
+    // if(VerifyIndex(shape[i]))shape[i].set_index(true);
+    // if(VerifyIndex(indices[i]))indices[i].set_index(true);
+
     Expr indice_cast = indices[i];
     optim::SimplifyCast(&indice_cast);
-    if (res.defined()) {
-      res = RampRelatedAdd(RampRelatedMul(res, shape[i]), indice_cast);
-      if (res.is_index_tmp()) {
-        res.set_index(true);
-        res = res.as_index().Normalize();
-      }
-    } else {
-      res = 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_tmp()) {
-        res.set_index(true);
-        res = MergeMulMod(&analyzer, res).Normalize();
+      if (res.is_index()) {
+        res = MergeMulMod(res).Normalize();
+      } else {
+        VLOG(8) << "**** expr is not index ****: " << res;
       }
     }
   }
@@ -660,10 +653,11 @@ bool ComparePriority(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs) {
     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 lhsLen = lhs.length();
   auto rhsLen = rhs.length();
   if (lhsLen < rhsLen) return false;
-  // Add < Mul < Div < Mod.
+  // Add < Mul < Div < Mod < Min < Max < Load.
   else if (lhsLen == rhsLen)
     return lhs.node_type() <= rhs.node_type();
   else
@@ -694,6 +688,10 @@ bool IsSumPartialBySymbol(const ir::IndexExpr &expr,
       return IsSumPartialBySymbol(expr.operand(0), symbol);
     }
     case ir::IrNodeTy::Mod:
+    case ir::IrNodeTy::Min:
+    case ir::IrNodeTy::Max:
+    case ir::IrNodeTy::Load:
+    case ir::IrNodeTy::Cast:
       return false;
     default:
       PADDLE_THROW(::common::errors::InvalidArgument(
@@ -771,6 +769,11 @@ bool IsDivisiblieBySymbol(const ir::IndexExpr &expr,
       if (ty != expr.node_type()) return false;
       return IsDivisiblieBySymbol(expr.operand(0), symbol, expr.node_type());
     }
+    case ir::IrNodeTy::Min:
+    case ir::IrNodeTy::Max:
+    case ir::IrNodeTy::Load:
+    case ir::IrNodeTy::Cast:
+      return false;
     default:
       PADDLE_THROW(::common::errors::InvalidArgument(
           "Unsupported type of expr in IsDivisiblieBySymbol which is: %s",
@@ -831,5 +834,34 @@ bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
   }
   return false;
 }
+
+bool VerifyIndex(const ir::Expr &expr) {
+  switch (expr.node_type()) {
+    case ir::IrNodeTy::_Var_:
+    case ir::IrNodeTy::IntImm: {
+      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:
+    case ir::IrNodeTy::Div:
+    case ir::IrNodeTy::Mod:
+    case ir::IrNodeTy::Max:
+    case ir::IrNodeTy::Min:
+      return VerifyIndex(expr->operand(0)) && VerifyIndex(expr->operand(1));
+  }
+  return false;
+}
 }  // namespace common
 }  // namespace cinn
@@ -38,8 +38,7 @@ Expr PrecedingAxisToAbsOffset(const std::vector<Expr> &shape,
 
 Expr CastIfNeeded(Expr body, Type type);
 
-ir::IndexExpr MergeMulMod(SymbolicExprAnalyzer *analyzer,
-                          const ir::IndexExpr &base);
+ir::IndexExpr MergeMulMod(const ir::IndexExpr &base);
 
 //! Substitute vars to other expressions.
 //! @param expr The expression to do modification.
@@ -337,5 +336,13 @@ bool ProveDivisible(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
 bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
                         ir::IndexExpr &expr);  // NOLINT
 
+/*!
+ * \brief Judge type of `expr` is valid type of `IndexExpr` or not.
+ * \param expr The expression to be checked.
+ * \return A boolean value indicating whether the type of `expr` is valid
+ * IndexExpr type.
+ */
+bool VerifyIndex(const ir::Expr &expr);
+
 }  // namespace common
 }  // namespace cinn
@@ -61,7 +61,6 @@ 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());
@@ -89,7 +88,7 @@ void IterMapRewriter::Visit(const ir::_Var_* op, Expr* expr) {
 void IterMapRewriter::Visit(const ir::Add* op, Expr* expr) {
   auto a = op->a();
   auto b = op->b();
-
+  VLOG(10) << "in visit add: " << a << " " << b;
   Visit(&a);
   Visit(&b);
 
@@ -103,7 +102,6 @@ void IterMapRewriter::Visit(const ir::Add* op, Expr* expr) {
 
   Expr ret = ir::ir_utils::IRCopy(ToIterSum(a));
   ir::IterSum* ret_sum = ret.As<ir::IterSum>();
-
   if (auto b_sum = b.As<ir::IterSum>()) {
     AddToLhs(ret_sum, *b_sum, 1);
   } else if (auto b_split = b.As<ir::IterSplit>()) {
@@ -112,12 +110,13 @@ void IterMapRewriter::Visit(const ir::Add* op, Expr* expr) {
     ret_sum->base = ret_sum->base + b.as_index();
   }
   *expr = ret;
+  VLOG(10) << "out visit add";
 }
 
 void IterMapRewriter::Visit(const ir::Sub* op, Expr* expr) {
   auto a = op->a();
   auto b = op->b();
-
+  VLOG(10) << "in visit sub: " << a << " " << b;
   Visit(&a);
   Visit(&b);
 
@@ -139,12 +138,13 @@ void IterMapRewriter::Visit(const ir::Sub* op, Expr* expr) {
   }
 
   *expr = ret;
+  VLOG(10) << "out visit sub";
 }
 
 void IterMapRewriter::Visit(const ir::Mul* op, Expr* expr) {
   auto a = op->a();
   auto b = op->b();
-
+  VLOG(10) << "in visit mul: " << a << " " << b;
   Visit(&a);
   Visit(&b);
 
@@ -176,12 +176,14 @@ void IterMapRewriter::Visit(const ir::Mul* op, Expr* expr) {
   }
 
   *expr = ret;
+  VLOG(10) << "out visit mul";
 }
 
 void IterMapRewriter::Visit(const ir::Div* op, Expr* expr) {
   auto a = op->a();
   auto b = op->b();
 
+  VLOG(10) << "in visit div: " << a << " " << b;
   Visit(&a);
   Visit(&b);
 
@@ -197,21 +199,19 @@ void IterMapRewriter::Visit(const ir::Div* op, Expr* expr) {
         "Division of iter and iter is not supported"));
     return;
   }
-
   auto ret = ir::ir_utils::IRCopy(a);
-
   auto preprocessed = PreprocessDividend(ret);
   auto preprocessed_sum = preprocessed.As<ir::IterSum>();
 
   ret = SplitDivConst(preprocessed_sum->args[0], preprocessed_sum->base, b);
-
   *expr = ret;
+  VLOG(10) << "out visit div";
 }
 
 void IterMapRewriter::Visit(const ir::Mod* op, Expr* expr) {
   auto a = op->a();
   auto b = op->b();
-
+  VLOG(10) << "in visit mod: " << a << " " << b;
   Visit(&a);
   Visit(&b);
 
@@ -236,6 +236,7 @@ void IterMapRewriter::Visit(const ir::Mod* op, Expr* expr) {
   ret = SplitModConst(preprocessed_sum->args[0], preprocessed_sum->base, b);
 
   *expr = ret;
+  VLOG(10) << "out visit mod";
 }
 
 Expr IterMapRewriter::PreprocessDividend(const Expr& dividend) {
@@ -471,7 +472,6 @@ std::optional<Expr> IterMapRewriter::TryFuse(const Expr& expr) {
       return opt.value();
     }
   }
-
   // Select iter with smallest scale as base iter.
   std::vector<bool> visited(iter_sum->args.size(), false);
   int base_index = FindBaseSplit(*iter_sum, visited, Expr(), -1);
@@ -484,7 +484,6 @@ std::optional<Expr> IterMapRewriter::TryFuse(const Expr& expr) {
   ir::IndexExpr expected_scale = base_scale;
   int first_possible_unit_extent_pos =
       FindFirstPossibleUnitExtentIndex(*iter_sum);
-
   // Find iter with same scale as expected_scale and update expected_scale.
   // e.g. i * 32 + j * 8 + k * 1, Extent(i, j, k) = 2, 4, 8.
   // first base_index = 2, expected_scale = 1. means select k as base iter.
@@ -493,7 +492,7 @@ std::optional<Expr> IterMapRewriter::TryFuse(const Expr& expr) {
   // finally matched_pos = 0, expected_scale = 32 * 2 = 64. means match i.
   // if match failed, indicates that expr is illegal and cannot be merged.
   for (size_t i = 0; i < iter_sum->args.size(); ++i) {
-    ir::IndexExpr matched_scale{nullptr};
+    ir::IndexExpr matched_scale;
     int matched_pos =
         i == 0 ? base_index
                : FindSplitWithExactScale(*iter_sum,
 
@@ -182,6 +182,7 @@ void SubstitudeTargetExprWithDestExpr(const ir::Expr& source,
   VLOG(4) << "SubstitideExpr Start";
   VLOG(5) << "Substitide Body : " << *body;
   ir::Expr new_dest = dest;
+  optim::Simplify(&new_dest);
   if (source.type() != dest.type()) {
     VLOG(4) << "Cast the dest" << dest << " to type" << source.type();
     new_dest = ir::Cast::Make(source.type(), dest);
Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@ struct DimExprToIrExprVisitor {`
`37`	`37`	`dim_expr,`
`38`	`38`	`/* is_reduce = */ false,`
`39`	`39`	`/* is_symbolic_constant = */ true);`
`40`		`- return x;`
	`40`	`+ return x.set_index(true);`
`41`	`41`	`}`
`42`	`42`
`43`	`43`	`ir::Expr operator()(const Negative<DimExpr>& dim_expr) {`