[SYCL] Refactor program build and cache (#8104)

In order to help device_global find associated programs in the cache,
the way the program manager and the cache interacts needs a refactoring.
This commit makes this refactoring.

Additionally this adds the notion of a "common key" (CommonProgramKeyT)
representing the common parts of the full cache key shared between
builds of the same program on the same device, excluding specializations
such as specialization constants and compilation options.

---------

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/source/detail/kernel_program_cache.cpp

@@ -14,7 +14,7 @@ namespace sycl {
 __SYCL_INLINE_VER_NAMESPACE(_V1) {
 namespace detail {
 KernelProgramCache::~KernelProgramCache() {
-  for (auto &ProgIt : MCachedPrograms) {
+  for (auto &ProgIt : MCachedPrograms.Cache) {
     ProgramWithBuildStateT &ProgWithState = ProgIt.second;
     PiProgramT *ToBeDeleted = ProgWithState.Ptr.load();
 

sycl/source/detail/kernel_program_cache.hpp

@@ -39,13 +39,16 @@ class KernelProgramCache {
     bool isFilledIn() const { return !Msg.empty(); }
   };
 
+  /// Denotes the state of a build.
+  enum BuildState { BS_InProgress, BS_Done, BS_Failed };
+
   /// Denotes pointer to some entity with its general state and build error.
   /// The pointer is not null if and only if the entity is usable.
   /// State of the entity is provided by the user of cache instance.
   /// Currently there is only a single user - ProgramManager class.
   template <typename T> struct BuildResult {
     std::atomic<T *> Ptr;
-    std::atomic<int> State;
+    std::atomic<BuildState> State;
     BuildError Error;
 
     /// Condition variable to signal that build result is ready.
@@ -62,15 +65,23 @@ class KernelProgramCache {
     /// A mutex to be employed along with MBuildCV.
     std::mutex MBuildResultMutex;
 
-    BuildResult(T *P, int S) : Ptr{P}, State{S}, Error{"", 0} {}
+    BuildResult(T *P, BuildState S) : Ptr{P}, State{S}, Error{"", 0} {}
   };
 
   using PiProgramT = std::remove_pointer<RT::PiProgram>::type;
   using PiProgramPtrT = std::atomic<PiProgramT *>;
   using ProgramWithBuildStateT = BuildResult<PiProgramT>;
   using ProgramCacheKeyT = std::pair<std::pair<SerializedObj, std::uintptr_t>,
                                      std::pair<RT::PiDevice, std::string>>;
-  using ProgramCacheT = std::map<ProgramCacheKeyT, ProgramWithBuildStateT>;
+  using CommonProgramKeyT = std::pair<std::uintptr_t, RT::PiDevice>;
+
+  struct ProgramCache {
+    std::map<ProgramCacheKeyT, ProgramWithBuildStateT> Cache;
+    std::multimap<CommonProgramKeyT, ProgramCacheKeyT> KeyMap;
+
+    size_t size() const noexcept { return Cache.size(); }
+  };
+
   using ContextPtr = context_impl *;
 
   using PiKernelT = std::remove_pointer<RT::PiKernel>::type;
@@ -91,21 +102,66 @@ class KernelProgramCache {
 
   void setContextPtr(const ContextPtr &AContext) { MParentContext = AContext; }
 
-  Locked<ProgramCacheT> acquireCachedPrograms() {
+  Locked<ProgramCache> acquireCachedPrograms() {
     return {MCachedPrograms, MProgramCacheMutex};
   }
 
   Locked<KernelCacheT> acquireKernelsPerProgramCache() {
     return {MKernelsPerProgramCache, MKernelsPerProgramCacheMutex};
   }
 
+  std::pair<ProgramWithBuildStateT *, bool>
+  getOrInsertProgram(const ProgramCacheKeyT &CacheKey) {
+    auto LockedCache = acquireCachedPrograms();
+    auto &ProgCache = LockedCache.get();
+    auto Inserted = ProgCache.Cache.emplace(
+        std::piecewise_construct, std::forward_as_tuple(CacheKey),
+        std::forward_as_tuple(nullptr, BS_InProgress));
+    if (Inserted.second) {
+      // Save reference between the common key and the full key.
+      CommonProgramKeyT CommonKey =
+          std::make_pair(CacheKey.first.second, CacheKey.second.first);
+      ProgCache.KeyMap.emplace(std::piecewise_construct,
+                               std::forward_as_tuple(CommonKey),
+                               std::forward_as_tuple(CacheKey));
+    }
+    return std::make_pair(&Inserted.first->second, Inserted.second);
+  }
+
+  std::pair<KernelWithBuildStateT *, bool>
+  getOrInsertKernel(RT::PiProgram Program, const std::string &KernelName) {
+    auto LockedCache = acquireKernelsPerProgramCache();
+    auto &Cache = LockedCache.get()[Program];
+    auto Inserted = Cache.emplace(
+        std::piecewise_construct, std::forward_as_tuple(KernelName),
+        std::forward_as_tuple(nullptr, BS_InProgress));
+    return std::make_pair(&Inserted.first->second, Inserted.second);
+  }
+
   template <typename T, class Predicate>
   void waitUntilBuilt(BuildResult<T> &BR, Predicate Pred) const {
     std::unique_lock<std::mutex> Lock(BR.MBuildResultMutex);
 
     BR.MBuildCV.wait(Lock, Pred);
   }
 
+  template <typename ExceptionT, typename RetT>
+  RetT *waitUntilBuilt(BuildResult<RetT> *BuildResult) {
+    // Any thread which will find nullptr in cache will wait until the pointer
+    // is not null anymore.
+    waitUntilBuilt(*BuildResult, [BuildResult]() {
+      int State = BuildResult->State.load();
+      return State == BuildState::BS_Done || State == BuildState::BS_Failed;
+    });
+
+    if (BuildResult->Error.isFilledIn()) {
+      const BuildError &Error = BuildResult->Error;
+      throw ExceptionT(Error.Msg, Error.Code);
+    }
+
+    return BuildResult->Ptr.load();
+  }
+
   template <typename T> void notifyAllBuild(BuildResult<T> &BR) const {
     BR.MBuildCV.notify_all();
   }
@@ -132,7 +188,7 @@ class KernelProgramCache {
   ///
   /// This member function should only be used in unit tests.
   void reset() {
-    MCachedPrograms = ProgramCacheT{};
+    MCachedPrograms = ProgramCache{};
     MKernelsPerProgramCache = KernelCacheT{};
     MKernelFastCache = KernelFastCacheT{};
   }
@@ -141,7 +197,7 @@ class KernelProgramCache {
   std::mutex MProgramCacheMutex;
   std::mutex MKernelsPerProgramCacheMutex;
 
-  ProgramCacheT MCachedPrograms;
+  ProgramCache MCachedPrograms;
   KernelCacheT MKernelsPerProgramCache;
   ContextPtr MParentContext;