[SYCL] Optimize kernel name based cache lookup (#18081)

Remove the kernel name string related operations during the cache lookup
by taking advantage of the kernel name type uniqueness.

---------

Signed-off-by: Sergei Vinogradov <[email protected]>
Co-authored-by: Sergei Vinogradov <[email protected]>

Author: sergey-semenov

sycl/include/sycl/detail/kernel_name_based_cache.hpp

@@ -0,0 +1,29 @@
+//==--------------------- kernel_name_based_cache.hpp ----------------------==//
+//
+// 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 <sycl/detail/export.hpp>
+
+namespace sycl {
+inline namespace _V1 {
+namespace detail {
+
+struct KernelNameBasedCacheT;
+__SYCL_EXPORT KernelNameBasedCacheT *createKernelNameBasedCache();
+
+// Retrieves a cache pointer unique to a kernel name type that can be used to
+// avoid kernel name based lookup in the runtime.
+template <typename KernelName>
+KernelNameBasedCacheT *getKernelNameBasedCache() {
+  static KernelNameBasedCacheT *Instance = createKernelNameBasedCache();
+  return Instance;
+}
+
+} // namespace detail
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/handler.hpp

@@ -18,6 +18,7 @@
 #include <sycl/detail/id_queries_fit_in_int.hpp>
 #include <sycl/detail/impl_utils.hpp>
 #include <sycl/detail/kernel_desc.hpp>
+#include <sycl/detail/kernel_name_based_cache.hpp>
 #include <sycl/detail/kernel_name_str_t.hpp>
 #include <sycl/detail/reduction_forward.hpp>
 #include <sycl/detail/string.hpp>
@@ -813,6 +814,7 @@ class __SYCL_EXPORT handler {
       // later during finalize.
       setArgsToAssociatedAccessors();
     }
+    setKernelNameBasedCachePtr(detail::getKernelNameBasedCache<KernelName>());
 
     // If the kernel lambda is callable with a kernel_handler argument, manifest
     // the associated kernel handler.
@@ -3789,6 +3791,9 @@ class __SYCL_EXPORT handler {
       sycl::handler &h, size_t size,
       const ext::oneapi::experimental::memory_pool &pool);
 
+  void setKernelNameBasedCachePtr(
+      detail::KernelNameBasedCacheT *KernelNameBasedCachePtr);
+
 protected:
   /// Registers event dependencies in this command group.
   void depends_on(const detail::EventImplPtr &Event);

sycl/source/CMakeLists.txt

@@ -273,6 +273,7 @@ set(SYCL_COMMON_SOURCES
     "detail/kernel_compiler/kernel_compiler_opencl.cpp"
     "detail/kernel_compiler/kernel_compiler_sycl.cpp"
     "detail/kernel_impl.cpp"
+    "detail/kernel_name_based_cache.cpp"
     "detail/kernel_program_cache.cpp"
     "detail/memory_manager.cpp"
     "detail/pipes.cpp"

sycl/source/detail/cg.hpp

@@ -257,6 +257,7 @@ class CGExecKernel : public CG {
   std::shared_ptr<detail::kernel_bundle_impl> MKernelBundle;
   std::vector<ArgDesc> MArgs;
   KernelNameStrT MKernelName;
+  KernelNameBasedCacheT *MKernelNameBasedCachePtr;
   std::vector<std::shared_ptr<detail::stream_impl>> MStreams;
   std::vector<std::shared_ptr<const void>> MAuxiliaryResources;
   /// Used to implement ext_oneapi_graph dynamic_command_group. Stores the list
@@ -272,6 +273,7 @@ class CGExecKernel : public CG {
                std::shared_ptr<detail::kernel_bundle_impl> KernelBundle,
                CG::StorageInitHelper CGData, std::vector<ArgDesc> Args,
                KernelNameStrT KernelName,
+               KernelNameBasedCacheT *KernelNameBasedCachePtr,
                std::vector<std::shared_ptr<detail::stream_impl>> Streams,
                std::vector<std::shared_ptr<const void>> AuxiliaryResources,
                CGType Type, ur_kernel_cache_config_t KernelCacheConfig,
@@ -281,7 +283,9 @@ class CGExecKernel : public CG {
         MNDRDesc(std::move(NDRDesc)), MHostKernel(std::move(HKernel)),
         MSyclKernel(std::move(SyclKernel)),
         MKernelBundle(std::move(KernelBundle)), MArgs(std::move(Args)),
-        MKernelName(std::move(KernelName)), MStreams(std::move(Streams)),
+        MKernelName(std::move(KernelName)),
+        MKernelNameBasedCachePtr(KernelNameBasedCachePtr),
+        MStreams(std::move(Streams)),
         MAuxiliaryResources(std::move(AuxiliaryResources)),
         MAlternativeKernels{}, MKernelCacheConfig(std::move(KernelCacheConfig)),
         MKernelIsCooperative(KernelIsCooperative),

sycl/source/detail/global_handler.cpp

@@ -14,6 +14,7 @@
 #include <detail/adapter.hpp>
 #include <detail/config.hpp>
 #include <detail/global_handler.hpp>
+#include <detail/kernel_name_based_cache_t.hpp>
 #include <detail/platform_impl.hpp>
 #include <detail/program_manager/program_manager.hpp>
 #include <detail/scheduler/scheduler.hpp>
@@ -252,6 +253,15 @@ ThreadPool &GlobalHandler::getHostTaskThreadPool() {
   return TP;
 }
 
+KernelNameBasedCacheT *GlobalHandler::createKernelNameBasedCache() {
+  static std::vector<std::unique_ptr<KernelNameBasedCacheT>>
+      &KernelNameBasedCaches = getOrCreate(MKernelNameBasedCaches);
+  LockGuard LG{MKernelNameBasedCaches.Lock};
+  return KernelNameBasedCaches
+      .emplace_back(std::make_unique<KernelNameBasedCacheT>())
+      .get();
+}
+
 void GlobalHandler::releaseDefaultContexts() {
   // Release shared-pointers to SYCL objects.
   // Note that on Windows the destruction of the default context

sycl/source/detail/global_handler.hpp

@@ -26,6 +26,7 @@ class Adapter;
 class ods_target_list;
 class XPTIRegistry;
 class ThreadPool;
+struct KernelNameBasedCacheT;
 
 using ContextImplPtr = std::shared_ptr<context_impl>;
 using AdapterPtr = std::shared_ptr<Adapter>;
@@ -73,7 +74,7 @@ class GlobalHandler {
   ods_target_list &getOneapiDeviceSelectorTargets(const std::string &InitValue);
   XPTIRegistry &getXPTIRegistry();
   ThreadPool &getHostTaskThreadPool();
-
+  KernelNameBasedCacheT *createKernelNameBasedCache();
   static void registerStaticVarShutdownHandler();
 
   bool isOkToDefer() const;
@@ -129,6 +130,8 @@ class GlobalHandler {
   InstWithLock<XPTIRegistry> MXPTIRegistry;
   // Thread pool for host task and event callbacks execution
   InstWithLock<ThreadPool> MHostTaskThreadPool;
+  InstWithLock<std::vector<std::unique_ptr<KernelNameBasedCacheT>>>
+      MKernelNameBasedCaches;
 };
 } // namespace detail
 } // namespace _V1

sycl/source/detail/graph_impl.cpp

@@ -813,8 +813,8 @@ exec_graph_impl::enqueueNodeDirect(sycl::context Ctx,
                                         CGExec->MLine, CGExec->MColumn);
     std::tie(CmdTraceEvent, InstanceID) = emitKernelInstrumentationData(
         StreamID, CGExec->MSyclKernel, CodeLoc, CGExec->MIsTopCodeLoc,
-        CGExec->MKernelName.data(), nullptr, CGExec->MNDRDesc,
-        CGExec->MKernelBundle, CGExec->MArgs);
+        CGExec->MKernelName.data(), CGExec->MKernelNameBasedCachePtr, nullptr,
+        CGExec->MNDRDesc, CGExec->MKernelBundle, CGExec->MArgs);
     if (CmdTraceEvent)
       sycl::detail::emitInstrumentationGeneral(
           StreamID, InstanceID, CmdTraceEvent, xpti::trace_task_begin, nullptr);
@@ -1503,7 +1503,8 @@ void exec_graph_impl::populateURKernelUpdateStructs(
     ur_program_handle_t UrProgram = nullptr;
     std::tie(UrKernel, std::ignore, EliminatedArgMask, UrProgram) =
         sycl::detail::ProgramManager::getInstance().getOrCreateKernel(
-            ContextImpl, DeviceImpl, ExecCG.MKernelName);
+            ContextImpl, DeviceImpl, ExecCG.MKernelName,
+            ExecCG.MKernelNameBasedCachePtr);
     BundleObjs = std::make_pair(UrProgram, UrKernel);
   }
 

sycl/source/detail/handler_impl.hpp

@@ -205,6 +205,9 @@ class handler_impl {
   detail::kernel_param_desc_t (*MKernelParamDescGetter)(int) = nullptr;
   bool MKernelIsESIMD = false;
   bool MKernelHasSpecialCaptures = true;
+
+  // A pointer to a kernel name based cache retrieved on the application side.
+  KernelNameBasedCacheT *MKernelNameBasedCachePtr;
 };
 
 } // namespace detail

sycl/source/detail/kernel_name_based_cache.cpp

@@ -0,0 +1,22 @@
+//==--------------------- kernel_name_based_cache.cpp ----------------------==//
+//
+// 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 <detail/global_handler.hpp>
+#include <sycl/detail/kernel_name_based_cache.hpp>
+
+namespace sycl {
+inline namespace _V1 {
+namespace detail {
+
+KernelNameBasedCacheT *createKernelNameBasedCache() {
+  return GlobalHandler::instance().createKernelNameBasedCache();
+}
+
+} // namespace detail
+} // namespace _V1
+} // namespace sycl

sycl/source/detail/kernel_name_based_cache_t.hpp

@@ -0,0 +1,43 @@
+//==-------------------- kernel_name_based_cache_t.hpp ---------------------==//
+//
+// 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 <detail/kernel_arg_mask.hpp>
+#include <sycl/detail/spinlock.hpp>
+#include <sycl/detail/ur.hpp>
+
+#include <mutex>
+
+#include <boost/unordered/unordered_flat_map.hpp>
+
+namespace sycl {
+inline namespace _V1 {
+namespace detail {
+using FastKernelCacheKeyT = std::pair<ur_device_handle_t, ur_context_handle_t>;
+using FastKernelCacheValT =
+    std::tuple<ur_kernel_handle_t, std::mutex *, const KernelArgMask *,
+               ur_program_handle_t>;
+using FastKernelSubcacheMapT =
+    ::boost::unordered_flat_map<FastKernelCacheKeyT, FastKernelCacheValT>;
+
+using FastKernelSubcacheMutexT = SpinLock;
+using FastKernelSubcacheReadLockT = std::lock_guard<FastKernelSubcacheMutexT>;
+using FastKernelSubcacheWriteLockT = std::lock_guard<FastKernelSubcacheMutexT>;
+
+struct FastKernelSubcacheT {
+  FastKernelSubcacheMapT Map;
+  FastKernelSubcacheMutexT Mutex;
+};
+
+struct KernelNameBasedCacheT {
+  FastKernelSubcacheT FastKernelSubcache;
+};
+
+} // namespace detail
+} // namespace _V1
+} // namespace sycl

sycl/source/detail/kernel_program_cache.cpp

@@ -16,6 +16,11 @@ namespace detail {
 const AdapterPtr &KernelProgramCache::getAdapter() {
   return MParentContext->getAdapter();
 }
+
+ur_context_handle_t KernelProgramCache::getURContext() const {
+  return MParentContext->getHandleRef();
+}
+
 } // namespace detail
 } // namespace _V1
 } // namespace sycl