optimize_dependece pass enhance

pull/560/head
huanghui 5 years ago
parent 53b3d187b9
commit 2ed4ad0f2a

@ -28,8 +28,7 @@ namespace mindspore {
namespace opt {
constexpr auto kSingleInputIndex = 1;
namespace {
AnfNodePtr GetReplaceNode(const FuncGraphPtr &func_graph, const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(func_graph);
AnfNodePtr GetReplaceNode(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
if (!node->isa<CNode>()) {
return nullptr;
@ -41,15 +40,6 @@ AnfNodePtr GetReplaceNode(const FuncGraphPtr &func_graph, const AnfNodePtr &node
if (op_name != kTransDataOpName && op_name != prim::kPrimCast->name()) {
return nullptr;
}
auto manager = func_graph->manager();
MS_EXCEPTION_IF_NULL(manager);
// Check whether the node has only one output node.
if (manager->node_users().find(cnode) == manager->node_users().end()) {
MS_LOG(EXCEPTION) << "The node should be used by at least another node's input";
}
if (manager->node_users()[cnode].size() > 1) {
return nullptr;
}
CheckCNodeInputSize(cnode, kSingleInputIndex + 1);
return cnode->input(kSingleInputIndex);
}
@ -63,7 +53,7 @@ bool ReplaceMakeTuple(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
std::vector<AnfNodePtr> new_make_tuple_inputs;
bool need_update = false;
for (const auto &input : cnode->inputs()) {
AnfNodePtr replace_input = GetReplaceNode(func_graph, input);
AnfNodePtr replace_input = GetReplaceNode(input);
// If replace input is not null, it will be the input of the TransData or Cast.
if (replace_input == nullptr) {
new_make_tuple_inputs.push_back(input);
@ -119,7 +109,7 @@ const AnfNodePtr OptimizeDependence::Process(const FuncGraphPtr &func_graph, con
if (ReplaceMakeTuple(func_graph, replacing_cnode)) {
return nullptr;
}
AnfNodePtr replace_node = GetReplaceNode(func_graph, replacing_cnode);
AnfNodePtr replace_node = GetReplaceNode(replacing_cnode);
if (replace_node == nullptr) {
MS_LOG(DEBUG) << "Can not find the TransData or Cast with single output node. Depend node: " << node->DebugString();
return nullptr;

Loading…
Cancel
Save