diff --git a/ge/single_op/single_op.cc b/ge/single_op/single_op.cc index 4b3f17cf..e326f3e0 100755 --- a/ge/single_op/single_op.cc +++ b/ge/single_op/single_op.cc @@ -361,6 +361,37 @@ Status DynamicSingleOp::SetHostTensorValue(const std::vector &input_desc, + const vector &input_buffers) { + for (auto &tensor_map : tensor_with_hostmem_) { + auto index = tensor_map.first; + if (index >= input_desc.size() || index >= input_buffers.size()) { + GELOGE(INTERNAL_ERROR, "[Check][Size]Index %d should smaller then input desc size %zu " + "and input buffers size %zu.", index, input_desc.size(), input_buffers.size()); + return INTERNAL_ERROR; + } + auto ge_tensor_desc = input_desc[index]; + // reconstruct GeTensor by DataBuffer + GeTensorPtr ge_tensor = MakeShared(ge_tensor_desc); + GE_CHECK_NOTNULL(ge_tensor); + GELOGD("The %d tensor input type is host, desc data type is %d, input buffer addr is %p, size is %ld.", + index, ge_tensor_desc.GetDataType(), input_buffers[index].data, input_buffers[index].length); + if (ge_tensor->SetData(reinterpret_cast(input_buffers[index].data), + static_cast(input_buffers[index].length)) != SUCCESS) { + GELOGE(INTERNAL_ERROR, "[Set][Data]Failed to set data of ge tensor."); + return INTERNAL_ERROR; + } + for (auto &tensor_desc : tensor_map.second) { + GE_CHECK_NOTNULL(tensor_desc); + if (!AttrUtils::SetTensor(tensor_desc, ATTR_NAME_VALUE, ge_tensor)) { + GELOGE(FAILED, "[Set][ATTR_NAME_VALUE]Failed to set ATTR_NAME_VALUE."); + return FAILED; + } + } + } + return SUCCESS; +} + Status DynamicSingleOp::ExecuteAsync(const vector &input_desc, const vector &input_buffers, vector &output_desc, @@ -374,6 +405,7 @@ Status DynamicSingleOp::ExecuteAsync(const vector &input_desc, if (!inputs_size.empty()) { StreamResource *stream_resource = SingleOpManager::GetInstance().GetResource(resource_id_, stream_); GE_CHK_STATUS_RET_NOLOG(UpdateInputsBufferAddr(stream_resource, stream_, inputs_size, update_buffers)); + GE_CHK_STATUS_RET_NOLOG(SetHostTensorValue(input_desc, input_buffers)); } if (hybrid_model_executor_ != nullptr) { diff --git a/ge/single_op/single_op.h b/ge/single_op/single_op.h index 01d6dfc0..deb4532e 100755 --- a/ge/single_op/single_op.h +++ b/ge/single_op/single_op.h @@ -81,9 +81,12 @@ class DynamicSingleOp { std::vector &outputs) const; Status SetHostTensorValue(const std::vector> &inputs_size, const vector &input_desc, const std::vector &input_buffers); + Status SetHostTensorValue(const vector &input_desc, const vector &input_buffers); std::unique_ptr op_task_; std::unique_ptr hybrid_model_; std::unique_ptr hybrid_model_executor_; + std::map> tensor_with_hostmem_; + uintptr_t resource_id_ = 0; std::mutex *stream_mutex_; rtStream_t stream_ = nullptr; diff --git a/ge/single_op/single_op_model.cc b/ge/single_op/single_op_model.cc index a4135999..6959c6b3 100755 --- a/ge/single_op/single_op_model.cc +++ b/ge/single_op/single_op_model.cc @@ -235,6 +235,13 @@ Status SingleOpModel::LoadAllNodes() { if (op_type == DATA_TYPE || op_type == AIPP_DATA_TYPE) { data_ops_.emplace_back(op_desc); + auto tensor = op_desc->MutableInputDesc(0); + if (AttrUtils::HasAttr(tensor, ATTR_NAME_VALUE)) { + int32_t index = 0; + (void) AttrUtils::GetInt(op_desc, ATTR_NAME_INDEX, index); + GELOGD("Node %s, index %d, has host mem.", node->GetName().c_str(), index); + op_with_hostmem_[index] = node; + } continue; } @@ -616,6 +623,7 @@ Status SingleOpModel::BuildDynamicOp(StreamResource &resource, DynamicSingleOp & if (need_hybrid_model) { GELOGD("Build single op HybridModel."); GE_CHK_STATUS_RET_NOLOG(hybrid::NodeExecutorManager::GetInstance().EnsureInitialized()); + GE_CHK_STATUS(SetHostMemTensor(single_op), "[Init][HostMem]Failed."); auto root_model = model_helper_.GetGeRootModel(); GE_CHECK_NOTNULL(root_model); root_model->SetRootGraph(GraphUtils::GetComputeGraph(ge_model->GetGraph())); @@ -634,4 +642,28 @@ Status SingleOpModel::BuildDynamicOp(StreamResource &resource, DynamicSingleOp & } return BuildTaskListForDynamicOp(&resource, single_op); } + +Status SingleOpModel::SetHostMemTensor(DynamicSingleOp &single_op) { + for (auto &node_map : op_with_hostmem_) { + auto node = node_map.second; + auto out_anchor = node->GetOutDataAnchor(0); + GE_CHECK_NOTNULL(out_anchor); + auto in_anchors = out_anchor->GetPeerInDataAnchors(); + vector tensor_descs; + auto idx = node_map.first; + for (auto anchor : in_anchors) { + GE_CHECK_NOTNULL(anchor); + auto output_node = anchor->GetOwnerNode(); + GE_CHECK_NOTNULL(output_node); + auto op_desc = output_node->GetOpDesc(); + GE_CHECK_NOTNULL(op_desc); + auto tensor_desc = op_desc->MutableInputDesc(anchor->GetIdx()); + tensor_descs.emplace_back(tensor_desc); + GELOGD("Get %d th input tensor desc of %s by %d data node: %s.", anchor->GetIdx(), + output_node->GetName().c_str(), idx, node->GetName().c_str()); + } + single_op.tensor_with_hostmem_[idx] = tensor_descs; + } + return SUCCESS; +} } // namespace ge diff --git a/ge/single_op/single_op_model.h b/ge/single_op/single_op_model.h index d900f09f..e7d07ee0 100755 --- a/ge/single_op/single_op_model.h +++ b/ge/single_op/single_op_model.h @@ -77,6 +77,7 @@ class SingleOpModel { static void ParseOpModelParams(ModelHelper &model_helper, SingleOpModelParam ¶m); void ParseArgTable(OpTask *task, SingleOp &op); Status InitHybridModelExecutor(const StreamResource &resource, const GeModelPtr &ge_model, SingleOp &single_op); + Status SetHostMemTensor(DynamicSingleOp &single_op); std::string model_name_; uint32_t model_id_ = 0; @@ -86,6 +87,7 @@ class SingleOpModel { ModelHelper model_helper_; map op_list_; + map op_with_hostmem_; SingleOpModelParam model_params_; std::vector input_offset_list_; diff --git a/tests/ut/ge/single_op/single_op_model_unittest.cc b/tests/ut/ge/single_op/single_op_model_unittest.cc index dadabaf6..f5d1a83c 100644 --- a/tests/ut/ge/single_op/single_op_model_unittest.cc +++ b/tests/ut/ge/single_op/single_op_model_unittest.cc @@ -27,6 +27,7 @@ #include "single_op/task/tbe_task_builder.h" #undef private #undef protected +#include "graph/passes/graph_builder_utils.h" using namespace std; using namespace testing; @@ -223,3 +224,19 @@ TEST_F(UtestSingleOpModel, test_build_dynamic_op) { model.BuildDynamicOp(res, dynamic_single_op); } +TEST_F(UtestSingleOpModel, test_host_mem) { + string model_data_str = "123456789"; + SingleOpModel model("model", model_data_str.c_str(), model_data_str.size()); + + // make graph + ut::GraphBuilder builder = ut::GraphBuilder("graph"); + auto data = builder.AddNode("Data", "Data", 0, 1); + auto netoutput = builder.AddNode("Netoutput", "NetOutput", 1, 0); + builder.AddDataEdge(data, 0, netoutput, 0); + auto graph = builder.GetGraph(); + model.op_with_hostmem_[0] = data; + + std::mutex stream_mu_; + DynamicSingleOp single_op(0, &stream_mu_, nullptr); + ASSERT_EQ(model.SetHostMemTensor(single_op), SUCCESS); +} diff --git a/tests/ut/ge/single_op/single_op_unittest.cc b/tests/ut/ge/single_op/single_op_unittest.cc index 8c2f6e51..3519811b 100644 --- a/tests/ut/ge/single_op/single_op_unittest.cc +++ b/tests/ut/ge/single_op/single_op_unittest.cc @@ -160,4 +160,23 @@ TEST_F(UtestSingleOp, test_singleop_execute_async2) { EXPECT_EQ(single_op.running_param_->mem_base, nullptr); EXPECT_EQ(single_op.tasks_.size(), 0); EXPECT_EQ(single_op.ExecuteAsync(input_buffers, output_buffers), PARAM_INVALID); -} \ No newline at end of file +} + +TEST_F(UtestSingleOp, test_set_host_mem) { + std::mutex stream_mu_; + DynamicSingleOp single_op(0, &stream_mu_, nullptr); + + vector input_buffers; + DataBuffer data_buffer; + input_buffers.emplace_back(data_buffer); + + vector input_descs; + GeTensorDesc tensor_desc1; + input_descs.emplace_back(tensor_desc1); + + vector op_input_descs; + auto tensor_desc2 = std::make_shared(); + op_input_descs.emplace_back(tensor_desc2); + single_op.tensor_with_hostmem_[0] = op_input_descs; + EXPECT_EQ(single_op.SetHostTensorValue(input_descs, input_buffers), SUCCESS); +}