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 "sycl/exception.hpp"
#include <detail/kernel_arg_mask.hpp>
#include <detail/platform_impl.hpp>
#include <sycl/detail/common.hpp>
#include <sycl/detail/locked.hpp>
#include <sycl/detail/os_util.hpp>
#include <sycl/detail/ur.hpp>
#include <sycl/detail/util.hpp>
 
#include <atomic>
#include <condition_variable>
#include <mutex>
#include <type_traits>
 
#include <boost/unordered/unordered_flat_map.hpp>
#include <boost/unordered_map.hpp>
 
// For testing purposes
class MockKernelProgramCache;
 
namespace sycl {
inline namespace _V1 {
namespace detail {
class context_impl;
class KernelProgramCache {
public:
  struct BuildError {
    std::string Msg;
    int32_t Code;
 
    bool isFilledIn() const { return !Msg.empty(); }
  };
 
  enum class BuildState { BS_Initial, BS_InProgress, BS_Done, BS_Failed };
 
  template <typename T> struct BuildResult {
    T Val;
    std::atomic<BuildState> State{BuildState::BS_Initial};
    BuildError Error{"", 0};
 
    std::condition_variable MBuildCV;
    std::mutex MBuildResultMutex;
 
    BuildState
    waitUntilTransition(BuildState From = BuildState::BS_InProgress) {
      BuildState To;
      std::unique_lock<std::mutex> Lock(MBuildResultMutex);
      MBuildCV.wait(Lock, [&] {
        To = State;
        return State != From;
      });
      return To;
    }
 
    void updateAndNotify(BuildState DesiredState) {
      {
        std::lock_guard<std::mutex> Lock(MBuildResultMutex);
        State.store(DesiredState);
      }
      MBuildCV.notify_all();
    }
  };
 
  struct ProgramBuildResult : public BuildResult<ur_program_handle_t> {
    PluginPtr Plugin;
    ProgramBuildResult(const PluginPtr &Plugin) : Plugin(Plugin) {
      Val = nullptr;
    }
    ProgramBuildResult(const PluginPtr &Plugin, BuildState InitialState)
        : Plugin(Plugin) {
      Val = nullptr;
      this->State.store(InitialState);
    }
    ~ProgramBuildResult() {
      if (Val) {
        ur_result_t Err = Plugin->call_nocheck(urProgramRelease, Val);
        __SYCL_CHECK_OCL_CODE_NO_EXC(Err);
      }
    }
  };
  using ProgramBuildResultPtr = std::shared_ptr<ProgramBuildResult>;
 
  /* Drop LinkOptions and CompileOptions from CacheKey since they are only used
   * when debugging environment variables are set and we can just ignore them
   * since all kernels will have their build options overridden with the same
   * string*/
  using ProgramCacheKeyT =
      std::pair<std::pair<SerializedObj, std::uintptr_t>, ur_device_handle_t>;
  using CommonProgramKeyT = std::pair<std::uintptr_t, ur_device_handle_t>;
 
  struct ProgramCache {
    ::boost::unordered_map<ProgramCacheKeyT, ProgramBuildResultPtr> Cache;
    ::boost::unordered_multimap<CommonProgramKeyT, ProgramCacheKeyT> KeyMap;
 
    size_t size() const noexcept { return Cache.size(); }
  };
 
  using ContextPtr = context_impl *;
 
  using KernelArgMaskPairT =
      std::pair<ur_kernel_handle_t, const KernelArgMask *>;
  struct KernelBuildResult : public BuildResult<KernelArgMaskPairT> {
    PluginPtr Plugin;
    KernelBuildResult(const PluginPtr &Plugin) : Plugin(Plugin) {
      Val.first = nullptr;
    }
    ~KernelBuildResult() {
      if (Val.first) {
        ur_result_t Err = Plugin->call_nocheck(urKernelRelease, Val.first);
        __SYCL_CHECK_OCL_CODE_NO_EXC(Err);
      }
    }
  };
  using KernelBuildResultPtr = std::shared_ptr<KernelBuildResult>;
 
  using KernelByNameT =
      ::boost::unordered_map<std::string, KernelBuildResultPtr>;
  using KernelCacheT =
      ::boost::unordered_map<ur_program_handle_t, KernelByNameT>;
 
  using KernelFastCacheKeyT =
      std::tuple<SerializedObj, ur_device_handle_t, std::string, std::string>;
  using KernelFastCacheValT =
      std::tuple<ur_kernel_handle_t, std::mutex *, const KernelArgMask *,
                 ur_program_handle_t>;
  // This container is used as a fast path for retrieving cached kernels.
  // unordered_flat_map is used here to reduce lookup overhead.
  // The slow path is used only once for each newly created kernel, so the
  // higher overhead of insertion that comes with unordered_flat_map is more
  // of an issue there. For that reason, those use regular unordered maps.
  using KernelFastCacheT =
      ::boost::unordered_flat_map<KernelFastCacheKeyT, KernelFastCacheValT>;
 
  ~KernelProgramCache() = default;
 
  void setContextPtr(const ContextPtr &AContext) { MParentContext = AContext; }
 
  Locked<ProgramCache> acquireCachedPrograms() {
    return {MCachedPrograms, MProgramCacheMutex};
  }
 
  Locked<KernelCacheT> acquireKernelsPerProgramCache() {
    return {MKernelsPerProgramCache, MKernelsPerProgramCacheMutex};
  }
 
  std::pair<ProgramBuildResultPtr, bool>
  getOrInsertProgram(const ProgramCacheKeyT &CacheKey) {
    auto LockedCache = acquireCachedPrograms();
    auto &ProgCache = LockedCache.get();
    auto [It, DidInsert] = ProgCache.Cache.try_emplace(CacheKey, nullptr);
    if (DidInsert) {
      It->second = std::make_shared<ProgramBuildResult>(getPlugin());
      // Save reference between the common key and the full key.
      CommonProgramKeyT CommonKey =
          std::make_pair(CacheKey.first.second, CacheKey.second);
      ProgCache.KeyMap.emplace(CommonKey, CacheKey);
    }
    return std::make_pair(It->second, DidInsert);
  }
 
  // Used in situation where you have several cache keys corresponding to the
  // same program. An example would be a multi-device build, or use of virtual
  // functions in kernels.
  //
  // Returns whether or not an insertion took place.
  bool insertBuiltProgram(const ProgramCacheKeyT &CacheKey,
                          ur_program_handle_t Program) {
    auto LockedCache = acquireCachedPrograms();
    auto &ProgCache = LockedCache.get();
    auto [It, DidInsert] = ProgCache.Cache.try_emplace(CacheKey, nullptr);
    if (DidInsert) {
      It->second = std::make_shared<ProgramBuildResult>(getPlugin(),
                                                        BuildState::BS_Done);
      It->second->Val = Program;
      // Save reference between the common key and the full key.
      CommonProgramKeyT CommonKey =
          std::make_pair(CacheKey.first.second, CacheKey.second);
      ProgCache.KeyMap.emplace(CommonKey, CacheKey);
    }
    return DidInsert;
  }
 
  std::pair<KernelBuildResultPtr, bool>
  getOrInsertKernel(ur_program_handle_t Program,
                    const std::string &KernelName) {
    auto LockedCache = acquireKernelsPerProgramCache();
    auto &Cache = LockedCache.get()[Program];
    auto [It, DidInsert] = Cache.try_emplace(KernelName, nullptr);
    if (DidInsert)
      It->second = std::make_shared<KernelBuildResult>(getPlugin());
    return std::make_pair(It->second, DidInsert);
  }
 
  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, 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() {
    std::lock_guard<std::mutex> L1(MProgramCacheMutex);
    std::lock_guard<std::mutex> L2(MKernelsPerProgramCacheMutex);
    std::lock_guard<std::mutex> L3(MKernelFastCacheMutex);
    MCachedPrograms = ProgramCache{};
    MKernelsPerProgramCache = KernelCacheT{};
    MKernelFastCache = KernelFastCacheT{};
  }
 
  template <errc Errc, typename GetCachedBuildFT, typename BuildFT>
  auto getOrBuild(GetCachedBuildFT &&GetCachedBuild, BuildFT &&Build) {
    using BuildState = KernelProgramCache::BuildState;
    constexpr size_t MaxAttempts = 2;
    for (size_t AttemptCounter = 0;; ++AttemptCounter) {
      auto Res = GetCachedBuild();
      auto &BuildResult = Res.first;
      BuildState Expected = BuildState::BS_Initial;
      BuildState Desired = BuildState::BS_InProgress;
      if (!BuildResult->State.compare_exchange_strong(Expected, Desired)) {
        // no insertion took place, thus some other thread has already inserted
        // smth in the cache
        BuildState NewState = BuildResult->waitUntilTransition();
 
        // Build succeeded.
        if (NewState == BuildState::BS_Done)
          return BuildResult;
 
        // Build failed, or this is the last attempt.
        if (NewState == BuildState::BS_Failed ||
            AttemptCounter + 1 == MaxAttempts) {
          if (BuildResult->Error.isFilledIn())
            throw detail::set_ur_error(
                exception(make_error_code(Errc), BuildResult->Error.Msg),
                BuildResult->Error.Code);
          else
            throw exception();
        }
 
        // NewState == BuildState::BS_Initial
        // Build state was set back to the initial state,
        // which means to go back to the beginning of the
        // loop and try again.
        continue;
      }
 
      // only the building thread will run this
      try {
        BuildResult->Val = Build();
 
        BuildResult->updateAndNotify(BuildState::BS_Done);
        return BuildResult;
      } catch (const exception &Ex) {
        BuildResult->Error.Msg = Ex.what();
        BuildResult->Error.Code = detail::get_ur_error(Ex);
        if (Ex.code() == errc::memory_allocation ||
            BuildResult->Error.Code == UR_RESULT_ERROR_OUT_OF_RESOURCES ||
            BuildResult->Error.Code == UR_RESULT_ERROR_OUT_OF_HOST_MEMORY) {
          reset();
          BuildResult->updateAndNotify(BuildState::BS_Initial);
          continue;
        }
 
        BuildResult->updateAndNotify(BuildState::BS_Failed);
        std::rethrow_exception(std::current_exception());
      } catch (...) {
        BuildResult->updateAndNotify(BuildState::BS_Initial);
        std::rethrow_exception(std::current_exception());
      }
    }
  }
 
private:
  std::mutex MProgramCacheMutex;
  std::mutex MKernelsPerProgramCacheMutex;
 
  ProgramCache MCachedPrograms;
  KernelCacheT MKernelsPerProgramCache;
  ContextPtr MParentContext;
 
  std::mutex MKernelFastCacheMutex;
  KernelFastCacheT MKernelFastCache;
  friend class ::MockKernelProgramCache;
 
  const PluginPtr &getPlugin();
};
} // namespace detail
} // namespace _V1
} // namespace sycl