1 /** 2 * Copyright 2022 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_PIPELINE_JIT_COMPILE_CACHE_MANAGER_H_ 18 #define MINDSPORE_CCSRC_PIPELINE_JIT_COMPILE_CACHE_MANAGER_H_ 19 20 #include <string> 21 #include <memory> 22 #include "pybind11/pybind11.h" 23 #include "include/backend/kernel_graph.h" 24 #include "ir/func_graph.h" 25 #include "load_mindir/load_model.h" 26 27 namespace mindspore { 28 namespace pipeline { 29 namespace py = pybind11; 30 31 // A class for loading and caching the func_graph. 32 class CompileCacheManager { 33 public: CompileCacheManager(size_t compile_cache_id)34 explicit CompileCacheManager(size_t compile_cache_id) : compile_cache_id_(compile_cache_id) {} 35 CompileCacheManager() = default; 36 37 ~CompileCacheManager() = default; 38 39 // Get cached data queue name 40 static std::string GetCachedDataQueueName(const std::string &dataset_phase); 41 // Get the hash of dependent files when compiling graph. 42 void InitCompileCacheHash(const py::list &compile_cache_dep_files); 43 // Init group checkpoint file path for parallel mode. 44 void InitParallelGroupCkptSaveFile(); 45 // Compare the dependency files hash. 46 bool CanLoadCache(); 47 // Load the cached func_graph from mindir file. 48 FuncGraphPtr GetCachedFuncGraph(const FuncGraphManagerPtr &manager, const py::dict &weights, 49 const std::string &queue_name); 50 // Export the func_graph to mindir file. 51 void CacheFuncGraph(const FuncGraphPtr &fg, const FuncGraphPtr &layout_fg); 52 layout_map()53 const LayoutMap &layout_map() const { return layout_map_; } 54 SetCompileCacheDir(const std::string & dir)55 void SetCompileCacheDir(const std::string &dir) { compile_cache_dir_ = dir; } CompileCacheDir()56 std::string CompileCacheDir() const { return compile_cache_dir_; } 57 static size_t data_queue_num_; 58 59 private: 60 size_t compile_cache_id_; 61 std::string compile_cache_dep_files_hash_; 62 LayoutMap layout_map_; 63 std::string compile_cache_dir_; 64 }; 65 using CompileCacheManagerPtr = std::shared_ptr<CompileCacheManager>; 66 } // namespace pipeline 67 } // namespace mindspore 68 #endif // MINDSPORE_CCSRC_PIPELINE_JIT_COMPILE_CACHE_MANAGER_H_ 69