Skip to content

Commit

Permalink
Refactor2023: Introduce CompiledKernelData & KernelCompiler & KernelC…
Browse files Browse the repository at this point in the history 
…ompilationManager, impl them for spirv/gfx backends and re-impl VulkanProgramImpl::compile() (fixed allways dump
  • Loading branch information
@PGZXB
PGZXB committed Dec 18, 2022
1 parent ace30eb commit c8e6285
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 33 deletions.
45 changes: 27 additions & 18 deletions taichi/cache/kernel_compilation_manager.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@ const CompiledKernelData &KernelCompilationManager::load_or_compile(
const CompileConfig &compile_config,
const DeviceCapabilityConfig &device_caps,
const Kernel &kernel_def) {
auto cache_mode = compile_config.offline_cache ? MemAndDiskCache : MemCache;
auto cache_mode = compile_config.offline_cache ? CacheData::MemAndDiskCache
: CacheData::MemCache;
const auto gukk = make_gukk(compile_config, kernel_def);
auto cached_kernel =
try_load_cached_kernel(kernel_def, gukk, compile_config.arch, cache_mode);
Expand All @@ -39,30 +40,32 @@ void KernelCompilationManager::dump_with_merging() {
auto lock_path = join_path(config_.offline_cache_path, kMetadataLockName);
if (lock_with_file(lock_path)) {
auto _ = make_unlocker(lock_path);
OfflineCacheData data;
CacheData data;
data.version[0] = TI_VERSION_MAJOR;
data.version[1] = TI_VERSION_MINOR;
data.version[2] = TI_VERSION_PATCH;
auto &kernels = data.kernels;
// Load old cached data
offline_cache::load_metadata_with_checking(data, filepath);
// Update the cached data
for (const auto *e : updated_data_) {
auto iter = data.kernels.find(e->kernel_key);
if (iter != data.kernels.end()) {
auto iter = kernels.find(e->kernel_key);
if (iter != kernels.end()) {
iter->second.last_used_at = e->last_used_at;
}
}
// Add new data
for (auto &[gukk, kernel] : caching_kernels_) {
auto [iter, ok] = data.kernels.insert({gukk, std::move(kernel)});
// data.kernels.insert()
if (ok) {
data.size += iter->second.size;
// TODO: Mangle iter->second.compiled_kernel_data->mangle_names(gukk);
if (kernel.cache_mode == CacheData::MemAndDiskCache) {
auto [iter, inserted] = kernels.insert({gukk, std::move(kernel)});
if (inserted) {
data.size += iter->second.size;
// TODO: iter->second.compiled_kernel_data->mangle_names(gukk);
}
}
}
// Dump
for (const auto &[_, k] : data.kernels) {
// Dump cache files
for (const auto &[_, k] : kernels) {
auto cache_filename = join_path(
config_.offline_cache_path,
fmt::format(kCacheFilenameFormat, k.kernel_key, arch_name(k.arch)));
Expand All @@ -72,7 +75,10 @@ void KernelCompilationManager::dump_with_merging() {
k.compiled_kernel_data->dump(fs);
}
}
write_to_binary_file(data, filepath);
// Dump offline cache metadata
if (!kernels.empty()) {
write_to_binary_file(data, filepath);
}
}
}
}
Expand Down Expand Up @@ -108,7 +114,7 @@ const CompiledKernelData *KernelCompilationManager::try_load_cached_kernel(
const Kernel &kernel_def,
const std::string &gukk,
Arch arch,
CacheMode cache_mode) {
CacheData::CacheMode cache_mode) {
{ // Find in memory-cache (caching_kernels_)
const auto &kernels = caching_kernels_;
auto iter = kernels.find(gukk);
Expand All @@ -119,11 +125,12 @@ const CompiledKernelData *KernelCompilationManager::try_load_cached_kernel(
}
}
// Find in disk-cache (cached_data_)
if (cache_mode == MemAndDiskCache) {
if (cache_mode == CacheData::MemAndDiskCache) {
auto &kernels = cached_data_.kernels;
auto iter = kernels.find(gukk);
if (iter != kernels.end()) {
auto &k = iter->second;
TI_ASSERT(k.arch == arch);
if (auto loaded = load_cache_file(gukk, arch)) {
TI_DEBUG("Create kernel '{}' from cache (key='{}')",
kernel_def.get_name(), gukk);
Expand All @@ -142,17 +149,19 @@ const CompiledKernelData &KernelCompilationManager::compile_and_cache_kernel(
const CompileConfig &compile_config,
const DeviceCapabilityConfig &device_caps,
const Kernel &kernel_def) {
auto cache_mode = compile_config.offline_cache ? MemAndDiskCache : MemCache;
TI_DEBUG_IF(cache_mode == MemAndDiskCache, "Cache kernel '{}' (key='{}')",
kernel_def.get_name(), gukk);
auto cache_mode = compile_config.offline_cache ? CacheData::MemAndDiskCache
: CacheData::MemCache;
TI_DEBUG_IF(cache_mode == CacheData::MemAndDiskCache,
"Cache kernel '{}' (key='{}')", kernel_def.get_name(), gukk);
TI_ASSERT(caching_kernels_.find(gukk) == caching_kernels_.end());
OfflineCacheKernelData k;
KernelCacheData k;
k.arch = compile_config.arch;
k.kernel_key = gukk;
k.created_at = k.last_used_at = std::time(nullptr);
k.compiled_kernel_data =
compile_kernel(compile_config, device_caps, kernel_def);
k.size = k.compiled_kernel_data->size();
k.cache_mode = cache_mode;
const auto &kernel_data = (caching_kernels_[gukk] = std::move(k));
return *kernel_data.compiled_kernel_data;
}
Expand Down
33 changes: 18 additions & 15 deletions taichi/cache/kernel_compilation_manager.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,10 @@

namespace taichi::lang {

struct OfflineCacheData {
struct CacheData {
enum CacheMode { MemCache, MemAndDiskCache };
using Version = std::uint16_t[3];

struct KernelData {
Arch arch;
std::string kernel_key;
Expand All @@ -22,10 +25,13 @@ struct OfflineCacheData {

std::unique_ptr<lang::CompiledKernelData> compiled_kernel_data;

TI_IO_DEF(kernel_key, size, created_at, last_used_at);
// Dump the kernel to disk if `cache_mode` == `MemAndDiskCache`
CacheMode cache_mode{MemCache};

TI_IO_DEF(arch, kernel_key, size, created_at, last_used_at);
};

offline_cache::Version version{};
Version version{};
std::size_t size{0};
std::unordered_map<std::string, KernelData> kernels;

Expand All @@ -38,12 +44,8 @@ class KernelCompilationManager final {
static constexpr char kCacheFilenameFormat[] = "{}-{}.tic";
static constexpr char kMetadataLockName[] = "ticache.lock";

using OfflineCacheData = OfflineCacheData;
using OfflineCacheKernelData = OfflineCacheData::KernelData;
using CachingKernels =
std::unordered_map<std::string, OfflineCacheKernelData>;

enum CacheMode { MemCache, MemAndDiskCache };
using KernelCacheData = CacheData::KernelData;
using CachingKernels = std::unordered_map<std::string, KernelCacheData>;

public:
struct Config {
Expand Down Expand Up @@ -76,10 +78,11 @@ class KernelCompilationManager final {
std::string make_gukk(const CompileConfig &compile_config,
const Kernel &kernel_def) const;

const CompiledKernelData *try_load_cached_kernel(const Kernel &kernel_def,
const std::string &gukk,
Arch arch,
CacheMode cache_mode);
const CompiledKernelData *try_load_cached_kernel(
const Kernel &kernel_def,
const std::string &gukk,
Arch arch,
CacheData::CacheMode cache_mode);

const CompiledKernelData &compile_and_cache_kernel(
const std::string &gukk,
Expand All @@ -92,8 +95,8 @@ class KernelCompilationManager final {

Config config_;
CachingKernels caching_kernels_;
OfflineCacheData cached_data_;
std::vector<OfflineCacheKernelData *> updated_data_;
CacheData cached_data_;
std::vector<KernelCacheData *> updated_data_;
};

} // namespace taichi::lang

0 comments on commit c8e6285

Please sign in to comment.