[SYCL] Intoduce thread-safety for kernels and programs cache (#866)

Signed-off-by: Sergey Kanaev <[email protected]>

Author: s-kanaev

sycl/include/CL/sycl/detail/context_impl.hpp

@@ -9,6 +9,7 @@
 #pragma once
 #include <CL/sycl/detail/common.hpp>
 #include <CL/sycl/detail/device_impl.hpp>
+#include <CL/sycl/detail/kernel_program_cache.hpp>
 #include <CL/sycl/detail/os_util.hpp>
 #include <CL/sycl/detail/pi.hpp>
 #include <CL/sycl/detail/platform_impl.hpp>
@@ -110,21 +111,6 @@ class context_impl {
     return MDevices;
   }
 
-  /// Gets cached programs.
-  ///
-  /// @return a map of cached programs.
-  std::map<KernelSetId, RT::PiProgram> &getCachedPrograms() {
-    return MCachedPrograms;
-  }
-
-  /// Gets cached kernels.
-  ///
-  /// @return a map of cached kernels.
-  std::map<RT::PiProgram, std::map<string_class, RT::PiKernel>> &
-  getCachedKernels() {
-    return MCachedKernels;
-  }
-
   /// Gets USM dispatcher.
   ///
   /// @return a pointer to USM dispatcher.
@@ -146,17 +132,18 @@ class context_impl {
     return MCachedLibPrograms;
   }
 
+  KernelProgramCache &getKernelProgramCache() const;
+
 private:
   async_handler MAsyncHandler;
   vector_class<device> MDevices;
   RT::PiContext MContext;
   PlatformImplPtr MPlatform;
   bool MPluginInterop;
   bool MHostContext;
-  std::map<KernelSetId, RT::PiProgram> MCachedPrograms;
-  std::map<RT::PiProgram, std::map<string_class, RT::PiKernel>> MCachedKernels;
   std::shared_ptr<usm::USMDispatcher> MUSMDispatch;
   std::map<DeviceLibExt, RT::PiProgram> MCachedLibPrograms;
+  mutable KernelProgramCache MKernelProgramCache;
 };
 
 } // namespace detail

sycl/include/CL/sycl/detail/kernel_program_cache.hpp

@@ -0,0 +1,84 @@
+//==--- kernel_program_cache.hpp - Cache for kernel and program -*- C++-*---==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <CL/sycl/detail/locked.hpp>
+#include <CL/sycl/detail/os_util.hpp>
+#include <CL/sycl/detail/pi.hpp>
+
+#include <atomic>
+#include <condition_variable>
+#include <map>
+#include <mutex>
+#include <type_traits>
+
+namespace cl {
+namespace sycl {
+namespace detail {
+class KernelProgramCache {
+public:
+  /// Denotes pointer to some entity with its state.
+  /// 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 EntityWithState {
+    std::atomic<T *> Ptr;
+    std::atomic<int> State;
+
+    EntityWithState(T* P, int S)
+      : Ptr{P}, State{S}
+    {}
+  };
+
+  using PiProgramT = std::remove_pointer<RT::PiProgram>::type;
+  using PiProgramPtrT = std::atomic<PiProgramT *>;
+  using ProgramWithBuildStateT = EntityWithState<PiProgramT>;
+  using ProgramCacheT = std::map<OSModuleHandle, ProgramWithBuildStateT>;
+
+  using PiKernelT = std::remove_pointer<RT::PiKernel>::type;
+  using PiKernelPtrT = std::atomic<PiKernelT *>;
+  using KernelWithBuildStateT = EntityWithState<PiKernelT>;
+  using KernelByNameT = std::map<string_class, KernelWithBuildStateT>;
+  using KernelCacheT = std::map<RT::PiProgram, KernelByNameT>;
+
+  ~KernelProgramCache();
+
+  Locked<ProgramCacheT> acquireCachedPrograms() {
+    return {MCachedPrograms, MProgramCacheMutex};
+  }
+
+  Locked<KernelCacheT> acquireKernelsPerProgramCache() {
+    return {MKernelsPerProgramCache, MKernelsPerProgramCacheMutex};
+  }
+
+  template<class Predicate>
+  void waitUntilBuilt(Predicate Pred) const {
+    std::unique_lock<std::mutex> Lock(MBuildCVMutex);
+
+    MBuildCV.wait(Lock, Pred);
+  }
+
+  void notifyAllBuild() const {
+    MBuildCV.notify_all();
+  }
+
+private:
+  std::mutex MProgramCacheMutex;
+  std::mutex MKernelsPerProgramCacheMutex;
+
+  mutable std::condition_variable MBuildCV;
+  mutable std::mutex MBuildCVMutex;
+
+  ProgramCacheT MCachedPrograms;
+  KernelCacheT MKernelsPerProgramCache;
+};
+}
+}
+}

