[SYCL] Reduce the time to get a kernel from cache (#4186)

Signed-off-by: Alexander Flegontov <[email protected]>

Author: alexanderfle

sycl/source/detail/kernel_program_cache.hpp

@@ -21,6 +21,9 @@
 #include <mutex>
 #include <type_traits>
 
+// For testing purposes
+class MockKernelProgramCache;
+
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
 namespace detail {
@@ -79,6 +82,13 @@ class KernelProgramCache {
   using KernelByNameT = std::map<std::string, KernelWithBuildStateT>;
   using KernelCacheT = std::map<RT::PiProgram, KernelByNameT>;
 
+  using KernelFastCacheKeyT =
+      std::tuple<SerializedObj, OSModuleHandle, RT::PiDevice, std::string,
+                 std::string>;
+  using KernelFastCacheValT =
+      std::tuple<RT::PiKernel, std::mutex *, RT::PiProgram>;
+  using KernelFastCacheT = std::map<KernelFastCacheKeyT, KernelFastCacheValT>;
+
   ~KernelProgramCache();
 
   void setContextPtr(const ContextPtr &AContext) { MParentContext = AContext; }
@@ -102,13 +112,35 @@ class KernelProgramCache {
     BR.MBuildCV.notify_all();
   }
 
+  template <typename KeyT>
+  KernelFastCacheValT tryToGetKernelFast(KeyT &&CacheKey) {
+    std::unique_lock<std::mutex> Lock(MKernelFastCacheMutex);
+    auto It = MKernelFastCache.find(CacheKey);
+    if (It != MKernelFastCache.end()) {
+      return It->second;
+    }
+    return std::make_tuple(nullptr, nullptr, nullptr);
+  }
+
+  template <typename KeyT, typename ValT>
+  void saveKernel(KeyT &&CacheKey, ValT &&CacheVal) {
+    std::unique_lock<std::mutex> Lock(MKernelFastCacheMutex);
+    // if no insertion took place, thus some other thread has already inserted
+    // smth in the cache
+    MKernelFastCache.emplace(CacheKey, CacheVal);
+  }
+
 private:
   std::mutex MProgramCacheMutex;
   std::mutex MKernelsPerProgramCacheMutex;
 
   ProgramCacheT MCachedPrograms;
   KernelCacheT MKernelsPerProgramCache;
   ContextPtr MParentContext;
+
+  std::mutex MKernelFastCacheMutex;
+  KernelFastCacheT MKernelFastCache;
+  friend class ::MockKernelProgramCache;
 };
 } // namespace detail
 } // namespace sycl

sycl/source/detail/program_impl.cpp

