diff --git a/src/CHIPDriver.hh b/src/CHIPDriver.hh index 9126bd5a3..ed8ae8efa 100644 --- a/src/CHIPDriver.hh +++ b/src/CHIPDriver.hh @@ -237,6 +237,7 @@ private: int L0CollectEventsTimeout_ = 0; bool OCLDisableQueueProfiling_ = false; std::optional OclUseAllocStrategy_; + std::optional ModuleCacheDir_; public: EnvVars() { @@ -264,6 +265,9 @@ public: const std::optional &getOclUseAllocStrategy() const noexcept { return OclUseAllocStrategy_; } + const std::optional &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 \ No newline at end of file diff --git a/src/backend/Level0/CHIPBackendLevel0.cc b/src/backend/Level0/CHIPBackendLevel0.cc index 44afa7dcb..306594244 100644 --- a/src/backend/Level0/CHIPBackendLevel0.cc +++ b/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) { diff --git a/src/backend/OpenCL/CHIPBackendOpenCL.cc b/src/backend/OpenCL/CHIPBackendOpenCL.cc index f062708c4..704834c83 100644 --- a/src/backend/OpenCL/CHIPBackendOpenCL.cc +++ b/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 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 &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 &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 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) {