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_FRONTEND_PARALLEL_GRAPH_UTIL_PIPELINE_SPLIT_UTILS_H_ 18 #define MINDSPORE_CCSRC_FRONTEND_PARALLEL_GRAPH_UTIL_PIPELINE_SPLIT_UTILS_H_ 19 20 #include <utility> 21 #include <vector> 22 #include <string> 23 #include "ir/anf.h" 24 #include "ir/manager.h" 25 26 namespace mindspore { 27 namespace parallel { 28 using PipelinePair = std::pair<std::vector<AnfNodePtr>, std::vector<AnfNodePtr>>; 29 AnfNodePtr FindAccuGrad(const CNodePtr &cnode); 30 bool IsLastStage(); 31 void InsertVirtualAssignAdd(const std::pair<AnfNodePtr, int> &node_user, const FuncGraphManagerPtr &manager, 32 const AnfNodePtr &accu_parameter); 33 void InsertVirtualAccuGrad(const AnfNodePtr &recv, const FuncGraphManagerPtr &manager, const AnfNodePtr ¶m); 34 AnfNodePtr FindGradAccuParameter(const std::vector<AnfNodePtr> ¶meters, const std::string &name); 35 void HandleReceiveParam(const FuncGraphPtr &root, const std::vector<AnfNodePtr> &all_nodes); 36 void AddVirtualAssignAdd(const FuncGraphPtr &root); 37 bool CompFunc(const AnfNodePtr &node1, const AnfNodePtr &node2); 38 void ReorderForForward(const std::vector<AnfNodePtr> &forward_start, const std::vector<AnfNodePtr> &forward_end, 39 const FuncGraphPtr &root); 40 void ReorderForBackward(const PipelinePair &forward_start_pair, const PipelinePair &forward_end_pair, 41 const PipelinePair &backward_start_pair, const PipelinePair &backward_end_pair, 42 const PipelinePair &forward_end_before_pair, const FuncGraphPtr &root); 43 void ReorderForParams(const std::vector<AnfNodePtr> &backward_params, const std::vector<AnfNodePtr> &forward_params, 44 const std::vector<AnfNodePtr> &allreduce_params, const PipelinePair &forward_params_pair, 45 const PipelinePair &backward_params_pair, const std::vector<AnfNodePtr> &backward_end, 46 const PipelinePair &forward_start_pair, const FuncGraphPtr &root); 47 int64_t GetMicroBatch(const AnfNodePtr &node); 48 void InsertDepend(const AnfNodePtr &prior_node, const AnfNodePtr &post_node, const FuncGraphManagerPtr &manager, 49 const FuncGraphPtr &root); 50 PipelinePair Deduplicate(const std::vector<AnfNodePtr> &node_vector, const FuncGraphPtr &root, int64_t micro_max); 51 AnfNodePtr GetActualOp(const AnfNodePtr &node); 52 void GetBorderNode(std::vector<AnfNodePtr> *forward_start, std::vector<AnfNodePtr> *forward_end, 53 std::vector<AnfNodePtr> *backward_start, std::vector<AnfNodePtr> *backward_end, 54 std::vector<AnfNodePtr> *forward_params, std::vector<AnfNodePtr> *backward_params, 55 std::vector<AnfNodePtr> *allreduce_params, const FuncGraphPtr &root); 56 void Reorder(const FuncGraphPtr &root); 57 void ReorderForPredict(const FuncGraphPtr &root, const FuncGraphManagerPtr &manager); 58 void HandleMicroBatch(const std::vector<AnfNodePtr> &all_nodes, const FuncGraphManagerPtr &manager); 59 void BroadCastMicroBatch(const CNodePtr &node, NodeUsersMap *node_users_map, const ValuePtr &value, size_t max_depth); 60 void LabelNeedGrad(const FuncGraphManagerPtr &manager, const FuncGraphPtr &root); 61 void BroadCastNeedGrad(const AnfNodePtr &node, NodeUsersMap *node_user_map, const FuncGraphPtr &root); 62 AnfNodePtr GetPreNode(const AnfNodePtr &node); 63 void LastStageEndNode(const std::vector<AnfNodePtr> &all_nodes, const FuncGraphManagerPtr &manager, 64 const FuncGraphPtr &root); 65 void SetStridedSliceStrategy(const AnfNodePtr &node); 66 void ParameterStartNode(const std::vector<AnfNodePtr> &all_nodes, const FuncGraphManagerPtr &manager); 67 ValuePtr Micro(const CNodePtr &cnode, NodeUsersMap *node_users_map, size_t max_depth); 68 void CheckBorderNode(const PipelinePair &forward_start_pair, const PipelinePair &forward_end_pair, 69 const PipelinePair &backward_start_pair, const PipelinePair &backward_end_pair, size_t micro_size); 70 } // namespace parallel 71 } // namespace mindspore 72 73 #endif // MINDSPORE_CCSRC_FRONTEND_PARALLEL_GRAPH_UTIL_PIPELINE_SPLIT_UTILS_H_ 74