PaddlePaddle
diff --git a/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+64-63Lines changed: 64 additions & 63 deletions b/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+64-63Lines changed: 64 additions & 63 deletions
diff --git a/‎paddle/cinn/ast_gen_ius/ast_gen.h
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion b/‎paddle/cinn/ast_gen_ius/ast_gen.h
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
Copy file name to clipboard
+18-7Lines changed: 18 additions & 7 deletions b/‎paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
Copy file name to clipboard
+18-7Lines changed: 18 additions & 7 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/op_lowering_impl.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/hlir/framework/pir/op_lowering_impl.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/ir/ir_printer.cc
Copy file name to clipboard
+10-1Lines changed: 10 additions & 1 deletion b/‎paddle/cinn/ir/ir_printer.cc
Copy file name to clipboard
+10-1Lines changed: 10 additions & 1 deletion
diff --git a/‎paddle/cinn/ir/lowered_func.cc
Copy file name to clipboard
+33Lines changed: 33 additions & 0 deletions b/‎paddle/cinn/ir/lowered_func.cc
Copy file name to clipboard
+33Lines changed: 33 additions & 0 deletions
@@ -21,58 +21,58 @@
 #include "paddle/cinn/lang/compute.h"
 #include "paddle/cinn/optim/replace_var_with_expr.h"
 
+using cinn::ir::stmt::BlockRef;
+using cinn::ir::stmt::For;
+using cinn::ir::stmt::Schedule;
+using cinn::ir::stmt::StmtRef;
+using cinn::ir::stmt::Store;
+
 namespace cinn {
 namespace ast_gen_ius {
 
 bool IsReduceBool(const ir::Expr& lhs, const ir::Expr& rhs) {
   return lhs.type().is_bool() || rhs.type().is_bool();
 }
 
-ir::Expr ConvertReduceBody(ir::Expr body,
-                           ir::Tensor tensor,
-                           const std::vector<Expr>& axis_exprs) {
+StmtRef ConvertReduceBody(ir::Expr body,
+                          ir::Tensor tensor,
+                          const std::vector<Expr>& axis_exprs) {
   ir::Reduce* reduce_node = body.As<ir::Reduce>();
   if (!reduce_node) {
-    return ir::Store::Make(tensor, body, axis_exprs);
+    return Store(tensor, body, axis_exprs);
   }
 
   switch (reduce_node->reduce_type) {
     case ir::Reduce::kSum:
       if (IsReduceBool(tensor(axis_exprs), reduce_node->body)) {
-        return ir::Store::Make(
+        return Store(
             tensor, tensor(axis_exprs) || reduce_node->body, axis_exprs);
       }
-      return ir::Store::Make(
-          tensor, tensor(axis_exprs) + reduce_node->body, axis_exprs);
+      return Store(tensor, tensor(axis_exprs) + reduce_node->body, axis_exprs);
     case ir::Reduce::kMul:
       if (IsReduceBool(tensor(axis_exprs), reduce_node->body)) {
-        return ir::Store::Make(
+        return Store(
             tensor, tensor(axis_exprs) && reduce_node->body, axis_exprs);
       }
-      return ir::Store::Make(
-          tensor, tensor(axis_exprs) * reduce_node->body, axis_exprs);
+      return Store(tensor, tensor(axis_exprs) * reduce_node->body, axis_exprs);
     case ir::Reduce::kMax:
-      return ir::Store::Make(
-          tensor,
-          ir::Max::Make(tensor(axis_exprs), reduce_node->body),
-          axis_exprs);
+      return Store(tensor,
+                   ir::Max::Make(tensor(axis_exprs), reduce_node->body),
+                   axis_exprs);
     case ir::Reduce::kMin:
-      return ir::Store::Make(
-          tensor,
-          ir::Min::Make(tensor(axis_exprs), reduce_node->body),
-          axis_exprs);
+      return Store(tensor,
+                   ir::Min::Make(tensor(axis_exprs), reduce_node->body),
+                   axis_exprs);
     case ir::Reduce::kAll:
-      return ir::Store::Make(
-          tensor, tensor(axis_exprs) && reduce_node->body, axis_exprs);
+      return Store(tensor, tensor(axis_exprs) && reduce_node->body, axis_exprs);
     case ir::Reduce::kAny:
-      return ir::Store::Make(
-          tensor, tensor(axis_exprs) || reduce_node->body, axis_exprs);
+      return Store(tensor, tensor(axis_exprs) || reduce_node->body, axis_exprs);
     default:
       CINN_NOT_IMPLEMENTED
   }
 }
 
-ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
+StmtRef AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
   const std::vector<ir::Var>& axis = tensor->axis();
   const std::vector<ir::Expr>& shape = tensor->shape;
   size_t axis_len = axis.size();
@@ -99,7 +99,7 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     tensor_group->Insert(init_tensor);
     tensor_group->MarkShareMemBuffer(tensor, init_tensor);
     tensor_group->CtrlDepend(tensor, init_tensor);
-    Expr init_body = ir::Store::Make(init_tensor, init_value, axis_exprs);
+    StmtRef init_body = Store(init_tensor, init_value, axis_exprs);
     // create schedule block itervars, i0,i1...
     std::vector<ir::Var> block_vars;
     std::vector<ir::Expr> iter_values;
@@ -113,21 +113,22 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
                                shape[i],
                                cinn::UniqName("i" + std::to_string(i)),
                                /*is_reduce = */ false));
-      optim::ReplaceVarWithExpr(&init_body, axis[i], block_vars.back());
+      optim::ReplaceVarWithExpr(init_body, axis[i], block_vars.back());
       axis_vars[i]->is_reduce_axis = false;
       iter_values.push_back(axis_vars[i]);
       ir::TryElevateInt32ToInt64({ir::Expr(axis_vars[i]), shape[i]});
     }
     VLOG(4) << "iter_value.size() and block_vars.size() is "
             << iter_values.size() << " " << block_vars.size();
-    init_body = ir::ScheduleBlockRealize::Make(
-        iter_values,
-        ir::ScheduleBlock::Make(
-            block_vars, {}, {}, reduce_init_name, init_body));
+    init_body = Schedule(block_vars,
+                         iter_values,
+                         {},
+                         {},
+                         reduce_init_name,
+                         BlockRef({init_body}));
 
     // For the remaining reduce axis, make reduce body
-    ir::Expr reduce_body =
-        ConvertReduceBody(tensor->body(), tensor, axis_exprs);
+    StmtRef reduce_body = ConvertReduceBody(tensor->body(), tensor, axis_exprs);
 
     VLOG(4) << "ast gen: reduce body is " << reduce_body;
 
@@ -169,7 +170,7 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     }();
     for (int i = 0; i < non_reduce_axis_vars.size(); ++i) {
       optim::ReplaceVarWithExpr(
-          &reduce_body, non_reduce_axis_vars[i], reduce_block_vars[i]);
+          reduce_body, non_reduce_axis_vars[i], reduce_block_vars[i]);
       ++non_zero_axis_size;
     }
 
