perf(fastrun): cache persistent cache

GitOrigin-RevId: 21f7d4c19d

imperative/python/src/utils.cpp

@@ -237,15 +237,67 @@ void init_utils(py::module m) {
         mgb::sys::TimedFuncInvoker::ins().fork_exec_impl_mainloop(user_data.c_str());
     });
     using mgb::PersistentCache;
-    class PyPersistentCache: public mgb::PersistentCache{
+    class PyPersistentCache: public mgb::PersistentCache {
+    private:
+        using KeyPair = std::pair<std::string, std::string>;
+        using BlobPtr = std::unique_ptr<Blob, void(*)(Blob*)>;
+
+        static size_t hash_key_pair(const KeyPair& kp) {
+            std::hash<std::string> hasher;
+            return hasher(kp.first) ^ hasher(kp.second);
+        }
+
+        std::string blob_to_str(const Blob& key) {
+            return std::string(reinterpret_cast<const char*>(key.ptr), key.size);
+        }
+
+        BlobPtr copy_blob(const Blob& blob) {
+            auto blob_deleter = [](Blob* blob){
+                if (blob) {
+                    std::free(const_cast<void*>(blob->ptr));
+                    delete blob;
+                }
+            };
+            auto blob_ptr = BlobPtr{ new Blob(), blob_deleter };
+            blob_ptr->ptr = std::malloc(blob.size);
+            std::memcpy(const_cast<void*>(blob_ptr->ptr), blob.ptr, blob.size);
+            blob_ptr->size = blob.size;
+            return blob_ptr;
+        }
+
+        BlobPtr str_to_blob(const std::string& str) {
+            auto blob = Blob{ str.data(), str.size() };
+            return copy_blob(blob);
+        }
+
+        std::unique_ptr<Blob, void(*)(Blob*)> empty_blob() {
+            return BlobPtr{ nullptr, [](Blob* blob){} };
+        }
     public:
         mgb::Maybe<Blob> get(const std::string& category, const Blob& key) override {
-            PYBIND11_OVERLOAD_PURE(mgb::Maybe<Blob>, PersistentCache, get, category, key);
+            thread_local std::unordered_map<KeyPair, BlobPtr, mgb::pairhash> m_local_cache;
+            auto py_get = [this](const std::string& category, const Blob& key) -> mgb::Maybe<Blob> {
+                PYBIND11_OVERLOAD_PURE(mgb::Maybe<Blob>, PersistentCache, get, category, key);
+            };
+            KeyPair kp = { category, blob_to_str(key) };
+            auto iter = m_local_cache.find(kp);
+            if (iter == m_local_cache.end()) {
+                auto py_ret = py_get(category, key);
+                if (!py_ret.valid()) {
+                    iter = m_local_cache.insert({kp, empty_blob()}).first;
+                } else {
+                    iter = m_local_cache.insert({kp, copy_blob(py_ret.val())}).first;
+                }
+            }
+            if (iter->second) {
+                return *iter->second;
+            } else {
+                return {};
+            }
         }
         void put(const std::string& category, const Blob& key, const Blob& value) override {
             PYBIND11_OVERLOAD_PURE(void, PersistentCache, put, category, key, value);
         }
-
     };
     py::class_<PersistentCache, PyPersistentCache, std::shared_ptr<PersistentCache>>(m, "PersistentCache")
         .def(py::init<>())