[refactor] Remove dependencies on Callable::program in lang::get_hash…

…ed_offline_cache_key (taichi-dev#7287)

Issue: taichi-dev#7286
* taichi-dev#7286: Part of taichi-dev#7002

---------

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

taichi/analysis/gen_offline_cache_key.cpp

@@ -53,8 +53,7 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
   using IRVisitor::visit;
 
  public:
-  ASTSerializer(Program *prog, std::ostream *os)
-      : ExpressionVisitor(true), prog_(prog), os_(os) {
+  explicit ASTSerializer(std::ostream *os) : ExpressionVisitor(true), os_(os) {
     // TODO(PGZXB): Set allow_undefined_visitor as false. (blocked by
     // constant-folding)
     this->allow_undefined_visitor = true;
@@ -414,8 +413,8 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
     emit(stmt->outputs);
   }
 
-  static void run(Program *prog, IRNode *ast, std::ostream *os) {
-    ASTSerializer serializer(prog, os);
+  static void run(IRNode *ast, std::ostream *os) {
+    ASTSerializer serializer(os);
     ast->accept(&serializer);
     serializer.emit_dependencies();
   }
@@ -434,7 +433,7 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
     // Serialize snode_trees(Temporary: using offline-cache-key of SNode)
     // Note: The result of serializing snode_tree_roots_ is not parsable now
     emit(static_cast<std::size_t>(snode_tree_roots_.size()));
-    for (auto *snode : snode_tree_roots_) {
+    for (const auto *snode : snode_tree_roots_) {
       auto key = get_hashed_offline_cache_key_of_snode(snode);
       emit_bytes(key.c_str(), key.size());
     }
@@ -517,12 +516,11 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
     }
   }
 
-  void emit(SNode *snode) {
-    TI_ASSERT(prog_);
+  void emit(const SNode *snode) {
     if (snode) {
       emit(static_cast<std::size_t>(snode->get_snode_tree_id()));
       emit(static_cast<std::size_t>(snode->id));
-      auto *root = prog_->get_snode_root(snode->get_snode_tree_id());
+      const auto *root = snode->get_root();
       snode_tree_roots_.insert(root);
     } else {
       emit(std::numeric_limits<std::size_t>::max());
@@ -643,17 +641,16 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
 
 #undef DEFINE_EMIT_ENUM
 
-  Program *prog_{nullptr};
   std::ostream *os_{nullptr};
-  std::unordered_set<SNode *> snode_tree_roots_;
+  std::unordered_set<const SNode *> snode_tree_roots_;
   std::unordered_map<Function *, std::size_t> real_funcs_;
   std::vector<char> string_pool_;
 };
 
 }  // namespace
 
-void gen_offline_cache_key(Program *prog, IRNode *ast, std::ostream *os) {
-  ASTSerializer::run(prog, ast, os);
+void gen_offline_cache_key(IRNode *ast, std::ostream *os) {
+  ASTSerializer::run(ast, os);
 }
 
 }  // namespace taichi::lang

taichi/analysis/offline_cache_util.cpp

@@ -68,7 +68,7 @@ static std::vector<std::uint8_t> get_offline_cache_key_of_compile_config(
 }
 
 static void get_offline_cache_key_of_snode_impl(
-    SNode *snode,
+    const SNode *snode,
     BinaryOutputSerializer &serializer,
     std::unordered_set<int> &visited) {
   if (auto iter = visited.find(snode->id); iter != visited.end()) {
@@ -122,7 +122,7 @@ static void get_offline_cache_key_of_snode_impl(
   serializer(snode->get_snode_tree_id());
 }
 
-std::string get_hashed_offline_cache_key_of_snode(SNode *snode) {
+std::string get_hashed_offline_cache_key_of_snode(const SNode *snode) {
   TI_ASSERT(snode);
 
   BinaryOutputSerializer serializer;
@@ -145,7 +145,7 @@ std::string get_hashed_offline_cache_key(const CompileConfig &config,
   std::string kernel_ast_string;
   if (kernel) {
     std::ostringstream oss;
-    gen_offline_cache_key(kernel->program, kernel->ir.get(), &oss);
+    gen_offline_cache_key(kernel->ir.get(), &oss);
     kernel_ast_string = oss.str();
   }
 

taichi/analysis/offline_cache_util.h

@@ -12,9 +12,9 @@ class IRNode;
 class SNode;
 class Kernel;
 
-std::string get_hashed_offline_cache_key_of_snode(SNode *snode);
+std::string get_hashed_offline_cache_key_of_snode(const SNode *snode);
 std::string get_hashed_offline_cache_key(const CompileConfig &config,
                                          Kernel *kernel);
-void gen_offline_cache_key(Program *prog, IRNode *ast, std::ostream *os);
+void gen_offline_cache_key(IRNode *ast, std::ostream *os);
 
 }  // namespace taichi::lang

taichi/ir/snode.cpp

@@ -330,4 +330,11 @@ int SNode::get_snode_tree_id() const {
   return snode_tree_id_;
 }
 
+const SNode *SNode::get_root() const {
+  if (!parent) {  // root->parent == nullptr
+    return this;
+  }
+  return parent->get_root();
+}
+
 }  // namespace taichi::lang

taichi/ir/snode.h

@@ -327,6 +327,8 @@ class SNode {
 
   int get_snode_tree_id() const;
 
+  const SNode *get_root() const;
+
   static void reset_counter() {
     counter = 0;
   }

taichi/program/function.cpp

@@ -19,7 +19,7 @@ void Function::set_function_body(const std::function<void()> &func) {
 
   if (program->compile_config().offline_cache) {  // For generating AST-Key
     std::ostringstream oss;
-    gen_offline_cache_key(program, ir.get(), &oss);
+    gen_offline_cache_key(ir.get(), &oss);
     ast_serialization_data_ = oss.str();
   }
 }