!1852 Fix mem leak and recursive depth protection

From: @zhao_zhixuan
Reviewed-by: 
Signed-off-by:

ge/common/ge/tbe_plugin_manager.cc

@@ -104,7 +104,15 @@ void TBEPluginManager::ProcessSoFullName(vector<string> &file_list, string &caff
   }
 }
 
-void TBEPluginManager::FindParserSo(const string &path, vector<string> &file_list, string &caffe_parser_path) {
+void TBEPluginManager::FindParserSo(const string &path, vector<string> &file_list,
+                                    string &caffe_parser_path, uint32_t recursive_depth) {
+  static const uint32_t max_recursive_depth = 20; // For recursive depth protection
+
+  if (recursive_depth >= max_recursive_depth) {
+    GELOGW("Recursive depth is become %u, Please check input!", recursive_depth);
+    return;
+  }
+
   // Path, change to absolute path
   string real_path = RealPath(path.c_str());
   // Plugin path does not exist
@@ -138,7 +146,7 @@ void TBEPluginManager::FindParserSo(const string &path, vector<string> &file_lis
           ProcessSoFullName(file_list, caffe_parser_path, full_name, caffe_parser_so_suff, aicpu_so_suff,
                             aicpu_host_so_suff);
       } else {
-          FindParserSo(full_name, file_list, caffe_parser_path);
+          FindParserSo(full_name, file_list, caffe_parser_path, recursive_depth + 1);
       }
   }
   mmScandirFree(entries, ret);

ge/common/ge/tbe_plugin_manager.h

@@ -57,7 +57,8 @@ class TBEPluginManager {
   static void ProcessSoFullName(vector<string> &file_list, string &caffe_parser_path, string &full_name,
                                 const string &caffe_parser_so_suff, const string &aicpu_so_suff,
                                 const string &aicpu_host_so_suff);
-  static void FindParserSo(const string &path, vector<string> &file_list, string &caffe_parser_path);
+  static void FindParserSo(const string &path, vector<string> &file_list, string &caffe_parser_path,
+                           uint32_t recursive_depth = 0);
   static void GetPluginSoFileList(const string &path, vector<string> &file_list, string &caffe_parser_path);
   static void GetCustomOpPath(std::string &customop_path);
   void LoadCustomOpLib();

ge/hybrid/node_executor/task_context.cc

@@ -43,6 +43,7 @@ TaskContext::~TaskContext() {
       output_tensor->Destroy();
     }
   }
+  ReleaseWorkspace();
 }
 
 void TaskContext::ReleaseWorkspace() {
@@ -50,6 +51,7 @@ void TaskContext::ReleaseWorkspace() {
   for (auto ws_addr : workspaces_) {
     execution_context_->allocator->Deallocate(ws_addr);
   }
+  workspaces_.clear();
 }
 
 std::unique_ptr<TaskContext> TaskContext::Create(NodeState *node_state, SubgraphContext *subgraph_context) {

ge/single_op/single_op_model.cc

@@ -376,11 +376,10 @@ Status SingleOpModel::BuildTaskList(StreamResource *stream_resource, SingleOp &s
     } else if (task_type == RT_MODEL_TASK_KERNEL_EX) {
       GELOGD("Building AICPU_TF task");
       AiCpuTask *aicpu_task = nullptr;
-      bool depend_compute_flag = false;
       uint64_t singleop_kernel_id = aicpu_kernel_id++;
       GELOGI("Build singleOp TfTask, kernel_id = %lu", singleop_kernel_id);
       GE_CHK_STATUS_RET_NOLOG(
-          BuildKernelExTask(task_def.kernel_ex(), &aicpu_task, false, depend_compute_flag, singleop_kernel_id));
+          BuildKernelExTask(task_def.kernel_ex(), &aicpu_task, singleop_kernel_id));
       aicpu_task->SetModelArgs(model_name_, model_id_);
       ParseArgTable(aicpu_task, single_op);
       single_op.tasks_.emplace_back(aicpu_task);
@@ -457,8 +456,7 @@ Status SingleOpModel::BuildKernelTask(const domi::TaskDef &task_def, TbeOpTask *
   return SUCCESS;
 }
 
-Status SingleOpModel::BuildKernelExTask(const domi::KernelExDef &kernel_def, AiCpuTask **task,
-                                        bool dynamic_flag, bool& depend_compute_flag, uint64_t kernel_id) {
+Status SingleOpModel::BuildKernelExTask(const domi::KernelExDef &kernel_def, AiCpuTask **task, uint64_t kernel_id) {
   auto iter = op_list_.find(kernel_def.op_index());
   if (iter == op_list_.end()) {
     GELOGE(ACL_ERROR_GE_INTERNAL_ERROR, 
@@ -476,12 +474,11 @@ Status SingleOpModel::BuildKernelExTask(const domi::KernelExDef &kernel_def, AiC
     return ACL_ERROR_GE_MEMORY_ALLOCATION;
   }
   auto builder = AiCpuTaskBuilder(iter->second->GetOpDesc(), kernel_def);
-  auto ret = builder.BuildTask(*aicpu_task, model_params_, dynamic_flag, kernel_id);
+  auto ret = builder.BuildTask(*aicpu_task, model_params_, kernel_id);
   if (ret != SUCCESS) {
     GELOGE(ret, "[Build][Task] failed, kernel_id:%lu.", kernel_id);
     return ret;
   }
-  depend_compute_flag = (aicpu_task->GetUnknownType() == DEPEND_COMPUTE);
 
   *task = aicpu_task.release();
   return SUCCESS;
@@ -628,12 +625,10 @@ Status SingleOpModel::BuildTaskListForDynamicOp(StreamResource *stream_resource,
       }
       GELOGD("Building AICPU_TF task");
       AiCpuTask *aicpu_task = nullptr;
-      bool depend_compute_flag = false;
       uint64_t dynamic_singleop_kernel_id = aicpu_kernel_id++;
       GELOGI("Build dynamic singleOp TfTask, kernel_id = %lu", dynamic_singleop_kernel_id);
-      GE_CHK_STATUS_RET_NOLOG(BuildKernelExTask(task_def.kernel_ex(), &aicpu_task, true,
-                                                depend_compute_flag, dynamic_singleop_kernel_id));
-      if (depend_compute_flag) {
+      GE_CHK_STATUS_RET_NOLOG(BuildKernelExTask(task_def.kernel_ex(), &aicpu_task, dynamic_singleop_kernel_id));
+      if (aicpu_task->GetUnknownType() == DEPEND_COMPUTE) {
         if (i >= tasks.size() - 1) {
           GELOGE(ACL_ERROR_GE_PARAM_INVALID, "[Check][Task]The copy task of the fourth operator was not found.");
           REPORT_INNER_ERROR("E19999", "The copy task of the fourth operator was not found.");

ge/single_op/single_op_model.h

@@ -69,8 +69,7 @@ class SingleOpModel {
   Status BuildTaskList(StreamResource *stream_resource, SingleOp &single_op);
   Status BuildTaskListForDynamicOp(StreamResource *stream_resource, DynamicSingleOp &dynamic_single_op);
   Status BuildKernelTask(const domi::TaskDef &task_def, TbeOpTask **task);
-  Status BuildKernelExTask(const domi::KernelExDef &kernel_def, AiCpuTask **task,
-                           bool dynamic_flag, bool& depend_compute_flag, uint64_t kernel_id);
+  Status BuildKernelExTask(const domi::KernelExDef &kernel_def, AiCpuTask **task, uint64_t kernel_id);
   Status BuildCpuKernelTask(const domi::KernelDef &kernel_def, OpTask **task, uint64_t kernel_id);
   Status BuildModelTaskKernel(StreamResource *stream_resource, const domi::TaskDef &task_def,
                               DynamicSingleOp &single_op);

ge/single_op/task/aicpu_task_builder.cc

@@ -63,7 +63,7 @@ namespace ge {
     return SUCCESS;
   }
 
-  Status AiCpuTaskBuilder::InitWorkspaceAndIO(AiCpuTask &task, const SingleOpModelParam &param, bool dynamic_flag) {
+  Status AiCpuTaskBuilder::InitWorkspaceAndIO(AiCpuTask &task, const SingleOpModelParam &param) {
     if (kernel_def_.args_size() > sizeof(STR_FWK_OP_KERNEL)) {
       GELOGE(ACL_ERROR_GE_PARAM_INVALID, "[Check][Size]sizeof STR_FWK_OP_KERNEL is: %lu, but args_size is: %d",
           sizeof(STR_FWK_OP_KERNEL), kernel_def_.args_size());
@@ -83,9 +83,8 @@ namespace ge {
     return SUCCESS;
   }
 
-  Status AiCpuTaskBuilder::BuildTask(ge::AiCpuTask &task, const SingleOpModelParam &param,
-                                     bool dynamic_flag, uint64_t kernel_id) {
-    GE_CHK_STATUS_RET_NOLOG(InitWorkspaceAndIO(task, param, dynamic_flag));
+  Status AiCpuTaskBuilder::BuildTask(ge::AiCpuTask &task, const SingleOpModelParam &param, uint64_t kernel_id) {
+    GE_CHK_STATUS_RET_NOLOG(InitWorkspaceAndIO(task, param));
 
     STR_FWK_OP_KERNEL fwk_op_kernel = {0};
     auto ret = SetFmkOpKernel(task.io_addr_, task.workspace_addr_, fwk_op_kernel);
@@ -124,7 +123,6 @@ namespace ge {
     task.arg_size_ = sizeof(STR_FWK_OP_KERNEL);
     task.op_type_ = op_desc_->GetName();
     task.task_info_ = kernel_def_.task_info();
-    task.dynamic_flag_ = dynamic_flag;
     task.kernel_id_ = kernel_id;
 
     auto debug_info = BuildTaskUtils::GetTaskInfo(op_desc_);

ge/single_op/task/aicpu_task_builder.h

@@ -29,12 +29,12 @@ namespace ge {
     AiCpuTaskBuilder(const OpDescPtr &op_desc, const domi::KernelExDef &kernel_def);
     ~AiCpuTaskBuilder() = default;
 
-    Status BuildTask(AiCpuTask &task, const SingleOpModelParam &param, bool dynamic_flag, uint64_t kernel_id);
+    Status BuildTask(AiCpuTask &task, const SingleOpModelParam &param, uint64_t kernel_id);
 
   private:
     static Status SetKernelArgs(void **args, STR_FWK_OP_KERNEL &kernel);
     Status SetFmkOpKernel(void *io_addr, void *ws_addr, STR_FWK_OP_KERNEL &kernel);
-    Status InitWorkspaceAndIO(AiCpuTask &task, const SingleOpModelParam &param, bool dynamic_flag);
+    Status InitWorkspaceAndIO(AiCpuTask &task, const SingleOpModelParam &param);
 
     const OpDescPtr op_desc_;
     const domi::KernelExDef &kernel_def_;

ge/single_op/task/op_task.cc

@@ -621,9 +621,7 @@ Status AiCpuBaseTask::UpdateIoAddr(const vector<DataBuffer> &inputs, const vecto
 AiCpuTask::~AiCpuTask() {
   FreeHbm(args_);
   FreeHbm(io_addr_);
-  if (dynamic_flag_) {
-    FreeHbm(workspace_addr_);
-  }
+  FreeHbm(workspace_addr_);
   FreeHbm(copy_workspace_buf_);
   FreeHbm(copy_ioaddr_dev_);
   FreeHbm(copy_input_release_flag_dev_);

ge/single_op/task/op_task.h

@@ -192,7 +192,6 @@ class AiCpuTask : public AiCpuBaseTask {
   // host addr
   std::vector<void *> io_addr_host_;
 
-  bool dynamic_flag_ = false;
   // for copy task
   void *copy_task_args_buf_ = nullptr;
   void *copy_workspace_buf_ = nullptr;

tests/ut/ge/CMakeLists.txt

@@ -811,6 +811,7 @@ set(MULTI_PARTS_TEST_FILES
     "session/inner_session_unittest.cc"
     "session/session_manager_unittest.cc"
     "common/host_cpu_engine_unittest.cc"
+    "common/tbe_plugin_manager_unittest.cc"
 )
 
 set(GENERATOR_TEST_FILES

tests/ut/ge/common/tbe_plugin_manager_unittest.cc

@@ -0,0 +1,40 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gtest/gtest.h>
+
+#define protected public
+#define private public
+#include "common/ge/tbe_plugin_manager.h"
+#undef private
+#undef protected
+
+namespace ge {
+class UtestTBEPluginManager: public testing::Test {
+ protected:
+  void SetUp() {}
+  void TearDown() {}
+};
+
+TEST_F(UtestTBEPluginManager, CheckFindParserSo) {
+  string path = "";
+  vector<string> file_list = {};
+  string caffe_parser_path = "";
+  TBEPluginManager::Instance().FindParserSo(path, file_list, caffe_parser_path);
+  path = "/lib64";
+  TBEPluginManager::Instance().FindParserSo(path, file_list, caffe_parser_path);
+}
+}  // namespace ge

tests/ut/ge/single_op/single_op_model_unittest.cc

@@ -310,3 +310,27 @@ TEST_F(UtestSingleOpModel, BuildTaskList) {
   MemcpyAsyncTask mem_task;
   ASSERT_EQ(mem_task.LaunchKernel(0), SUCCESS);
 }
+
+TEST_F(UtestSingleOpModel, build_aicpu_task) {
+  ComputeGraphPtr graph = make_shared<ComputeGraph>("single_op");
+  GeModelPtr ge_model = make_shared<GeModel>();
+  ge_model->SetGraph(GraphUtils::CreateGraphFromComputeGraph(graph));
+  shared_ptr<domi::ModelTaskDef> model_task_def = make_shared<domi::ModelTaskDef>();
+  ge_model->SetModelTaskDef(model_task_def);
+
+  domi::TaskDef *task_def = model_task_def->add_task();
+  task_def->set_type(RT_MODEL_TASK_KERNEL_EX);
+
+  string model_data_str = "123456789";
+  SingleOpModel model("model", model_data_str.c_str(), model_data_str.size());
+  std::mutex stream_mu;
+  rtStream_t stream = nullptr;
+  rtStreamCreate(&stream, 0);
+  DynamicSingleOp single_op(0, &stream_mu, stream);
+  model.model_helper_.model_ = ge_model;
+  auto op_desc = std::make_shared<ge::OpDesc>("add", "Add");
+  NodePtr node = graph->AddNode(op_desc); 
+  model.op_list_[0] = node;
+  StreamResource *res = new (std::nothrow) StreamResource(1);
+  ASSERT_EQ(model.BuildTaskListForDynamicOp(res, single_op), SUCCESS);
+}