【CINN】Add load,min,max,cast for IndexExpr (#70450)

liuruyan · web-flow · commit d3665db2effe · 2024-12-30T10:41:58.000+08:00
* add load/min/max/cast

* fix bug

* fix bug

* fix ci

* fix ci
diff --git a/paddle/cinn/common/const_fold.h b/paddle/cinn/common/const_fold.h
@@ -107,5 +107,28 @@ 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
diff --git a/paddle/cinn/common/ir_util.cc b/paddle/cinn/common/ir_util.cc
@@ -478,7 +478,7 @@ bool ComparePriority(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs) {
   auto lhsLen = lhs.length();
   auto rhsLen = rhs.length();
   if (lhsLen < rhsLen) return false;
-  // Add < Mul < Div < Mod.
+  // Add < Mul < Div < Mod < Min < Max < Cast < Load.
   else if (lhsLen == rhsLen)
     return lhs.node_type() <= rhs.node_type();
   else
@@ -509,6 +509,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(
@@ -586,6 +590,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",
@@ -647,26 +656,38 @@ bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
   return false;
 }
 
-bool VerifyIndex(const ir::Expr &expr) {
+IndexType VerifyIndex(const ir::Expr &expr) {
   switch (expr.node_type()) {
     case ir::IrNodeTy::_Var_:
-    case ir::IrNodeTy::IntImm:
+    case ir::IrNodeTy::IntImm: {
+      return expr.type().is_index_type() ? IndexType::kValid
+                                         : IndexType::kInvalid;
+    }
     case ir::IrNodeTy::Load: {
-      if (expr.type().is_index_type()) return true;
-      return false;
+      return expr.type().is_index_type() ? IndexType::kLoad
+                                         : IndexType::kInvalid;
+    }
+    case ir::IrNodeTy::Cast: {
+      IndexType result = VerifyIndex(expr->operand(0));
+      return result == IndexType::kValid && expr.type().is_index_type()
+                 ? IndexType::kCast
+                 : IndexType::kInvalid;
     }
-    case ir::IrNodeTy::Cast:
-      return VerifyIndex(expr->operand(0));
     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));
+    case ir::IrNodeTy::Min: {
+      IndexType left = VerifyIndex(expr->operand(0));
+      IndexType right = VerifyIndex(expr->operand(1));
+      if (left == IndexType::kInvalid || right == IndexType::kInvalid)
+        return IndexType::kInvalid;
+      return std::max(left, right);
+    }
   }
-  return false;
+  return IndexType::kInvalid;
 }
 }  // namespace common
 }  // namespace cinn
diff --git a/paddle/cinn/common/ir_util.h b/paddle/cinn/common/ir_util.h
@@ -339,5 +339,22 @@ bool ProveDivisible(const ir::IndexExpr &lhs, const ir::IndexExpr &rhs);
 bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
                         ir::IndexExpr &expr);  // NOLINT
 
+enum IndexType {
+  kInvalid = 0,  // invalid expr
+  kValid = 1,    // valid expr
+  kLoad = 2,     // exist Load
+  kCast = 3      // exist cast
+};
+
+/*!
+ * \brief Judge type of `expr` is valid type of `IndexExpr` or not.
+ * \param expr The expression to be checked.
+ * \return A enum IndexType value indicating whether the type of `expr` is valid
+ * IndexExpr.
+ *
+ * Note: Although load and cast are also legal IndexExpr, we need to know this
+ * information in some scenarios.
+ */
+IndexType VerifyIndex(const ir::Expr &expr);
 }  // namespace common
 }  // namespace cinn
diff --git a/paddle/cinn/common/iter_simplify.cc b/paddle/cinn/common/iter_simplify.cc
@@ -493,7 +493,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,
diff --git a/paddle/cinn/ir/ir.cc b/paddle/cinn/ir/ir.cc
@@ -55,7 +55,35 @@ Expr Cast::Make(Type t, Expr v) {
                         "The expression is not defined. "
                         "A defined expression is required for casting."));
 
