!31677 support the dynamic shape infer in the sub graph partition

Merge pull request !31677 from limingqi107/bug_fix4
This commit is contained in:
i-robot 2022-03-24 10:01:03 +00:00 committed by Gitee
commit 56183ab741
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
8 changed files with 60 additions and 44 deletions

View File

@ -1054,7 +1054,7 @@ void AnfRuntimeAlgorithm::InferShape(const CNodePtr &node, std::map<uint32_t, te
}
}
}
common::AnfAlgo::AddArgList(&args_spec_list, cnode_input, real_input);
common::AnfAlgo::AddArgList(&args_spec_list, real_input, input_with_index.second);
}
auto eval_result = opt::CppInferShape(primitive, args_spec_list);
node->set_abstract(eval_result);

View File

@ -168,9 +168,7 @@ class COMMON_EXPORT AnfAlgo {
static std::vector<int64_t> GetOutputMaxShape(const AnfNodePtr &anf_node, size_t index);
static std::vector<int64_t> 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<AnfNodePtr> *result,
std::set<AnfNodePtr> *visited);

View File

@ -57,20 +57,26 @@ void KernelMod::InferShape() {
AbstractBasePtrList args_spec_list;
auto primitive = GetValueNode<PrimitivePtr>(inputs[0]);
auto input_size = common::AnfAlgo::GetInputTensorNum(cnode);
bool skip_nop_node = !context->get_param<bool>(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<bool>(MS_CTX_ENABLE_MINDRT);
auto output_addr = AnfAlgo::GetPrevNodeMutableOutputAddr(cnode, i, skip_nop_node);
std::vector<int64_t> 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<tensor::Tensor>(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<abstract::AbstractTensor>()) {
real_abs->set_value(out_tensor);
} else if (real_abs->isa<abstract::AbstractTuple>()) {
auto tuple_get_item_index = common::AnfAlgo::GetTupleGetItemOutIndex(cnode_input->cast<CNodePtr>());
auto abstract_tuple = real_abs->cast<abstract::AbstractTuplePtr>();
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);

View File

@ -20,6 +20,7 @@
#include <memory>
#include <map>
#include <set>
#include <utility>
#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<AddressPtr> inputs_addr_;
std::vector<AddressPtr> workspaces_addr_;
std::vector<AddressPtr> outputs_addr_;
// HashMap <input_index, pair<pre_node, pre_node_output_index>> 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<size_t, std::pair<AnfNodeWeakPtr, size_t>> real_input_nodes_;
};
using KernelModPtr = std::shared_ptr<KernelMod>;
} // namespace kernel

View File

@ -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);

View File

@ -246,7 +246,7 @@ void HostQueueDataSourceActor::OnMemoryAllocFinish(OpContext<DeviceTensor> *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;

View File

@ -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<CNode>()) {
auto kernel_mod = AnfAlgo::GetKernelMod(to_kernel_with_input_idx.first->cast<CNodePtr>());
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<Parameter>() || front_node->isa<CNode>()) &&
(!parser->IsRootGraphPersistentDeviceTensor(front_node)))) {
if ((front_node == nullptr) ||
(!front_node->isa<ValueNode>() && !parser->IsRootGraphPersistentDeviceTensor(front_node))) {
continue;
}

View File

@ -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<abstract::TupleShape>()) {
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<abstract::AbstractTuplePtr>();
MS_EXCEPTION_IF_NULL(abs);
auto tuple_get_item_indexk = AnfAlgo::GetTupleGetItemOutIndex(cnode_input->cast<CNodePtr>());
auto abs_i = abs->elements()[tuple_get_item_indexk];
(void)args_spec_list->emplace_back(abs_i);
} else if (cnode_input->isa<CNode>() && 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<abstract::AbstractTuple>()) {
auto abs_tuple = real_abs->Clone()->cast<abstract::AbstractTuplePtr>();
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());
}
}