From e9bab262b736d4d56e3d78c8ed921fa811b752bd Mon Sep 17 00:00:00 2001 From: zhaozhixuan Date: Mon, 28 Jun 2021 16:43:47 +0800 Subject: [PATCH 1/4] Fix bug of aicore input const. --- ge/single_op/task/op_task.cc | 54 +++++++++++++++++++++--- ge/single_op/task/op_task.h | 2 + tests/ut/ge/single_op/single_op_task_unittest.cc | 37 ++++++++++++++-- 3 files changed, 83 insertions(+), 10 deletions(-) diff --git a/ge/single_op/task/op_task.cc b/ge/single_op/task/op_task.cc index b6a78f9e..92d1e325 100755 --- a/ge/single_op/task/op_task.cc +++ b/ge/single_op/task/op_task.cc @@ -345,6 +345,53 @@ Status TbeOpTask::AllocateWorkspaces(const vector &workspace_sizes) { return SUCCESS; } +Status TbeOpTask::UpdateIoAddr(std::vector &args, const std::vector &inputs, + const std::vector &outputs) { + uintptr_t *arg_base = nullptr; + size_t arg_num = 0; + GetIoAddr(arg_base, arg_num); + + const vector v_is_input_const = op_desc_->GetIsInputConst(); + size_t non_const_index = 0; + for (size_t i = 0; i < op_desc_->GetAllInputsSize(); ++i) { + const GeTensorDescPtr tensor_desc = op_desc_->MutableInputDesc(static_cast(i)); + if (tensor_desc == nullptr) { + GELOGD("SingleOp: %s, Index: %zu, has no input", op_desc_->GetName().c_str(), i); + continue; + } + if (i < v_is_input_const.size() && v_is_input_const[i]) { + if (i >= arg_num) { + GELOGE(ACL_ERROR_GE_PARAM_INVALID, "[Check][Size] Args size is %zu, but get index is %zu.", arg_num, i); + REPORT_INNER_ERROR("E19999", "[Check][Size] Args size is %zu, but get index is %zu.", arg_num, i); + return ACL_ERROR_GE_PARAM_INVALID; + } + auto addr = reinterpret_cast(arg_base[i]); + GELOGD("SingleOp: %s, Index: %zu, input is const, addr = %p", op_desc_->GetName().c_str(), i, addr); + args.emplace_back(addr); + continue; + } + if (non_const_index >= inputs.size()) { + GELOGE(ACL_ERROR_GE_PARAM_INVALID, "[Check][Size] Input size is %zu, but get non_const_index is %zu", + inputs.size(), non_const_index); + REPORT_INNER_ERROR("E19999", "[Check][Size] Input size is %zu, but get non_const_index is %zu", + inputs.size(), non_const_index); + return ACL_ERROR_GE_PARAM_INVALID; + } + auto addr = inputs[non_const_index].data; + GELOGD("SingleOp: %s, input[%zu], addr = %p", op_desc_->GetName().c_str(), i, addr); + args.emplace_back(addr); + non_const_index++; + } + + for (size_t i = 0; i < outputs.size(); ++i) { + auto addr = outputs[i].data; + GELOGD("SingleOp: %s, output[%zu] addr = %p", op_desc_->GetName().c_str(), i, addr); + args.emplace_back(addr); + } + + return SUCCESS; +} + Status TbeOpTask::LaunchKernel(const vector &input_desc, const vector &input_buffers, vector &output_desc, @@ -355,12 +402,7 @@ Status TbeOpTask::LaunchKernel(const vector &input_desc, GE_CHK_STATUS_RET_NOLOG(UpdateRunInfo()); GE_CHK_STATUS_RET(AllocateWorkspaces(run_info_workspaces_), "[Allocate][Workspaces] failed."); std::vector args; - for (auto &buffer : input_buffers) { - args.emplace_back(buffer.data); - } - for (auto &buffer : output_buffers) { - args.emplace_back(buffer.data); - } + GE_CHK_STATUS_RET(UpdateIoAddr(args, input_buffers, output_buffers), "[Update][IoAddr] failed."); for (auto &buffer : workspaces_) { args.emplace_back(buffer); } diff --git a/ge/single_op/task/op_task.h b/ge/single_op/task/op_task.h index 19320bc0..0cbc1a29 100644 --- a/ge/single_op/task/op_task.h +++ b/ge/single_op/task/op_task.h @@ -101,6 +101,8 @@ class TbeOpTask : public OpTask { const vector &output_desc); Status AllocateWorkspaces(const std::vector &workspace_sizes); Status DoLaunchKernel(rtStream_t stream); + Status UpdateIoAddr(std::vector &args, const std::vector &inputs, + const std::vector &outputs); const void *stub_func_ = nullptr; std::unique_ptr args_; diff --git a/tests/ut/ge/single_op/single_op_task_unittest.cc b/tests/ut/ge/single_op/single_op_task_unittest.cc index a17c9012..472a88c3 100644 --- a/tests/ut/ge/single_op/single_op_task_unittest.cc +++ b/tests/ut/ge/single_op/single_op_task_unittest.cc @@ -91,8 +91,9 @@ TEST_F(UtestSingleOpTask, test_build_kernel_task) { TbeOpTask task_tmp; TbeOpTask *task = &task_tmp; ASSERT_EQ(model.BuildKernelTask(task_def, &task), SUCCESS); + ge::DataBuffer data_buffer; vector input_desc; - vector input_buffers; + vector input_buffers = { data_buffer }; vector output_desc; vector output_buffers; task->node_ = node; @@ -110,8 +111,36 @@ TEST_F(UtestSingleOpTask, test_build_kernel_task) { task->args_.reset(&task_args); ASSERT_EQ(task->LaunchKernel(input_desc, input_buffers, output_desc, output_buffers, stream_), SUCCESS); - char handle_tmp = '0'; - char *handle = &handle_tmp; + char *handle = "00"; task->SetHandle(handle); ASSERT_EQ(task->LaunchKernel(input_desc, input_buffers, output_desc, output_buffers, stream_), SUCCESS); -} \ No newline at end of file +} + +TEST_F(UtestSingleOpTask, test_update_ioaddr) { + auto graph = make_shared("graph"); + auto op_desc = make_shared("Add", "Add"); + + GeTensorDesc desc; + op_desc->AddInputDesc(desc); + op_desc->AddInputDesc(desc); + op_desc->AddOutputDesc(desc); + vector is_input_const = { true, false }; + op_desc->SetIsInputConst(is_input_const); + auto node = graph->AddNode(op_desc); + + TbeOpTask task; + task.op_desc_ = op_desc; + task.args_.reset(new (std::nothrow) uint8_t[sizeof(void *) * 3]); + + vector args; + vector inputs; + vector outputs; + ASSERT_EQ(task.UpdateIoAddr(args, inputs, outputs), ACL_ERROR_GE_PARAM_INVALID); + task.arg_size_ = sizeof(void *) * 3; + ASSERT_EQ(task.UpdateIoAddr(args, inputs, outputs), ACL_ERROR_GE_PARAM_INVALID); + + ge::DataBuffer data_buffer; + inputs = { data_buffer }; + ASSERT_EQ(task.UpdateIoAddr(args, inputs, outputs), SUCCESS); +} + From ae348cf55a6cb7c5729dd6b4b5fb700723e8362b Mon Sep 17 00:00:00 2001 From: zhaozhixuan Date: Tue, 29 Jun 2021 15:42:36 +0800 Subject: [PATCH 2/4] Fix ut. --- ge/single_op/task/op_task.cc | 120 ++++++++++++----------- ge/single_op/task/op_task.h | 6 +- ge/single_op/task/tbe_task_builder.cc | 1 + tests/ut/ge/single_op/single_op_task_unittest.cc | 23 +++-- tests/ut/ge/single_op/single_op_unittest.cc | 2 +- 5 files changed, 84 insertions(+), 68 deletions(-) diff --git a/ge/single_op/task/op_task.cc b/ge/single_op/task/op_task.cc index 92d1e325..632cd4d8 100755 --- a/ge/single_op/task/op_task.cc +++ b/ge/single_op/task/op_task.cc @@ -345,14 +345,46 @@ Status TbeOpTask::AllocateWorkspaces(const vector &workspace_sizes) { return SUCCESS; } -Status TbeOpTask::UpdateIoAddr(std::vector &args, const std::vector &inputs, - const std::vector &outputs) { - uintptr_t *arg_base = nullptr; - size_t arg_num = 0; - GetIoAddr(arg_base, arg_num); +Status TbeOpTask::UpdateTilingArgs(rtStream_t stream) { + size_t args_size = op_desc_->GetInputsSize() + op_desc_->GetOutputsSize() + workspaces_.size(); + if (tiling_buffer_ != nullptr) { + args_size++; + } + size_t temp_size = args_size * sizeof(void *); + if (arg_size_ < temp_size) { + GELOGD("Need to reset size of args_ from %zu to %zu.", arg_size_, temp_size); + std::unique_ptr args(new (std::nothrow) uint8_t[temp_size]()); + GE_CHECK_NOTNULL(args); + if (memcpy_s(args.get(), temp_size, args_.get(), arg_size_) != EOK) { + GELOGE(ACL_ERROR_GE_MEMORY_OPERATE_FAILED, "[Update][KernelArgs] failed for [%s].", node_->GetName().c_str()); + REPORT_INNER_ERROR("E19999", "update kernel args failed for %s.", node_->GetName().c_str()); + return ACL_ERROR_GE_MEMORY_OPERATE_FAILED; + } + args_.reset(args.release()); + arg_size_ = temp_size; + } + + uintptr_t *arg_base = reinterpret_cast(args_.get()); + size_t arg_index = op_desc_->GetInputsSize() + op_desc_->GetOutputsSize(); + for (size_t i = 0; i < workspaces_.size(); ++i) { + arg_base[arg_index++] = reinterpret_cast(workspaces_[i]); + } + + if (tiling_buffer_ != nullptr) { + GELOGD("[%s] Start to copy tiling info. size = %zu", node_->GetName().c_str(), tiling_data_.size()); + GE_CHK_RT_RET(rtMemcpyAsync(tiling_buffer_, max_tiling_size_, tiling_data_.data(), tiling_data_.size(), + RT_MEMCPY_HOST_TO_DEVICE_EX, stream)); + arg_base[arg_index] = reinterpret_cast(tiling_buffer_); + } + + return SUCCESS; +} + +Status TbeOpTask::SetArgIndex() { + arg_index_.clear(); const vector v_is_input_const = op_desc_->GetIsInputConst(); - size_t non_const_index = 0; + size_t input_index = 0; for (size_t i = 0; i < op_desc_->GetAllInputsSize(); ++i) { const GeTensorDescPtr tensor_desc = op_desc_->MutableInputDesc(static_cast(i)); if (tensor_desc == nullptr) { @@ -360,33 +392,33 @@ Status TbeOpTask::UpdateIoAddr(std::vector &args, const std::vector= arg_num) { - GELOGE(ACL_ERROR_GE_PARAM_INVALID, "[Check][Size] Args size is %zu, but get index is %zu.", arg_num, i); - REPORT_INNER_ERROR("E19999", "[Check][Size] Args size is %zu, but get index is %zu.", arg_num, i); - return ACL_ERROR_GE_PARAM_INVALID; - } - auto addr = reinterpret_cast(arg_base[i]); - GELOGD("SingleOp: %s, Index: %zu, input is const, addr = %p", op_desc_->GetName().c_str(), i, addr); - args.emplace_back(addr); + GELOGD("SingleOp: %s, Index: %zu, input is const", op_desc_->GetName().c_str(), i); + input_index++; continue; } - if (non_const_index >= inputs.size()) { - GELOGE(ACL_ERROR_GE_PARAM_INVALID, "[Check][Size] Input size is %zu, but get non_const_index is %zu", - inputs.size(), non_const_index); - REPORT_INNER_ERROR("E19999", "[Check][Size] Input size is %zu, but get non_const_index is %zu", - inputs.size(), non_const_index); - return ACL_ERROR_GE_PARAM_INVALID; - } - auto addr = inputs[non_const_index].data; - GELOGD("SingleOp: %s, input[%zu], addr = %p", op_desc_->GetName().c_str(), i, addr); - args.emplace_back(addr); - non_const_index++; + arg_index_.emplace_back(input_index); + input_index++; } + return SUCCESS; +} - for (size_t i = 0; i < outputs.size(); ++i) { - auto addr = outputs[i].data; - GELOGD("SingleOp: %s, output[%zu] addr = %p", op_desc_->GetName().c_str(), i, addr); - args.emplace_back(addr); +Status TbeOpTask::UpdateIoAddr(const vector &inputs, const vector &outputs) { + if (arg_index_.size() != inputs.size()) { + GELOGE(ACL_ERROR_GE_PARAM_INVALID, "[Check][Size] Args size is %zu, but get input size is %zu.", + arg_index_.size(), inputs.size()); + REPORT_INNER_ERROR("E19999", "[Check][Size] Args size is %zu, but get input size is %zu.", + arg_index_.size(), inputs.size()); + return ACL_ERROR_GE_PARAM_INVALID; + } + + uintptr_t *arg_base = reinterpret_cast(args_.get()); + for (size_t i = 0; i < arg_index_.size(); ++i) { + arg_base[arg_index_[i]] = reinterpret_cast(inputs[i].data); + } + + size_t input_size = op_desc_->GetInputsSize(); + for (size_t i = 0; i < op_desc_->GetOutputsSize(); ++i) { + arg_base[input_size + i] = reinterpret_cast(outputs[i].data); } return SUCCESS; @@ -398,37 +430,11 @@ Status TbeOpTask::LaunchKernel(const vector &input_desc, vector &output_buffers, rtStream_t stream) { GELOGD("[%s] Start to launch kernel", node_->GetName().c_str()); + GE_CHK_STATUS_RET(UpdateIoAddr(input_buffers, output_buffers), "[Update][IoAddr] failed."); GE_CHK_STATUS_RET_NOLOG(UpdateNodeByShape(input_desc, output_desc)); GE_CHK_STATUS_RET_NOLOG(UpdateRunInfo()); GE_CHK_STATUS_RET(AllocateWorkspaces(run_info_workspaces_), "[Allocate][Workspaces] failed."); - std::vector args; - GE_CHK_STATUS_RET(UpdateIoAddr(args, input_buffers, output_buffers), "[Update][IoAddr] failed."); - for (auto &buffer : workspaces_) { - args.emplace_back(buffer); - } - - if (tiling_buffer_ != nullptr) { - GELOGD("[%s] Start to copy tiling info. size = %zu", node_->GetName().c_str(), tiling_data_.size()); - GE_CHK_RT_RET(rtMemcpyAsync(tiling_buffer_, max_tiling_size_, tiling_data_.data(), tiling_data_.size(), - RT_MEMCPY_HOST_TO_DEVICE_EX, stream)); - - args.emplace_back(tiling_buffer_); - } - - GELOGD("Dst size is %zu, src size is %zu.", arg_size_, args.size() * sizeof(void *)); - // node with workspace: build can not get size of workspace, need to update arg_size_ when execute - if (arg_size_ < (args.size() * sizeof(void *))) { - size_t temp_size = args.size() * sizeof(void *); - GELOGD("Need to reset size of args_ from %zu to %zu.", arg_size_, temp_size); - args_.reset(new(std::nothrow) uint8_t[temp_size]()); - GE_CHECK_NOTNULL(args_); - arg_size_ = temp_size; - } - if (memcpy_s(args_.get(), arg_size_, args.data(), args.size() * sizeof(void *)) != EOK) { - GELOGE(ACL_ERROR_GE_MEMORY_OPERATE_FAILED, "[Update][KernelArgs] failed for [%s].", node_->GetName().c_str()); - REPORT_INNER_ERROR("E19999", "update kernel args failed for %s.", node_->GetName().c_str()); - return ACL_ERROR_GE_MEMORY_OPERATE_FAILED; - } + GE_CHK_STATUS_RET(UpdateTilingArgs(stream), "[Update][TilingArgs] failed."); GELOGD("[%s] Start to invoke rtKernelLaunch", node_->GetName().c_str()); GE_CHK_STATUS_RET(DoLaunchKernel(stream), "Failed to do launch kernel."); diff --git a/ge/single_op/task/op_task.h b/ge/single_op/task/op_task.h index 0cbc1a29..d3e8383d 100644 --- a/ge/single_op/task/op_task.h +++ b/ge/single_op/task/op_task.h @@ -85,6 +85,7 @@ class TbeOpTask : public OpTask { const OpDescPtr &op_desc, const domi::KernelDefWithHandle& kernel_def_with_handle); Status UpdateRunInfo() override; + Status SetArgIndex(); const void *GetArgs() const; size_t GetArgSize() const; @@ -100,9 +101,9 @@ class TbeOpTask : public OpTask { Status UpdateNodeByShape(const vector &input_desc, const vector &output_desc); Status AllocateWorkspaces(const std::vector &workspace_sizes); + Status UpdateTilingArgs(rtStream_t stream); Status DoLaunchKernel(rtStream_t stream); - Status UpdateIoAddr(std::vector &args, const std::vector &inputs, - const std::vector &outputs); + Status UpdateIoAddr(const vector &inputs, const vector &outputs); const void *stub_func_ = nullptr; std::unique_ptr args_; @@ -122,6 +123,7 @@ class TbeOpTask : public OpTask { void* handle_ = nullptr; std::string original_kernel_key_; std::string node_info_; + std::vector arg_index_; }; class AiCpuBaseTask : public OpTask { diff --git a/ge/single_op/task/tbe_task_builder.cc b/ge/single_op/task/tbe_task_builder.cc index c7ff13d1..e5206ea6 100644 --- a/ge/single_op/task/tbe_task_builder.cc +++ b/ge/single_op/task/tbe_task_builder.cc @@ -387,6 +387,7 @@ Status TbeTaskBuilder::BuildTask(TbeOpTask &task, const SingleOpModelParam ¶ } task.SetStubFunc(stub_name_, stub_func); } + GE_CHK_STATUS_RET(task.SetArgIndex(), "[Set][ArgTable] failed."); return SUCCESS; } diff --git a/tests/ut/ge/single_op/single_op_task_unittest.cc b/tests/ut/ge/single_op/single_op_task_unittest.cc index 472a88c3..020efc23 100644 --- a/tests/ut/ge/single_op/single_op_task_unittest.cc +++ b/tests/ut/ge/single_op/single_op_task_unittest.cc @@ -95,7 +95,7 @@ TEST_F(UtestSingleOpTask, test_build_kernel_task) { vector input_desc; vector input_buffers = { data_buffer }; vector output_desc; - vector output_buffers; + vector output_buffers = { data_buffer }; task->node_ = node; OpTilingFunc op_tiling_func = [](const TeOpParas &, const OpCompileInfo &, OpRunInfo &) -> bool {return true;}; OpTilingRegistryInterf("Add", op_tiling_func); @@ -107,8 +107,7 @@ TEST_F(UtestSingleOpTask, test_build_kernel_task) { task->max_tiling_size_ = 64; task->tiling_data_ = "tiling_data"; task->arg_size_ = 64; - uint8_t task_args{0}; - task->args_.reset(&task_args); + task->args_.reset(new (std::nothrow) uint8_t[sizeof(void *) * 3]); ASSERT_EQ(task->LaunchKernel(input_desc, input_buffers, output_desc, output_buffers, stream_), SUCCESS); char *handle = "00"; @@ -130,17 +129,25 @@ TEST_F(UtestSingleOpTask, test_update_ioaddr) { TbeOpTask task; task.op_desc_ = op_desc; - task.args_.reset(new (std::nothrow) uint8_t[sizeof(void *) * 3]); + task.node_ = node; + ASSERT_EQ(task.SetArgIndex(), SUCCESS); + task.arg_size_ = sizeof(void *) * 4; + task.args_.reset(new (std::nothrow) uint8_t[task.arg_size_]); + task.arg_index_ = {0}; vector args; vector inputs; vector outputs; - ASSERT_EQ(task.UpdateIoAddr(args, inputs, outputs), ACL_ERROR_GE_PARAM_INVALID); - task.arg_size_ = sizeof(void *) * 3; - ASSERT_EQ(task.UpdateIoAddr(args, inputs, outputs), ACL_ERROR_GE_PARAM_INVALID); + ASSERT_EQ(task.UpdateIoAddr(inputs, outputs), ACL_ERROR_GE_PARAM_INVALID); ge::DataBuffer data_buffer; inputs = { data_buffer }; - ASSERT_EQ(task.UpdateIoAddr(args, inputs, outputs), SUCCESS); + outputs = { data_buffer }; + ASSERT_EQ(task.UpdateIoAddr(inputs, outputs), SUCCESS); + + task.tiling_buffer_ = (void *)0x0001; + task.workspaces_ = { (void *)0x0002 }; + ASSERT_EQ(task.UpdateTilingArgs(nullptr), SUCCESS); + task.tiling_buffer_ = nullptr; } diff --git a/tests/ut/ge/single_op/single_op_unittest.cc b/tests/ut/ge/single_op/single_op_unittest.cc index db3de7ec..09aac153 100644 --- a/tests/ut/ge/single_op/single_op_unittest.cc +++ b/tests/ut/ge/single_op/single_op_unittest.cc @@ -103,7 +103,7 @@ TEST_F(UtestSingleOp, test_dynamic_singleop_execute_async1) { EXPECT_EQ(desc_ptr->AddInputDesc("x", GeTensorDesc(GeShape({2}), FORMAT_NCHW)), GRAPH_SUCCESS); dynamic_single_op.op_task_->op_desc_ = desc_ptr; // UpdateRunInfo failed - EXPECT_EQ(dynamic_single_op.ExecuteAsync(input_desc, input_buffers, output_desc, output_buffers), PARAM_INVALID); + EXPECT_EQ(dynamic_single_op.ExecuteAsync(input_desc, input_buffers, output_desc, output_buffers), ACL_ERROR_GE_PARAM_INVALID); } From 813f2fe4a2eb234153ce0ad5a5801c98bdb9be4f Mon Sep 17 00:00:00 2001 From: zhaozhixuan Date: Tue, 29 Jun 2021 19:11:56 +0800 Subject: [PATCH 3/4] Fix ut. --- ge/single_op/task/op_task.cc | 10 ++++------ ge/single_op/task/op_task.h | 4 +++- ge/single_op/task/tbe_task_builder.cc | 2 ++ 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/ge/single_op/task/op_task.cc b/ge/single_op/task/op_task.cc index 632cd4d8..28ec7f64 100755 --- a/ge/single_op/task/op_task.cc +++ b/ge/single_op/task/op_task.cc @@ -346,7 +346,7 @@ Status TbeOpTask::AllocateWorkspaces(const vector &workspace_sizes) { } Status TbeOpTask::UpdateTilingArgs(rtStream_t stream) { - size_t args_size = op_desc_->GetInputsSize() + op_desc_->GetOutputsSize() + workspaces_.size(); + size_t args_size = input_num_ + output_num_ + workspaces_.size(); if (tiling_buffer_ != nullptr) { args_size++; } @@ -361,12 +361,12 @@ Status TbeOpTask::UpdateTilingArgs(rtStream_t stream) { return ACL_ERROR_GE_MEMORY_OPERATE_FAILED; } - args_.reset(args.release()); + args_ = std::move(args); arg_size_ = temp_size; } uintptr_t *arg_base = reinterpret_cast(args_.get()); - size_t arg_index = op_desc_->GetInputsSize() + op_desc_->GetOutputsSize(); + size_t arg_index = input_num_ + output_num_; for (size_t i = 0; i < workspaces_.size(); ++i) { arg_base[arg_index++] = reinterpret_cast(workspaces_[i]); } @@ -382,7 +382,6 @@ Status TbeOpTask::UpdateTilingArgs(rtStream_t stream) { } Status TbeOpTask::SetArgIndex() { - arg_index_.clear(); const vector v_is_input_const = op_desc_->GetIsInputConst(); size_t input_index = 0; for (size_t i = 0; i < op_desc_->GetAllInputsSize(); ++i) { @@ -416,9 +415,8 @@ Status TbeOpTask::UpdateIoAddr(const vector &inputs, const vector(inputs[i].data); } - size_t input_size = op_desc_->GetInputsSize(); for (size_t i = 0; i < op_desc_->GetOutputsSize(); ++i) { - arg_base[input_size + i] = reinterpret_cast(outputs[i].data); + arg_base[input_num_ + i] = reinterpret_cast(outputs[i].data); } return SUCCESS; diff --git a/ge/single_op/task/op_task.h b/ge/single_op/task/op_task.h index d3e8383d..f93e031a 100644 --- a/ge/single_op/task/op_task.h +++ b/ge/single_op/task/op_task.h @@ -123,7 +123,9 @@ class TbeOpTask : public OpTask { void* handle_ = nullptr; std::string original_kernel_key_; std::string node_info_; - std::vector arg_index_; + std::vector arg_index_; // data index in args + size_t input_num_; // Include const input + size_t output_num_; }; class AiCpuBaseTask : public OpTask { diff --git a/ge/single_op/task/tbe_task_builder.cc b/ge/single_op/task/tbe_task_builder.cc index e5206ea6..db8ecfe2 100644 --- a/ge/single_op/task/tbe_task_builder.cc +++ b/ge/single_op/task/tbe_task_builder.cc @@ -388,6 +388,8 @@ Status TbeTaskBuilder::BuildTask(TbeOpTask &task, const SingleOpModelParam ¶ task.SetStubFunc(stub_name_, stub_func); } GE_CHK_STATUS_RET(task.SetArgIndex(), "[Set][ArgTable] failed."); + task.input_num_ = op_desc_->GetInputsSize(); + task.output_num_ = op_desc_->GetOutputsSize(); return SUCCESS; } From 4f3f54d56407d045fa37bf02508cc101a46e1bef Mon Sep 17 00:00:00 2001 From: zhaozhixuan Date: Tue, 29 Jun 2021 19:19:08 +0800 Subject: [PATCH 4/4] Fix ut. --- ge/single_op/task/op_task.h | 2 +- tests/ut/ge/single_op/single_op_task_unittest.cc | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/ge/single_op/task/op_task.h b/ge/single_op/task/op_task.h index f93e031a..97d8a342 100644 --- a/ge/single_op/task/op_task.h +++ b/ge/single_op/task/op_task.h @@ -124,7 +124,7 @@ class TbeOpTask : public OpTask { std::string original_kernel_key_; std::string node_info_; std::vector arg_index_; // data index in args - size_t input_num_; // Include const input + size_t input_num_; // include const input size_t output_num_; }; diff --git a/tests/ut/ge/single_op/single_op_task_unittest.cc b/tests/ut/ge/single_op/single_op_task_unittest.cc index 020efc23..b0c98205 100644 --- a/tests/ut/ge/single_op/single_op_task_unittest.cc +++ b/tests/ut/ge/single_op/single_op_task_unittest.cc @@ -134,6 +134,8 @@ TEST_F(UtestSingleOpTask, test_update_ioaddr) { task.arg_size_ = sizeof(void *) * 4; task.args_.reset(new (std::nothrow) uint8_t[task.arg_size_]); task.arg_index_ = {0}; + task.input_num_ = 2; + task.output_num_ = 1; vector args; vector inputs;