-  if (v.node_type() != ir::IrNodeTy::_Var_ && v.is_index() && t == Int(64)) {
+#define __CAST_TO_TYPE(type__)                  \
+  if (auto *i = v.As<ir::IntImm>()) {           \
+    return Expr(static_cast<type__>(i->value)); \
+  } else if (auto *u = v.As<ir::UIntImm>()) {   \
+    return Expr(static_cast<type__>(u->value)); \
+  }
+
+  if (v.is_constant()) {
+    if (t == type_of<int8_t>()) {
+      __CAST_TO_TYPE(int8_t)
+    } else if (t == type_of<int16_t>()) {
+      __CAST_TO_TYPE(int16_t)
+    } else if (t == type_of<int32_t>()) {
+      __CAST_TO_TYPE(int32_t)
+    } else if (t == type_of<int64_t>()) {
+      __CAST_TO_TYPE(int64_t)
+    } else if (t == type_of<uint8_t>()) {
+      __CAST_TO_TYPE(uint8_t)
+    } else if (t == type_of<uint16_t>()) {
+      __CAST_TO_TYPE(uint16_t)
+    } else if (t == type_of<uint32_t>()) {
+      __CAST_TO_TYPE(uint32_t)
+    } else if (t == type_of<uint64_t>()) {
+      __CAST_TO_TYPE(uint64_t)
+    }
+  }
+#undef __CAST_TO_TYPE
+
+  if (v.node_type() != ir::IrNodeTy::Load && v.is_index() && t == Int(64)) {
     v->convert_int32_to_int64();
     return v;
   }
diff --git a/paddle/cinn/ir/ir_base.cc b/paddle/cinn/ir/ir_base.cc
@@ -285,30 +285,26 @@ double Expr::get_constant() const {
 bool Expr::is_var() const { return As<_Var_>(); }
 
 bool Expr::is_index() const {
-  // Temporarily use `is_index_tmp`. because `get_index` depends on marking
+  // Temporarily use `VerifyIndex`. because `get_index` depends on marking
   // `indexExpr` in For::make and sch
-  return is_index_tmp();
+  return VerifyIndex(*this);
   // return get()->get_index();
 }
 
-bool Expr::is_index_tmp() const {
-  switch (node_type()) {
-    case ir::IrNodeTy::_Var_:
-    case ir::IrNodeTy::IntImm: {
-      if (type().is_index_type()) return true;
-      return false;
-    }
-    case ir::IrNodeTy::Add:
-    case ir::IrNodeTy::Sub:
-    case ir::IrNodeTy::Mul:
-    case ir::IrNodeTy::Div:
-    case ir::IrNodeTy::Mod:
-      return p_->operand(0).is_index_tmp() && p_->operand(1).is_index_tmp();
+Expr &Expr::set_index(bool flag) {
+  if (flag && !VerifyIndex(*this)) {
+    PADDLE_THROW(::common::errors::InvalidType(
+        "Expr: %s is not IndexExpr! cannot be set as IndexExpr.", *this));
   }
-  return false;
+  get()->set_index(flag);
+  return *this;
 }
 
-Expr &Expr::set_index(bool flag) {
+const Expr &Expr::set_index(bool flag) const {
+  if (flag && !VerifyIndex(*this)) {
+    PADDLE_THROW(::common::errors::InvalidType(
+        "Expr: %s is not IndexExpr! cannot be set as IndexExpr.", *this));
+  }
   get()->set_index(flag);
   return *this;
 }
@@ -410,6 +406,10 @@ const IndexExpr IndexExpr::operand(int32_t i) const {
 int64_t IndexExpr::GetLargestMultiplyPart() const {
   switch (node_type()) {
     case cinn::ir::IrNodeTy::_Var_:
+    case ir::IrNodeTy::Min:
+    case ir::IrNodeTy::Max:
+    case ir::IrNodeTy::Load:
+    case ir::IrNodeTy::Cast:
       return 1;
     case cinn::ir::IrNodeTy::Div: {
       if (operand(1).type().is_index_type()) {
@@ -441,15 +441,21 @@ int32_t IndexExpr::length() const {
   switch (node_type()) {
     case ir::IrNodeTy::_Var_:
     case ir::IrNodeTy::IntImm:
+    case ir::IrNodeTy::Load:
       return 1;
     case ir::IrNodeTy::Add:
     case ir::IrNodeTy::Mul:
     case ir::IrNodeTy::Div:
-    case ir::IrNodeTy::Mod: {
+    case ir::IrNodeTy::Mod:
+    case ir::IrNodeTy::Min:
+    case ir::IrNodeTy::Max: {
       int lhs_count = operand(0).length();
       int rhs_count = operand(1).length();
       return lhs_count + rhs_count + 1;
     }
+    case ir::IrNodeTy::Cast: {
+      return operand(0).length() + 1;
+    }
     default:
       PADDLE_THROW(::common::errors::InvalidArgument(
           "Unsupported type in length, which is: %s", node_type()));
@@ -460,13 +466,18 @@ bool IndexExpr::IsDynamic() const {
   switch (node_type()) {
     case ir::IrNodeTy::_Var_:
       return as_var()->name.at(0) == 'S';
-    case ir::IrNodeTy::IntImm: {
+    case ir::IrNodeTy::Load:
+      return true;
+    case ir::IrNodeTy::IntImm:
       return false;
-    }
+    case ir::IrNodeTy::Cast:
+      return operand(0).IsDynamic();
     case ir::IrNodeTy::Add:
     case ir::IrNodeTy::Mul:
     case ir::IrNodeTy::Div:
-    case ir::IrNodeTy::Mod: {
+    case ir::IrNodeTy::Mod:
+    case ir::IrNodeTy::Min:
+    case ir::IrNodeTy::Max: {
       auto lFlag = operand(0).IsDynamic();
       auto rFlag = operand(1).IsDynamic();
       return lFlag || rFlag;
@@ -590,6 +601,7 @@ void IrNode::replace(Expr old_op, Expr new_op) {
 
 bool IrNode::get_index() const { return is_index_; }
 void IrNode::set_index(bool flag) {
+  if (is_index_ == flag) return;
   is_index_ = flag;
   if (flag) {
     for (Expr &operand : operands) {
@@ -623,8 +635,9 @@ void IrNode::convert_int64_to_int32() {
                         true,
                         ::common::errors::InvalidArgument(
                             "Current only support convert int64_t "
-                            "to int32_t, but get type is: %s",
-                            type_));
+                            "to int32_t, but get type is: %s, node type is: %s",
+                            type_,
+                            node_type()));
 
   if (node_type() == IrNodeTy::IntImm) {
     auto *int_imm = static_cast<IntImm *>(this);
diff --git a/paddle/cinn/ir/ir_base.h b/paddle/cinn/ir/ir_base.h
@@ -474,13 +474,11 @@ struct Expr : public IrNodeRef {
 
   bool is_index() const;
 
-  //! `is_index_tmp` is used to judge whether the expr is a index temporary.
-  bool is_index_tmp() const;
-
   IndexExpr as_index();
   const IndexExpr as_index() const;
 
   Expr& set_index(bool flag);
+  const Expr& set_index(bool flag) const;
 
   operator Var();
 
@@ -491,8 +489,8 @@ struct IndexExpr : public IrNodeRef {
  public:
   IndexExpr() = default;
   IndexExpr(const IndexExpr& other) : IrNodeRef(other.ptr()) {}
-  IndexExpr(IrNode* p) : IrNodeRef(p) {}  // NOLINT
-  IndexExpr(const Expr& e);               // NOLINT
+  IndexExpr(IrNode* p) : IrNodeRef(p) { p->set_index(true); }  // NOLINT
+  IndexExpr(const Expr& e);                                    // NOLINT
 
   explicit IndexExpr(int32_t x) : IrNodeRef(new IntImm(Int(32), x)) {}
   explicit IndexExpr(int64_t x) : IrNodeRef(new IntImm(Int(64), x)) {}
diff --git a/paddle/cinn/ir/ir_visitor.cc b/paddle/cinn/ir/ir_visitor.cc
@@ -97,15 +97,35 @@ bool operator==(IndexExpr a, IndexExpr b) {
     case ir::IrNodeTy::_Var_: {
       return a.as_var()->name == b.as_var()->name;
     }
+    case ir::IrNodeTy::Cast: {
+      auto lhs = a.As<ir::Cast>();
+      auto rhs = b.As<ir::Cast>();
+      return lhs->type() == rhs->type() && lhs->v() == rhs->v();
+    }
+
+    case ir::IrNodeTy::Load: {
+      auto lhs = a.As<ir::Load>();
+      auto rhs = b.As<ir::Load>();
+      bool equal = lhs->tensor == rhs->tensor;
+      for (int32_t i = 0; i < lhs->indices.size(); ++i) {
+        equal &= (lhs->indices[i] == rhs->indices[i]);
+      }
+      return equal;
+    }
     case ir::IrNodeTy::Div:
     case ir::IrNodeTy::Mod: {
-      return a->operand(0) == b->operand(0) && a->operand(1) == b->operand(1);
+      return a.operand(0) == b.operand(0) && a.operand(1) == b.operand(1);
     }
     case ir::IrNodeTy::Add:
       return CompareExpressions<ir::Add>(a, b);
     case ir::IrNodeTy::Mul:
       return CompareExpressions<ir::Mul>(a, b);
+    case ir::IrNodeTy::Min:
+      return CompareExpressions<ir::Min>(a, b);
+    case ir::IrNodeTy::Max:
+      return CompareExpressions<ir::Max>(a, b);
   }
+  return false;
 }
 
 bool operator!=(IndexExpr a, IndexExpr b) { return !(a == b); }
