JitCache.hpp

Go to the documentation of this file.

//===-- JitCache.hpp -- JIT in-memory code cache header implementation --===//
//
// Part of the Proteus 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
//
//===----------------------------------------------------------------------===//
//
//===----------------------------------------------------------------------===//
 
#ifndef PROTEUS_JITCACHE_HPP
#define PROTEUS_JITCACHE_HPP
 
#include <cstdint>
#include <iostream>
#include <mutex>
 
#include "proteus/CompilerInterfaceTypes.h"
#include "proteus/Hashing.hpp"
#include "proteus/TimeTracing.hpp"
#include "proteus/Utils.h"
 
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/ADT/StringRef.h>
#include <llvm/Config/llvm-config.h>
 
namespace proteus {
 
using namespace llvm;
 
template <typename Function_t> class JitCache {
public:
  Function_t lookup(HashT &HashValue) {
    TIMESCOPE("lookup");
    Accesses++;
 
    auto It = CacheMap.find(HashValue);
    if (It == CacheMap.end())
      return nullptr;
 
    It->second.NumExecs++;
    It->second.NumHits++;
    Hits++;
    return It->second.FunctionPtr;
  }
 
  void insert(HashT &HashValue, Function_t FunctionPtr,
              [[maybe_unused]] StringRef FnName,
              [[maybe_unused]] ArrayRef<RuntimeConstant> RCArray) {
#if PROTEUS_ENABLE_DEBUG
    if (CacheMap.count(HashValue))
      PROTEUS_FATAL_ERROR("JitCache collision detected");
#endif
 
    auto &CacheEntry = CacheMap[HashValue];
    CacheEntry.FunctionPtr = FunctionPtr;
    CacheEntry.NumExecs = 1;
    CacheEntry.NumHits = 0;
 
#if PROTEUS_ENABLE_DEBUG
    CacheEntry.FnName = FnName.str();
    CacheEntry.RCVector = SmallVector<RuntimeConstant>{RCArray};
#endif
  }
 
  void printStats() {
    // outs() << "JitCache hits " << Hits << " total " << Accesses << "\n";
    // Use printf to avoid re-ordering outputs by outs() in HIP.
    printf("JitCache hits %lu total %lu\n", Hits, Accesses);
    for (const auto &[HashValue, JCE] : CacheMap) {
      std::cout << "HashValue " << HashValue.toString() << " NumExecs "
                << JCE.NumExecs << " NumHits " << JCE.NumHits;
#if PROTEUS_ENABLE_DEBUG
      // outs() << " FnName " << JCE.FnName << " RCs [";
      printf(" FnName %s RCs [", JCE.FnName.c_str());
      for (auto &RC : JCE.RCVector)
        // outs() << RC.Int64Val << ", ";
        printf("%ld, ", RC.Value.Int64Val);
      // outs() << "]";
      printf("]");
#endif
      // outs() << "\n";
      printf("\n");
    }
  }
 
  JitCache() {}
 
private:
  struct JitCacheEntry {
    Function_t FunctionPtr;
    uint64_t NumExecs;
    uint64_t NumHits;
#if PROTEUS_ENABLE_DEBUG
    std::string FnName;
    SmallVector<RuntimeConstant> RCVector;
#endif
  };
 
  std::unordered_map<HashT, JitCacheEntry> CacheMap;
  // Use the executable binary path when hashing to differentiate between
  // same-named kernels generated by other executables.
  uint64_t Hits = 0;
  uint64_t Accesses = 0;
};
 
} // namespace proteus
 
#endif