PaddlePaddle
diff --git a/‎paddle/cinn/backends/codegen_c.cc
Copy file name to clipboard
+185Lines changed: 185 additions & 0 deletions b/‎paddle/cinn/backends/codegen_c.cc
Copy file name to clipboard
+185Lines changed: 185 additions & 0 deletions
diff --git a/‎paddle/cinn/backends/codegen_c.h
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions b/‎paddle/cinn/backends/codegen_c.h
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions
diff --git a/‎paddle/cinn/backends/codegen_device_util.cc
Copy file name to clipboard
+26-13Lines changed: 26 additions & 13 deletions b/‎paddle/cinn/backends/codegen_device_util.cc
Copy file name to clipboard
+26-13Lines changed: 26 additions & 13 deletions
@@ -270,6 +270,77 @@ void CodeGenC::Visit(const ir::For *op) {
     str_ += "}";
   }
 }
+
+void CodeGenC::VisitStmt(const ir::stmt::For &stmt) {
+  Expr extent = stmt->extent();
+  Expr min = stmt->min();
+  int num_task = 1;
+  if (stmt->is_parallel()) {
+    str_ += "int num_task = max_concurrency();\n";
+    DoIndent();
+    str_ += "omp_set_num_threads(num_task);\n";
+    DoIndent();
+    str_ += "auto flambda = [=](int task_id, int num_task) -> int {\n";
+    IncIndent();
+    DoIndent();
+    str_ += "int n_per_task = ";
+    Expr num_task_var = Var("num_task");
+    IrPrinter::Visit((stmt->extent() + num_task_var - 1) / num_task_var);
+    str_ += ";\n";
+    PADDLE_ENFORCE_EQ(min.as_int32(),
+                      0,
+                      ::common::errors::InvalidArgument(
+                          "The min of the for loop should be 0"));
+    auto task_id = Var("task_id");
+    auto n_per_task = Var("n_per_task");
+    min = task_id * n_per_task;
+    extent = (task_id + 1) * n_per_task;
+    DoIndent();
+  }
+  str_ += "for (";
+  str_ += GetTypeRepr(Int(32));
+  str_ += " ";
+  str_ += stmt->loop_var()->name;
+  str_ += " = ";
+  IrPrinter::Visit(min);
+  str_ += "; ";
+  str_ += stmt->loop_var()->name;
+  str_ += " < ";
+  IrPrinter::Visit(stmt->extent());
+  if (stmt->is_parallel()) {
+    str_ += " && ";
+    str_ += stmt->loop_var()->name;
+    str_ += " < ";
+    IrPrinter::Visit(extent);
+  }
+  str_ += "; ";
+
+  str_ += stmt->loop_var()->name;
+  str_ += " += 1";
+  str_ += ") ";
+
+  VisitBlock(stmt->body());
+  if (stmt->is_parallel()) {
+    str_ += "\n";
+    DoIndent();
+    str_ += "return 0;\n";
+    DecIndent();
+    DoIndent();
+    str_ += "};\n";
+    str_ += "#pragma omp parallel num_threads(num_task)\n";
+    DoIndent();
+    str_ += "{\n";
+    IncIndent();
+    DoIndent();
+    str_ += "int task_id = omp_get_thread_num();\n";
+    DoIndent();
+    str_ += "flambda(task_id, num_task);\n";
+    DecIndent();
+    DoIndent();
+    str_ += "}";
+  }
+}
+
 void CodeGenC::Visit(const ir::PolyFor *op) {
   str_ += "for (";
   str_ += GetTypeRepr(Int(32));
@@ -310,6 +381,20 @@ void CodeGenC::Visit(const ir::IfThenElse *op) {
     IrPrinter::Visit(op->false_case);
   }
 }
