【CINN】Fix bug of nvrtc compile (#72457)

liuruyan · web-flow · commit 1d979da8457d · 2025-04-27T20:46:04.000+08:00
* fix bug

* fix_dpa3_bug
diff --git a/paddle/cinn/common/ir_util.cc b/paddle/cinn/common/ir_util.cc
@@ -512,16 +512,6 @@ void OpDataTypePromote(Expr *expr) {
       node->body = promote_args.at(1);
       IRMutator::Visit(op, expr);
     }
-
-    void Visit(const ir::For *op, ir::Expr *expr) {
-      auto node = expr->As<ir::For>();
-      auto promote_args = std::move(ir::TryElevateInt32ToInt64(
-          {node->loop_var, node->min, node->extent}));
-      node->loop_var = promote_args.at(0);
-      node->min = promote_args.at(1);
-      node->extent = promote_args.at(2);
-      IRMutator::Visit(op, expr);
-    }
   };
 
   TypePromote visitor;
diff --git a/paddle/cinn/ir/ir.cc b/paddle/cinn/ir/ir.cc
@@ -352,6 +352,11 @@ Expr Let::Make(Expr symbol, Expr body) {
   body = promote_args.at(1);
 
   n->symbol = symbol;
+
+  if (n->symbol.is_var()) {
+    n->symbol.as_var()->is_let_symbol = true;
+  }
+
   n->body = body;
   n->set_type(n->symbol->type());
   return Expr(n);
@@ -395,13 +400,15 @@ Expr _Var_::Make(Expr lower_bound,
                  const std::string &name,
                  bool is_reduce_axis,
                  bool is_symbolic_constant,
-                 bool is_keepdim) {
+                 bool is_keepdim,
+                 bool is_let_symbol) {
   auto *n = make_shared<_Var_>();
   n->lower_bound = lower_bound;
   n->upper_bound = upper_bound;
   n->is_reduce_axis = is_reduce_axis;
   n->is_keepdim = is_keepdim;
   n->is_symbolic_constant = is_symbolic_constant;
+  n->is_let_symbol = is_let_symbol;
   n->name = name;
   n->set_type(lower_bound.type());
   return Expr(n);
@@ -412,6 +419,8 @@ Expr _Var_::Copy() const {
   n->name = name;
   n->is_reduce_axis = is_reduce_axis;
   n->is_keepdim = is_keepdim;
+  n->is_symbolic_constant = is_symbolic_constant;
+  n->is_let_symbol = is_let_symbol;
   n->set_index(get_index());
   n->lower_bound = lower_bound;
   n->upper_bound = upper_bound;
diff --git a/paddle/cinn/ir/ir.h b/paddle/cinn/ir/ir.h
@@ -393,6 +393,7 @@ struct _Var_ : public ExprNode<_Var_> {
   bool is_reduce_axis{false};
   bool is_keepdim{false};
   bool is_symbolic_constant{false};
+  bool is_let_symbol{false};
   //! Lower bound and upper bound of a axis.
   // @{
   Expr lower_bound;
@@ -413,7 +414,8 @@ struct _Var_ : public ExprNode<_Var_> {
                    const std::string& name,
                    bool is_reduce,
                    bool is_symbolic_constant = false,
-                   bool is_keepdim = false);
+                   bool is_keepdim = false,
+                   bool is_let_symbol = false);
 
   void Verify() const override;
 
diff --git a/paddle/cinn/ir/ir_base.cc b/paddle/cinn/ir/ir_base.cc
@@ -641,6 +641,9 @@ void IrNode::convert_int32_to_int64() {
       }
     } else if (operand->node_type() == IrNodeTy::Load) {
       operand = ir::Cast::Make(type(), operand);
+    } else if (operand->node_type() == IrNodeTy::_Var_ &&
+               operand.as_var()->is_let_symbol) {
+      operand = ir::Cast::Make(type(), operand);
     }
   }
 }
@@ -676,6 +679,9 @@ void IrNode::convert_int64_to_int32() {
       }
     } else if (operand->node_type() == IrNodeTy::Load) {
       operand = ir::Cast::Make(type(), operand);
+    } else if (operand->node_type() == IrNodeTy::_Var_ &&
+               operand.as_var()->is_let_symbol) {
+      operand = ir::Cast::Make(type(), operand);
     }
   }
 }
