diff --git a/ge/graph/load/model_manager/task_info/kernel_ex_task_info.cc b/ge/graph/load/model_manager/task_info/kernel_ex_task_info.cc index b9644946..400b51bf 100644 --- a/ge/graph/load/model_manager/task_info/kernel_ex_task_info.cc +++ b/ge/graph/load/model_manager/task_info/kernel_ex_task_info.cc @@ -29,6 +29,10 @@ #include "hybrid/node_executor/aicpu/aicpu_ext_info.h" #include "framework/common/debug/log.h" +namespace { +const char *const kAicpuAllshape = "_AllShape"; +} // namespace + namespace ge { Status KernelExTaskInfo::InitTaskExtInfo(const std::string &ext_info, const OpDescPtr &op_desc) { if (ext_info.empty()) { @@ -50,6 +54,24 @@ Status KernelExTaskInfo::InitTaskExtInfo(const std::string &ext_info, const OpDe GE_CHK_STATUS_RET(ext_handle->UpdateExecuteMode(true), "UpdateExecuteMode failed."); GELOGD("Update aicpu_task ext_info bit_map execute mode to 1."); + bool all_shape = false; + (void)AttrUtils::GetBool(op_desc, kAicpuAllshape, all_shape); + if (all_shape) { + for (uint32_t i = 0; i < num_inputs; i++) { + auto input_desc = op_desc->MutableInputDesc(i); + GE_CHECK_NOTNULL(input_desc); + GE_CHK_STATUS_RET(ext_handle->UpdateInputShapeAndType(i, *input_desc), + "Input[%zu] update input shape failed.", i); + } + if (unknown_type = DEPEND_COMPUTE) { + for (uint32_t j = 0; j < num_outputs; j++) { + auto output_desc = op_desc->MutableOutputDesc(j); + GE_CHECK_NOTNULL(output_desc); + GE_CHK_STATUS_RET(ext_handle->UpdateOutputShapeAndType(j, *output_desc), + "Output[%zu] update output shape failed.", j); + } + } + } auto rt_ret = rtMalloc(&ext_info_addr_, ext_handle->GetExtInfoLen(), RT_MEMORY_HBM); GE_IF_BOOL_EXEC(rt_ret != RT_ERROR_NONE, GELOGE(RT_FAILED, "rtMalloc ext_info error: 0x%X, size=%zu", rt_ret, ext_info.size()); diff --git a/ge/graph/load/model_manager/task_info/kernel_task_info.cc b/ge/graph/load/model_manager/task_info/kernel_task_info.cc index 1d2a74cc..54dfc2b9 100755 --- a/ge/graph/load/model_manager/task_info/kernel_task_info.cc +++ b/ge/graph/load/model_manager/task_info/kernel_task_info.cc @@ -43,6 +43,7 @@ constexpr int64_t kInvalidGroupKey = -1; constexpr uint32_t kSKTSingleSize = 1; const char *kIsLastNode = "is_last_node"; const char *kIsFirstNode = "is_first_node"; +const char *const kAicpuAllshape = "_AllShape"; const int64_t kCloseSkt = 100; const uint32_t kAddrLen = sizeof(void *); const int kBaseInt = 10; @@ -985,6 +986,22 @@ Status KernelTaskInfo::InitAicpuTaskExtInfo(const std::string &ext_info) { GE_CHK_STATUS_RET(ext_handle->UpdateExecuteMode(true), "UpdateExecuteMode failed."); GELOGD("Update aicpu_task ext_info bit_map execute mode to 1."); + bool all_shape = false; + (void)AttrUtils::GetBool(op_desc_, kAicpuAllshape, all_shape); + if (all_shape) { + for (uint32_t i = 0; i < num_inputs; i++) { + auto input_desc = op_desc_->MutableInputDesc(i); + GE_CHECK_NOTNULL(input_desc); + GE_CHK_STATUS_RET(ext_handle->UpdateInputShapeAndType(i, *input_desc), + "Input[%zu] update input shape failed.", i); + } + for (uint32_t j = 0; j < num_outputs; j++) { + auto output_desc = op_desc_->MutableOutputDesc(j); + GE_CHECK_NOTNULL(output_desc); + GE_CHK_STATUS_RET(ext_handle->UpdateOutputShapeAndType(j, *output_desc), + "Output[%zu] update output shape failed.", j); + } + } auto rt_ret = rtMalloc(&aicpu_ext_info_addr_, ext_handle->GetExtInfoLen(), RT_MEMORY_HBM); if (rt_ret != RT_ERROR_NONE) { GELOGE(RT_FAILED, "rtMalloc ext_info error: 0x%X, size=%zu", rt_ret, ext_info.size()); diff --git a/ge/hybrid/model/node_item.cc b/ge/hybrid/model/node_item.cc index acc250ef..c37bdb57 100644 --- a/ge/hybrid/model/node_item.cc +++ b/ge/hybrid/model/node_item.cc @@ -222,8 +222,8 @@ void NodeItem::ResolveUnknownShapeType() { Status NodeItem::Init() { GE_CHK_STATUS_RET_NOLOG(InitInputsAndOutputs()); GE_CHK_STATUS_RET_NOLOG(ResolveDynamicState()); + ResolveUnknownShapeType(); if (is_dynamic) { - ResolveUnknownShapeType(); GE_CHK_STATUS_RET_NOLOG(ResolveStaticInputsAndOutputs()); GE_CHK_STATUS_RET(ParseFusedSubgraph(*this), "[%s] Failed to parse fused subgraph", node_name.c_str()); } diff --git a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc index ca9e908b..b342151d 100755 --- a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc +++ b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc @@ -28,6 +28,7 @@ namespace hybrid { namespace { // mem need release constexpr uint64_t kReleaseFlag = 1; +const char *const kAicpuAllshape = "_AllShape"; } REGISTER_NODE_EXECUTOR_BUILDER(NodeExecutorManager::ExecutorType::AICPU_TF, AiCpuNodeExecutor); REGISTER_NODE_EXECUTOR_BUILDER(NodeExecutorManager::ExecutorType::AICPU_CUSTOM, AiCpuNodeExecutor); @@ -60,6 +61,7 @@ Status AicpuNodeTaskBase::InitExtInfo(const std::string &kernel_ext_info, int64_ GELOGD("To update aicpu_task ext_info session_info session_id to %lu", session_id); GE_CHK_STATUS_RET(aicpu_ext_handle_.UpdateSessionInfoSessionId(session_id), "UpdateSessionInfoSessionId failed."); + GE_CHK_STATUS_RET(aicpu_ext_handle_.UpdateExecuteMode(!node_item_->is_dynamic), "UpdateExecuteMode failed."); // copy task args buf GE_CHK_STATUS_RET(AllocTensorBuffer(aicpu_ext_handle_.GetExtInfoLen(), ext_info_addr_dev_), @@ -136,7 +138,6 @@ Status AicpuNodeTaskBase::UpdateExtInfo() { return SUCCESS; } - GE_CHK_STATUS_RET(aicpu_ext_handle_.UpdateExecuteMode(false), "UpdateExecuteMode failed."); for (auto i = 0; i < node_item_->num_inputs; ++i) { auto input_desc = node_item_->MutableInputDesc(i); GE_CHECK_NOTNULL(input_desc); @@ -176,10 +177,14 @@ Status AicpuNodeTaskBase::UpdateArgs(TaskContext &context) { } GE_CHK_STATUS_RET(UpdateIoAddr(context), "Node[%s] update io addr failed.", node_name_.c_str()); - if (node_item_->is_dynamic) { - // dynamic node need update ext info. + bool all_shape = false; + const OpDescPtr op_desc = node_item_->GetOpDesc(); + (void)AttrUtils::GetBool(op_desc, kAicpuAllshape, all_shape); + if (node_item_->is_dynamic || all_shape) { + // dynamic node and all_shape kernel need update ext info. GE_CHK_STATUS_RET(UpdateExtInfo(), "Node[%s] update ext info failed.", node_name_.c_str()); } + GELOGD("Node[%s] update args end.", node_name_.c_str()); return SUCCESS; } diff --git a/tests/ut/ge/graph/load/kernel_ex_task_info_unittest.cc b/tests/ut/ge/graph/load/kernel_ex_task_info_unittest.cc index 850a29ad..4c47163c 100644 --- a/tests/ut/ge/graph/load/kernel_ex_task_info_unittest.cc +++ b/tests/ut/ge/graph/load/kernel_ex_task_info_unittest.cc @@ -140,6 +140,7 @@ TEST_F(UtestKernelExTaskInfo, kernel_ex_task_info_calculate_args) { TEST_F(UtestKernelExTaskInfo, kernel_ex_task_ext_info) { const string ext_info = {1, 1, 1, 1, 0, 0, 0, 0}; const OpDescPtr op_desc = CreateOpDesc("FrameworkOp", "FrameworkOp"); + (void)AttrUtils::GetBool(op_desc, "_Allshape", true); KernelExTaskInfo kernel_ex_task_info; EXPECT_EQ(kernel_ex_task_info.InitTaskExtInfo(ext_info, op_desc), SUCCESS); diff --git a/tests/ut/ge/graph/load/kernel_task_info_unittest.cc b/tests/ut/ge/graph/load/kernel_task_info_unittest.cc index 65783ac9..5776ddfe 100644 --- a/tests/ut/ge/graph/load/kernel_task_info_unittest.cc +++ b/tests/ut/ge/graph/load/kernel_task_info_unittest.cc @@ -390,6 +390,7 @@ TEST_F(UtestKernelTaskInfo, init_kernel_taskInfo_with_aicpu_kernel_type_fail) { rtStreamCreate(&stream, 0); model.stream_list_ = { stream }; model.op_list_[0] = CreateOpDesc("FrameworkOp", "FrameworkOp"); + (void)AttrUtils::GetBool(op_list_[0], "_Allshape", true); domi::TaskDef task_def; KernelTaskInfo kernel_task_info;