You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
120 lines
4.9 KiB
120 lines
4.9 KiB
/**
|
|
* 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 "graph/passes/variable_format_pass.h"
|
|
#include <map>
|
|
#include <set>
|
|
#include <string>
|
|
#include "framework/common/debug/ge_log.h"
|
|
|
|
namespace ge {
|
|
Status VariableFormatPass::Run(ge::ComputeGraphPtr graph) {
|
|
GE_CHECK_NOTNULL(graph);
|
|
|
|
for (auto &node : graph->GetDirectNode()) {
|
|
GE_IF_BOOL_EXEC(node->GetOpDesc() == nullptr, continue);
|
|
GE_IF_BOOL_EXEC(node->GetOpDesc()->GetType() != VARIABLE, continue);
|
|
|
|
ge::NodePtr use_node = nullptr;
|
|
if (GetApplyMomentumOpByVariableInput(node, use_node)) {
|
|
GE_CHK_STATUS_RET(UpdateVariableOutFormat(node, use_node), "update variable out format failed");
|
|
GE_CHK_STATUS_RET(UpdateApplyMomentumInputFormat(use_node), "update apply momentum input format failed");
|
|
}
|
|
}
|
|
|
|
return domi::SUCCESS;
|
|
}
|
|
|
|
bool VariableFormatPass::GetApplyMomentumOpByVariableInput(const ge::NodePtr &var_node, ge::NodePtr &use_node) {
|
|
GE_IF_BOOL_EXEC(var_node == nullptr, return false);
|
|
|
|
std::map<std::string, std::set<int>> confirm_ops = {{"ApplyMomentum", {1}}};
|
|
for (auto &out_anchor : var_node->GetAllOutDataAnchors()) {
|
|
for (auto &in_anchor : out_anchor->GetPeerInDataAnchors()) {
|
|
GE_IF_BOOL_EXEC(ConfirmUseOpAndIndexByAnchor(in_anchor, confirm_ops, use_node), return true);
|
|
}
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
bool VariableFormatPass::ConfirmUseOpAndIndexByAnchor(const ge::InDataAnchorPtr &in_anchor,
|
|
const map<string, std::set<int>> &confirm_ops,
|
|
ge::NodePtr &use_node) {
|
|
GE_IF_BOOL_EXEC(in_anchor == nullptr, return false);
|
|
ge::NodePtr dst_node = in_anchor->GetOwnerNode();
|
|
ge::OpDescPtr dst_op_desc = dst_node->GetOpDesc();
|
|
GE_IF_BOOL_EXEC(dst_op_desc == nullptr, return false);
|
|
const string &dst_type = dst_op_desc->GetType();
|
|
int input_index = in_anchor->GetIdx();
|
|
|
|
GELOGD("ConfirmUseOpAndIndex, var name %s, dst_type = %s, input index %d", dst_node->GetName().c_str(),
|
|
dst_type.c_str(), input_index);
|
|
|
|
GE_IF_BOOL_EXEC(confirm_ops.count(dst_type) > 0,
|
|
GE_IF_BOOL_EXEC(confirm_ops.at(dst_type).count(input_index) > 0, use_node = dst_node; return true););
|
|
return false;
|
|
}
|
|
|
|
Status VariableFormatPass::UpdateVariableOutFormat(const ge::NodePtr &var_node, ge::NodePtr &use_node) {
|
|
GE_CHECK_NOTNULL(var_node);
|
|
GE_CHECK_NOTNULL(use_node);
|
|
ge::OpDescPtr op_desc_ptr = use_node->GetOpDesc();
|
|
GE_CHECK_NOTNULL(op_desc_ptr);
|
|
GE_CHECK_NOTNULL(use_node->GetInDataAnchor(0));
|
|
GE_CHECK_NOTNULL(use_node->GetInDataAnchor(0)->GetPeerOutAnchor());
|
|
NodePtr in_node = use_node->GetInDataAnchor(0)->GetPeerOutAnchor()->GetOwnerNode();
|
|
if (in_node != nullptr) {
|
|
string in_op_type = in_node->GetType();
|
|
if ((in_op_type == VARIABLE) && (in_node->GetOpDesc() != nullptr) &&
|
|
(in_node->GetOpDesc()->MutableOutputDesc(0) != nullptr)) {
|
|
ge::Format format = in_node->GetOpDesc()->MutableOutputDesc(0)->GetFormat();
|
|
ge::OpDescPtr cur_op_desc_ptr = var_node->GetOpDesc();
|
|
if (cur_op_desc_ptr != nullptr) {
|
|
cur_op_desc_ptr->MutableOutputDesc(0)->SetFormat(format);
|
|
cur_op_desc_ptr->MutableOutputDesc(0)->SetOriginFormat(format);
|
|
}
|
|
}
|
|
}
|
|
return domi::SUCCESS;
|
|
}
|
|
|
|
Status VariableFormatPass::UpdateApplyMomentumInputFormat(const ge::NodePtr &node) {
|
|
GE_CHECK_NOTNULL(node);
|
|
ge::OpDescPtr op_desc_ptr = node->GetOpDesc();
|
|
GE_CHECK_NOTNULL(op_desc_ptr);
|
|
GE_CHECK_NOTNULL(node->GetInDataAnchor(0));
|
|
GE_CHECK_NOTNULL(node->GetInDataAnchor(0)->GetPeerOutAnchor());
|
|
GE_CHECK_NOTNULL(op_desc_ptr->MutableInputDesc(0));
|
|
GE_CHECK_NOTNULL(op_desc_ptr->MutableInputDesc(1));
|
|
GE_CHECK_NOTNULL(op_desc_ptr->MutableOutputDesc(0));
|
|
NodePtr in_node = node->GetInDataAnchor(0)->GetPeerOutAnchor()->GetOwnerNode();
|
|
if (in_node != nullptr) {
|
|
string in_op_type = in_node->GetType();
|
|
if ((in_op_type == VARIABLE) && (in_node->GetOpDesc() != nullptr)) {
|
|
ge::Format format = in_node->GetOpDesc()->MutableOutputDesc(0)->GetFormat();
|
|
op_desc_ptr->MutableInputDesc(0)->SetFormat(format);
|
|
op_desc_ptr->MutableInputDesc(0)->SetOriginFormat(format);
|
|
op_desc_ptr->MutableInputDesc(1)->SetFormat(format);
|
|
op_desc_ptr->MutableInputDesc(1)->SetOriginFormat(format);
|
|
op_desc_ptr->MutableOutputDesc(0)->SetFormat(format);
|
|
op_desc_ptr->MutableOutputDesc(0)->SetOriginFormat(format);
|
|
}
|
|
}
|
|
return domi::SUCCESS;
|
|
}
|
|
} // namespace ge
|