@@ -713,6 +719,9 @@ void TryElevateInt32ToInt64_(std::vector<Expr> &expr_vec) {  // NOLINT
         }
       } else if (expr->node_type() == IrNodeTy::Load) {
         expr = ir::Cast::Make(Int(64), expr);
+      } else if (expr->node_type() == IrNodeTy::_Var_ &&
+                 expr.as_var()->is_let_symbol) {
+        expr = ir::Cast::Make(Int(64), expr);
       }
     }
   }
@@ -748,6 +757,10 @@ void ElevateInt64ToInt32_(Expr &expr) {  // NOLINT
       }
     } else if (expr->node_type() == IrNodeTy::Load) {
       expr = ir::Cast::Make(Int(32), expr);
+    } else if (expr->node_type() == IrNodeTy::_Var_ &&
+               expr.as_var()->is_let_symbol) {
+      // symbol of `let` op should be use cast to convert to int32.
+      expr = ir::Cast::Make(Int(32), expr);
     }
   }
 }
diff --git a/paddle/cinn/ir/utils/ir_copy.cc b/paddle/cinn/ir/utils/ir_copy.cc
@@ -200,6 +200,7 @@ struct IRCopyVisitor : public ir::IRVisitorRequireReImpl<Expr>,
     n->name = op->name;
     n->is_reduce_axis = op->is_reduce_axis;
     n->is_symbolic_constant = op->is_symbolic_constant;
