1 /** 2 * Copyright 2019 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_RESOURCE_H_ 18 #define MINDSPORE_CCSRC_PIPELINE_JIT_RESOURCE_H_ 19 20 #include <iostream> 21 #include <vector> 22 #include <string> 23 #include <unordered_map> 24 #include <memory> 25 #include <unordered_set> 26 27 #include "pybind11/pybind11.h" 28 #include "pybind11/stl.h" 29 30 #include "utils/any.h" 31 #include "utils/profile.h" 32 #include "ir/manager.h" 33 34 #include "pipeline/jit/resource_base.h" 35 #include "pipeline/jit/static_analysis/prim.h" 36 #include "pipeline/jit/static_analysis/static_analysis.h" 37 38 namespace mindspore { 39 namespace pipeline { 40 41 namespace py = pybind11; 42 43 const char kBackend[] = "backend"; 44 const char kStepParallelGraph[] = "step_parallel"; 45 const char kOutput[] = "output"; 46 const char kPynativeGraphId[] = "graph_id"; 47 48 class InferenceResource; 49 50 using BuiltInTypeMap = std::unordered_map<int64_t, std::unordered_map<std::string, Any>>; 51 52 BuiltInTypeMap &GetMethodMap(); 53 54 BuiltInTypeMap &GetAttrMap(); 55 56 class Resource : public ResourceBase { 57 public: 58 explicit Resource(const py::object &obj = py::none()); 59 60 ~Resource() override; 61 engine()62 abstract::AnalysisEnginePtr engine() { return engine_; } 63 64 static bool IsTypeInBuiltInMap(const TypeId &type); 65 66 static Any GetMethodPtr(const TypeId &type, const std::string &name); 67 68 static Any GetAttrPtr(const TypeId &type, const std::string &name); 69 source_input()70 const py::object &source_input() const { return source_input_; } 71 func_graph()72 FuncGraphPtr func_graph() const { return func_graph_; } set_func_graph(const FuncGraphPtr & func_graph)73 void set_func_graph(const FuncGraphPtr &func_graph) { func_graph_ = func_graph; } 74 args_spec()75 const abstract::AbstractBasePtrList &args_spec() const { return args_spec_; } set_args_spec(const abstract::AbstractBasePtrList & args_spec)76 void set_args_spec(const abstract::AbstractBasePtrList &args_spec) { args_spec_ = args_spec; } 77 set_vm_loop(const bool & flag,const int64_t size)78 void set_vm_loop(const bool &flag, const int64_t size) { 79 vm_loop_flag_ = flag; 80 loop_size_ = size; 81 } set_is_load(bool flag)82 void set_is_load(bool flag) { is_load_ = flag; } is_load()83 bool is_load() { return is_load_; } vm_loop_flag()84 bool vm_loop_flag() { return vm_loop_flag_; } loop_size()85 int64_t loop_size() { return loop_size_; } 86 // Reclaim resource and clear the cache. 87 // GraphExecutorPy::Compile() can be called multiple times, so cache 88 // should be cleared. 89 void Clean(); 90 91 private: 92 abstract::AnalysisEnginePtr engine_; 93 FuncGraphPtr func_graph_; 94 abstract::AbstractBasePtrList args_spec_; 95 // The source obj to compile, usually a `Cell` or `ms_function` decorated function. 96 py::object source_input_; 97 bool is_cleaned_; 98 // The func_graph_ is loaded from mindir 99 bool is_load_{false}; 100 bool vm_loop_flag_{false}; 101 int64_t loop_size_{1}; 102 }; 103 104 using ResourcePtr = std::shared_ptr<pipeline::Resource>; 105 106 } // namespace pipeline 107 } // namespace mindspore 108 109 #endif // MINDSPORE_CCSRC_PIPELINE_JIT_RESOURCE_H_ 110