forked from huawei/mindspore2022
!31677 support the dynamic shape infer in the sub graph partition
Merge pull request !31677 from limingqi107/bug_fix4
This commit is contained in:
commit
56183ab741
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue