1 /** 2 * Copyright 2020-2023 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_LITE_SRC_RUNTIME_SCHEDULER_H_ 18 #define MINDSPORE_LITE_SRC_RUNTIME_SCHEDULER_H_ 19 20 #include <utility> 21 #include <vector> 22 #include <memory> 23 #include <map> 24 #include <deque> 25 #include <unordered_map> 26 #include <set> 27 #include <string> 28 #include "src/executor/sub_graph_kernel.h" 29 #include "src/litert/inner_context.h" 30 #include "include/model.h" 31 #include "src/litert/scheduler_cb.h" 32 #include "include/api/delegate.h" 33 #include "src/control_flow/control_flow_scheduler.h" 34 #include "src/litert/runtime_shape_fusion_pass.h" 35 36 namespace mindspore::lite { 37 constexpr int kDefaultDeviceType = -1; 38 class Scheduler { 39 public: 40 Scheduler(InnerContext *ctx, const mindspore::Context *ms_ctx, Model *src_model, std::vector<Tensor *> *src_tensors, 41 std::vector<Tensor *> *input_tensors, std::vector<Tensor *> *output_tensors, bool is_train_session, 42 int *is_infershape, bool *is_control_flow, bool *infer_along_running, 43 std::map<std::string, TypeId> *executions, std::shared_ptr<Delegate> delegate = nullptr, 44 int delegate_device_type = -1) context_(ctx)45 : context_(ctx), 46 ms_context_(ms_ctx), 47 src_model_(src_model), 48 src_tensors_(src_tensors), 49 inputs_(input_tensors), 50 outputs_(output_tensors), 51 is_train_session_(is_train_session), 52 is_control_flow_(is_control_flow), 53 infer_along_running_(infer_along_running), 54 is_infershape_(is_infershape), 55 delegate_(delegate), 56 delegate_device_type_(delegate_device_type), 57 execution_plan_(executions) {} 58 ~Scheduler() = default; 59 int Schedule(std::vector<kernel::KernelExec *> *dst_kernels); SetupSchedulerCb(std::unique_ptr<SchedulerCb> cb)60 void SetupSchedulerCb(std::unique_ptr<SchedulerCb> cb) { sched_cb_ = std::move(cb); } SetConfig(const std::map<std::string,std::map<std::string,std::string>> * config_info)61 void SetConfig(const std::map<std::string, std::map<std::string, std::string>> *config_info) { 62 config_info_ = config_info; 63 } 64 std::vector<kernel::KernelExec *> NonTailCallNodes(); 65 66 private: 67 int SchedulePreProcess(); 68 int CheckInputParam(const std::vector<kernel::KernelExec *> *dst_kernels) const; 69 void FindNodeInoutTensors(const LiteGraph::Node &node, std::vector<Tensor *> *inputs, std::vector<Tensor *> *outputs); 70 LiteGraph::Node *NodeInputIsPartial(const LiteGraph::Node *node); 71 int InferPartialShape(const LiteGraph::Node *node); 72 int InferCallShape(const LiteGraph::Node *node); 73 int InferNodeShape(const LiteGraph::Node *node); 74 void FreeOpParameters(); 75 int InferSubGraphShape(size_t subgraph_index); 76 // schedule a node to kernel according to context and kernels registered 77 int HandleBuildinCpuKernelWeight(const kernel::SubGraphType &belong_subgraph_type, const kernel::KernelExec *kernel); 78 kernel::KernelExec *FindBackendKernel(const std::vector<Tensor *> &in_tensors, 79 const std::vector<Tensor *> &out_tensors, const LiteGraph::Node *node, 80 TypeId prefer_data_type = kTypeUnknown); 81 int FindCpuKernel(const std::vector<Tensor *> &in_tensors, const std::vector<Tensor *> &out_tensors, 82 OpParameter *op_parameter, const kernel::KernelKey &desc, TypeId kernel_data_type, 83 kernel::KernelExec **kernel); 84 int CheckCpuValid(const std::vector<kernel::KernelExec *> *dst_kernels) const; 85 void ResetByExecutionPlan(std::string node_name, TypeId *data_type); 86 87 #ifdef GPU_OPENCL 88 int FindGpuKernel(const std::vector<Tensor *> &in_tensors, const std::vector<Tensor *> &out_tensors, 89 OpParameter *op_parameter, const kernel::KernelKey &desc, kernel::KernelExec **kernel, 90 TypeId prefer_data_type); 91 #endif 92 int FindProviderKernel(const std::vector<Tensor *> &in_tensors, const std::vector<Tensor *> &out_tensors, 93 const LiteGraph::Node *node, TypeId data_type, kernel::KernelExec **kernel); 94 95 int InitKernels(std::vector<kernel::KernelExec *> &&dst_kernels); 96 kernel::KernelExec *SchedulePartialToKernel(const lite::LiteGraph::Node *src_node); 97 // schedule a partial node to a subgraph_kernel 98 std::vector<kernel::KernelExec *> ScheduleSubGraphToSubGraphKernels(const int &subgraph_index); 99 // schedule a node to a kernel 100 kernel::KernelExec *ScheduleNodeToKernel(const LiteGraph::Node *src_node, TypeId prefer_data_type = kTypeUnknown); 101 // schedule a Model::Graph into a vector of subgraph_kernel 102 int ScheduleGraphToKernels(std::vector<kernel::KernelExec *> *dst_kernels, TypeId prefer_data_type = kTypeUnknown); 103 // schedule a LiteGraph::SubGraph into a vector of kernel and subgraph_kernel 104 int ScheduleSubGraphToKernels(size_t subgraph_index, std::vector<kernel::KernelExec *> *dst_kernels, 105 std::vector<lite::Tensor *> *in_tensors, std::vector<lite::Tensor *> *out_tensors, 106 TypeId prefer_data_type = kTypeUnknown); 107 // vector<KernelExec/SubGraphKernel> --> vector<SubGraphKernel> 108 int ConstructNormalSubGraphs(const std::vector<kernel::KernelExec *> &src_kernel, 109 std::vector<kernel::KernelExec *> *dst_kernel, 110 std::map<const kernel::KernelExec *, bool> *sinked_kernel_map); 111 112 int ConstructSubGraphs(std::vector<kernel::KernelExec *> *dst_kernel); 113 114 int ProcessSubGraphTranspose(std::vector<kernel::KernelExec *> *dst_kernels); 115 116 // create subgraph_kernel from a vector of kernel 117 std::vector<kernel::KernelExec *> ScheduleMainSubGraphToKernels(); 118 kernel::KernelExec *SchedulePartialToSubGraphKernel(const int &subgraph_index); 119 kernel::SubGraphType PartialSubGraphType(const std::vector<kernel::KernelExec *> &kernels); 120 121 // other methods 122 static TypeId GetFirstFp32Fp16OrInt8Type(const std::vector<Tensor *> &in_tensors); 123 int CopyPartialShapeToSubGraph(const lite::LiteGraph::Node *partial_node); 124 int RestoreSubGraphInput(const lite::LiteGraph::Node *partial_node); 125 126 bool IsControlFlowPattern(const lite::LiteGraph::Node &partial_node); 127 STATUS DelQuantDTypeCastKernel(std::vector<kernel::KernelExec *> *kernels); 128 #ifdef ENABLE_FP16 129 int SubGraphPreferDataType(const int &subgraph_index, TypeId *prefer_data_type); 130 #endif 131 int InferSwitchShape(const LiteGraph::Node *node); 132 LiteGraph::Node *NodeInputIsSwitchType(const LiteGraph::Node *node); 133 bool SubGraphHasScheduled(const int &index); 134 void SubGraphMarkScheduled(const int &index); 135 int ConstructControlFlowMainGraph(std::vector<kernel::KernelExec *> *kernels); 136 137 #ifndef DELEGATE_CLIP 138 /* delegate related */ 139 int ReplaceDelegateKernels(std::vector<kernel::KernelExec *> *dst_kernels); 140 int InitDelegateKernels(std::vector<kernel::KernelExec *> *dst_kernels); 141 #else InitDelegateKernels(std::vector<kernel::KernelExec * > * dst_kernels)142 int InitDelegateKernels(std::vector<kernel::KernelExec *> *dst_kernels) { return RET_OK; } 143 #endif 144 GetEnableGLTexture()145 bool GetEnableGLTexture() const { return context_->GetDeviceInfo(DT_GPU).gpu_device_info_.enable_gl_texture_; } GetGLContext()146 void *GetGLContext() const { return context_->GetDeviceInfo(DT_GPU).gpu_device_info_.gl_context_; } GetGLDisplay()147 void *GetGLDisplay() const { return context_->GetDeviceInfo(DT_GPU).gpu_device_info_.gl_display_; } 148 149 protected: 150 InnerContext *context_ = nullptr; 151 const mindspore::Context *ms_context_ = nullptr; 152 Model *src_model_ = nullptr; 153 std::vector<Tensor *> *src_tensors_; 154 std::vector<Tensor *> *inputs_; 155 std::vector<Tensor *> *outputs_; 156 std::vector<mindspore::MSTensor> ms_inputs_; 157 std::vector<mindspore::MSTensor> ms_outputs_; 158 std::vector<size_t> graph_output_node_indexes_; 159 std::map<int, OpParameter *> op_parameters_; 160 bool is_train_session_ = false; 161 bool *is_control_flow_ = nullptr; 162 bool *infer_along_running_ = nullptr; 163 int *is_infershape_ = nullptr; 164 std::unique_ptr<SchedulerCb> sched_cb_; 165 std::map<kernel::Kernel *, const schema::Primitive *> primitives_; 166 std::shared_ptr<Delegate> delegate_ = nullptr; 167 int delegate_device_type_ = -1; 168 std::deque<int> subgraphs_to_schedule_{}; 169 std::unordered_map<size_t, kernel::KernelExec *> subgraph_index_subgraph_kernel_map_{}; 170 std::set<int> scheduled_subgraph_index_{}; 171 std::unordered_map<kernel::KernelExec *, size_t> partial_kernel_subgraph_index_map_{}; 172 std::set<lite::LiteGraph::Node *> partial_cnode_inferred_{}; 173 ControlFlowSchedulerPtr control_flow_scheduler_ = nullptr; 174 std::map<std::string, TypeId> *execution_plan_ = nullptr; 175 const std::map<std::string, std::map<std::string, std::string>> *config_info_ = nullptr; 176 std::shared_ptr<ShapeFusionPass> shape_fusion_pass_ = nullptr; 177 std::vector<int> infer_subgraph_index_; 178 }; 179 } // namespace mindspore::lite 180 181 #endif // MINDSPORE_LITE_SRC_RUNTIME_SCHEDULER_H_ 182