From e038020dc50570f8bd8037536ad50176487e0435 Mon Sep 17 00:00:00 2001
From: Mingrui Zhang <33411325+erizmr@users.noreply.github.com>
Date: Thu, 19 Dec 2024 19:27:38 +0800
Subject: [PATCH] [lang] Accelerate compilation by improving key generation
 (#8604)

---
 taichi/analysis/gen_offline_cache_key.cpp | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/taichi/analysis/gen_offline_cache_key.cpp b/taichi/analysis/gen_offline_cache_key.cpp
index 62eca26d378ba..592908f6f239f 100644
--- a/taichi/analysis/gen_offline_cache_key.cpp
+++ b/taichi/analysis/gen_offline_cache_key.cpp
@@ -445,7 +445,14 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
     // Note: The result of serializing snode_tree_roots_ is not parsable now
     emit(static_cast<std::size_t>(snode_tree_roots_.size()));
     for (const auto *snode : snode_tree_roots_) {
-      auto key = get_hashed_offline_cache_key_of_snode(snode);
+      std::string key;
+      if (snode_key_cache_.find(snode) == snode_key_cache_.end()) {
+        key = get_hashed_offline_cache_key_of_snode(snode);
+        snode_key_cache_[snode] = key;
+      } else {
+        key = snode_key_cache_[snode];
+      }
+      snode_key_cache_[snode] = key;
       emit_bytes(key.c_str(), key.size());
     }
 
@@ -655,6 +662,7 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
 
   std::ostream *os_{nullptr};
   std::vector<const SNode *> snode_tree_roots_;
+  std::unordered_map<const SNode *, std::string> snode_key_cache_;
   std::map<Function *, std::size_t> real_funcs_;
   std::vector<char> string_pool_;
 };