[refactor] [ir] Refactor ExternalFuncCallExpression into a frontend stmt

taichi/inc/statements.inc.h

@@ -1,4 +1,5 @@
 // Frontend statements
+PER_STATEMENT(FrontendExternalFuncStmt)
 PER_STATEMENT(FrontendExprStmt)
 PER_STATEMENT(FrontendIfStmt)
 PER_STATEMENT(FrontendForStmt)

taichi/ir/frontend_ir.cpp

@@ -285,51 +285,6 @@ void InternalFuncCallExpression::flatten(FlattenContext *ctx) {
   stmt = ctx->back_stmt();
 }
 
-void ExternalFuncCallExpression::type_check() {
-  for (auto &arg : args) {
-    TI_ASSERT_TYPE_CHECKED(arg);
-    // no arg type compatibility check for now due to lack of specification
-  }
-  for (auto &output : outputs) {
-    TI_ASSERT_TYPE_CHECKED(output);
-    // no output type compatibility check for now due to lack of specification
-  }
-  // external func calls have no return type for now
-}
-
-void ExternalFuncCallExpression::flatten(FlattenContext *ctx) {
-  TI_ASSERT((int)(so_func != nullptr) + (int)(!asm_source.empty()) +
-                (int)(!bc_filename.empty()) ==
-            1)
-  std::vector<Stmt *> arg_statements, output_statements;
-  if (so_func != nullptr || !asm_source.empty()) {
-    for (auto &s : args) {
-      s.set(load_if_ptr(s));
-      s->flatten(ctx);
-      arg_statements.push_back(s->stmt);
-    }
-    for (auto &s : outputs) {
-      output_statements.push_back(s.cast<IdExpression>()->flatten_noload(ctx));
-    }
-    ctx->push_back(std::make_unique<ExternalFuncCallStmt>(
-        (so_func != nullptr) ? ExternalFuncCallStmt::SHARED_OBJECT
-                             : ExternalFuncCallStmt::ASSEMBLY,
-        so_func, asm_source, "", "", arg_statements, output_statements));
-    stmt = ctx->back_stmt();
-  } else {
-    for (auto &s : args) {
-      TI_ASSERT_INFO(
-          s.is<IdExpression>(),
-          "external func call via bitcode must pass in local variables.")
-      arg_statements.push_back(s.cast<IdExpression>()->flatten_noload(ctx));
-    }
-    ctx->push_back(std::make_unique<ExternalFuncCallStmt>(
-        ExternalFuncCallStmt::BITCODE, nullptr, "", bc_filename, bc_funcname,
-        arg_statements, output_statements));
-    stmt = ctx->back_stmt();
-  }
-}
-
 void ExternalTensorExpression::flatten(FlattenContext *ctx) {
   auto ptr = Stmt::make<ArgLoadStmt>(arg_id, dt, /*is_ptr=*/true);
   ctx->push_back(std::move(ptr));

taichi/ir/frontend_ir.h

@@ -13,6 +13,31 @@
 TLANG_NAMESPACE_BEGIN
 
 // Frontend Statements
+class FrontendExternalFuncStmt : public Stmt {
+ public:
+  void *so_func;
+  std::string asm_source;
+  std::string bc_filename;
+  std::string bc_funcname;
+  std::vector<Expr> args;
+  std::vector<Expr> outputs;
+
+  FrontendExternalFuncStmt(void *so_func,
+                           const std::string &asm_source,
+                           const std::string &bc_filename,
+                           const std::string &bc_funcname,
+                           const std::vector<Expr> &args,
+                           const std::vector<Expr> &outputs)
+      : so_func(so_func),
+        asm_source(asm_source),
+        bc_filename(bc_filename),
+        bc_funcname(bc_funcname),
+        args(args),
+        outputs(outputs) {
+  }
+
+  TI_DEFINE_ACCEPT
+};
 
 class FrontendExprStmt : public Stmt {
  public:
@@ -370,58 +395,7 @@ class InternalFuncCallExpression : public Expression {
   void flatten(FlattenContext *ctx) override;
 };
 
-class ExternalFuncCallExpression : public Expression {
- public:
-  void *so_func;
-  std::string asm_source;
-  std::string bc_filename;
-  std::string bc_funcname;
-  std::vector<Expr> args;
-  std::vector<Expr> outputs;
-
-  ExternalFuncCallExpression(void *so_func,
-                             const std::string &asm_source,
-                             const std::string &bc_filename,
-                             const std::string &bc_funcname,
-                             const std::vector<Expr> &args,
-                             const std::vector<Expr> &outputs)
-      : so_func(so_func),
-        asm_source(asm_source),
-        bc_filename(bc_filename),
-        bc_funcname(bc_funcname),
-        args(args),
-        outputs(outputs) {
-  }
-
-  void type_check() override;
-
-  void serialize(std::ostream &ss) override {
-    if (so_func != nullptr) {
-      ss << fmt::format("so {:x} (", (uint64)so_func);
-    } else if (!asm_source.empty()) {
-      ss << fmt::format("asm \"{}\" (", asm_source);
-    } else {
-      ss << fmt::format("bc {}:{} (", bc_filename, bc_funcname);
-    }
-
-    ss << "inputs=";
-
-    for (auto &s : args) {
-      s.serialize(ss);
-    }
-
-    ss << ", outputs=";
-
-    for (auto &s : outputs) {
-      s.serialize(ss);
-    }
-
-    ss << ')';
-  }
-
-  void flatten(FlattenContext *ctx) override;
-};
-
+// TODO: Make this a non-expr
 class ExternalTensorExpression : public Expression {
  public:
   DataType dt;
@@ -448,6 +422,7 @@ class ExternalTensorExpression : public Expression {
   void flatten(FlattenContext *ctx) override;
 };
 
+// TODO: Make this a non-expr
 class GlobalVariableExpression : public Expression {
  public:
   Identifier ident;
@@ -676,6 +651,7 @@ class SNodeOpExpression : public Expression {
   void flatten(FlattenContext *ctx) override;
 };
 
+// TODO: Make this a non-expr
 class LocalLoadExpression : public Expression {
  public:
   Expr ptr;
@@ -695,6 +671,7 @@ class LocalLoadExpression : public Expression {
   void flatten(FlattenContext *ctx) override;
 };
 
+// TODO: make this a non-expr
 class GlobalLoadExpression : public Expression {
  public:
   Expr ptr;

taichi/python/export_lang.cpp

@@ -657,11 +657,10 @@ void export_lang(py::module &m) {
         [](std::size_t func_addr, std::string source, std::string filename,
            std::string funcname, const ExprGroup &args,
            const ExprGroup &outputs) {
-          auto expr = Expr::make<ExternalFuncCallExpression>(
+          auto stmt = Stmt::make<FrontendExternalFuncStmt>(
               (void *)func_addr, source, filename, funcname, args.exprs,
               outputs.exprs);
-
-          current_ast_builder().insert(Stmt::make<FrontendExprStmt>(expr));
+          current_ast_builder().insert(std::move(stmt));
         });
 
   m.def("insert_is_active", [](SNode *snode, const ExprGroup &indices) {

taichi/transforms/frontend_type_check.cpp

@@ -28,6 +28,10 @@ class FrontendTypeCheck : public IRVisitor {
       stmt->accept(this);
   }
 
+  void visit(FrontendExternalFuncStmt *stmt) override {
+    // TODO: noop for now; add typechecking after we have type specification
+  }
+
   void visit(FrontendExprStmt *stmt) override {
     // Noop
   }

taichi/transforms/ir_printer.cpp

@@ -726,6 +726,25 @@ class IRPrinter : public IRVisitor {
   void visit(MeshPatchIndexStmt *stmt) override {
     print("{}{} = mesh patch idx", stmt->type_hint(), stmt->name());
   }
+
+  void visit(FrontendExternalFuncStmt *stmt) override {
+    if (stmt->so_func != nullptr) {
+      print("so {:x}", (uint64)stmt->so_func);
+    } else if (!stmt->asm_source.empty()) {
+      print("asm \"{}\"", stmt->asm_source);
+    } else {
+      print("bc {}:{}", stmt->bc_filename, stmt->bc_funcname);
+    }
+    print(" (inputs=");
+    for (auto &s : stmt->args) {
+      print(s.serialize());
+    }
+    print(", outputs=");
+    for (auto &s : stmt->outputs) {
+      print(s.serialize());
+    }
+    print(")");
+  }
 };
 
 }  // namespace

taichi/transforms/lower_ast.cpp

@@ -486,6 +486,42 @@ class LowerAST : public IRVisitor {
     stmt->parent->replace_with(stmt, std::move(fctx.stmts));
   }
 
+  void visit(FrontendExternalFuncStmt *stmt) override {
+    auto ctx = make_flatten_ctx();
+    TI_ASSERT((int)(stmt->so_func != nullptr) +
+                  (int)(!stmt->asm_source.empty()) +
+                  (int)(!stmt->bc_filename.empty()) ==
+              1)
+    std::vector<Stmt *> arg_statements, output_statements;
+    if (stmt->so_func != nullptr || !stmt->asm_source.empty()) {
+      for (auto &s : stmt->args) {
+        s.set(load_if_ptr(s));
+        s->flatten(&ctx);
+        arg_statements.push_back(s->stmt);
+      }
+      for (auto &s : stmt->outputs) {
+        output_statements.push_back(
+            s.cast<IdExpression>()->flatten_noload(&ctx));
+      }
+      ctx.push_back(std::make_unique<ExternalFuncCallStmt>(
+          (stmt->so_func != nullptr) ? ExternalFuncCallStmt::SHARED_OBJECT
+                                     : ExternalFuncCallStmt::ASSEMBLY,
+          stmt->so_func, stmt->asm_source, "", "", arg_statements,
+          output_statements));
+    } else {
+      for (auto &s : stmt->args) {
+        TI_ASSERT_INFO(
+            s.is<IdExpression>(),
+            "external func call via bitcode must pass in local variables.")
+        arg_statements.push_back(s.cast<IdExpression>()->flatten_noload(&ctx));
+      }
+      ctx.push_back(std::make_unique<ExternalFuncCallStmt>(
+          ExternalFuncCallStmt::BITCODE, nullptr, "", stmt->bc_filename,
+          stmt->bc_funcname, arg_statements, output_statements));
+    }
+    stmt->parent->replace_with(stmt, std::move(ctx.stmts));
+  }
+
   static void run(IRNode *node) {
     LowerAST inst(irpass::analysis::detect_fors_with_break(node));
     node->accept(&inst);