#include "printer.h"
#include "scoped_mutex_lock.h"
namespace Halide {
namespace Internal {
template<typename ContextT, typename ModuleStateT>
class GPUCompilationCache {
struct CachedCompilation {
ContextT context{};
ModuleStateT module_state{};
uint32_t kernel_id{};
uint32_t use_count{0};
CachedCompilation(ContextT context, ModuleStateT module_state,
uint32_t kernel_id, uint32_t use_count)
: context(context), module_state(module_state),
kernel_id(kernel_id), use_count(use_count) {
}
};
halide_mutex mutex;
static constexpr float kLoadFactor{.5f};
static constexpr int kInitialTableBits{7};
int log2_compilations_size{0}; // number of bits in index into compilations table.
CachedCompilation *compilations{nullptr};
int count{0};
static constexpr uint32_t kInvalidId{0};
static constexpr uint32_t kDeletedId{1};
uint32_t unique_id{2}; // zero is an invalid id
public:
static ALWAYS_INLINE uintptr_t kernel_hash(ContextT context, uint32_t id, uint32_t bits) {
uintptr_t addr = (uintptr_t)context + id;
// Fibonacci hashing. The golden ratio is 1.9E3779B97F4A7C15F39...
// in hexadecimal.
if (sizeof(uintptr_t) >= 8) {
return (addr * (uintptr_t)0x9E3779B97F4A7C15) >> (64 - bits);
} else {
return (addr * (uintptr_t)0x9E3779B9) >> (32 - bits);
}
}
HALIDE_MUST_USE_RESULT bool insert(const CachedCompilation &entry) {
if (log2_compilations_size == 0) {
if (!resize_table(kInitialTableBits)) {
return false;
}
}
if ((count + 1) > (1 << log2_compilations_size) * kLoadFactor) {
if (!resize_table(log2_compilations_size + 1)) {
return false;
}
}
count += 1;
uintptr_t index = kernel_hash(entry.context, entry.kernel_id, log2_compilations_size);
for (int i = 0; i < (1 << log2_compilations_size); i++) {
uintptr_t effective_index = (index + i) & ((1 << log2_compilations_size) - 1);
if (compilations[effective_index].kernel_id <= kDeletedId) {
compilations[effective_index] = entry;
return true;
}
}
// This is a logic error that should never occur. It means the table is
// full, but it should have been resized.
halide_debug_assert(nullptr, false);
return false;
}
HALIDE_MUST_USE_RESULT bool find_internal(ContextT context, uint32_t id,
ModuleStateT *&module_state, int increment) {
if (log2_compilations_size == 0) {
return false;
}
uintptr_t index = kernel_hash(context, id, log2_compilations_size);
for (int i = 0; i < (1 << log2_compilations_size); i++) {
uintptr_t effective_index = (index + i) & ((1 << log2_compilations_size) - 1);
if (compilations[effective_index].kernel_id == kInvalidId) {
return false;
}
if (compilations[effective_index].context == context &&
compilations[effective_index].kernel_id == id) {
module_state = &compilations[effective_index].module_state;
if (increment != 0) {
compilations[effective_index].use_count += increment;
}
return true;
}
}
return false;
}
HALIDE_MUST_USE_RESULT bool lookup(ContextT context, void *state_ptr, ModuleStateT &module_state) {
ScopedMutexLock lock_guard(&mutex);
uint32_t id = (uint32_t)(uintptr_t)state_ptr;
ModuleStateT *mod_ptr;
if (find_internal(context, id, mod_ptr, 0)) {
module_state = *mod_ptr;
return true;
}
return false;
}
HALIDE_MUST_USE_RESULT bool resize_table(int size_bits) {
if (size_bits != log2_compilations_size) {
int new_size = (1 << size_bits);
int old_size = (1 << log2_compilations_size);
CachedCompilation *new_table = (CachedCompilation *)malloc(new_size * sizeof(CachedCompilation));
if (new_table == nullptr) {
// signal error.
return false;
}
memset(new_table, 0, new_size * sizeof(CachedCompilation));
CachedCompilation *old_table = compilations;
compilations = new_table;
log2_compilations_size = size_bits;
if (count > 0) { // Mainly to catch empty initial table case
for (int32_t i = 0; i < old_size; i++) {
if (old_table[i].kernel_id != kInvalidId &&
old_table[i].kernel_id != kDeletedId) {
bool result = insert(old_table[i]);
halide_debug_assert(nullptr, result); // Resizing the table while resizing the table is a logic error.
(void)result;
}
}
}
free(old_table);
}
return true;
}
template<typename FreeModuleT>
void release_context(void *user_context, bool all, ContextT context, FreeModuleT &f) {
if (count == 0) {
return;
}
for (int i = 0; i < (1 << log2_compilations_size); i++) {
if (compilations[i].kernel_id > kInvalidId &&
(all || (compilations[i].context == context)) &&
compilations[i].use_count == 0) {
debug(user_context) << "Releasing cached compilation: " << compilations[i].module_state
<< " id " << compilations[i].kernel_id
<< " context " << compilations[i].context << "\n";
f(compilations[i].module_state);
compilations[i].module_state = nullptr;
compilations[i].kernel_id = kDeletedId;
count--;
}
}
}
template<typename FreeModuleT>
void delete_context(void *user_context, ContextT context, FreeModuleT &f) {
ScopedMutexLock lock_guard(&mutex);
release_context(user_context, false, context, f);
}
template<typename FreeModuleT>
void release_all(void *user_context, FreeModuleT &f) {
ScopedMutexLock lock_guard(&mutex);
release_context(user_context, true, nullptr, f);
// Some items may have been in use, so can't free.
if (count == 0) {
free(compilations);
compilations = nullptr;
log2_compilations_size = 0;
}
}
template<typename CompileModuleT, typename... Args>
HALIDE_MUST_USE_RESULT bool kernel_state_setup(void *user_context, void **state_ptr,
ContextT context, ModuleStateT &result,
CompileModuleT f,
Args... args) {
ScopedMutexLock lock_guard(&mutex);
uint32_t *id_ptr = (uint32_t *)state_ptr;
if (*id_ptr == 0) {
*id_ptr = unique_id++;
}
ModuleStateT *mod;
if (find_internal(context, *id_ptr, mod, 1)) {
result = *mod;
return true;
}
// TODO(zvookin): figure out the calling signature here...
ModuleStateT compiled_module = f(args...);
debug(user_context) << "Caching compiled kernel: " << compiled_module
<< " id " << *id_ptr << " context " << context << "\n";
if (compiled_module == nullptr) {
return false;
}
if (!insert({context, compiled_module, *id_ptr, 1})) {
return false;
}
result = compiled_module;
return true;
}
void release_hold(void *user_context, ContextT context, void *state_ptr) {
ModuleStateT *mod;
uint32_t id = (uint32_t)(uintptr_t)state_ptr;
bool result = find_internal(context, id, mod, -1);
halide_debug_assert(user_context, result); // Value must be in cache to be released
(void)result;
}
};
} // namespace Internal
} // namespace Halide