1#pragma once
2
3#include "taichi/aot/module_builder.h"
4#include "taichi/aot/module_loader.h"
5#include "taichi/runtime/gfx/runtime.h"
6#include "taichi/util/offline_cache.h"
7
8namespace taichi::lang {
9namespace gfx {
10
11struct OfflineCacheKernelMetadata {
12 std::string kernel_key;
13 std::size_t size{0}; // byte
14 std::time_t created_at{0}; // sec
15 std::time_t last_used_at{0}; // sec
16 std::size_t num_files{0};
17
18 TI_IO_DEF(kernel_key, size, created_at, last_used_at, num_files);
19};
20
21class CacheManager {
22 using CompiledKernelData = gfx::GfxRuntime::RegisterParams;
23
24 public:
25 using Metadata = offline_cache::Metadata<OfflineCacheKernelMetadata>;
26 enum Mode { NotCache, MemCache, MemAndDiskCache };
27
28 struct Params {
29 Arch arch;
30 Mode mode{MemCache};
31 std::string cache_path;
32 GfxRuntime *runtime{nullptr};
33 const CompileConfig *compile_config{nullptr};
34 DeviceCapabilityConfig caps{};
35 const std::vector<spirv::CompiledSNodeStructs> *compiled_structs;
36 };
37
38 explicit CacheManager(Params &&init_params);
39
40 CompiledKernelData load_or_compile(const CompileConfig &config,
41 Kernel *kernel);
42 void dump_with_merging() const;
43 void clean_offline_cache(offline_cache::CleanCachePolicy policy,
44 int max_bytes,
45 double cleaning_factor) const;
46
47 private:
48 std::optional<CompiledKernelData> try_load_cached_kernel(
49 Kernel *kernel,
50 const std::string &key);
51 CompiledKernelData compile_and_cache_kernel(const std::string &key,
52 Kernel *kernel);
53 std::string make_kernel_key(const CompileConfig &config,
54 Kernel *kernel) const;
55
56 Mode mode_{MemCache};
57 std::string path_;
58 GfxRuntime *runtime_{nullptr};
59 const CompileConfig &compile_config_;
60 const std::vector<spirv::CompiledSNodeStructs> &compiled_structs_;
61 Metadata offline_cache_metadata_;
62 std::unique_ptr<AotModuleBuilder> caching_module_builder_{nullptr};
63 std::unique_ptr<aot::Module> cached_module_{nullptr};
64};
65
66} // namespace gfx
67} // namespace taichi::lang
68