From 0769193185f50ea00be575f5d43cc209cbc1a00f Mon Sep 17 00:00:00 2001 From: He Wei Date: Tue, 23 Mar 2021 16:28:30 +0800 Subject: [PATCH] Fix a bug by revert GetCNodeTarget() for Load --- mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc | 4 ++++ mindspore/core/ir/anf.cc | 2 -- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc b/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc index 5cc124df9f..95cb264606 100644 --- a/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc +++ b/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc @@ -1221,9 +1221,13 @@ class AutoMonadConverter { } CNodePtr MakeLoad(const CNodePtr &cnode, const AnfNodePtr &ref, const AnfNodePtr &u) { + static const std::string primitive_target = "primitive_target"; // Create Load cnode. auto load_prim = NewValueNode(prim::kPrimLoad); auto load_cnode = func_graph_->NewCNode({load_prim, ref, u}); + // Set device target for Load CNode. + std::string target = GetCNodeTarget(cnode); + load_cnode->set_user_data(primitive_target, std::make_shared(target)); // Set load_cnode abstract to Tensor according the input Ref[Tensor]. auto ref_abs = dyn_cast(ref->abstract()); MS_EXCEPTION_IF_NULL(ref_abs); diff --git a/mindspore/core/ir/anf.cc b/mindspore/core/ir/anf.cc index 645b1833a1..379d0919a2 100644 --- a/mindspore/core/ir/anf.cc +++ b/mindspore/core/ir/anf.cc @@ -454,8 +454,6 @@ std::string GetCNodeTarget(const AnfNodePtr &node) { if (inputs.size() == 3 && !IsPrimitiveCNode(inputs[1], prim::kPrimMakeTuple)) { return GetCNodeTarget(inputs[1]); } - } else if (IsPrimitiveCNode(node, prim::kPrimLoad)) { - return GetCNodeTarget(cnode->input(1)); } else if (IsPrimitiveCNode(node, prim::kPrimMakeTuple)) { return GetMaketupleNodeTarget(cnode); } else if (IsPrimitiveCNode(node, prim::kPrimTupleGetItem)) {