[ir] Basic function definition/call instructions (#612)

* [skip ci] add FuncBodyStmt

* add FrontendFuncDefStmt & begin_func/end_func

* function using scope_stack & real func_call

* [skip ci] enforce code format

* Update TaichiCore.cmake

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>
Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>

taichi/inc/statements.inc.h

@@ -20,6 +20,9 @@ PER_STATEMENT(StructForStmt)
 PER_STATEMENT(IfStmt)
 PER_STATEMENT(WhileStmt)
 PER_STATEMENT(WhileControlStmt)
+PER_STATEMENT(FuncBodyStmt)
+PER_STATEMENT(FrontendFuncDefStmt)
+PER_STATEMENT(FuncCallStmt)
 
 PER_STATEMENT(ArgLoadStmt)
 PER_STATEMENT(ExternalPtrStmt)

taichi/ir/ir.h

@@ -1865,6 +1865,51 @@ class StructForStmt : public Stmt {
   DEFINE_ACCEPT
 };
 
+class FuncBodyStmt : public Stmt {
+ public:
+  std::string funcid;
+  std::unique_ptr<Block> body;
+
+  FuncBodyStmt(const std::string &funcid, std::unique_ptr<Block> &&body)
+      : funcid(funcid), body(std::move(body)) {
+  }
+
+  bool is_container_statement() const override {
+    return true;
+  }
+
+  DEFINE_ACCEPT
+};
+
+class FrontendFuncDefStmt : public Stmt {
+ public:
+  std::string funcid;
+  std::unique_ptr<Block> body;
+
+  FrontendFuncDefStmt(const std::string &funcid) : funcid(funcid) {
+  }
+
+  bool is_container_statement() const override {
+    return true;
+  }
+
+  DEFINE_ACCEPT
+};
+
+class FuncCallStmt : public Stmt {
+ public:
+  std::string funcid;
+
+  FuncCallStmt(const std::string &funcid) : funcid(funcid) {
+  }
+
+  bool is_container_statement() const override {
+    return false;
+  }
+
+  DEFINE_ACCEPT
+};
+
 class WhileStmt : public Stmt {
  public:
   Stmt *mask;

taichi/python/export_lang.cpp

@@ -297,6 +297,21 @@ void export_lang(py::module &m) {
     current_ast_builder().insert(Stmt::make<FrontendBreakStmt>());
   });
 
+  m.def("begin_func", [&](const std::string &funcid) {
+    auto stmt_unique = std::make_unique<FrontendFuncDefStmt>(funcid);
+    auto stmt = stmt_unique.get();
+    current_ast_builder().insert(std::move(stmt_unique));
+    scope_stack.push_back(current_ast_builder().create_scope(stmt->body));
+  });
+
+  m.def("end_func", [&](const std::string &funcid) { scope_stack.pop_back(); });
+
+  m.def("func_call", [&](const std::string &funcid) {
+    auto func = Stmt::make<FuncCallStmt>(
+        funcid);  // TODO: use FuncCallExpr with return values & args
+    current_ast_builder().insert(std::move(func));
+  });
+
   m.def("layout", layout);
 
   m.def("value_cast", static_cast<Expr (*)(const Expr &expr, DataType)>(cast));

taichi/transforms/ir_printer.cpp

@@ -195,6 +195,22 @@ class IRPrinter : public IRVisitor {
     print("while control {}, {}", stmt->mask->name(), stmt->cond->name());
   }
 
+  void visit(FuncCallStmt *stmt) override {
+    print("{}{} = call \"{}\"", stmt->type_hint(), stmt->name(), stmt->funcid);
+  }
+
+  void visit(FrontendFuncDefStmt *stmt) override {
+    print("function \"{}\" {{", stmt->funcid);
+    stmt->body->accept(this);
+    print("}}");
+  }
+
+  void visit(FuncBodyStmt *stmt) override {
+    print("func \"{}\" {{");
+    stmt->body->accept(this);
+    print("}}");
+  }
+
   void visit(WhileStmt *stmt) override {
     print("while true {{");
     stmt->body->accept(this);