[ir] Basic function definition/call instructions (taichi-dev#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 <[email protected]> Co-authored-by: Yuanming Hu <[email protected]>
archibate · Mar 31, 2020 · 49c893a · 49c893a
1 parent 6dfec4c
commit 49c893a
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 0 deletions.
diff --git a/taichi/inc/statements.inc.h b/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)

diff --git a/taichi/ir/ir.h b/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;

diff --git a/taichi/python/export_lang.cpp b/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));

diff --git a/taichi/transforms/ir_printer.cpp b/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);