Refactor2023: Remove dependencies on Program::this_thread_config() in…

… metal::CacheManager
taichi-dev · Dec 6, 2022 · e595513 · e595513
1 parent 0f2a0b5
commit e595513
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/taichi/cache/metal/cache_manager.cpp b/taichi/cache/metal/cache_manager.cpp
@@ -71,14 +71,14 @@ CacheManager::CompiledKernelData CacheManager::load_or_compile(
     const CompileConfig *compile_config,
     Kernel *kernel) {
   if (kernel->is_evaluator || config_.mode == NotCache) {
-    return compile_kernel(kernel);
+    return compile_kernel(compile_config, kernel);
   }
   TI_ASSERT(config_.mode > NotCache);
   const auto kernel_key = make_kernel_key(compile_config, kernel);
   if (auto opt = try_load_cached_kernel(kernel, kernel_key)) {
     return *opt;
   }
-  return compile_and_cache_kernel(kernel_key, kernel);
+  return compile_and_cache_kernel(kernel_key, compile_config, kernel);
 }
 
 void CacheManager::dump_with_merging() const {
@@ -141,8 +141,10 @@ void CacheManager::clean_offline_cache(offline_cache::CleanCachePolicy policy,
   }
 }
 
-CompiledKernelData CacheManager::compile_kernel(Kernel *kernel) const {
-  irpass::ast_to_ir(kernel->program->this_thread_config(), *kernel);
+CompiledKernelData CacheManager::compile_kernel(
+    const CompileConfig *compile_config,
+    Kernel *kernel) const {
+  irpass::ast_to_ir(*compile_config, *kernel);
   return run_codegen(config_.compiled_runtime_module_,
                      *config_.compiled_snode_trees_, kernel, nullptr);
 }
@@ -196,13 +198,14 @@ std::optional<CompiledKernelData> CacheManager::try_load_cached_kernel(
 
 CompiledKernelData CacheManager::compile_and_cache_kernel(
     const std::string &key,
+    const CompileConfig *compile_config,
     Kernel *kernel) {
   TI_DEBUG_IF(config_.mode == MemAndDiskCache, "Cache kernel '{}' (key='{}')",
               kernel->get_name(), key);
   OfflineCacheKernelMetadata k;
   k.kernel_key = key;
   k.created_at = k.last_used_at = std::time(nullptr);
-  k.compiled_kernel_data = compile_kernel(kernel);
+  k.compiled_kernel_data = compile_kernel(compile_config, kernel);
   k.size = k.compiled_kernel_data.source_code.size();
   const auto &kernel_data = (caching_kernels_[key] = std::move(k));
   return kernel_data.compiled_kernel_data;

diff --git a/taichi/cache/metal/cache_manager.h b/taichi/cache/metal/cache_manager.h
@@ -55,14 +55,17 @@ class CacheManager {
                            double cleaning_factor) const;
 
  private:
-  CompiledKernelData compile_kernel(Kernel *kernel) const;
+  CompiledKernelData compile_kernel(const CompileConfig *compile_config,
+                                    Kernel *kernel) const;
   std::string make_kernel_key(const CompileConfig *compile_config,
                               Kernel *kernel) const;
   std::optional<CompiledKernelData> try_load_cached_kernel(
       Kernel *kernel,
       const std::string &key);
-  CompiledKernelData compile_and_cache_kernel(const std::string &key,
-                                              Kernel *kernel);
+  CompiledKernelData compile_and_cache_kernel(
+      const std::string &key,
+      const CompileConfig *compile_config,
+      Kernel *kernel);
   bool load_kernel_source_code(OfflineCacheKernelMetadata &kernel_data);
 
   Params config_;