1 /** 2 * Copyright 2021 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_STATIC_ANALYSIS_STACK_FRAME_H_ 18 #define MINDSPORE_CCSRC_PIPELINE_JIT_STATIC_ANALYSIS_STACK_FRAME_H_ 19 20 #include <utility> 21 #include <memory> 22 #include <string> 23 #include <vector> 24 25 #include "pipeline/jit/static_analysis/evaluator.h" 26 27 namespace mindspore { 28 namespace abstract { 29 class StackFrame; 30 using StackFramePtr = std::shared_ptr<StackFrame>; 31 using EvaluatorWeakPtr = std::weak_ptr<Evaluator>; 32 using BaseFuncGraphEvaluatorPtr = std::shared_ptr<BaseFuncGraphEvaluator>; 33 34 class StackFrame : public Base { 35 public: StackFrame(const EvaluatorPtr & evaluator,const FuncGraphPtr & func_graph,const AnalysisContextPtr & current_context,const AnalysisContextPtr & parent_context)36 StackFrame(const EvaluatorPtr &evaluator, const FuncGraphPtr &func_graph, const AnalysisContextPtr ¤t_context, 37 const AnalysisContextPtr &parent_context) 38 : evaluator_(EvaluatorWeakPtr(evaluator)), 39 func_graph_(func_graph), 40 current_context_(current_context), 41 parent_context_(parent_context), 42 slot_index_(0), 43 done_(false) { 44 Load(); 45 } 46 virtual ~StackFrame() = default; 47 Load()48 void Load() { 49 node_slots_ = TopoSort(func_graph_->get_return(), SuccIncoming, [this](const AnfNodePtr &node) -> IncludeType { 50 if (node->isa<ValueNode>() || node->isa<Parameter>()) { 51 return EXCLUDE; 52 } 53 return FOLLOW; 54 }); 55 slot_index_ = 0; 56 args_abs_list_.clear(); 57 } 58 59 // Check if we need branch to another func graph. 60 StackFramePtr Jump(const AnalysisEnginePtr &engine); 61 // Run one step in current func graph. 62 EvalResultPtr Step(const AnalysisEnginePtr &engine); 63 // Return back from branch func graph. 64 void Back(const AnalysisEnginePtr &engine, const StackFramePtr &last_stack_frame, const EvalResultPtr &eval_result); 65 Done()66 bool Done() { return done_; } 67 CurrentNode()68 AnfNodePtr &CurrentNode() { 69 if (slot_index_ >= node_slots_.size()) { 70 MS_LOG(EXCEPTION) << "The stack frame of " << func_graph_->ToString() 71 << " is invalid. Try to access frame sequence by index " << slot_index_ 72 << ", while the size is " << node_slots_.size() << "."; 73 } 74 return node_slots_[slot_index_]; 75 } 76 NextNode()77 AnfNodePtr &NextNode() { 78 auto ¤t_node = CurrentNode(); 79 // Set `done_` true, if the stack frames is being exhausted. 80 if (current_node == func_graph_->get_return()) { 81 done_ = true; 82 } 83 // Move cursor to next node. 84 slot_index_++; 85 return current_node; 86 } 87 evaluator()88 EvaluatorPtr evaluator() const { return evaluator_.lock(); } func_graph()89 FuncGraphPtr func_graph() const { return func_graph_; } current_context()90 AnalysisContextPtr current_context() const { return current_context_; } parent_context()91 AnalysisContextPtr parent_context() const { return parent_context_; } 92 args_abs_list()93 const AbstractBasePtrList &args_abs_list() { return args_abs_list_; } set_args_abs_list(const AbstractBasePtrList && args_abs_list)94 void set_args_abs_list(const AbstractBasePtrList &&args_abs_list) { args_abs_list_ = args_abs_list; } 95 ToString()96 std::string ToString() const override { 97 MS_EXCEPTION_IF_NULL(func_graph_); 98 std::ostringstream buffer; 99 buffer << "StackFrame: " << this << ", " << func_graph_->ToString(); 100 if (slot_index_ < node_slots_.size()) { 101 auto current_node = node_slots_[slot_index_]; 102 buffer << "(#" << slot_index_ << " / Running " << current_node->DebugString() << ")"; 103 } else { 104 buffer << "(Exhausted..)"; 105 } 106 buffer << ", parent: "; 107 auto parent_graph = parent_context_->func_graph(); 108 if (parent_graph != nullptr) { 109 buffer << parent_graph << "/" << parent_graph->ToString(); 110 } else { 111 buffer << "NULL"; 112 } 113 return buffer.str(); 114 } 115 116 friend std::ostream &operator<<(std::ostream &os, const StackFramePtr &frame) { 117 MS_EXCEPTION_IF_NULL(frame); 118 os << frame->ToString(); 119 return os; 120 } 121 122 private: 123 AbstractBasePtrList GenerateArgsAbsList(const AnalysisEnginePtr &engine, const EvaluatorPtr &evaluator, 124 const CNodePtr current_cnode); 125 AnalysisContextPtr GetParentContext(const BaseFuncGraphEvaluatorPtr &fg_evaluator, 126 const AbstractFunctionPtr &graph_func); 127 StackFramePtr DoJump(const AnalysisEnginePtr &engine, const CNodePtr current_cnode, 128 const AbstractFunctionPtr &graph_func); 129 130 EvaluatorWeakPtr evaluator_; 131 FuncGraphPtr func_graph_; 132 AnalysisContextPtr current_context_; 133 AnalysisContextPtr parent_context_; 134 AbstractBasePtrList args_abs_list_; 135 std::vector<AnfNodePtr> node_slots_; 136 size_t slot_index_; 137 bool done_; 138 }; 139 } // namespace abstract 140 } // namespace mindspore 141 #endif // MINDSPORE_CCSRC_PIPELINE_JIT_STATIC_ANALYSIS_STACK_FRAME_H_ 142