Implement CHIP_MODULE_CACHE_DIR

src/CHIPDriver.hh

@@ -237,6 +237,7 @@ private:
   int L0CollectEventsTimeout_ = 0;
   bool OCLDisableQueueProfiling_ = false;
   std::optional<std::string> OclUseAllocStrategy_;
+  std::optional<std::string> ModuleCacheDir_;
 
 public:
   EnvVars() {
@@ -264,6 +265,9 @@ public:
   const std::optional<std::string> &getOclUseAllocStrategy() const noexcept {
     return OclUseAllocStrategy_;
   }
+  const std::optional<std::string> &getModuleCacheDir() const {
+    return ModuleCacheDir_;
+  }
 
 private:
   void parseEnvironmentVariables() {
@@ -299,6 +303,11 @@ private:
         readEnvVar("CHIP_OCL_USE_ALLOC_STRATEGY", value, true)
             ? value
             : OclUseAllocStrategy_;
+    if (readEnvVar("CHIP_MODULE_CACHE_DIR", value, true)) {
+      ModuleCacheDir_ = value; // If set (even if empty), use the value
+    } else {
+      ModuleCacheDir_ = "/tmp"; // If not set, default to "/tmp"
+    }
   }
 
   int parseInt(const std::string &value) {
@@ -343,9 +352,11 @@ private:
     logInfo("CHIP_OCL_USE_ALLOC_STRATEGY={}", OclUseAllocStrategy_.has_value()
                                                   ? OclUseAllocStrategy_.value()
                                                   : "off");
+    logInfo("CHIP_MODULE_CACHE_DIR={}",
+            ModuleCacheDir_.has_value() ? ModuleCacheDir_.value() : "off");
   }
 };
 
 extern EnvVars ChipEnvVars;
 
-#endif
+#endif

src/backend/Level0/CHIPBackendLevel0.cc

@@ -2371,7 +2371,14 @@ void save(const ze_module_desc_t &desc, const ze_module_handle_t &module) {
   }
 
   size_t hash = hasher(combinedInput);
-  std::string fullPath = "/tmp/chipstar_module_cache_" + std::to_string(hash);
+
+  if (!ChipEnvVars.getModuleCacheDir().has_value()) {
+    logTrace("Module caching is disabled");
+    return;
+  }
+
+  std::string cacheDir = ChipEnvVars.getModuleCacheDir().value();
+  std::string fullPath = cacheDir + "/chipstar_module_cache_" + std::to_string(hash);
 
   size_t binarySize;
   zeStatus = zeModuleGetNativeBinary(module, &binarySize, nullptr);
@@ -2410,7 +2417,13 @@ bool load(ze_module_desc_t &desc) {
   }
 
   size_t hash = hasher(combinedInput);
-  std::string fullPath = "/tmp/chipstar_module_cache_" + std::to_string(hash);
+
+  if (!ChipEnvVars.getModuleCacheDir().has_value()) {
+    return false;
+  }
+
+  std::string cacheDir = ChipEnvVars.getModuleCacheDir().value();
+  std::string fullPath = cacheDir + "/chipstar_module_cache_" + std::to_string(hash);
   // Open the binary file
   std::ifstream inFile(fullPath, std::ios::in | std::ios::binary);
   if (!inFile) {

src/backend/OpenCL/CHIPBackendOpenCL.cc

@@ -915,6 +915,14 @@ static void appendRuntimeObjects(cl::Context Ctx, CHIPDeviceOpenCL &ChipDev,
 }
 
 static void save(const cl::Program &program, const std::string &cacheName) {
+  if (!ChipEnvVars.getModuleCacheDir().has_value()) {
+    logTrace("Module caching is disabled");
+    return;
+  }
+
+  std::string cacheDir = ChipEnvVars.getModuleCacheDir().value();
+  std::string fullPath = cacheDir + "/" + cacheName;
+
   // Step 1: Get the sizes of the binaries for each device
   std::vector<size_t> binarySizes;
   program.getInfo(CL_PROGRAM_BINARY_SIZES, &binarySizes);
@@ -951,7 +959,7 @@ static void save(const cl::Program &program, const std::string &cacheName) {
   }
 
   // Step 4: Write the binaries to the output file
-  std::ofstream outFile(cacheName, std::ios::out | std::ios::binary);
+  std::ofstream outFile(fullPath, std::ios::out | std::ios::binary);
   if (!outFile) {
     logError("Failed to open file for writing kernel binary");
     // Clean up allocated memory
@@ -980,7 +988,7 @@ static void save(const cl::Program &program, const std::string &cacheName) {
   outFile.close();
 
   // Step 5: Verify the file size
-  std::ifstream inFile(cacheName,
+  std::ifstream inFile(fullPath,
                        std::ios::in | std::ios::binary | std::ios::ate);
   if (!inFile) {
     logError("Failed to open file for reading to verify size");
@@ -1003,7 +1011,7 @@ static void save(const cl::Program &program, const std::string &cacheName) {
     return;
   }
 
-  logTrace("Kernel binary cached as {}", cacheName);
+  logTrace("Kernel binary cached as {}", fullPath);
   logTrace("Number of binaries: {}", numDevices);
 
   // Step 6: Clean up allocated memory
@@ -1014,9 +1022,16 @@ static void save(const cl::Program &program, const std::string &cacheName) {
 
 static bool load(cl::Context &context, const std::vector<cl::Device> &devices,
                  const std::string &cacheName, cl::Program &program) {
-  logTrace("Loading kernel binary from cache at {}", cacheName);
+  if (!ChipEnvVars.getModuleCacheDir().has_value()) {
+    return false;
+  }
+
+  std::string cacheDir = ChipEnvVars.getModuleCacheDir().value();
+  std::string fullPath = cacheDir + "/" + cacheName;
+
+  logTrace("Loading kernel binary from cache at {}", fullPath);
 
-  std::ifstream inFile(cacheName,
+  std::ifstream inFile(fullPath,
                        std::ios::in | std::ios::binary | std::ios::ate);
   if (!inFile) {
     return false;
@@ -1092,14 +1107,14 @@ static bool load(cl::Context &context, const std::vector<cl::Device> &devices,
     return false;
   }
 
-  logTrace("Kernel binary loaded from cache as {}", cacheName);
+  logTrace("Kernel binary loaded from cache as {}", fullPath);
   return true;
 }
 
 std::string generateCacheName(const std::string strIn) {
   std::hash<std::string> hasher;
   size_t hash = hasher(strIn);
-  return "/tmp/chipstar_kernel_cache_" + std::to_string(hash);
+  return "chipstar_kernel_cache_" + std::to_string(hash);
 }
 
 void CHIPModuleOpenCL::compile(chipstar::Device *ChipDev) {