text: /**
   * Copyright 2019-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 "graph/passes/memcpy_addr_async_pass.h"
  
  #include "common/ge/ge_util.h"
  #include "framework/common/debug/log.h"
  #include "graph/utils/node_utils.h"
  #include "graph/utils/op_desc_utils.h"
  #include "graph/utils/tensor_utils.h"
  
  namespace ge {
  Status MemcpyAddrAsyncPass::Run(ComputeGraphPtr graph) {
    GE_CHECK_NOTNULL(graph);
    for (auto &node : graph->GetAllNodes()) {
      auto op_desc = node->GetOpDesc();
      GE_IF_BOOL_EXEC(op_desc == nullptr, continue);
  
      if (op_desc->GetType() == STREAMSWITCHN || op_desc->GetType() == STREAMMERGE) {
        Status ret = AddMemcpyAddrAsyncNode(graph, node);
        if (ret != SUCCESS) {
          GELOGE(ret, "AddMemcpyAddrAsyncNode failed.");
          return ret;
        }
      }
      // handle data->netoutput, const->netoutput in root graph, use mem_addr_async to improve performance
      if (op_desc->GetType() == NETOUTPUT) {
        // check this netoutput is on root graph
        if (node->GetOwnerComputeGraph()->GetParentNode() == nullptr) {
          Status ret = InsertMemAddrAsyncNodeBeforeNetoutput(node->GetOwnerComputeGraph(), node);
          if (ret != SUCCESS) {
            GELOGE(ret, "AddMemcpyAddrAsyncNode failed.");
            return ret;
          }
        }
      }
    }
    return SUCCESS;
  }
  
  Status MemcpyAddrAsyncPass::AddMemcpyAddrAsyncNode(const ComputeGraphPtr &graph, const NodePtr &node) {
    GELOGI("Start AddMemcpyAddrAsyncNode for %s.", node->GetName().c_str());
    for (InDataAnchorPtr &in_data_anchor : node->GetAllInDataAnchors()) {
      OutDataAnchorPtr peer_out_anchor = in_data_anchor->GetPeerOutAnchor();
      GE_IF_BOOL_EXEC(peer_out_anchor == nullptr, continue);
      NodePtr in_node = peer_out_anchor->GetOwnerNode();
  
      if (in_node->GetType() == DATA) {
        ComputeGraphPtr owner_graph = in_node->GetOwnerComputeGraph();
        GE_CHECK_NOTNULL(owner_graph);
        // Data is in parent_graph
        if (owner_graph->GetParentGraph() == nullptr) {
          GELOGI("Need to insert MemcpyAddrAsync directly when data in parent graph.");
          NodePtr memcpy_addr_async_node = CreateMemcpyAddrAsyncNode(graph, peer_out_anchor, node);
          GE_IF_BOOL_EXEC(memcpy_addr_async_node == nullptr, GELOGE(INTERNAL_ERROR, "CreateMemcpyAddrAsyncNode failed.");
                          return INTERNAL_ERROR);
  
          Status ret = InsertMemcpyAddrAsyncNode(peer_out_anchor, in_data_anchor, memcpy_addr_async_node);
          GE_IF_BOOL_EXEC(ret != SUCCESS, GELOGE(ret, "InsertMemcpyAddrAsyncNode failed."); return ret);
        } else {
          uint32_t parent_index = 0;
          if (!AttrUtils::GetInt(in_node->GetOpDesc(), ATTR_NAME_PARENT_NODE_INDEX, parent_index)) {
            GELOGE(INTERNAL_ERROR, "Failed to get parent index of %s", in_node->GetName().c_str());
            return INTERNAL_ERROR;
          }
          // Data is in sub_graph
          GELOGI("Need to find data in parent graph, then insert MemcpyAddrAsync.");
          NodePtr parent_node = owner_graph->GetParentNode();
          user_data_for_known_ = in_node;
          out_of_user_data_for_known_ = node;
          peer_out_anchor_for_known_ = peer_out_anchor;
          in_anchor_for_known_ = in_data_anchor;
          FindUserData(parent_node, parent_index);
          if (find_user_data_) {
            GELOGI("Insert memcpy_addr_async for non_dynamic.");
            GE_CHECK_NOTNULL(peer_out_anchor_);
            NodePtr memcpy_addr_async_node = CreateMemcpyAddrAsyncNode(graph, peer_out_anchor_, out_of_user_data_);
            GE_IF_BOOL_EXEC(memcpy_addr_async_node == nullptr,
                            GELOGE(INTERNAL_ERROR, "CreateMemcpyAddrAsyncNode failed.");
                            return INTERNAL_ERROR);
  
            Status ret = InsertMemcpyAddrAsyncNode(peer_out_anchor_, in_anchor_, memcpy_addr_async_node);
            GE_IF_BOOL_EXEC(ret != SUCCESS, GELOGE(ret, "InsertMemcpyAddrAsyncNode failed."); return ret);
          }
          if (find_user_data_for_known_) {
            GELOGI("Insert memcpy_addr_async for known graph.");
            auto sub_graph = user_data_for_known_->GetOwnerComputeGraph();
            NodePtr memcpy_addr_async_node =
              CreateMemcpyAddrAsyncNode(sub_graph, peer_out_anchor_for_known_, out_of_user_data_for_known_);
            GE_IF_BOOL_EXEC(memcpy_addr_async_node == nullptr,
                            GELOGE(INTERNAL_ERROR, "CreateMemcpyAddrAsyncNode for known failed.");
                            return INTERNAL_ERROR);
  
            Status ret =
              InsertMemcpyAddrAsyncNode(peer_out_anchor_for_known_, in_anchor_for_known_, memcpy_addr_async_node);
            GE_IF_BOOL_EXEC(ret != SUCCESS, GELOGE(ret, "InsertMemcpyAddrAsyncNode for known failed."); return ret);
          }
        }
      }
    }
    return SUCCESS;
  }
  
  void MemcpyAddrAsyncPass::FindUserDataForKnown(const NodePtr &parent_node, uint32_t &parent_index) {
    GELOGI("Start FindUserDataForKnown of %s.", parent_node->GetName().c_str());
    if (user_data_for_known_->GetOpDesc() == nullptr) {
      GELOGI("Cannot get op_desc of %s.", user_data_for_known_->GetName().c_str());
      return;
    }
    string src_var_name;
    if (ge::AttrUtils::GetStr(user_data_for_known_->GetOpDesc(), REF_VAR_SRC_VAR_NAME, src_var_name)) {
      GELOGI("The data in known graph is variable, no need to insert memcpy_addr_async.");
      find_user_data_for_known_ = false;
      return;
    } else {
      find_user_data_for_known_ = true;
    }
  }
  
  void MemcpyAddrAsyncPass::FindUserDataForNonDynamic(const ge::NodePtr &parent_node, uint32_t &parent_index) {
    GELOGI("Start to FindUserDataForNonDynamic of %s.", parent_node->GetName().c_str());
    InDataAnchorPtr in_data_anchor = parent_node->GetInDataAnchor(parent_index);
    OutDataAnchorPtr out_anchor = in_data_anchor->GetPeerOutAnchor();
    GE_IF_BOOL_EXEC(out_anchor == nullptr,
                    GELOGE(INTERNAL_ERROR, "Cannot find out_anchor of %s.", parent_node->GetName().c_str());
                    return );
    NodePtr in_node = out_anchor->GetOwnerNode();
    GELOGI("in_node of parent_node is %s.", in_node->GetName().c_str());
    if (in_node->GetType() == DATA) {
      if (in_node->GetOwnerComputeGraph()->GetParentGraph() != nullptr) {
        // DATA is in sub graph again, update user_data of known firstly
        user_data_for_known_ = in_node;
        out_of_user_data_for_known_ = parent_node;
        peer_out_anchor_for_known_ = out_anchor;
        in_anchor_for_known_ = in_data_anchor;
        NodePtr pre_in_node = in_node->GetOwnerComputeGraph()->GetParentNode();
        if (!AttrUtils::GetInt(in_node->GetOpDesc(), ATTR_NAME_PARENT_NODE_INDEX, parent_index)) {
          GELOGE(INTERNAL_ERROR, "Failed to refresh parent index of %s", in_node->GetName().c_str());
          return;
        }
        FindUserData(pre_in_node, parent_index);
      } else {
        // DATA is in parent graph and not has input
        user_data_ = in_node;
        out_of_user_data_ = parent_node;
        peer_out_anchor_ = out_anchor;
        in_anchor_ = in_data_anchor;
        find_user_data_ = true;
        GELOGI("%s connect with %s, will insert memcpyaddr.", user_data_->GetName().c_str(),
               out_of_user_data_->GetName().c_str());
      }
    } else if (in_node->GetType() == IF || in_node->GetType() == WHILE || in_node->GetType() == CASE) {
      if (!AttrUtils::GetInt(parent_node->GetOpDesc(), ATTR_NAME_PARENT_NODE_INDEX, parent_index)) {
        GELOGE(INTERNAL_ERROR, "Failed to refresh parent index of %s", in_node->GetName().c_str());
        return;
      }
      FindUserData(in_node, parent_index);
    } else {
      GELOGI("%s connect with %s, which is not user_data.", parent_node->GetName().c_str(), in_node->GetName().c_str());
      find_user_data_ = false;
    }
  }
  
  void MemcpyAddrAsyncPass::FindUserData(const NodePtr &parent_node, uint32_t &parent_index) {
    auto parent_op_desc = parent_node->GetOpDesc();
    if (parent_op_desc == nullptr) {
      GELOGI("Cannot get op_desc of %s.", parent_node->GetName().c_str());
      return;
    }
    bool is_unknown_shape = false;
    if (parent_node->GetType() == PARTITIONEDCALL &&
        AttrUtils::GetBool(parent_op_desc, ATTR_NAME_IS_UNKNOWN_SHAPE, is_unknown_shape) && !is_unknown_shape) {
      FindUserDataForKnown(parent_node, parent_index);
    } else {
      FindUserDataForNonDynamic(parent_node, parent_index);
    }
  }
  
  NodePtr MemcpyAddrAsyncPass::CreateMemcpyAddrAsyncNode(const ComputeGraphPtr &graph,
                                                         const OutDataAnchorPtr &out_data_anchor,
                                                         const NodePtr &out_of_user_data) {
    GELOGD("Start CreateMemcpyAddrAsyncNode.");
    OpDescPtr pre_op_desc = out_data_anchor->GetOwnerNode()->GetOpDesc();
    GE_CHK_BOOL_EXEC(pre_op_desc != nullptr, return nullptr, "Op_desc of pre node is invalid.");
    std::string node_name = pre_op_desc->GetName() + "_" + MEMCPYADDRASYNC;
  
    OpDescPtr op_desc = MakeShared<OpDesc>(node_name, MEMCPYADDRASYNC);
    GE_CHECK_NOTNULL_EXEC(op_desc, return nullptr);
  
    if (op_desc->AddInputDesc(pre_op_desc->GetOutputDesc(out_data_anchor->GetIdx())) != GRAPH_SUCCESS) {
      GELOGE(INTERNAL_ERROR, "Add memcpy_addr_async input desc failed.");
      return nullptr;
    }
  
    if (op_desc->AddOutputDesc(pre_op_desc->GetOutputDesc(out_data_anchor->GetIdx())) != GRAPH_SUCCESS) {
      GELOGE(INTERNAL_ERROR, "Add memcpy_addr_async output desc failed.");
      return nullptr;
    }
  
    int64_t stream_id = out_of_user_data->GetOpDesc()->GetStreamId();
    op_desc->SetStreamId(stream_id);
    GELOGI("SetStreamId: Node %s assign stream is %ld.", op_desc->GetName().c_str(), stream_id);
    bool labeled_input = false;
    (void)ge::AttrUtils::GetBool(out_of_user_data->GetOpDesc(), ATTR_NAME_NODE_CONNECT_INPUT, labeled_input);
    if (labeled_input) {
      if (!ge::AttrUtils::SetBool(out_of_user_data->GetOpDesc(), ATTR_NAME_NODE_CONNECT_INPUT, false)) {
        GELOGE(FAILED, "Failed to unset attr %s for node %s.", ATTR_NAME_NODE_CONNECT_INPUT.c_str(),
               out_of_user_data->GetName().c_str());
        return nullptr;
      }
      if (!ge::AttrUtils::SetBool(op_desc, ATTR_NAME_NODE_CONNECT_INPUT, true)) {
        GELOGE(FAILED, "Failed to set attr %s for node %s.", ATTR_NAME_NODE_CONNECT_INPUT.c_str(),
               op_desc->GetName().c_str());
        return nullptr;
      }
    }
  
    NodePtr memcpy_addr_async_node = graph->AddNodeAfter(op_desc, out_data_anchor->GetOwnerNode());
    GE_CHECK_NOTNULL_EXEC(memcpy_addr_async_node, return nullptr);
  
    return memcpy_addr_async_node;
  }
  
  Status MemcpyAddrAsyncPass::InsertMemcpyAddrAsyncNode(const OutDataAnchorPtr &out_anchor,
                                                        const InDataAnchorPtr &in_anchor, const NodePtr &node) {
    // insert memcpy_addr of each user_data and out_of_user_data
    if (GraphUtils::RemoveEdge(out_anchor, in_anchor) != GRAPH_SUCCESS) {
      GELOGE(INTERNAL_ERROR, "Remove edge of %s and %s failed.", out_anchor->GetOwnerNode()->GetName().c_str(),
             in_anchor->GetOwnerNode()->GetName().c_str());
      return INTERNAL_ERROR;
    }
    if (GraphUtils::AddEdge(out_anchor, node->GetInDataAnchor(0)) != GRAPH_SUCCESS) {
      GELOGE(INTERNAL_ERROR, "Add edge of %s and %s failed.", out_anchor->GetOwnerNode()->GetName().c_str(),
             node->GetName().c_str());
      return INTERNAL_ERROR;
    }
    if (GraphUtils::AddEdge(node->GetOutDataAnchor(0), in_anchor) != GRAPH_SUCCESS) {
      GELOGE(INTERNAL_ERROR, "Add edge of %s and %s failed.", node->GetName().c_str(),
             in_anchor->GetOwnerNode()->GetName().c_str());
      return INTERNAL_ERROR;
    }
    return SUCCESS;
  }
  
  Status MemcpyAddrAsyncPass::InsertMemAddrAsyncNodeBeforeNetoutput(const ComputeGraphPtr &graph, const NodePtr &node) {
    GELOGI("Start AddMemcpyAddrAsyncNode for %s.", node->GetName().c_str());
    for (const auto &in_data_anchor : node->GetAllInDataAnchors()) {
      auto in_node = NodeUtils::GetInDataNodeByIndex(*node, in_data_anchor->GetIdx());
      GE_CHECK_NOTNULL(in_node);
      auto peer_out_anchor = in_data_anchor->GetPeerOutAnchor();
      if ((in_node->GetType() != CONSTANT) && (in_node->GetType() != CONSTANTOP) && (in_node->GetType() != DATA)) {
        continue;
      }
      auto desc = in_node->GetOpDesc();
      GE_CHECK_NOTNULL(desc);
      if (IsEmptyTenor(desc->GetOutputDesc(peer_out_anchor->GetIdx()).GetShape())) {
        continue;
      }
      GELOGI("Need to insert MemcpyAddrAsync before netoutput on parent graph.");
      NodePtr memcpy_addr_async_node = CreateMemcpyAddrAsyncNode(graph, peer_out_anchor, in_node);
      GE_IF_BOOL_EXEC(memcpy_addr_async_node == nullptr, GELOGE(INTERNAL_ERROR, "CreateMemcpyAddrAsyncNode failed.");
                      return INTERNAL_ERROR);
  
      Status ret = InsertMemcpyAddrAsyncNode(peer_out_anchor, in_data_anchor, memcpy_addr_async_node);
      GE_IF_BOOL_EXEC(ret != SUCCESS, GELOGE(ret, "InsertMemcpyAddrAsyncNode failed."); return ret);
      GELOGI("Insert mem_addr_async node %s success between %s and %s.", memcpy_addr_async_node->GetName().c_str(),
             in_node->GetName().c_str(), node->GetName().c_str());
      // if src node is const, need to update attr and offset here because this pass process is after offset set.
      if ((in_node->GetType() == CONSTANT) || (in_node->GetType() == CONSTANTOP)) {
        NodeUtils::UpdateIsInputConst(memcpy_addr_async_node);
        auto output_desc = node->GetOpDesc();
        GE_CHECK_NOTNULL(output_desc);
        auto output_tensor_desc = output_desc->MutableInputDesc(static_cast<uint32_t>(in_data_anchor->GetIdx()));
        int64_t data_offset = 0;
        (void)TensorUtils::GetDataOffset(*output_tensor_desc, data_offset);
        auto input_tensor = memcpy_addr_async_node->GetOpDesc()->MutableInputDesc(0);
        GELOGI("Need update const Offset %ld to op [%s]", data_offset, memcpy_addr_async_node->GetName().c_str());
        TensorUtils::SetDataOffset(*input_tensor, data_offset);
        TensorUtils::SetDataOffset(*output_tensor_desc, 0);
      }
    }
    NodeUtils::UpdateIsInputConst(node);
    return SUCCESS;
  }
  
  bool MemcpyAddrAsyncPass::IsEmptyTenor(const GeShape &shape) const {
    for (const auto dim : shape.GetDims()) {
      if (dim == 0) {
        return true;
      }
    }
    return false;
  }
  }  // namespace ge
  