xxq250
/
mindspore
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[auto-monad] Fix primitive_target and IsFeatureMapOutput check for Load

This commit is contained in:
He Wei 2021-03-17 17:23:03 +08:00
parent ddffb61c62
commit ce690a5489
3 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mindspore/ccsrc/backend/session/anf_runtime_algorithm.cc
View File

@ -1176,9 +1176,12 @@ AnfNodePtr AnfRuntimeAlgorithm::GetInputNode(const CNodePtr &node, size_t index)
bool AnfRuntimeAlgorithm::IsFeatureMapOutput(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
if (node->isa<ValueNode>() || IsPrimitiveCNode(node, prim::kPrimLoad)) {
if (node->isa<ValueNode>()) {
return false;
}
if (IsPrimitiveCNode(node, prim::kPrimLoad)) {
return IsFeatureMapOutput(node->cast<CNodePtr>()->input(1));
}
auto kernel_info = static_cast<const device::KernelInfo *>(node->kernel_info());
MS_EXCEPTION_IF_NULL(kernel_info);
return kernel_info->is_feature_map();

4
mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc
View File

@ -1220,13 +1220,9 @@ 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<std::string>(target));
// Set load_cnode abstract to Tensor according the input Ref[Tensor].
auto ref_abs = dyn_cast<abstract::AbstractRef>(ref->abstract());
MS_EXCEPTION_IF_NULL(ref_abs);

2
mindspore/core/ir/anf.cc
View File

@ -454,6 +454,8 @@ 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)) {