[SYCL] Add per-kernel mutex to fix the race when setting kernel parameters in parallel

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

Author: alexanderfle

sycl/source/detail/kernel_program_cache.hpp

@@ -58,8 +58,15 @@ class KernelProgramCache {
   using ContextPtr = context_impl *;
 
   using PiKernelT = std::remove_pointer<RT::PiKernel>::type;
+
+  struct BuildResultKernel : public BuildResult<PiKernelT> {
+    std::mutex MKernelMutex;
+
+    BuildResultKernel(PiKernelT *P, int S) : BuildResult(P, S) {}
+  };
+
   using PiKernelPtrT = std::atomic<PiKernelT *>;
-  using KernelWithBuildStateT = BuildResult<PiKernelT>;
+  using KernelWithBuildStateT = BuildResultKernel;
   using KernelByNameT = std::map<string_class, KernelWithBuildStateT>;
   using KernelCacheT = std::map<RT::PiProgram, KernelByNameT>;
 

sycl/source/detail/program_impl.cpp

@@ -396,8 +396,9 @@ RT::PiKernel program_impl::get_pi_kernel(const string_class &KernelName) const {
   RT::PiKernel Kernel;
 
   if (is_cacheable()) {
-    Kernel = ProgramManager::getInstance().getOrCreateKernel(
-        MProgramModuleHandle, get_context(), KernelName, this);
+    std::tie(Kernel, std::ignore) =
+        ProgramManager::getInstance().getOrCreateKernel(
+            MProgramModuleHandle, get_context(), KernelName, this);
     getPlugin().call<PiApiKind::piKernelRetain>(Kernel);
   } else {
     const detail::plugin &Plugin = getPlugin();

sycl/source/detail/program_manager/program_manager.cpp

@@ -167,8 +167,9 @@ RetT *waitUntilBuilt(KernelProgramCache &Cache,
 ///         cache. Accepts nothing. Return pointer to built entity.
 template <typename RetT, typename ExceptionT, typename KeyT, typename AcquireFT,
           typename GetCacheFT, typename BuildFT>
-RetT *getOrBuild(KernelProgramCache &KPCache, KeyT &&CacheKey,
-                 AcquireFT &&Acquire, GetCacheFT &&GetCache, BuildFT &&Build) {
+KernelProgramCache::BuildResult<RetT> *
+getOrBuild(KernelProgramCache &KPCache, KeyT &&CacheKey, AcquireFT &&Acquire,
+           GetCacheFT &&GetCache, BuildFT &&Build) {
   bool InsertionTookPlace;
   KernelProgramCache::BuildResult<RetT> *BuildResult;
 
@@ -190,7 +191,7 @@ RetT *getOrBuild(KernelProgramCache &KPCache, KeyT &&CacheKey,
       RetT *Result = waitUntilBuilt<ExceptionT>(KPCache, BuildResult);
 
       if (Result)
-        return Result;
+        return BuildResult;
 
       // Previous build is failed. There was no SYCL exception though.
       // We might try to build once more.
@@ -220,7 +221,7 @@ RetT *getOrBuild(KernelProgramCache &KPCache, KeyT &&CacheKey,
 
     KPCache.notifyAllBuild();
 
-    return Desired;
+    return BuildResult;
   } catch (const exception &Ex) {
     BuildResult->Error.Msg = Ex.what();
     BuildResult->Error.Code = Ex.get_cl_code();
@@ -400,14 +401,15 @@ RT::PiProgram ProgramManager::getBuiltPIProgram(OSModuleHandle M,
   if (Prg)
     Prg->stableSerializeSpecConstRegistry(SpecConsts);
 
-  return getOrBuild<PiProgramT, compile_program_error>(
+  auto BuildResult = getOrBuild<PiProgramT, compile_program_error>(
       Cache, KeyT(std::move(SpecConsts), KSId), AcquireF, GetF, BuildF);
+  return BuildResult->Ptr.load();
 }
 
-RT::PiKernel ProgramManager::getOrCreateKernel(OSModuleHandle M,
-                                               const context &Context,
-                                               const string_class &KernelName,
-                                               const program_impl *Prg) {
+std::pair<RT::PiKernel, std::mutex *>
+ProgramManager::getOrCreateKernel(OSModuleHandle M, const context &Context,
+                                  const string_class &KernelName,
+                                  const program_impl *Prg) {
   if (DbgProgMgr > 0) {
     std::cerr << ">>> ProgramManager::getOrCreateKernel(" << M << ", "
               << getRawSyclObjImpl(Context) << ", " << KernelName << ")\n";
@@ -441,8 +443,10 @@ RT::PiKernel ProgramManager::getOrCreateKernel(OSModuleHandle M,
     return Result;
   };
 
-  return getOrBuild<PiKernelT, invalid_object_error>(Cache, KernelName,
-                                                     AcquireF, GetF, BuildF);
+  auto BuildResult = static_cast<KernelProgramCache::BuildResultKernel *>(
+      getOrBuild<PiKernelT, invalid_object_error>(Cache, KernelName, AcquireF,
+                                                  GetF, BuildF));
+  return std::make_pair(BuildResult->Ptr.load(), &(BuildResult->MKernelMutex));
 }
 
 RT::PiProgram

sycl/source/detail/program_manager/program_manager.hpp

@@ -77,9 +77,9 @@ class ProgramManager {
   RT::PiProgram getBuiltPIProgram(OSModuleHandle M, const context &Context,
                                   const string_class &KernelName,
                                   const program_impl *Prg = nullptr);
-  RT::PiKernel getOrCreateKernel(OSModuleHandle M, const context &Context,
-                                 const string_class &KernelName,
-                                 const program_impl *Prg);
+  std::pair<RT::PiKernel, std::mutex *>
+  getOrCreateKernel(OSModuleHandle M, const context &Context,
+                    const string_class &KernelName, const program_impl *Prg);
   RT::PiProgram getPiProgramFromPiKernel(RT::PiKernel Kernel,
                                          const ContextImplPtr Context);
 

sycl/source/detail/scheduler/commands.cpp

@@ -1799,15 +1799,19 @@ cl_int ExecCGCommand::enqueueImp() {
     sycl::context Context = MQueue->get_context();
     const detail::plugin &Plugin = MQueue->getPlugin();
     RT::PiKernel Kernel = nullptr;
+    std::mutex *KernelMutex = nullptr;
 
     if (nullptr != ExecKernel->MSyclKernel) {
       assert(ExecKernel->MSyclKernel->get_info<info::kernel::context>() ==
              Context);
       Kernel = ExecKernel->MSyclKernel->getHandleRef();
-    } else
-      Kernel = detail::ProgramManager::getInstance().getOrCreateKernel(
-          ExecKernel->MOSModuleHandle, Context, ExecKernel->MKernelName,
-          nullptr);
+    } else {
+      std::tie(Kernel, KernelMutex) =
+          detail::ProgramManager::getInstance().getOrCreateKernel(
+              ExecKernel->MOSModuleHandle, Context, ExecKernel->MKernelName,
+              nullptr);
+      KernelMutex->lock();
+    }
 
     for (ArgDesc &Arg : ExecKernel->MArgs) {
       switch (Arg.MType) {
@@ -1863,6 +1867,9 @@ cl_int ExecCGCommand::enqueueImp() {
         &NDRDesc.GlobalSize[0], HasLocalSize ? &NDRDesc.LocalSize[0] : nullptr,
         RawEvents.size(), RawEvents.empty() ? nullptr : &RawEvents[0], &Event);
 
+    if (KernelMutex != nullptr)
+      KernelMutex->unlock();
+
     if (PI_SUCCESS != Error) {
       // If we have got non-success error code, let's analyze it to emit nice
       // exception explaining what was wrong

sycl/source/detail/scheduler/scheduler.cpp

@@ -80,7 +80,10 @@ EventImplPtr Scheduler::addCG(std::unique_ptr<detail::CG> CommandGroup,
     default:
       NewCmd = MGraphBuilder.addCG(std::move(CommandGroup), std::move(Queue));
     }
+  }
 
+  {
+    std::shared_lock<std::shared_timed_mutex> Lock(MGraphLock);
     // TODO: Check if lazy mode.
     EnqueueResultT Res;
     bool Enqueued = GraphProcessor::enqueueCommand(NewCmd, Res);