Added a program cache (per-context) next to the per-device binary cache

1 files changed, 44 insertions, 4 deletions

diff --git a/src/cache.cc b/src/cache.cc
index 18731a51..4dbdb711 100644
--- a/src/cache.cc
+++ b/src/cache.cc
@@ -7,7 +7,7 @@
 // Author(s):
 //   Cedric Nugteren <www.cedricnugteren.nl>
 //
-// This file implements the caching functionality of compiled binaries.
+// This file implements the caching functionality of compiled binaries and programs.
 //
 // =================================================================================================
 
@@ -22,13 +22,21 @@ namespace cache {
 // =================================================================================================
 
 // Stores the compiled binary or IR in the cache
-void StoreBinaryToCache(const std::string& binary, const std::string &device_name,
+void StoreBinaryToCache(const std::string &binary, const std::string &device_name,
                         const Precision &precision, const std::string &routine_name) {
   binary_cache_mutex_.lock();
-  binary_cache_.push_back({binary, device_name, precision, routine_name});
+  binary_cache_.push_back(BinaryCache{binary, device_name, precision, routine_name});
   binary_cache_mutex_.unlock();
 }
 
+// Stores the compiled program in the cache
+void StoreProgramToCache(const Program &program, const Context &context,
+                         const Precision &precision, const std::string &routine_name) {
+  program_cache_mutex_.lock();
+  program_cache_.push_back(ProgramCache{program, context.pointer(), precision, routine_name});
+  program_cache_mutex_.unlock();
+}
+
 // Queries the cache and retrieves a matching binary. Assumes that the match is available, throws
 // otherwise.
 const std::string& GetBinaryFromCache(const std::string &device_name, const Precision &precision,
@@ -44,6 +52,21 @@ const std::string& GetBinaryFromCache(const std::string &device_name, const Prec
   throw std::runtime_error("Internal CLBlast error: Expected binary in cache, but found none.");
 }
 
+// Queries the cache and retrieves a matching program. Assumes that the match is available, throws
+// otherwise.
+const Program& GetProgramFromCache(const Context &context, const Precision &precision,
+                                   const std::string &routine_name) {
+  program_cache_mutex_.lock();
+  for (auto &cached_program: program_cache_) {
+    if (cached_program.MatchInCache(context.pointer(), precision, routine_name)) {
+      program_cache_mutex_.unlock();
+      return cached_program.program;
+    }
+  }
+  program_cache_mutex_.unlock();
+  throw std::runtime_error("Internal CLBlast error: Expected program in cache, but found none.");
+}
+
 // Queries the cache to see whether or not the compiled kernel is already there
 bool BinaryIsInCache(const std::string &device_name, const Precision &precision,
                      const std::string &routine_name) {
@@ -58,13 +81,30 @@ bool BinaryIsInCache(const std::string &device_name, const Precision &precision,
   return false;
 }
 
+// Queries the cache to see whether or not the compiled kernel is already there
+bool ProgramIsInCache(const Context &context, const Precision &precision,
+                      const std::string &routine_name) {
+  program_cache_mutex_.lock();
+  for (auto &cached_program: program_cache_) {
+    if (cached_program.MatchInCache(context.pointer(), precision, routine_name)) {
+      program_cache_mutex_.unlock();
+      return true;
+    }
+  }
+  program_cache_mutex_.unlock();
+  return false;
+}
+
 // =================================================================================================
 
-// Clears the cache of stored binaries
+// Clears the cache of stored binaries and programs
 StatusCode ClearCache() {
   binary_cache_mutex_.lock();
   binary_cache_.clear();
   binary_cache_mutex_.unlock();
+  program_cache_mutex_.lock();
+  program_cache_.clear();
+  program_cache_mutex_.unlock();
   return StatusCode::kSuccess;
 }