CINN Add more Simplify scenario: Select2MinMax, BoundSimplify, PowerLog2BitOp. (PaddlePaddle#74292)

Author: YuhanXu

paddle/cinn/common/integer_set.cc

@@ -164,7 +164,8 @@ cas_intervals_t CollectVarIntervalsOfExprs(const std::vector<ir::Expr>& exprs,
           lower_bound = ir::Expr(1);
         }
         var_intervals.insert(
-            {var->name, CasInterval(lower_bound, upper_bound)});
+            {var->name,
+             CasInterval(lower_bound, NormalizeUpperBound(upper_bound))});
       }
       return false;
     });
@@ -572,14 +573,21 @@ class BoundReplacer : public ir::IRMutator<> {
 ir::Expr SymbolicExprAnalyzer::LowerBound(const ir::Expr& expr) const {
   BoundReplacer bound_replacer(var_intervals_, true);
   ir::Expr bound = ir::ir_utils::IRCopy(expr);
+  if (bound.is_index()) {
+    bound = bound.as_index().Normalize(ir::IndexExpr::OptLevel::kLevel3);
+  }
   bound_replacer(&bound);
   return optim::ArithSimplify(bound);
 }
 
 ir::Expr SymbolicExprAnalyzer::UpperBound(const ir::Expr& expr) const {
   BoundReplacer bound_replacer(var_intervals_, false);
   ir::Expr bound = ir::ir_utils::IRCopy(expr);
+  if (bound.is_index()) {
+    bound = bound.as_index().Normalize(ir::IndexExpr::OptLevel::kLevel3);
+  }
   bound_replacer(&bound);
+
   return optim::ArithSimplify(bound);
 }
 
@@ -709,7 +717,8 @@ SingleIntervalIntSet::SingleIntervalIntSet(const ir::Expr& min,
                                    ? x->as_var()->upper_bound
                                    : SymbolicExprLimit::positive_inf;
         var_intervals_.insert(
-            {x->as_var()->name, CasInterval(lower_bound, upper_bound)});
+            {x->as_var()->name,
+             CasInterval(lower_bound, NormalizeUpperBound(upper_bound))});
       }
       return false;
     };

paddle/cinn/common/ir_util.cc

@@ -270,6 +270,16 @@ bool is_zero(Expr v) {
   return false;
 }
 
+Expr NormalizeUpperBound(Expr upper_bound, bool minus_one /* = true */) {
+  if (upper_bound == SymbolicExprLimit::positive_inf) {
+    return upper_bound;
+  }
+  if (minus_one) {
+    return upper_bound - ir::Expr(1);  // [lower, upper) to [lower, upper]
+  }
+  return upper_bound + ir::Expr(1);  // (lower, upper] to [lower, upper)
+}
+
 Expr CastIfNeeded(Expr body, Type type) {
   if (body.type() == type) return body;
   return ir::Cast::Make(type, body);

paddle/cinn/common/ir_util.h

@@ -91,6 +91,8 @@ std::vector<std::string> GatherItersToTensorProducer(
 
 bool is_zero(Expr v);
 
+Expr NormalizeUpperBound(Expr upper_bound, bool minus_one = true);
+
 bool MathEqual(const Expr &a, const Expr &b);
 
 //! helper function to get a ir::Select node.

paddle/cinn/ir/group_schedule/tactic/arrange_storage_tactic.cc

@@ -141,7 +141,8 @@ IntSet Evaluate(Expr expr,
                 const std::unordered_map<ir::Var, IntSet>& var_domain) {
   Expr copy_for_upper_bound = ir::ir_utils::IRCopy(expr);
   Expr copy_for_lower_bound = ir::ir_utils::IRCopy(expr);
-  common::cas_intervals_t var_intervals;
+  common::cas_intervals_t
+      var_intervals;  // variable name -> CasIntervals[lower_bound, upper_bound]
   std::vector<ir::Expr> var_vec = ir::ir_utils::CollectIRNodesWithoutTensor(
       expr, [](const ir::Expr* x) { return x->as_var(); });
   for (Expr var_expr : var_vec) {
@@ -150,7 +151,9 @@ IntSet Evaluate(Expr expr,
       const ir::Var& fixed_var = fixed.at(var);
       var_intervals.emplace(
           fixed_var->name,
-          common::CasInterval(fixed_var->lower_bound, fixed_var->upper_bound));
+          common::CasInterval(
+              fixed_var->lower_bound,
+              cinn::common::NormalizeUpperBound(fixed_var->upper_bound)));
       optim::ReplaceVarWithExpr(&copy_for_lower_bound, var, Expr(fixed_var));
       optim::ReplaceVarWithExpr(&copy_for_upper_bound, var, Expr(fixed_var));
     } else if (var_domain.count(var) != 0) {
@@ -172,7 +175,8 @@ IntSet Evaluate(Expr expr,
           ::common::errors::InvalidArgument(
               "The 'upper_bound' of the variable must be defined."));
       optim::ReplaceVarWithExpr(&copy_for_lower_bound, var, var->lower_bound);
-      optim::ReplaceVarWithExpr(&copy_for_upper_bound, var, var->upper_bound);
+      optim::ReplaceVarWithExpr(
+          &copy_for_upper_bound, var, NormalizeUpperBound(var->upper_bound));
     }
   }
   ir::Expr lower_bound = optim::ArithSimplify(copy_for_lower_bound);

paddle/cinn/ir/ir.h

@@ -421,6 +421,7 @@ struct _Var_ : public ExprNode<_Var_> {
 };
 
 //! A named variable.
+// i ∈ [lower_bound, upper_bound)
 struct Var : public IrNodeRef {
   Var() = default;
   explicit Var(IrNode* n) : IrNodeRef(n) {}
@@ -846,6 +847,7 @@ struct For : public ExprNode<For>, public ForBase {
   //! The minimum value of the iteration.
   Expr min;
   //! The extent of the iteration.
+  // loop_var ∈ [min, min + extent)
   Expr extent;
 
   Expr body;

paddle/cinn/ir/ir_analyzer/ir_analyzer.cc

@@ -621,7 +621,8 @@ std::vector<ir::Var> IndicesToVars(const std::vector<ir::Expr>& indices) {
     if (e.is_constant()) {
       std::string var_name =
           cinn::UniqName("constant" + static_cast<int>(e.get_constant()));
-      result.emplace_back(e, e, var_name, /* is_reduce = */ false);
+      result.emplace_back(
+          e, NormalizeUpperBound(e, false), var_name, /* is_reduce = */ false);
     } else if (e.As<ir::_Var_>() != nullptr) {
       ir::Expr copy_e = ir::ir_utils::IRCopy(e);
       ir::_Var_* var_ref = copy_e.As<ir::_Var_>();
@@ -635,14 +636,17 @@ std::vector<ir::Var> IndicesToVars(const std::vector<ir::Expr>& indices) {
           ir::Var var = x->as_var_ref();
           var_intervals.insert(
               {var->name,
-               common::CasInterval{var->lower_bound, var->upper_bound}});
+               common::CasInterval{var->lower_bound,
+                                   NormalizeUpperBound(var->upper_bound)}});
           if (var->is_reduce_axis) is_reduce = true;
         }
         return false;
       });
       common::SymbolicExprAnalyzer analyzer(var_intervals);
-      result.emplace_back(
-          analyzer.LowerBound(e), analyzer.UpperBound(e), var_name, is_reduce);
+      result.emplace_back(analyzer.LowerBound(e),
+                          NormalizeUpperBound(analyzer.UpperBound(e), false),
+                          var_name,
+                          is_reduce);
     }
   }
   return result;