[refactor] Add ASTSerializer and use it to generate offline-cache-key (…

…#4863)

* Add ASTSerializer, using it to generate offline-cache-key

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

taichi/analysis/offline_cache_util.cpp

@@ -152,8 +152,9 @@ std::string get_hashed_offline_cache_key(CompileConfig *config,
                                          Kernel *kernel) {
   std::string kernel_ast_string;
   if (kernel) {
-    irpass::gen_offline_cache_key(kernel->program, kernel->ir.get(),
-                                  &kernel_ast_string);
+    std::ostringstream oss;
+    gen_offline_cache_key(kernel->program, kernel->ir.get(), &oss);
+    kernel_ast_string = oss.str();
   }
 
   std::vector<std::uint8_t> compile_config_key;

taichi/analysis/offline_cache_util.h

@@ -6,11 +6,14 @@ namespace taichi {
 namespace lang {
 
 struct CompileConfig;
-class Kernel;
+class Program;
+class IRNode;
 class SNode;
+class Kernel;
 
 std::string get_hashed_offline_cache_key_of_snode(SNode *snode);
 std::string get_hashed_offline_cache_key(CompileConfig *config, Kernel *kernel);
+void gen_offline_cache_key(Program *prog, IRNode *ast, std::ostream *os);
 
 }  // namespace lang
 }  // namespace taichi

taichi/inc/frontend_statements.inc.h

@@ -0,0 +1,15 @@
+PER_STATEMENT(FrontendExternalFuncStmt)
+PER_STATEMENT(FrontendExprStmt)
+PER_STATEMENT(FrontendIfStmt)
+PER_STATEMENT(FrontendForStmt)
+PER_STATEMENT(FrontendPrintStmt)
+PER_STATEMENT(FrontendWhileStmt)
+PER_STATEMENT(FrontendBreakStmt)
+PER_STATEMENT(FrontendContinueStmt)
+PER_STATEMENT(FrontendAllocaStmt)
+PER_STATEMENT(FrontendAssignStmt)
+PER_STATEMENT(FrontendEvalStmt)
+PER_STATEMENT(FrontendSNodeOpStmt)  // activate, deactivate, append, clear
+PER_STATEMENT(FrontendAssertStmt)
+PER_STATEMENT(FrontendFuncDefStmt)
+PER_STATEMENT(FrontendReturnStmt)

taichi/inc/statements.inc.h

@@ -1,19 +1,5 @@
 // Frontend statements
-PER_STATEMENT(FrontendExternalFuncStmt)
-PER_STATEMENT(FrontendExprStmt)
-PER_STATEMENT(FrontendIfStmt)
-PER_STATEMENT(FrontendForStmt)
-PER_STATEMENT(FrontendPrintStmt)
-PER_STATEMENT(FrontendWhileStmt)
-PER_STATEMENT(FrontendBreakStmt)
-PER_STATEMENT(FrontendContinueStmt)
-PER_STATEMENT(FrontendAllocaStmt)
-PER_STATEMENT(FrontendAssignStmt)
-PER_STATEMENT(FrontendEvalStmt)
-PER_STATEMENT(FrontendSNodeOpStmt)  // activate, deactivate, append, clear
-PER_STATEMENT(FrontendAssertStmt)
-PER_STATEMENT(FrontendFuncDefStmt)
-PER_STATEMENT(FrontendReturnStmt)
+#include "frontend_statements.inc.h"
 
 // Middle-end statement
 

taichi/ir/expression_printer.h

@@ -18,9 +18,8 @@ class ExpressionPrinter : public ExpressionVisitor {
     os_ = os;
   }
 
-  std::ostream &get_ostream() {
-    TI_ASSERT(os_);
-    return *os_;
+  std::ostream *get_ostream() {
+    return os_;
   }
 
  private:
@@ -227,7 +226,8 @@ class ExpressionHumanFriendlyPrinter : public ExpressionPrinter {
  protected:
   template <typename... Args>
   void emit(Args &&...args) {
-    (this->get_ostream() << ... << std::forward<Args>(args));
+    TI_ASSERT(this->get_ostream());
+    (*this->get_ostream() << ... << std::forward<Args>(args));
   }
 
   template <typename T>

taichi/ir/transforms.h

@@ -47,7 +47,6 @@ void full_simplify(IRNode *root,
                    const CompileConfig &config,
                    const FullSimplifyPass::Args &args);
 void print(IRNode *root, std::string *output = nullptr);
-void gen_offline_cache_key(Program *program, IRNode *root, std::string *output);
 void frontend_type_check(IRNode *root);
 void lower_ast(IRNode *root);
 void type_check(IRNode *root, const CompileConfig &config);

taichi/transforms/ir_printer.cpp

@@ -805,12 +805,6 @@ void print(IRNode *root, std::string *output) {
   return IRPrinter::run(&expr_printer, root, output);
 }
 
-void gen_offline_cache_key(Program *prog, IRNode *root, std::string *output) {
-  irpass::re_id(root);
-  ExpressionOfflineCacheKeyGenerator cache_key_generator(prog);
-  return IRPrinter::run(&cache_key_generator, root, output);
-}
-
 }  // namespace irpass
 
 TLANG_NAMESPACE_END