text: /**
   * Copyright 2020 Huawei Technologies Co., Ltd
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  #include "frontend/parallel/allreduce_fusion/allreduce_graph.h"
  #include <algorithm>
  #include <functional>
  #include "ir/anf.h"
  #include "frontend/parallel/allreduce_fusion/allreduce_node.h"
  #include "frontend/parallel/ops_info/ops_utils.h"
  #include "utils/log_adapter.h"
  
  namespace mindspore {
  namespace parallel {
  Status AllreduceGraph::AddNode(const CNodePtr &node, const AnfNodePtr &para) {
    AllreduceNodePtr arnode;
    auto cnode_emplace_return = cnode_set_.emplace(node);
    if (!cnode_emplace_return.second) {
      MS_LOG(INFO) << "node: " << node->DebugString() << " has already been added!";
      auto cnode_arnode_pair = cnode_arnode_map_.find(node);
      if (cnode_arnode_pair == cnode_arnode_map_.end()) {
        MS_LOG(EXCEPTION) << "node is not in cnode_arnode_map_!";
      }
      arnode = cnode_arnode_pair->second;
    } else {
      arnode = std::make_shared<AllreduceNode>(AllreduceNode());
    }
  
    if (arnode->Init(node) != SUCCESS) {
      MS_LOG(ERROR) << "AllreduceNode Init failed";
      return FAILED;
    }
    if (arnode->AddPara(para) != SUCCESS) {
      MS_LOG(ERROR) << "AllreduceNode AddPara failed";
      return FAILED;
    }
    cnode_arnode_map_[node] = arnode;
  
    auto arnode_emplace_return = arnode_set_.insert(arnode);
    if (!arnode_emplace_return.second) {
      MS_LOG(INFO) << "node: " << node->DebugString() << "'s arnode has already been added!";
    }
    cnode_emplace_return = para_cnodeset_map_[para].emplace(node);
    if (!cnode_emplace_return.second) {
      MS_LOG(INFO) << "node: " << node->DebugString() << " already in para: " << para->fullname_with_scope()
                   << "'s cnodeset!";
    }
    auto para_emplace_return = cnode_paraset_map_[node].emplace(para);
    if (!para_emplace_return.second) {
      MS_LOG(INFO) << "para: " << para->fullname_with_scope() << " already in node: " << node->DebugString()
                   << "'s paraset!";
    }
    return SUCCESS;
  }
  
  Status AllreduceGraph::AddEdge(const CNodePtr &from, const CNodePtr &to, double dist) {
    auto from_arnode_iter = cnode_arnode_map_.find(from);
    if (from_arnode_iter == cnode_arnode_map_.end()) {
      MS_LOG(ERROR) << "cnode from: " << from->DebugString() << "has not been added";
      PrintCNodeSet();
      return FAILED;
    }
    auto to_arnode_iter = cnode_arnode_map_.find(to);
    if (to_arnode_iter == cnode_arnode_map_.end()) {
      MS_LOG(ERROR) << "cnode to: " << to->DebugString() << "has not been added";
      PrintCNodeSet();
      return FAILED;
    }
    auto from_arnode = from_arnode_iter->second;
    auto to_arnode = to_arnode_iter->second;
    if (from_arnode->AddNext(to_arnode) != SUCCESS) {
      MS_LOG(ERROR) << "from_arnode AddNext failed";
      return FAILED;
    }
    if (to_arnode->AddPrev(from_arnode, dist, &max_) != SUCCESS) {
      MS_LOG(ERROR) << "to_arnode AddPrev failed";
      return FAILED;
    }
    max_ = std::max(max_, to_arnode->depend_feat_size());
    MS_LOG(DEBUG) << "from " << from->DebugString() << ", to " << to->DebugString();
    MS_LOG(DEBUG) << "from depend_feat_size: " << from_arnode->depend_feat_size()
                  << ", to depend_feat_size: " << to_arnode->depend_feat_size();
    return SUCCESS;
  }
  
  bool AllreduceGraph::NodeInGraph(const CNodePtr &node) const {
    auto cnode_iter = cnode_set_.find(node);
    return !(cnode_iter == cnode_set_.end());
  }
  
  std::vector<AnfNodePtr> AllreduceGraph::GetParaByCost(double from, double to) {
    std::vector<AnfNodePtr> nodes;
    for (auto &cnode_arnode : cnode_arnode_map_) {
      MS_LOG(DEBUG) << "cnode: " << cnode_arnode.first->DebugString()
                    << ", depend_feat_size: " << cnode_arnode.second->depend_feat_size()
                    << " curr_para_size: " << cnode_arnode.second->curr_para_size();
      if ((cnode_arnode.second->depend_feat_size() <= to) && (cnode_arnode.second->depend_feat_size() > from)) {
        (void)nodes.insert(nodes.end(), cnode_paraset_map_[cnode_arnode.first].begin(),
                           cnode_paraset_map_[cnode_arnode.first].end());
      }
    }
    return nodes;
  }
  
  std::pair<std::vector<AnfNodePtr>, double> AllreduceGraph::GetParaByParaSize(double to, double para_size) {
    std::vector<AnfNodePtr> nodes;
    double cur_para_size = 0;
    double from = to;
    for (auto &arnode : arnode_vec_) {
      if ((arnode.depend_feat_size() - max_ <= EPS) && arnode.depend_feat_size() >= to) {
        continue;
      }
      if (para_size > 0 && cur_para_size >= para_size && arnode.depend_feat_size() < from) {
        return std::make_pair(nodes, from);
      }
      (void)nodes.insert(nodes.end(), arnode.paras().begin(), arnode.paras().end());
      cur_para_size += arnode.curr_para_size();
      from = arnode.depend_feat_size();
    }
    MS_LOG(INFO) << "GetParaByParaSize has reached head node! para_size: " << para_size
                 << " cur_para_size: " << cur_para_size << " from: " << from;
    return std::make_pair(nodes, from);
  }
  
  void AllreduceGraph::PrintCNodeSet() const {
    MS_LOG(INFO) << "CNodeSet:";
    for (auto &cnode : cnode_set_) {
      MS_LOG(INFO) << cnode->DebugString();
    }
  }
  
  void AllreduceGraph::PrintAllredueGraphInfo() const {
    MS_LOG(INFO) << "max: " << max_;
    for (auto &cnode_arnode : cnode_arnode_map_) {
      MS_LOG(INFO) << "cnode: " << cnode_arnode.first->DebugString();
      MS_LOG(INFO) << "arnode info: ";
      cnode_arnode.second->ToString();
    }
  }
  
  void AllreduceGraph::PrintArnodeVec() const {
    MS_LOG(INFO) << "ArnodeVec:";
    for (auto &arnode : arnode_vec_) {
      arnode.ToString();
    }
  }
  
  void AllreduceGraph::PrintArnodeSet() const {
    MS_LOG(INFO) << "ArnodeSet:";
    for (auto &arnode : arnode_set_) {
      arnode->ToString();
    }
  }
  
  void AllreduceGraph::SortArnode() {
    arnode_vec_.clear();
    for (auto &node : arnode_set_) {
      arnode_vec_.emplace_back(*node);
    }
    std::sort(arnode_vec_.begin(), arnode_vec_.end(), std::greater<>());
  }
  
  Status AllreduceGraph::RemoveExtraParas() {
    std::unordered_set<AnfNodePtr> para_map;
    for (auto &node : arnode_vec_) {
      for (auto &para : node.paras()) {
        auto emplac_result = para_map.emplace(para);
        if (!emplac_result.second) {
          MS_LOG(DEBUG) << "parameter: " << para->fullname_with_scope() << "in arnode";
          if (node.RemovePara(para) != SUCCESS) {
            MS_LOG(ERROR) << "remove para failed";
            return FAILED;
          }
        }
      }
    }
    return SUCCESS;
  }
  
  Status AllreduceGraph::set_head_cnode(const CNodePtr &node) {
    auto arnode = std::make_shared<AllreduceNode>(AllreduceNode());
    if (arnode->Init(node) != SUCCESS) {
      MS_LOG(ERROR) << "AllreduceNode Init failed";
    }
    head_cnode_ = node;
    cnode_arnode_map_[node] = arnode;
    auto arnode_emplace_return = arnode_set_.insert(arnode);
    if (!arnode_emplace_return.second) {
      MS_LOG(WARNING) << "node: " << node->DebugString() << "'s arnode has already been added!";
    }
    auto cnode_emplace_return = cnode_set_.emplace(node);
    if (!cnode_emplace_return.second) {
      MS_LOG(WARNING) << "node: " << node->DebugString() << " has already been added!";
    }
    return SUCCESS;
  }
  }  // namespace parallel
  }  // namespace mindspore
  