diff --git a/mindspore/ccsrc/backend/common/session/anf_runtime_algorithm.cc b/mindspore/ccsrc/backend/common/session/anf_runtime_algorithm.cc index e0aad4c4f05..f9d913e26cb 100644 --- a/mindspore/ccsrc/backend/common/session/anf_runtime_algorithm.cc +++ b/mindspore/ccsrc/backend/common/session/anf_runtime_algorithm.cc @@ -1054,7 +1054,7 @@ void AnfRuntimeAlgorithm::InferShape(const CNodePtr &node, std::mapset_abstract(eval_result); diff --git a/mindspore/ccsrc/include/common/utils/anfalgo.h b/mindspore/ccsrc/include/common/utils/anfalgo.h index 8e867ee0e99..8989009060d 100644 --- a/mindspore/ccsrc/include/common/utils/anfalgo.h +++ b/mindspore/ccsrc/include/common/utils/anfalgo.h @@ -168,9 +168,7 @@ class COMMON_EXPORT AnfAlgo { static std::vector GetOutputMaxShape(const AnfNodePtr &anf_node, size_t index); static std::vector GetOutputMinShape(const AnfNodePtr &anf_node, size_t index); static bool IsHostKernel(const CNodePtr &node); - // return true if use cnode_input's abstract, false if use real_input's abstract - static void AddArgList(AbstractBasePtrList *args_spec_list, const AnfNodePtr &cnode_input, - const AnfNodePtr &real_input); + static void AddArgList(AbstractBasePtrList *args_spec_list, const AnfNodePtr &real_input, size_t real_input_index); // Find real input nodes. static void GetAllFatherRealNode(const AnfNodePtr &anf_node, std::vector *result, std::set *visited); diff --git a/mindspore/ccsrc/kernel/kernel.cc b/mindspore/ccsrc/kernel/kernel.cc index 9bba03a5c44..c196d33c31b 100644 --- a/mindspore/ccsrc/kernel/kernel.cc +++ b/mindspore/ccsrc/kernel/kernel.cc @@ -57,20 +57,26 @@ void KernelMod::InferShape() { AbstractBasePtrList args_spec_list; auto primitive = GetValueNode(inputs[0]); auto input_size = common::AnfAlgo::GetInputTensorNum(cnode); + bool skip_nop_node = !context->get_param(MS_CTX_ENABLE_MINDRT); for (size_t i = 0; i < input_size; i++) { - auto input_node_with_index = common::AnfAlgo::GetPrevNodeOutput(cnode, i); - auto real_input = input_node_with_index.first; + AnfNodePtr real_input = nullptr; + size_t real_input_index = 0; + if (real_input_nodes_.count(i) > 0) { + real_input = real_input_nodes_[i].first.lock(); + real_input_index = real_input_nodes_[i].second; + } else { + auto input_node_with_index = common::AnfAlgo::GetPrevNodeOutput(cnode, i); + real_input = input_node_with_index.first; + real_input_index = input_node_with_index.second; + } MS_EXCEPTION_IF_NULL(real_input); - auto cnode_input = cnode->input(i + 1); - MS_EXCEPTION_IF_NULL(cnode_input); - InferShapeForNopNode(real_input); + if (skip_nop_node) { + InferShapeForNopNode(real_input); + } if (depend_list_.find(i) != depend_list_.end()) { - auto pre_node_with_index = common::AnfAlgo::GetPrevNodeOutput(cnode, i); - bool skip_nop_node = !context->get_param(MS_CTX_ENABLE_MINDRT); - auto output_addr = AnfAlgo::GetPrevNodeMutableOutputAddr(cnode, i, skip_nop_node); - std::vector shapes = - trans::GetRuntimePaddingShape(pre_node_with_index.first, pre_node_with_index.second); - auto host_type = common::AnfAlgo::GetOutputInferDataType(pre_node_with_index.first, pre_node_with_index.second); + auto output_addr = AnfAlgo::GetMutableOutputAddr(real_input, real_input_index, skip_nop_node); + auto shapes = trans::GetRuntimePaddingShape(real_input, real_input_index); + auto host_type = common::AnfAlgo::GetOutputInferDataType(real_input, real_input_index); auto out_tensor = std::make_shared(host_type, shapes); MS_EXCEPTION_IF_NULL(out_tensor); // The second parameter must be false, otherwise the device address cannot be released and allocated, and the @@ -88,14 +94,14 @@ void KernelMod::InferShape() { if (real_abs->isa()) { real_abs->set_value(out_tensor); } else if (real_abs->isa()) { - auto tuple_get_item_index = common::AnfAlgo::GetTupleGetItemOutIndex(cnode_input->cast()); auto abstract_tuple = real_abs->cast(); MS_EXCEPTION_IF_NULL(abstract_tuple); - auto tuple_elements = abstract_tuple->elements()[tuple_get_item_index]; + MS_EXCEPTION_IF_CHECK_FAIL((real_input_index < abstract_tuple->elements().size()), "Index is out of range."); + auto tuple_elements = abstract_tuple->elements()[real_input_index]; tuple_elements->set_value(out_tensor); } } - common::AnfAlgo::AddArgList(&args_spec_list, cnode_input, real_input); + common::AnfAlgo::AddArgList(&args_spec_list, real_input, real_input_index); } auto eval_result = opt::CppInferShape(primitive, args_spec_list); cnode->set_abstract(eval_result); diff --git a/mindspore/ccsrc/kernel/kernel.h b/mindspore/ccsrc/kernel/kernel.h index 1b8877aa0a1..4e9d426404e 100644 --- a/mindspore/ccsrc/kernel/kernel.h +++ b/mindspore/ccsrc/kernel/kernel.h @@ -20,6 +20,7 @@ #include #include #include +#include #include "nlohmann/json.hpp" #include "ir/anf.h" #include "ir/dtype.h" @@ -222,6 +223,10 @@ class KernelMod { // set true if need to update output's shape after launch in dynamic_shape, like Unique virtual bool IsNeedUpdateOp() { return is_need_updateop_; } + void InsertRealInputNode(const AnfNodePtr &pre_node, size_t pre_node_out_index, size_t input_index) { + real_input_nodes_[input_index] = {pre_node, pre_node_out_index}; + } + protected: void InferShape(); void GetDepndLists(const CNodePtr &cnode); @@ -249,6 +254,11 @@ class KernelMod { std::vector inputs_addr_; std::vector workspaces_addr_; std::vector outputs_addr_; + + // HashMap > is used to record the real input node to infer the + // dynamic shape information of the nodes located at the boundary of the graph partition, such as heterogeneous + // scenario and so on. + mindspore::HashMap> real_input_nodes_; }; using KernelModPtr = std::shared_ptr; } // namespace kernel diff --git a/mindspore/ccsrc/runtime/device/executor/dynamic_kernel.cc b/mindspore/ccsrc/runtime/device/executor/dynamic_kernel.cc index 178def51085..99634b933b8 100644 --- a/mindspore/ccsrc/runtime/device/executor/dynamic_kernel.cc +++ b/mindspore/ccsrc/runtime/device/executor/dynamic_kernel.cc @@ -105,7 +105,7 @@ void DynamicKernel::InferShape() { tuple_elements->set_value(out_tensor); } } - common::AnfAlgo::AddArgList(&args_spec_list, cnode_input, real_input); + common::AnfAlgo::AddArgList(&args_spec_list, real_input, input_node_with_index.second); } auto eval_result = opt::CppInferShape(primitive, args_spec_list); cnode->set_abstract(eval_result); diff --git a/mindspore/ccsrc/runtime/graph_scheduler/actor/data_source_actor.cc b/mindspore/ccsrc/runtime/graph_scheduler/actor/data_source_actor.cc index 89b86e1a914..9e02c844735 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/actor/data_source_actor.cc +++ b/mindspore/ccsrc/runtime/graph_scheduler/actor/data_source_actor.cc @@ -246,7 +246,7 @@ void HostQueueDataSourceActor::OnMemoryAllocFinish(OpContext *cons if (tensor_device_address.get() == device_tensor) { continue; } - if ((!Copy(device_tensor, tensor_device_address.get()))) { + if (!Copy(device_tensor, tensor_device_address.get())) { SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), "Copy data failed."); } continue; diff --git a/mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc b/mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc index b676f037882..68446ed2b9b 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc +++ b/mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc @@ -1182,6 +1182,15 @@ void GraphScheduler::LinkDataArrowForInternalParameter(AbstractActor *const, Abs kernel_type = actor_pair.first->type_; } + // Update the real input node. + MS_EXCEPTION_IF_NULL(to_kernel_with_input_idx.first); + if (to_kernel_with_input_idx.first->isa()) { + auto kernel_mod = AnfAlgo::GetKernelMod(to_kernel_with_input_idx.first->cast()); + MS_EXCEPTION_IF_NULL(kernel_mod); + kernel_mod->InsertRealInputNode(real_from_kernel_with_output_idx.first, real_from_kernel_with_output_idx.second, + to_kernel_with_input_idx.second); + } + if (kKernelTypeToLinkFunc.count(kernel_type) == 0) { MS_LOG(EXCEPTION) << "Invalid internal parameter:" << internal_parameter->DebugString() << ", type:" << kernel_type; } @@ -2036,8 +2045,8 @@ void GraphScheduler::PersistDeviceTensor(const GraphCompilerInfo &graph_compiler front_node = FetchFrontNodeByBackendNode(input_node, graph); } // The front node may be value node in the heterogeneous scene, needs to handle. - if ((front_node == nullptr) || ((front_node->isa() || front_node->isa()) && - (!parser->IsRootGraphPersistentDeviceTensor(front_node)))) { + if ((front_node == nullptr) || + (!front_node->isa() && !parser->IsRootGraphPersistentDeviceTensor(front_node))) { continue; } diff --git a/mindspore/ccsrc/utils/anfalgo.cc b/mindspore/ccsrc/utils/anfalgo.cc index b32039b5922..725d815a286 100644 --- a/mindspore/ccsrc/utils/anfalgo.cc +++ b/mindspore/ccsrc/utils/anfalgo.cc @@ -1421,29 +1421,22 @@ bool AnfAlgo::IsHostKernel(const CNodePtr &kernel_node) { return true; } -void AnfAlgo::AddArgList(AbstractBasePtrList *args_spec_list, const AnfNodePtr &cnode_input, - const AnfNodePtr &real_input) { - if (AnfAlgo::CheckPrimitiveType(cnode_input, prim::kPrimTupleGetItem)) { - // cppcheck-suppress unreadVariable - auto lock = AnfUtils::GetAbstractLock(real_input.get()); - auto base_shape = real_input->Shape(); - if (!base_shape->isa()) { - MS_LOG(EXCEPTION) << "Node input is a tuple_get_item but real input node shape is not a TupleShape. trace: " - << trace::DumpSourceLines(real_input); - } - auto abs = real_input->abstract()->Clone()->cast(); - MS_EXCEPTION_IF_NULL(abs); - auto tuple_get_item_indexk = AnfAlgo::GetTupleGetItemOutIndex(cnode_input->cast()); - auto abs_i = abs->elements()[tuple_get_item_indexk]; - (void)args_spec_list->emplace_back(abs_i); - } else if (cnode_input->isa() && AnfAlgo::GetCNodeName(cnode_input) == prim::kPrimReshape->name()) { - // cppcheck-suppress unreadVariable - auto lock = AnfUtils::GetAbstractLock(cnode_input.get()); - (void)args_spec_list->emplace_back(cnode_input->abstract()->Clone()); +void AnfAlgo::AddArgList(AbstractBasePtrList *args_spec_list, const AnfNodePtr &real_input, size_t real_input_index) { + MS_EXCEPTION_IF_NULL(args_spec_list); + MS_EXCEPTION_IF_NULL(real_input); + + // cppcheck-suppress unreadVariable + auto lock = AnfUtils::GetAbstractLock(real_input.get()); + auto real_abs = real_input->abstract(); + MS_EXCEPTION_IF_NULL(real_abs); + if (real_abs->isa()) { + auto abs_tuple = real_abs->Clone()->cast(); + MS_EXCEPTION_IF_NULL(abs_tuple); + MS_EXCEPTION_IF_CHECK_FAIL((real_input_index < abs_tuple->elements().size()), "Index is out of range."); + auto abs_index = abs_tuple->elements()[real_input_index]; + (void)args_spec_list->emplace_back(abs_index); } else { - // cppcheck-suppress unreadVariable - auto lock = AnfUtils::GetAbstractLock(real_input.get()); - (void)args_spec_list->emplace_back(real_input->abstract()->Clone()); + (void)args_spec_list->emplace_back(real_abs->Clone()); } }