From d763aaf30491dc8cb58b95c0c8f97ad3e9190261 Mon Sep 17 00:00:00 2001
From: leavelet <leavelet@163.com>
Date: Wed, 19 Nov 2025 14:16:45 +0000
Subject: [PATCH] feat: add kernel preload

---
 csrc/jit/cache.hpp    | 43 +++++++++++++++++++++++++++++++++++++++++++
 csrc/jit/compiler.hpp |  7 +++++++
 deep_gemm/__init__.py | 11 +++++++++++
 3 files changed, 61 insertions(+)

diff --git a/csrc/jit/cache.hpp b/csrc/jit/cache.hpp
index 1e8659f..c063d5c 100644
--- a/csrc/jit/cache.hpp
+++ b/csrc/jit/cache.hpp
@@ -5,6 +5,7 @@
 #include <unordered_map>
 
 #include "kernel_runtime.hpp"
+#include "../utils/system.hpp"
 
 namespace deep_gemm {
 
@@ -24,6 +25,48 @@ public:
             return cache[dir_path] = std::make_shared<KernelRuntime>(dir_path);
         return nullptr;
     }
+
+    // Preload all cached kernels from disk into memory
+    void preload_all(const std::filesystem::path& cache_dir) {
+        const auto kernel_cache_dir = cache_dir / "cache";
+        
+        // If cache directory doesn't exist, return early
+        if (!std::filesystem::exists(kernel_cache_dir))
+            return;
+            
+        int loaded_count = 0;
+        int total_count = 0;
+        
+        // Iterate through all kernel subdirectories in the cache
+        for (const auto& entry : std::filesystem::directory_iterator(kernel_cache_dir)) {
+            if (!entry.is_directory())
+                continue;
+                
+            total_count++;
+            const auto& dir_path = entry.path();
+            
+            // Skip if already in memory cache
+            if (cache.find(dir_path) != cache.end())
+                continue;
+            
+            // Check validity and load the kernel
+            if (KernelRuntime::check_validity(dir_path)) {
+                try {
+                    cache[dir_path] = std::make_shared<KernelRuntime>(dir_path);
+                    loaded_count++;
+                } catch (const std::exception& e) {
+                    // Skip failed loads
+                    if (get_env<int>("DG_JIT_DEBUG"))
+                        printf("Warning: Failed to preload kernel from %s: %s\n", 
+                               dir_path.c_str(), e.what());
+                }
+            }
+        }
+        
+        if (get_env<int>("DG_JIT_DEBUG") || loaded_count > 0)
+            printf("DeepGEMM: Preloaded %d/%d cached kernels from %s\n", 
+                   loaded_count, total_count, kernel_cache_dir.c_str());
+    }
 };
 
 static auto kernel_runtime_cache = std::make_shared<KernelRuntimeCache>();
diff --git a/csrc/jit/compiler.hpp b/csrc/jit/compiler.hpp
index 55bcc60..52405ab 100644
--- a/csrc/jit/compiler.hpp
+++ b/csrc/jit/compiler.hpp
@@ -42,6 +42,13 @@ public:
         Compiler::library_version = get_library_version();
     }
 
+    static std::filesystem::path get_cache_dir() {
+        std::filesystem::path cache_dir = std::filesystem::path(get_env<std::string>("HOME")) / ".deep_gemm";
+        if (const auto& env_cache_dir = get_env<std::string>("DG_JIT_CACHE_DIR"); !env_cache_dir.empty())
+            cache_dir = env_cache_dir;
+        return cache_dir;
+    }
+
     std::string signature, flags;
     std::filesystem::path cache_dir_path;
 
diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py
index df01daf..f088511 100644
--- a/deep_gemm/__init__.py
+++ b/deep_gemm/__init__.py
@@ -137,3 +137,14 @@ _ensure_initialized()
 
 if __debug__:
     _verify_ops_loaded()
+    
+# Preload cached kernels if enabled via environment variable
+if os.environ.get('DG_PRELOAD_KERNELS', '0') == '1':
+    try:
+        deep_gemm_cpp.preload_kernels()
+    except Exception as e:
+        # Preloading failure should not block initialization
+        import warnings
+        warnings.warn(f"Failed to preload DeepGEMM kernels: {e}")
+
+__version__ = '2.1.2'
\ No newline at end of file