sycl/include/CL/sycl/detail/locked.hpp

@@ -0,0 +1,37 @@
+//==---------------- locked.hpp - Reference with lock -----------*- C++-*---==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <functional>
+#include <mutex>
+
+namespace cl {
+namespace sycl {
+namespace detail {
+  /// Represents a reference to value with appropriate lock acquired.
+  /// Employed for acquire/release logic. Acquire action is construction
+  /// of instance of locked<>. Release action is destruction of instance of
+  /// locked<>.
+  template <typename T>
+  class Locked {
+    std::reference_wrapper<T> m_Value;
+    std::unique_lock<std::mutex> m_Lock;
+
+  public:
+    Locked(T &v, std::mutex &mutex)
+      : m_Value{v}, m_Lock{mutex}
+    {}
+
+    T& get() const {
+      return m_Value.get();
+    }
+  };
+}
+}
+}

sycl/source/CMakeLists.txt

@@ -61,6 +61,7 @@ set(SYCL_SOURCES
     "detail/image_impl.cpp"
     "detail/kernel_impl.cpp"
     "detail/kernel_info.cpp"
+    "detail/kernel_program_cache.cpp"
     "detail/memory_manager.cpp"
     "detail/platform_impl.cpp"
     "detail/platform_info.cpp"

sycl/source/detail/context_impl.cpp

@@ -83,13 +83,6 @@ context_impl::~context_impl() {
     // TODO catch an exception and put it to list of asynchronous exceptions
     PI_CALL(piContextRelease)(MContext);
   }
-  // Release all programs and kernels created with this context
-  for (auto ProgIt : MCachedPrograms) {
-    RT::PiProgram ToBeDeleted = ProgIt.second;
-    for (auto KernIt : MCachedKernels[ToBeDeleted])
-      PI_CALL(piKernelRelease)(KernIt.second);
-    PI_CALL(piProgramRelease)(ToBeDeleted);
-  }
   for (auto LibProg : MCachedLibPrograms) {
     assert(LibProg.second && "Null program must not be kept in the cache");
     PI_CALL(piProgramRelease)(LibProg.second);
@@ -123,6 +116,10 @@ std::shared_ptr<usm::USMDispatcher> context_impl::getUSMDispatch() const {
   return MUSMDispatch;
 }
 
+KernelProgramCache &context_impl::getKernelProgramCache() const {
+  return MKernelProgramCache;
+}
+
 } // namespace detail
 } // namespace sycl
 } // namespace cl

sycl/source/detail/kernel_program_cache.cpp

@@ -0,0 +1,40 @@
+//==--- kernel_program_cache.cpp - Cache for kernel and program -*- C++-*---==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <CL/sycl/detail/kernel_program_cache.hpp>
+
+namespace cl {
+namespace sycl {
+namespace detail {
+KernelProgramCache::~KernelProgramCache() {
+  for (auto &ProgIt : MCachedPrograms) {
+    ProgramWithBuildStateT &ProgWithState = ProgIt.second;
+    PiProgramT *ToBeDeleted = ProgWithState.Ptr.load();
+
+    if (!ToBeDeleted)
+      continue;
+
+    auto KernIt = MKernelsPerProgramCache.find(ToBeDeleted);
+
+    if (KernIt == MKernelsPerProgramCache.end())
+      continue;
+
+    for (auto &p : KernIt->second) {
+      KernelWithBuildStateT &KernelWithState = p.second;
+      PiKernelT *Kern = KernelWithState.Ptr.load();
+
+      if (Kern)
+        PI_CALL(piKernelRelease)(Kern);
+    }
+
+    PI_CALL(piProgramRelease)(ToBeDeleted);
+  }
+}
+}
+}
+}