+    n->is_let_symbol = op->is_let_symbol;
     n->set_type(op->type());
 
     if (op->lower_bound.defined()) {
diff --git a/paddle/cinn/optim/longlong2int_pass.cc b/paddle/cinn/optim/longlong2int_pass.cc
@@ -141,9 +141,10 @@ class CastLonglong2IntMutator : public ir::IRMutator<> {
     // min(min(S0, 1ll), 1ll) ==> min(min(S0, 1), 1)
     // min(V[S0, S1], 1ll)    ==> min(V[S0, S1], 1ll)
     // min(S0 + 1ll, 1ll)     ==> max(S0 + 1, 1)
-    // IndexType::kValid means expr only has +-*/%, Const, Symbol, Min, Max.
+    // min(V[0], S0)          ==> min((int32)V[0], S1)
+    // min(var_local, S0)     ==> min((int32)var_local, S0)
     // IsDynamic == true means expr has Symbol.
-    if (optim::VerifyIndex(*expr) == ir::IndexExpr::IndexType::kValid &&
+    if (optim::VerifyIndex(*expr) != ir::IndexExpr::IndexType::kInvalid &&
         expr->as_index().IsDynamic()) {
       ir::ElevateInt64ToInt32_((*expr)->operands);
     } else {
@@ -153,14 +154,28 @@ class CastLonglong2IntMutator : public ir::IRMutator<> {
   }
   void Visit(const ir::Max* op, Expr* expr) override {
     auto node = expr->As<ir::Max>();
-    if (optim::VerifyIndex(*expr) == ir::IndexExpr::IndexType::kValid &&
+    if (optim::VerifyIndex(*expr) != ir::IndexExpr::IndexType::kInvalid &&
         expr->as_index().IsDynamic()) {
       ir::ElevateInt64ToInt32_((*expr)->operands);
     } else {
       ir::IRMutator<>::Visit(&node->a(), &node->a());
       ir::IRMutator<>::Visit(&node->b(), &node->b());
     }
   }
+  void Visit(const ir::Call* op, Expr* expr) override {
+    auto node = expr->As<ir::Call>();
+    if (op->name == "CINN_ENTAIL_LOOP_CONDITION") {
+      // args of CINN_ENTAIL_LOOP_CONDITION is [loop_var, condition, stride],
+      // loop_var type is equal to stride type, so we only need to elevate
+      // condition and stride to int32.
+      ir::ElevateInt64ToInt32_(node->read_args[1]->operands);
+      ir::ElevateInt64ToInt32_(node->read_args[2]);
+    } else {
+      for (auto& expr : node->read_args) {
+        ir::IRMutator<>::Visit(&expr, &expr);
+      }
+    }
+  }
 };
 
 class LongLong2IntStmtPass : public StmtPass {
diff --git a/paddle/cinn/optim/simplify_util.cc b/paddle/cinn/optim/simplify_util.cc
@@ -224,7 +224,14 @@ bool IsNegatedIndexExpr(const ir::IndexExpr &candidate,
 
 ir::IndexExpr::IndexType VerifyIndex(const ir::Expr &expr) {
   switch (expr.node_type()) {
-    case ir::IrNodeTy::_Var_:
+    case ir::IrNodeTy::_Var_: {
+      if (expr.type().is_index_type()) {
+        return expr.as_var()->is_let_symbol ? ir::IndexExpr::IndexType::kLoad
+                                            : ir::IndexExpr::IndexType::kValid;
+      } else {
+        return ir::IndexExpr::IndexType::kInvalid;
+      }
+    }
     case ir::IrNodeTy::IntImm: {
       return expr.type().is_index_type() ? ir::IndexExpr::IndexType::kValid
                                          : ir::IndexExpr::IndexType::kInvalid;
diff --git a/paddle/cinn/pybind/ir/ir_api.cc b/paddle/cinn/pybind/ir/ir_api.cc
@@ -384,6 +384,7 @@ void BindIrIr(py::module *m) {
                                     const std::string &,
                                     bool,
                                     bool,
+                                    bool,
                                     bool>(&ir::_Var_::Make))
       .def("copy", &ir::_Var_::Copy);
 

Original file line number	Diff line number	Diff line change
`@@ -641,6 +641,9 @@ void IrNode::convert_int32_to_int64() {`
`641`	`641`	`}`
`642`	`642`	`} else if (operand->node_type() == IrNodeTy::Load) {`
`643`	`643`	`operand = ir::Cast::Make(type(), operand);`
	`644`	`+ } else if (operand->node_type() == IrNodeTy::_Var_ &&`
	`645`	`+ operand.as_var()->is_let_symbol) {`
	`646`	`+ operand = ir::Cast::Make(type(), operand);`
`644`	`647`	`}`
`645`	`648`	`}`
`646`	`649`	`}`
`@@ -676,6 +679,9 @@ void IrNode::convert_int64_to_int32() {`
`676`	`679`	`}`
`677`	`680`	`} else if (operand->node_type() == IrNodeTy::Load) {`
`678`	`681`	`operand = ir::Cast::Make(type(), operand);`
	`682`	`+ } else if (operand->node_type() == IrNodeTy::_Var_ &&`
	`683`	`+ operand.as_var()->is_let_symbol) {`
	`684`	`+ operand = ir::Cast::Make(type(), operand);`
`679`	`685`	`}`
`680`	`686`	`}`
`681`	`687`	`}`
`@@ -713,6 +719,9 @@ void TryElevateInt32ToInt64_(std::vector<Expr> &expr_vec) { // NOLINT`
`713`	`719`	`}`
`714`	`720`	`} else if (expr->node_type() == IrNodeTy::Load) {`
`715`	`721`	`expr = ir::Cast::Make(Int(64), expr);`
	`722`	`+ } else if (expr->node_type() == IrNodeTy::_Var_ &&`
	`723`	`+ expr.as_var()->is_let_symbol) {`
	`724`	`+ expr = ir::Cast::Make(Int(64), expr);`
`716`	`725`	`}`
`717`	`726`	`}`
`718`	`727`	`}`
`@@ -748,6 +757,10 @@ void ElevateInt64ToInt32_(Expr &expr) { // NOLINT`
`748`	`757`	`}`
`749`	`758`	`} else if (expr->node_type() == IrNodeTy::Load) {`
`750`	`759`	`expr = ir::Cast::Make(Int(32), expr);`
	`760`	`+ } else if (expr->node_type() == IrNodeTy::_Var_ &&`
	`761`	`+ expr.as_var()->is_let_symbol) {`
	`762`	+ // symbol of `let` op should be use cast to convert to int32.
	`763`	`+ expr = ir::Cast::Make(Int(32), expr);`
`751`	`764`	`}`
`752`	`765`	`}`
`753`	`766`	`}`