@@ -183,66 +184,66 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     }
     VLOG(4) << "before replace body: " << reduce_body;
     for (int i = non_zero_axis_size; i < reduce_block_vars.size(); ++i) {
-      optim::ReplaceVarWithExpr(&reduce_body,
+      optim::ReplaceVarWithExpr(reduce_body,
                                 reduce_axis[i - non_zero_axis_size],
                                 reduce_block_vars[i]);
     }
 
-    reduce_body = ir::ScheduleBlockRealize::Make(
-        reduce_iter_values,
-        ir::ScheduleBlock::Make(
-            reduce_block_vars, {}, {}, tensor->name, reduce_body));
+    reduce_body = Schedule(reduce_block_vars,
+                           reduce_iter_values,
+                           {},
+                           {},
+                           tensor->name,
+                           BlockRef({reduce_body}));
     for (int i = static_cast<int>(reduce_axis.size()) - 1; i >= 0; --i) {
       ir::TryElevateInt32ToInt64({reduce_axis[i],
                                   reduce_axis[i]->lower_bound,
                                   reduce_axis[i]->upper_bound});
-      reduce_body = ir::For::Make(reduce_axis[i],
-                                  reduce_axis[i]->lower_bound,
-                                  reduce_axis[i]->upper_bound,
-                                  ir::ForType::Serial,
-                                  ir::DeviceAPI::Host,
-                                  ir::Block::Make({reduce_body}));
+      reduce_body = For(reduce_axis[i],
+                        reduce_axis[i]->lower_bound,
+                        reduce_axis[i]->upper_bound,
+                        ir::ForType::Serial,
+                        ir::DeviceAPI::Host,
+                        BlockRef({reduce_body}));
     }
 
     // Put the two parts together
-    ir::Expr body = ir::Block::Make({init_body, reduce_body});
+    std::vector<StmtRef> block_body{init_body, reduce_body};
     for (int i = static_cast<int>(axis_len) - 1; i >= 0; --i) {
       ir::Var loop_var = axis[i];
       ir::Expr loop_extent = shape[i];
-      body = ir::For::Make(
-          loop_var,
-          Expr(0),
-          loop_extent,
-          ir::ForType::Serial,
-          ir::DeviceAPI::Host,
-          i == static_cast<int>(axis_len) - 1 ? body : ir::Block::Make({body}));
+      block_body = std::vector<StmtRef>{For(loop_var,
+                                            Expr(0),
+                                            loop_extent,
+                                            ir::ForType::Serial,
+                                            ir::DeviceAPI::Host,
+                                            BlockRef(block_body))};
     }
-    return body;
+    return block_body[0];
   } else {
-    ir::Expr body = ir::Store::Make(tensor, tensor->body(), axis_exprs);
+    StmtRef body = Store(tensor, tensor->body(), axis_exprs);
     // create schedule block itervars, i0,i1...
     std::vector<ir::Var> block_vars;
     std::vector<ir::Expr> iter_values;
     std::vector<Var> axis_vars = cinn::common::GenDefaultAxis(axis_len);
     for (int i = 0; i < shape.size(); ++i) {
       block_vars.push_back(Var(
           Expr(0), shape[i], cinn::UniqName("i" + std::to_string(i)), false));
-      optim::ReplaceVarWithExpr(&body, axis[i], block_vars[i]);
+      optim::ReplaceVarWithExpr(body, axis[i], block_vars[i]);
       axis_vars[i]->is_reduce_axis = false;
       iter_values.push_back(axis_vars[i]);
     }
-    body = ir::ScheduleBlockRealize::Make(
-        iter_values,
-        ir::ScheduleBlock::Make(block_vars, {}, {}, tensor->name, body));
+    body = Schedule(
+        block_vars, iter_values, {}, {}, tensor->name, BlockRef({body}));
     for (int i = static_cast<int>(axis_len) - 1; i >= 0; --i) {
       ir::Var loop_var = axis[i];
       ir::Expr loop_extent = shape[i];
-      body = ir::For::Make(loop_var,
-                           Expr(0),
-                           loop_extent,
-                           ir::ForType::Serial,
-                           ir::DeviceAPI::Host,
-                           ir::Block::Make({body}));
+      body = For(loop_var,
+                 Expr(0),
+                 loop_extent,
+                 ir::ForType::Serial,
+                 ir::DeviceAPI::Host,
+                 BlockRef({body}));
     }
     return body;
   }
 