+
+void CodeGenC::VisitStmt(const ir::stmt::IfThenElse &stmt) {
+  str_ += "if (";
+  IrPrinter::Visit(stmt->condition());
+  str_ += ") ";
+
+  VisitBlock(stmt->true_case());
+
+  if (!stmt->false_case()->stmts().empty()) {
+    str_ += " else ";
+    VisitBlock(stmt->false_case());
+  }
+}
+
 void CodeGenC::Visit(const ir::Block *op) {
   str_ += "{\n";
 
@@ -332,6 +417,29 @@ void CodeGenC::Visit(const ir::Block *op) {
   DoIndent();
   str_ += "}";
 }
+void CodeGenC::VisitBlock(const ir::stmt::BlockRef &stmt) {
+  str_ += "{\n";
+
+  IncIndent();
+
+  // Note: size_t (0 - 1) = 18446744073709551615
+  if (stmt->stmts().size() >= 1) {
+    for (int i = 0; i < stmt->stmts().size() - 1; i++) {
+      DoIndent();
+      IrPrinter::VisitStmt(stmt->stmts()[i]);
+      str_ += ";\n";
+    }
+    DoIndent();
+    IrPrinter::VisitStmt(stmt->stmts().back());
+    str_ += ";";
+  }
+
+  DecIndent();
+  str_ += "\n";
+  DoIndent();
+  str_ += "}";
+}
+
 void CodeGenC::Visit(const ir::Call *op) {
   if (op->name == runtime::intrinsic::buffer_malloc) {
     PrintCall_buffer_malloc(op);
@@ -527,6 +635,30 @@ void CodeGenC::Visit(const ir::Load *op) {
   }
 }
 
+void CodeGenC::VisitStmt(const ir::stmt::Store &stmt) {
+  PADDLE_ENFORCE_EQ(
+      stmt->is_addr_tensor(),
+      true,
+      ::common::errors::InvalidArgument(
+          "The operation type is invalid. It must be an address tensor."));
+  ir::Expr offset = [&] {
+    if (store_stmt_to_offset_.count(stmt) == 0) {
+      store_stmt_to_offset_[stmt] = stmt->index();
+    }
+    return store_stmt_to_offset_.at(stmt);
+  }();
+  auto *tensor = stmt->tensor().As<ir::_Tensor_>();
+  PADDLE_ENFORCE_NOT_NULL(tensor,
+                          ::common::errors::InvalidArgument(
+                              "The tensor is null. It must not be null."));
+  str_ += tensor->name;
+  str_ += "[";
+  IrPrinter::Visit(offset);
+  str_ += "]";
+  str_ += " = ";
+  IrPrinter::Visit(stmt->value());
+}
+
 void CodeGenC::Visit(const ir::Store *op) {
   PADDLE_ENFORCE_EQ(
       op->is_addr_tensor(),
@@ -560,6 +692,16 @@ void CodeGenC::Visit(const ir::Alloc *op) {
   str_ += ")";
 }
 
+void CodeGenC::VisitStmt(const ir::stmt::Alloc &stmt) {
+  str_ += runtime::intrinsic::buffer_malloc;
+  str_ += "(";
+  str_ += "(void*)(0), ";
+
+  auto *buffer = stmt->destination().As<ir::_Buffer_>();
+  str_ += buffer->name;
+  str_ += ")";
+}
+
 void CodeGenC::Visit(const ir::Free *op) {
   str_ += runtime::intrinsic::buffer_free;
   str_ += "(";
@@ -570,6 +712,16 @@ void CodeGenC::Visit(const ir::Free *op) {
   str_ += ")";
 }
 
+void CodeGenC::VisitStmt(const ir::stmt::Free &stmt) {
+  str_ += runtime::intrinsic::buffer_free;
+  str_ += "(";
+  str_ += "(void*)(0), ";
+
+  auto *buffer = stmt->destination().As<ir::_Buffer_>();
+  str_ += buffer->name;
+  str_ += ")";
+}
+
 void CodeGenC::Visit(const ir::_Buffer_ *op) { str_ += op->name; }
 void CodeGenC::Visit(const ir::_Tensor_ *op) { str_ += op->buffer->name; }
 void CodeGenC::Visit(const ir::Let *op) {
@@ -602,6 +754,36 @@ void CodeGenC::Visit(const ir::Let *op) {
   }
 }
 
+void CodeGenC::VisitStmt(const ir::stmt::Let &stmt) {
+  bool is_vec = false;
+  PADDLE_ENFORCE_EQ(stmt->type().valid(),
+                    true,
+                    ::common::errors::InvalidArgument(
+                        "The operation type is invalid. It must be valid."));
+  if (stmt->body().defined() && stmt->body().As<ir::Broadcast>()) {
+    // broadcast's type is hard to decide, so use c++11 auto instead.
+    str_ += "auto";
+    is_vec = true;
+  } else {
+    str_ += GetTypeRepr(stmt->type());
+  }
+
+  str_ += " ";
+  IrPrinter::Visit(stmt->symbol());
+
+  // native C array.
+  if (stmt->type().lanes() > 1 && !is_vec) {
+    str_ += "[";
+    str_ += std::to_string(stmt->type().lanes());
+    str_ += "]";
+  }
+
+  if (stmt->body().defined()) {
+    str_ += " = ";
+    IrPrinter::Visit(stmt->body());
+  }
+}
+
 void CodeGenC::Visit(const ir::Reduce *op) {
   PADDLE_THROW(::common::errors::InvalidArgument(
       "Reduce IR is just for internal representation, should not be "
@@ -828,6 +1010,9 @@ void CodeGenC::PrintStackVecType(Type type, int lanes) {
 void CodeGenC::Visit(const ir::PrimitiveNode *op) { CINN_NOT_IMPLEMENTED }
 void CodeGenC::Visit(const ir::_BufferRange_ *op) { CINN_NOT_IMPLEMENTED }
 void CodeGenC::Visit(const ir::ScheduleBlock *op) { CINN_NOT_IMPLEMENTED }
+void CodeGenC::VisitStmt(const ir::stmt::Schedule &stmt) {
+  CINN_NOT_IMPLEMENTED
+}
 void CodeGenC::Visit(const ir::ScheduleBlockRealize *op) {
   CINN_NOT_IMPLEMENTED
 }
 
@@ -105,6 +105,15 @@ class CodeGenC : public ir::IrPrinter {
   void Visit(const ir::_Module_* op) override;
   void Visit(const ir::_LoweredFunc_* op) override;
 
+  void VisitBlock(const ir::stmt::BlockRef& block) override;
+  void VisitStmt(const ir::stmt::Let& stmt) override;
+  void VisitStmt(const ir::stmt::Store& stmt) override;
+  void VisitStmt(const ir::stmt::Alloc& stmt) override;
+  void VisitStmt(const ir::stmt::Free& stmt) override;
+  void VisitStmt(const ir::stmt::IfThenElse& stmt) override;
+  void VisitStmt(const ir::stmt::For& stmt) override;
+  void VisitStmt(const ir::stmt::Schedule& stmt) override;
+
 #define __DEFINE_VISIT(op__) \
   void Visit(const ir::intrinsics::op__* op) override;
   INTRINSIC_KIND_FOR_EACH(__DEFINE_VISIT)
@@ -116,11 +125,19 @@ class CodeGenC : public ir::IrPrinter {
 
   friend class ExternFunctionEmitter;
 
+  struct StoreHash {
+    size_t operator()(const ir::stmt::Store& stmt) const {
+      return std::hash<const Object*>()(stmt.get());
+    }
+  };
+
  protected:
   Target target_;
   std::stringstream ss_;
   bool inline_builtin_codes_{true};
   std::unordered_map<const ir::Store*, ir::Expr> store_to_offset_;
+  std::unordered_map<const ir::stmt::Store, ir::Expr, StoreHash>
+      store_stmt_to_offset_;
   std::unordered_map<const ir::Load*, ir::Expr> load_to_offset_;
 };
 
 
@@ -244,6 +244,7 @@ void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
       [&](common::HygonDCUArchSYCL) {
         call_kernel = runtime::intrinsic::call_sycl_kernel;
       });
+  // TODO(Dmovic): use new ir when backend update done.
   ir::Expr call_extern_api =
       ir::Call::Make(Void(),
                      call_kernel.value(),
@@ -264,7 +265,7 @@ void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
                      0);
 
   // create memset calls for temp_spaces if needed
-  std::vector<ir::Expr> call_kernel_stmts;
+  std::vector<ir::stmt::StmtRef> call_kernel_stmts;
   for (auto &temp_space : func_node->temp_spaces) {
     if (temp_space.need_zero_init()) {
       ir::Expr size = common::cast(temp_space.size(), common::UInt(64));
@@ -274,23 +275,26 @@ void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
       ir::Expr call_memset = lang::CallExtern(
           runtime::intrinsic::call_cuda_memset,
           {call_get_arg, ir::Expr(1), ir::Expr(0), size, kernel_stream_});
-      call_kernel_stmts.push_back(call_memset);
+      call_kernel_stmts.push_back(ir::stmt::Evaluate(call_memset));
     }
   }
-  call_kernel_stmts.push_back(call_extern_api);
-  call_extern_api = ir::Block::Make(call_kernel_stmts);
+  call_kernel_stmts.push_back(ir::stmt::Evaluate(call_extern_api));
+  auto call_extern_api_block = ir::stmt::BlockRef(call_kernel_stmts);
 
   if (buckets_.empty()) {
-    buckets_.emplace_back(ir::IfThenElse::Make(predicate, call_extern_api));
+    buckets_.emplace_back(
+        ir::stmt::IfThenElse(predicate, call_extern_api_block));
   } else {
     auto false_expr = buckets_.back();
     buckets_.pop_back();
-    buckets_.emplace_back(
-        ir::IfThenElse::Make(predicate, call_extern_api, false_expr));
+    buckets_.emplace_back(ir::stmt::IfThenElse(
+        predicate,
+        call_extern_api_block,
+        ir::stmt::BlockRef(std::vector<ir::stmt::StmtRef>{false_expr})));
   }
 
   // create infer shape calls for temp_spaces
-  std::vector<ir::Expr> temp_space_infer_shape_stmts;
+  std::vector<ir::stmt::StmtRef> temp_space_infer_shape_stmts;
   for (int i = 0; i < func_node->temp_spaces.size(); ++i) {
     ir::Var tensor_shape_args(TENSOR_SHAPE_ARGS, type_of<int64_t **>());
     ir::Expr size =
@@ -301,12 +305,20 @@ void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
                           ir::Expr(0),
                           size,
                           tensor_shape_args});
-    temp_space_infer_shape_stmts.push_back(call_set_value);
+    temp_space_infer_shape_stmts.push_back(ir::stmt::Evaluate(call_set_value));
   }
   if (!temp_space_infer_shape_stmts.empty()) {
-    ir::Expr if_body = ir::Block::Make(temp_space_infer_shape_stmts);
-    temp_space_infer_shape_body_ =
-        ir::IfThenElse::Make(predicate, if_body, temp_space_infer_shape_body_);
+    ir::stmt::BlockRef if_body =
+        ir::stmt::BlockRef(temp_space_infer_shape_stmts);
+    if (temp_space_infer_shape_body_.defined()) {
+      temp_space_infer_shape_body_ = ir::stmt::IfThenElse(
+          predicate,
+          if_body,
+          ir::stmt::BlockRef(
+              std::vector<ir::stmt::StmtRef>{temp_space_infer_shape_body_}));
+    } else {
+      temp_space_infer_shape_body_ = ir::stmt::IfThenElse(predicate, if_body);
+    }
   }
 }
 
@@ -325,7 +337,8 @@ void detail::CollectBucketStrategyHostFunctionVisitor::ProcessArgs(
                          0);
       ir::Expr let_symbol = ir::Expr(args[i].var_arg());
       let_symbol->set_type(type_of<int64_t>());
-      ir::Expr stmt = ir::Let::Make(let_symbol, call_get_value_in_kernel_args);
+      ir::stmt::StmtRef stmt =
+          ir::stmt::Let(let_symbol, call_get_value_in_kernel_args);
       arg_defs_.push_back(stmt);
     }
   }