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 | |
8 | namespace taichi::lang { |
9 | namespace gfx { |
10 | |
11 | struct 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 | |
21 | class 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 | |