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.
89 lines
3.4 KiB
89 lines
3.4 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/reshape_recovery_pass.h"
|
|
#include "common/ge/ge_util.h"
|
|
|
|
namespace ge {
|
|
namespace {
|
|
NodePtr CreateReshape(const ConstGeTensorDescPtr &src, const ConstGeTensorDescPtr &dst, const ComputeGraphPtr &graph) {
|
|
static std::atomic_long reshape_num(0);
|
|
auto next_num = reshape_num.fetch_add(1);
|
|
auto reshape = MakeShared<OpDesc>("Reshape_ReshapeRecoveryPass_" + std::to_string(next_num), RESHAPE);
|
|
if (reshape == nullptr) {
|
|
return nullptr;
|
|
}
|
|
auto ret = reshape->AddInputDesc("x", *src);
|
|
if (ret != GRAPH_SUCCESS) {
|
|
return nullptr;
|
|
}
|
|
ret = reshape->AddInputDesc("shape", GeTensorDesc(GeShape(), Format(), DT_INT32));
|
|
if (ret != GRAPH_SUCCESS) {
|
|
return nullptr;
|
|
}
|
|
ret = reshape->AddOutputDesc("y", *dst);
|
|
if (ret != GRAPH_SUCCESS) {
|
|
return nullptr;
|
|
}
|
|
|
|
return graph->AddNode(reshape);
|
|
}
|
|
|
|
Status InsertReshapeIfNeed(const NodePtr &node) {
|
|
GE_CHECK_NOTNULL(node);
|
|
GE_CHECK_NOTNULL(node->GetOpDesc());
|
|
for (auto src_anchor : node->GetAllOutDataAnchors()) {
|
|
auto src_tensor = node->GetOpDesc()->GetOutputDescPtr(src_anchor->GetIdx());
|
|
GE_CHECK_NOTNULL(src_tensor);
|
|
for (auto dst_anchor : src_anchor->GetPeerInDataAnchors()) {
|
|
auto dst_node = dst_anchor->GetOwnerNode();
|
|
GELOGD("Try insert reshape between %s[%d] and %s[%d] to keep the shape continues",
|
|
node->GetName().c_str(), src_anchor->GetIdx(), dst_node->GetName().c_str(), dst_anchor->GetIdx());
|
|
GE_CHECK_NOTNULL(dst_node);
|
|
GE_CHECK_NOTNULL(dst_node->GetOpDesc());
|
|
auto dst_tensor = dst_node->GetOpDesc()->GetInputDescPtr(dst_anchor->GetIdx());
|
|
GE_CHECK_NOTNULL(dst_tensor);
|
|
bool is_need_insert_reshape = src_tensor->GetShape().GetDims() != UNKNOWN_RANK &&
|
|
dst_tensor->GetShape().GetDims() != UNKNOWN_RANK &&
|
|
src_tensor->GetShape().GetDims() != dst_tensor->GetShape().GetDims();
|
|
if (is_need_insert_reshape) {
|
|
auto reshape = CreateReshape(src_tensor, dst_tensor, node->GetOwnerComputeGraph());
|
|
GE_CHECK_NOTNULL(reshape);
|
|
auto ret = GraphUtils::InsertNodeBetweenDataAnchors(src_anchor, dst_anchor, reshape);
|
|
if (ret != GRAPH_SUCCESS) {
|
|
GELOGE(INTERNAL_ERROR, "Failed to insert reshape between node %s and %s",
|
|
node->GetName().c_str(), dst_node->GetName().c_str());
|
|
return INTERNAL_ERROR;
|
|
}
|
|
GELOGI("Insert reshape between %s and %s to keep the shape continues",
|
|
node->GetName().c_str(), dst_node->GetName().c_str());
|
|
}
|
|
}
|
|
}
|
|
return SUCCESS;
|
|
}
|
|
} // namespace
|
|
|
|
Status ReshapeRecoveryPass::Run(ComputeGraphPtr graph) {
|
|
for (const auto &node : graph->GetDirectNode()) {
|
|
auto ret = InsertReshapeIfNeed(node);
|
|
if (ret != SUCCESS) {
|
|
return ret;
|
|
}
|
|
}
|
|
return SUCCESS;
|
|
}
|
|
} // namespace ge
|