@@ -439,7 +439,7 @@ RT::PiKernel program_impl::get_pi_kernel(const std::string &KernelName) const {
   RT::PiKernel Kernel = nullptr;
 
   if (is_cacheable()) {
-    std::tie(Kernel, std::ignore) =
+    std::tie(Kernel, std::ignore, std::ignore) =
         ProgramManager::getInstance().getOrCreateKernel(
             MProgramModuleHandle, detail::getSyclObjImpl(get_context()),
             detail::getSyclObjImpl(get_devices()[0]), KernelName, this);

sycl/source/detail/program_manager/program_manager.cpp

@@ -519,24 +519,41 @@ RT::PiProgram ProgramManager::getBuiltPIProgram(
   return BuildResult->Ptr.load();
 }
 
-std::pair<RT::PiKernel, std::mutex *> ProgramManager::getOrCreateKernel(
-    OSModuleHandle M, const ContextImplPtr &Context,
-    const DeviceImplPtr &Device, const std::string &KernelName,
-    const program_impl *Prg) {
+std::tuple<RT::PiKernel, std::mutex *, RT::PiProgram>
+ProgramManager::getOrCreateKernel(OSModuleHandle M,
+                                  const ContextImplPtr &ContextImpl,
+                                  const DeviceImplPtr &DeviceImpl,
+                                  const std::string &KernelName,
+                                  const program_impl *Prg) {
   if (DbgProgMgr > 0) {
     std::cerr << ">>> ProgramManager::getOrCreateKernel(" << M << ", "
-              << Context.get() << ", " << Device.get() << ", " << KernelName
-              << ")\n";
+              << ContextImpl.get() << ", " << DeviceImpl.get() << ", "
+              << KernelName << ")\n";
   }
 
-  RT::PiProgram Program =
-      getBuiltPIProgram(M, Context, Device, KernelName, Prg);
-
   using PiKernelT = KernelProgramCache::PiKernelT;
   using KernelCacheT = KernelProgramCache::KernelCacheT;
   using KernelByNameT = KernelProgramCache::KernelByNameT;
 
-  KernelProgramCache &Cache = Context->getKernelProgramCache();
+  KernelProgramCache &Cache = ContextImpl->getKernelProgramCache();
+
+  std::string CompileOpts, LinkOpts;
+  SerializedObj SpecConsts;
+  if (Prg) {
+    CompileOpts = Prg->get_build_options();
+    Prg->stableSerializeSpecConstRegistry(SpecConsts);
+  }
+  applyOptionsFromEnvironment(CompileOpts, LinkOpts);
+  const RT::PiDevice PiDevice = DeviceImpl->getHandleRef();
+
+  auto key = std::make_tuple(std::move(SpecConsts), M, PiDevice,
+                             CompileOpts + LinkOpts, KernelName);
+  auto ret_tuple = Cache.tryToGetKernelFast(key);
+  if (std::get<0>(ret_tuple))
+    return ret_tuple;
+
+  RT::PiProgram Program =
+      getBuiltPIProgram(M, ContextImpl, DeviceImpl, KernelName, Prg);
 
   auto AcquireF = [](KernelProgramCache &Cache) {
     return Cache.acquireKernelsPerProgramCache();
@@ -545,12 +562,12 @@ std::pair<RT::PiKernel, std::mutex *> ProgramManager::getOrCreateKernel(
       [&Program](const Locked<KernelCacheT> &LockedCache) -> KernelByNameT & {
     return LockedCache.get()[Program];
   };
-  auto BuildF = [&Program, &KernelName, &Context] {
+  auto BuildF = [&Program, &KernelName, &ContextImpl] {
     PiKernelT *Result = nullptr;
 
     // TODO need some user-friendly error/exception
     // instead of currently obscure one
-    const detail::plugin &Plugin = Context->getPlugin();
+    const detail::plugin &Plugin = ContextImpl->getPlugin();
     Plugin.call<PiApiKind::piKernelCreate>(Program, KernelName.c_str(),
                                            &Result);
 
@@ -564,8 +581,10 @@ std::pair<RT::PiKernel, std::mutex *> ProgramManager::getOrCreateKernel(
 
   auto BuildResult = getOrBuild<PiKernelT, invalid_object_error>(
       Cache, KernelName, AcquireF, GetF, BuildF);
-  return std::make_pair(BuildResult->Ptr.load(),
-                        &(BuildResult->MBuildResultMutex));
+  auto ret_val = std::make_tuple(BuildResult->Ptr.load(),
+                                 &(BuildResult->MBuildResultMutex), Program);
+  Cache.saveKernel(key, ret_val);
+  return ret_val;
 }
 
 RT::PiProgram

sycl/source/detail/program_manager/program_manager.hpp

@@ -106,7 +106,7 @@ class ProgramManager {
                                   const property_list &PropList,
                                   bool JITCompilationIsRequired = false);
 
-  std::pair<RT::PiKernel, std::mutex *>
+  std::tuple<RT::PiKernel, std::mutex *, RT::PiProgram>
   getOrCreateKernel(OSModuleHandle M, const ContextImplPtr &ContextImpl,
                     const DeviceImplPtr &DeviceImpl,
                     const std::string &KernelName, const program_impl *Prg);

sycl/source/detail/scheduler/commands.cpp

@@ -2063,21 +2063,18 @@ cl_int ExecCGCommand::enqueueImp() {
       Program = SyclProg->getHandleRef();
       if (SyclProg->is_cacheable()) {
         RT::PiKernel FoundKernel = nullptr;
-        std::tie(FoundKernel, KernelMutex) =
+        std::tie(FoundKernel, KernelMutex, std::ignore) =
             detail::ProgramManager::getInstance().getOrCreateKernel(
                 ExecKernel->MOSModuleHandle, ContextImpl, DeviceImpl,
                 ExecKernel->MKernelName, SyclProg.get());
         assert(FoundKernel == Kernel);
       } else
         KnownProgram = false;
     } else {
-      std::tie(Kernel, KernelMutex) =
+      std::tie(Kernel, KernelMutex, Program) =
           detail::ProgramManager::getInstance().getOrCreateKernel(
               ExecKernel->MOSModuleHandle, ContextImpl, DeviceImpl,
               ExecKernel->MKernelName, nullptr);
-      MQueue->getPlugin().call<PiApiKind::piKernelGetInfo>(
-          Kernel, PI_KERNEL_INFO_PROGRAM, sizeof(RT::PiProgram), &Program,
-          nullptr);
     }
 
     pi_result Error = PI_SUCCESS;

sycl/unittests/kernel-and-program/Cache.cpp

@@ -434,3 +434,146 @@ TEST_F(KernelAndProgramCacheTest, KernelNegativeSource) {
       CtxImpl->getKernelProgramCache().acquireKernelsPerProgramCache().get();
   EXPECT_EQ(Cache.size(), 0U) << "Expect empty cache for kernels";
 }
+
+typedef KernelAndProgramCacheTest KernelAndProgramFastCacheTest;
+
+class MockKernelProgramCache : public detail::KernelProgramCache {
+public:
+  static detail::KernelProgramCache::KernelFastCacheT &
+  getFastCache(detail::KernelProgramCache &cache) {
+    return (reinterpret_cast<MockKernelProgramCache &>(cache)).get();
+  }
+
+  detail::KernelProgramCache::KernelFastCacheT &get() {
+    return this->MKernelFastCache;
+  }
+};
+
+// Check that kernels built without options are cached.
+TEST_F(KernelAndProgramFastCacheTest, KernelPositive) {
+  if (Plt.is_host() || Plt.get_backend() != backend::opencl) {
+    return;
+  }
+
+  context Ctx{Plt};
+  auto CtxImpl = detail::getSyclObjImpl(Ctx);
+
+  globalCtx.reset(new TestCtx{CtxImpl->getHandleRef()});
+
+  program Prg{Ctx};
+
+  Prg.build_with_kernel_type<TestKernel>();
+  kernel Ker = Prg.get_kernel<TestKernel>();
+  detail::KernelProgramCache::KernelFastCacheT &Cache =
+      MockKernelProgramCache::getFastCache(CtxImpl->getKernelProgramCache());
+  EXPECT_EQ(Cache.size(), 1U) << "Expect non-empty cache for kernels";
+}
+
+// Check that kernels built with options are cached.
+TEST_F(KernelAndProgramFastCacheTest, KernelPositiveBuildOpts) {
+  if (Plt.is_host() || Plt.get_backend() != backend::opencl) {
+    return;
+  }
+
+  context Ctx{Plt};
+  auto CtxImpl = detail::getSyclObjImpl(Ctx);
+
+  globalCtx.reset(new TestCtx{CtxImpl->getHandleRef()});
+
+  program Prg{Ctx};
+
+  Prg.build_with_kernel_type<TestKernel>("-g");
+
+  kernel Ker = Prg.get_kernel<TestKernel>();
+  detail::KernelProgramCache::KernelFastCacheT &Cache =
+      MockKernelProgramCache::getFastCache(CtxImpl->getKernelProgramCache());
+  EXPECT_EQ(Cache.size(), 1U) << "Expect non-empty cache for kernels";
+}
+
+// Check that kernels built with compile options are not cached.
+TEST_F(KernelAndProgramFastCacheTest, KernelNegativeCompileOpts) {
+  if (Plt.is_host() || Plt.get_backend() != backend::opencl) {
+    return;
+  }
+
+  context Ctx{Plt};
+  auto CtxImpl = detail::getSyclObjImpl(Ctx);
+
+  globalCtx.reset(new TestCtx{CtxImpl->getHandleRef()});
+
+  program Prg{Ctx};
+
+  Prg.compile_with_kernel_type<TestKernel>("-g");
+  Prg.link();
+  kernel Ker = Prg.get_kernel<TestKernel>();
+  detail::KernelProgramCache::KernelFastCacheT &Cache =
+      MockKernelProgramCache::getFastCache(CtxImpl->getKernelProgramCache());
+  EXPECT_EQ(Cache.size(), 0U) << "Expect empty cache for kernels";
+}
+
+// Check that kernels built with link options are not cached.
+TEST_F(KernelAndProgramFastCacheTest, KernelNegativeLinkOpts) {
+  if (Plt.is_host() || Plt.get_backend() != backend::opencl) {
+    return;
+  }
+
+  context Ctx{Plt};
+  auto CtxImpl = detail::getSyclObjImpl(Ctx);
+
+  globalCtx.reset(new TestCtx{CtxImpl->getHandleRef()});
+
+  program Prg{Ctx};
+
+  Prg.compile_with_kernel_type<TestKernel>();
+  Prg.link("-g");
+  kernel Ker = Prg.get_kernel<TestKernel>();
+  detail::KernelProgramCache::KernelFastCacheT &Cache =
+      MockKernelProgramCache::getFastCache(CtxImpl->getKernelProgramCache());
+  EXPECT_EQ(Cache.size(), 0U) << "Expect empty cache for kernels";
+}
+
+// Check that kernels are not cached if program is created from multiple
+// programs.
+TEST_F(KernelAndProgramFastCacheTest, KernelNegativeLinkedProgs) {
+  if (Plt.is_host() || Plt.get_backend() != backend::opencl) {
+    return;
+  }
+
+  context Ctx{Plt};
+  auto CtxImpl = detail::getSyclObjImpl(Ctx);
+
+  globalCtx.reset(new TestCtx{CtxImpl->getHandleRef()});
+
+  program Prg1{Ctx};
+  program Prg2{Ctx};
+
+  Prg1.compile_with_kernel_type<TestKernel>();
+  Prg2.compile_with_kernel_type<TestKernel2>();
+  program Prg({Prg1, Prg2});
+  kernel Ker = Prg.get_kernel<TestKernel>();
+
+  detail::KernelProgramCache::KernelFastCacheT &Cache =
+      MockKernelProgramCache::getFastCache(CtxImpl->getKernelProgramCache());
+  EXPECT_EQ(Cache.size(), 0U) << "Expect empty cache for kernels";
+}
+
+// Check that kernels created from source are not cached.
+TEST_F(KernelAndProgramFastCacheTest, KernelNegativeSource) {
+  if (Plt.is_host() || Plt.get_backend() != backend::opencl) {
+    return;
+  }
+
+  context Ctx{Plt};
+  auto CtxImpl = detail::getSyclObjImpl(Ctx);
+
+  globalCtx.reset(new TestCtx{CtxImpl->getHandleRef()});
+
+  program Prg{Ctx};
+
+  Prg.build_with_source("");
+  kernel Ker = Prg.get_kernel("test");
+
+  detail::KernelProgramCache::KernelFastCacheT &Cache =
+      MockKernelProgramCache::getFastCache(CtxImpl->getKernelProgramCache());
+  EXPECT_EQ(Cache.size(), 0U) << "Expect empty cache for kernels";
+}