From 3f0839b26aabeb6e8feba68e9b623c7713e31753 Mon Sep 17 00:00:00 2001
From: chuxing <chuxing@huawei.com>
Date: Mon, 22 Mar 2021 12:06:13 +0800
Subject: [PATCH] load davinci model in load phase

---
 .../compiledsubgraph/known_node_executor.cc   | 40 +++++++++++--------
 .../compiledsubgraph/known_node_executor.h    |  5 +--
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
index bb96c275..918603f0 100755
--- a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
+++ b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
@@ -95,13 +95,6 @@ Status KnownNodeTask::UpdateArgs(TaskContext &context) {
 Status KnownNodeTask::Init(TaskContext &context) {
   // allocate output mem
   GE_CHK_STATUS_RET(context.AllocateOutputs(), "known node task allocate output failed.");
-
-  // init davinicmodel
-  if (!load_flag_) {
-    davinci_model_->InitRuntimeParams();
-    GE_CHK_STATUS_RET(davinci_model_->InitVariableMem(), "init variable mem failed.");
-  }
-
   // allocate mem base
   void *buffer = nullptr;
   if (davinci_model_->TotalMemSize() != 0) {
@@ -129,14 +122,6 @@ Status KnownNodeTask::Init(TaskContext &context) {
       void *global_step = context.GetExecutionContext()->global_step;
       davinci_model_->SetKnownShapeGlobalStep(global_step);
     }
-    int32_t device_id = 0;
-    rtError_t rt_ret = rtGetDevice(&device_id);
-    if (rt_ret != RT_ERROR_NONE || device_id < 0) {
-      GELOGE(rt_ret, "Call rtGetDevice failed, ret = 0x%X, device_id = %d.", rt_ret, device_id);
-      return RT_ERROR_TO_GE_STATUS(rt_ret);
-    }
-    davinci_model_->SetDeviceId(device_id);
-    GE_CHK_STATUS_RET(davinci_model_->Init(), "KnownNodeExecutor::InitDavinciModel failed.");
     load_flag_ = true;
   } else {
     GE_CHK_STATUS_RET(ModelManager::GetInstance()->DestroyAicpuKernel(davinci_model_->GetSessionId(),
@@ -146,6 +131,25 @@ Status KnownNodeTask::Init(TaskContext &context) {
   return SUCCESS;
 }
 
+Status KnownNodeTask::InitDavinciModel() {
+  davinci_model_->InitRuntimeParams();
+  GE_CHK_STATUS_RET(davinci_model_->InitVariableMem(), "init variable mem failed");
+  auto mem_size = davinci_model_->TotalMemSize();
+  if (mem_size > 0) {
+    auto buffer = TensorBuffer::Create(NpuMemoryAllocator::GetAllocator(), mem_size);
+    GE_CHECK_NOTNULL(buffer);
+    davinci_model_->UpdateMemBase(static_cast<uint8_t *>(buffer->GetData()));
+    GELOGI("KnownNodeTask::Init mem base is %p, size %lu.",
+           davinci_model_->GetRuntimeParam().mem_base, davinci_model_->GetRuntimeParam().mem_size);
+  }
+
+  int32_t device_id = 0;
+  GE_CHK_RT_RET(rtGetDevice(&device_id));
+  davinci_model_->SetDeviceId(static_cast<uint32_t>(device_id));
+  GE_CHK_STATUS_RET(davinci_model_->Init(), "[Init][Model] Failed to init davinci model.");
+  return SUCCESS;
+}
+
 Status KnownNodeExecutor::PrepareTask(NodeTask &task, TaskContext &context) const {
   GELOGD("[%s] KnownNodeExecutor::PrepareTask in.", context.GetNodeName());
   RECORD_EXECUTION_EVENT(context.GetExecutionContext(), context.GetNodeName(), "[KnownNodeExecutorPrepareTask] Start");
@@ -182,9 +186,11 @@ Status KnownNodeExecutor::LoadTask(const HybridModel &model, const NodePtr &node
 
   GE_CHK_STATUS_RET(davinci_model->Assign(ge_model), "KnownNodeExecutor::LoadTask davincimodel assign failed.");
 
-  task = MakeShared<KnownNodeTask>(davinci_model);
-  GE_CHECK_NOTNULL(task);
+  auto known_node_task = MakeShared<KnownNodeTask>(davinci_model);
+  GE_CHECK_NOTNULL(known_node_task);
+  GE_CHK_STATUS_RET_NOLOG(known_node_task->InitDavinciModel());
   GELOGI("[%s] KnownNodeExecutor::LoadTask success.", node->GetName().c_str());
+  task = std::move(known_node_task);
   return SUCCESS;
 }
 
diff --git a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h
index 6e9740ad..762e9991 100644
--- a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h
+++ b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h
@@ -31,11 +31,12 @@ class KnownNodeTask : public NodeTask {
       : davinci_model_(davinci_model)
     {}
 
-  ~KnownNodeTask() {}
+  ~KnownNodeTask() = default;
 
   Status UpdateArgs(TaskContext &context) override;
   Status ExecuteAsync(TaskContext &context, std::function<void()> done_callback) override;
   Status Init(TaskContext &context) override;
+  Status InitDavinciModel();
  private:
   std::shared_ptr<DavinciModel> davinci_model_ = nullptr;
   bool load_flag_ = false;
@@ -47,8 +48,6 @@ class KnownNodeExecutor : public NodeExecutor {
   Status PrepareTask(NodeTask &task, TaskContext &context) const;
   Status ExecuteTask(NodeTask &task, TaskContext &context, const std::function<void()> &callback) const;
   ~KnownNodeExecutor() {}
- private:
-  std::shared_ptr<DavinciModel> davinci_model_ = nullptr;
 };
 }  // namespace hybrid
 }  // namespace ge