@@ -16,14 +16,16 @@
 
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/ir_base.h"
+#include "paddle/cinn/ir/stmt.h"
 #include "paddle/cinn/ir/tensor.h"
 
 namespace cinn {
 namespace ast_gen_ius {
 
 class AstGen {
  public:
-  static ir::Expr Build(const ir::Tensor& tensor, TensorGroup* tensor_group);
+  static ir::stmt::StmtRef Build(const ir::Tensor& tensor,
+                                 TensorGroup* tensor_group);
 };
 
 }  // namespace ast_gen_ius
 
@@ -99,9 +99,14 @@ BucketLoweredFuncsWrapper OpLowererImpl::BucketLower(
   std::unordered_map<::pir::Value, ir::Tensor> tensor_map;
   // for some op, it will output more tmp value and regard as
   // XX_0, XX_1, so we log them in tmp_tensor_info;
-  std::vector<ir::Expr> func_bodies =
+  std::vector<ir::stmt::BlockRef> func_body_blocks =
       LowerOps(group, ops, &group_func_arg_tensors, &tensor_map);
-
+  // TODO(Hongqing-work): delete this convert after using new IR update
+  // schedule.
+  std::vector<ir::Expr> func_bodies;
+  for (const auto& body : func_body_blocks) {
+    func_bodies.emplace_back(ir::ConvertStmtBlockToExprBlock(body));
+  }
   if (FLAGS_cinn_check_tensor_buffer_map) {
     optim::CheckTensorBufferMap(func_bodies, "BucketLower LowerOps");
     VLOG(3) << "LowerOps tensor-buffer map check succeed";
@@ -410,13 +415,13 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
   return lowered_funcs;
 }
 
-std::vector<ir::Expr> OpLowererImpl::LowerOps(
+std::vector<ir::stmt::BlockRef> OpLowererImpl::LowerOps(
     const OpLoweringGroupPtr& group,
     const std::vector<::pir::Operation*>& ops,
     std::vector<ir::Tensor>* group_func_arg_tensors,
     std::unordered_map<::pir::Value, ir::Tensor>* tensor_map) {
   auto& strategy = Operator::GetAttrs<StrategyFunction>("CINNStrategy");
-  std::vector<Expr> func_bodies;
+  std::vector<ir::stmt::BlockRef> func_bodies;
 
   for (auto* op : ops) {
     VLOG(4) << "start lowering op:" << op->name() << " id: " << op->id();
@@ -460,7 +465,7 @@ std::vector<ir::Expr> OpLowererImpl::LowerOps(
         DoOpLower(op_impl, op, tensor_map, &op_func_arg_tensors);
 
     for (const ir::LoweredFunc& func : funcs) {
-      func_bodies.push_back(func->body);
+      func_bodies.push_back(func->body_block);
     }
   }
 
@@ -776,11 +781,17 @@ ir::Expr OpLowererImpl::LowerX86(const OpLoweringGroupPtr& group,
   this->target_ = common::DefaultHostTarget();
   cinn::runtime::CurrentTarget::SetCurrentTarget(this->target_);
 
-  std::vector<ir::Expr> func_bodies =
+  std::vector<ir::stmt::BlockRef> func_bodies =
       LowerOps(group, ops, &group_func_arg_tensors, &tensor_map);
   this->target_ = common::DefaultDeviceTarget();
   cinn::runtime::CurrentTarget::SetCurrentTarget(this->target_);
-  ir::ModuleExpr mod_expr(func_bodies);
+  // TODO(Hongqing-work): delete this convert after using new IR update
+  // schedule.
+  std::vector<ir::Expr> expr_func_bodies;
+  for (const auto& body : func_bodies) {
+    expr_func_bodies.emplace_back(ir::ConvertStmtBlockToExprBlock(body));
+  }
+  ir::ModuleExpr mod_expr(expr_func_bodies);
   ir::IRSchedule ir_sch(
       mod_expr, -1, false, cinn::utils::ErrorMessageLevel::kGeneral, true);
   ir_sch.MergeExprs();
 
@@ -90,7 +90,7 @@ class OpLowererImpl : public OpLowererImplBase<OpLoweringGroupPtr> {
    * @param tensor_map All tensors used for calculating the group.
    * @return The lowered func bodies of Op set.
    */
-  std::vector<ir::Expr> LowerOps(
+  std::vector<ir::stmt::BlockRef> LowerOps(
       const OpLoweringGroupPtr& group,
       const std::vector<::pir::Operation*>& ops,
       std::vector<ir::Tensor>* group_func_arg_tensors,
 
@@ -522,7 +522,16 @@ void IrPrinter::Visit(const _LoweredFunc_ *f) {
   str_ += utils::Join(arg_names, ", ");
   str_ += ")\n";
 
-  Visit(f->body);
+  if (f->body.defined()) {
+    Visit(f->body);
+  } else {
+    PADDLE_ENFORCE_EQ(
+        f->body_block.defined(),
+        true,
+        ::common::errors::InvalidArgument(
+            "Please ensure that `f->body` or `f->body_block` is defined."));
+    VisitBlock(f->body_block);
+  }
 }
 void IrPrinter::Visit(const Let *f) {
   PADDLE_ENFORCE_EQ(
 
@@ -66,6 +66,29 @@ LoweredFunc _LoweredFunc_::Make(const std::string& name,
   return LoweredFunc(n);
 }
 
+LoweredFunc _LoweredFunc_::Make(const std::string& name,
+                                const std::vector<Argument>& args,
+                                const stmt::BlockRef& body,
+                                const std::vector<ir::Buffer>& temp_bufs) {
+  auto* n = make_shared<_LoweredFunc_>();
+  n->name = name;
+  n->args = args;
+  n->body_block = body;
+  n->temp_bufs = temp_bufs;
+
+  n->CheckValid();
+  n->PrepareAllocOutputBufferExprs();
+  n->PrepareCreateTempBufferExprs();
+  n->PrepareAllocTempBufferExprs();
+  n->AllocTempBuffer();
+  bool with_expr_gen_tensor = false;
+  n->PrepareBufferCastExprs(with_expr_gen_tensor);
+  n->PrepareArgumentExprs();
+  n->PrepareDeallocTempBufferExprs();
+  n->PrepareDeallocOutputBufferExprs();
+  return LoweredFunc(n);
+}
+
 LoweredFunc _LoweredFunc_::Make(const std::string& name,
                                 const std::vector<Argument>& args,
                                 const Expr& body) {
@@ -76,6 +99,16 @@ LoweredFunc _LoweredFunc_::Make(const std::string& name,
   return LoweredFunc(n);
 }
 
+LoweredFunc _LoweredFunc_::Make(const std::string& name,
+                                const std::vector<Argument>& args,
+                                const stmt::BlockRef& body) {
+  auto* n = make_shared<_LoweredFunc_>();
+  n->name = name;
+  n->args = args;
+  n->body_block = body;
+  return LoweredFunc(n);
+}
+
 void _LoweredFunc_::CheckValid() const {
   // check there is at least one output
   int out_count = 0;