kernel_program_cache.hpp

Go to the documentation of this file.

//==--- 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 <detail/platform_impl.hpp>
#include <sycl/detail/common.hpp>
#include <sycl/detail/locked.hpp>
#include <sycl/detail/os_util.hpp>
#include <sycl/detail/pi.hpp>
#include <sycl/detail/util.hpp>
 
#include <atomic>
#include <condition_variable>
#include <map>
#include <mutex>
#include <type_traits>
 
// For testing purposes
class MockKernelProgramCache;
 
namespace sycl {
__SYCL_INLINE_VER_NAMESPACE(_V1) {
namespace detail {
class context_impl;
class KernelProgramCache {
public:
  struct BuildError {
    std::string Msg;
    pi_int32 Code;
 
    bool isFilledIn() const { return !Msg.empty(); }
  };
 
  enum BuildState { BS_InProgress, BS_Done, BS_Failed };
 
  template <typename T> struct BuildResult {
    std::atomic<T *> Ptr;
    std::atomic<BuildState> State;
    BuildError Error;
 
    std::condition_variable MBuildCV;
    std::mutex MBuildResultMutex;
 
    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 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;
 
  using PiKernelPtrT = std::atomic<PiKernelT *>;
  using KernelWithBuildStateT = BuildResult<PiKernelT>;
  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; }
 
  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();
  }
 
  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);
  }
 
  void reset() {
    MCachedPrograms = ProgramCache{};
    MKernelsPerProgramCache = KernelCacheT{};
    MKernelFastCache = KernelFastCacheT{};
  }
 
private:
  std::mutex MProgramCacheMutex;
  std::mutex MKernelsPerProgramCacheMutex;
 
  ProgramCache MCachedPrograms;
  KernelCacheT MKernelsPerProgramCache;
  ContextPtr MParentContext;
 
  std::mutex MKernelFastCacheMutex;
  KernelFastCacheT MKernelFastCache;
  friend class ::MockKernelProgramCache;
};
} // namespace detail
} // __SYCL_INLINE_VER_NAMESPACE(_V1)
} // namespace sycl