diff --git a/taichi/codegen/codegen_llvm.cpp b/taichi/codegen/codegen_llvm.cpp index 919e54a78a0b6..ee1c091850ba9 100644 --- a/taichi/codegen/codegen_llvm.cpp +++ b/taichi/codegen/codegen_llvm.cpp @@ -2380,7 +2380,7 @@ FunctionType CodeGenLLVM::gen() { std::string kernel_key; if (config.offline_cache && this->supports_offline_cache() && !kernel->is_evaluator) { - kernel_key = get_offline_cache_key_of_kernel(kernel); + kernel_key = get_offline_cache_key(&kernel->program->config, kernel); LlvmOfflineCacheFileReader reader(config.offline_cache_file_path); LlvmOfflineCache::KernelCacheData cache_data; diff --git a/taichi/llvm/llvm_offline_cache.cpp b/taichi/llvm/llvm_offline_cache.cpp index 071b0e9b737b5..3930d0ea2d78b 100644 --- a/taichi/llvm/llvm_offline_cache.cpp +++ b/taichi/llvm/llvm_offline_cache.cpp @@ -11,11 +11,85 @@ namespace taichi { namespace lang { -std::string get_offline_cache_key_of_kernel(Kernel *kernel) { - std::string res, kernel_ast_string; - irpass::re_id(kernel->ir.get()); - irpass::print(kernel->ir.get(), &kernel_ast_string); - picosha2::hash256_hex_string(kernel_ast_string, res); +static TI_FORCE_INLINE std::vector +get_offline_cache_key_of_compile_config(CompileConfig *config) { + TI_ASSERT(config); + BinaryOutputSerializer serializer; + serializer.initialize(); + serializer(config->arch); + serializer(config->debug); + serializer(config->cfg_optimization); + serializer(config->check_out_of_bound); + serializer(config->simd_width); + serializer(config->opt_level); + serializer(config->external_optimization_level); + serializer(config->max_vector_width); + serializer(config->packed); + serializer(config->serial_schedule); + serializer(config->move_loop_invariant_outside_if); + serializer(config->demote_dense_struct_fors); + serializer(config->advanced_optimization); + serializer(config->constant_folding); + serializer(config->use_llvm); + serializer(config->fast_math); + serializer(config->dynamic_index); + serializer(config->flatten_if); + serializer(config->make_thread_local); + serializer(config->make_block_local); + serializer(config->detect_read_only); + serializer(config->default_fp->to_string()); + serializer(config->default_ip.to_string()); + serializer(config->extra_flags); + if (arch_is_cpu(config->arch)) { + serializer(config->default_cpu_block_dim); + serializer(config->cpu_max_num_threads); + } else if (arch_is_gpu(config->arch)) { + serializer(config->default_gpu_block_dim); + serializer(config->gpu_max_reg); + serializer(config->saturating_grid_dim); + serializer(config->cpu_max_num_threads); + } + serializer(config->ad_stack_size); + serializer(config->default_ad_stack_size); + serializer(config->random_seed); + if (config->arch == Arch::cc) { + serializer(config->cc_compile_cmd); + serializer(config->cc_link_cmd); + } else if (config->arch == Arch::opengl) { + serializer(config->allow_nv_shader_extension); + serializer(config->use_gles); + } + serializer(config->make_mesh_block_local); + serializer(config->optimize_mesh_reordered_mapping); + serializer(config->mesh_localize_to_end_mapping); + serializer(config->mesh_localize_from_end_mapping); + serializer(config->mesh_localize_all_attr_mappings); + serializer(config->demote_no_access_mesh_fors); + serializer(config->experimental_auto_mesh_local); + serializer(config->auto_mesh_local_default_occupacy); + serializer.finalize(); + + return serializer.data; +} + +std::string get_offline_cache_key(CompileConfig *config, Kernel *kernel) { + std::string kernel_ast_string; + if (kernel) { + irpass::re_id(kernel->ir.get()); + irpass::print(kernel->ir.get(), &kernel_ast_string); + } + + std::vector compile_config_key; + if (config) { + compile_config_key = get_offline_cache_key_of_compile_config(config); + } + + picosha2::hash256_one_by_one hasher; + hasher.process(compile_config_key.begin(), compile_config_key.end()); + hasher.process(kernel_ast_string.begin(), kernel_ast_string.end()); + hasher.finish(); + + auto res = picosha2::get_hash_hex_string(hasher); res.insert(res.begin(), kernel->grad ? 'g' : 'n'); return res; } diff --git a/taichi/llvm/llvm_offline_cache.h b/taichi/llvm/llvm_offline_cache.h index 54da7bd7866a4..98589ce586069 100644 --- a/taichi/llvm/llvm_offline_cache.h +++ b/taichi/llvm/llvm_offline_cache.h @@ -8,7 +8,7 @@ namespace taichi { namespace lang { -std::string get_offline_cache_key_of_kernel(Kernel *kernel); +std::string get_offline_cache_key(CompileConfig *config, Kernel *kernel); struct LlvmOfflineCache { struct OffloadedTaskCacheData {