From 0a04e99c8abae6135ccdf472ef48a697f0b5a997 Mon Sep 17 00:00:00 2001
From: unknown <zhaozhixuan2@hisilicon.com>
Date: Thu, 29 Oct 2020 17:26:58 +0800
Subject: [PATCH] Modify the execution of single_op to async.

---
 ge/single_op/task/op_task.cc | 23 +++++++++++++----------
 ge/single_op/task/op_task.h  |  3 ++-
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/ge/single_op/task/op_task.cc b/ge/single_op/task/op_task.cc
index f1d293d5..8ddcba58 100755
--- a/ge/single_op/task/op_task.cc
+++ b/ge/single_op/task/op_task.cc
@@ -263,7 +263,8 @@ Status AiCpuBaseTask::SetExtInfoAndType(const std::string &kernel_ext_info) {
 }
 
 Status AiCpuBaseTask::UpdateExtInfo(const std::vector<GeTensorDesc> &input_desc, 
-                                    std::vector<GeTensorDesc> &output_desc) {
+                                    std::vector<GeTensorDesc> &output_desc,
+                                    rtStream_t stream) {
   GELOGI("Update ext info begin, unknown_type=%d.", unknown_type_);
   if (num_inputs_ == 0 && num_outputs_ == 0) {
     GELOGI("No input and output, no need update ext info.");
@@ -283,11 +284,12 @@ Status AiCpuBaseTask::UpdateExtInfo(const std::vector<GeTensorDesc> &input_desc,
     }
   }
 
-  GE_CHK_RT_RET(rtMemcpy(ext_info_addr_dev_,
-                         aicpu_ext_handle_->GetExtInfoLen(), // check size
-                         aicpu_ext_handle_->GetExtInfo(),
-                         aicpu_ext_handle_->GetExtInfoLen(),
-                         RT_MEMCPY_HOST_TO_DEVICE));
+  GE_CHK_RT_RET(rtMemcpyAsync(ext_info_addr_dev_,
+                              aicpu_ext_handle_->GetExtInfoLen(), // check size
+                              aicpu_ext_handle_->GetExtInfo(),
+                              aicpu_ext_handle_->GetExtInfoLen(),
+                              RT_MEMCPY_HOST_TO_DEVICE_EX,
+                              stream));
 
   GELOGI("Update ext info end.");
   return SUCCESS;
@@ -618,7 +620,7 @@ Status AiCpuTask::LaunchKernel(const std::vector<GeTensorDesc> &input_desc,
                                std::vector<GeTensorDesc> &output_desc,
                                std::vector<DataBuffer> &output_buffers,
                                rtStream_t stream) {
-  GE_CHK_STATUS_RET_NOLOG(UpdateExtInfo(input_desc, output_desc));
+  GE_CHK_STATUS_RET_NOLOG(UpdateExtInfo(input_desc, output_desc, stream));
   std::vector<void *> inputs;
   std::vector<void *> outputs;
   for (auto &buffer : input_buffers) {
@@ -629,11 +631,12 @@ Status AiCpuTask::LaunchKernel(const std::vector<GeTensorDesc> &input_desc,
   }
   GE_CHK_STATUS_RET_NOLOG(SetIO(inputs, outputs));
   GE_CHK_STATUS_RET_NOLOG(LaunchKernel(stream));
-  GE_CHK_RT_RET(rtStreamSynchronize(stream));
 
   if (unknown_type_ == DEPEND_SHAPE_RANGE) {
+    GE_CHK_RT_RET(rtStreamSynchronize(stream));
     GE_CHK_STATUS_RET_NOLOG(UpdateOutputShape(output_desc));
   } else if (unknown_type_ == DEPEND_COMPUTE) {
+    GE_CHK_RT_RET(rtStreamSynchronize(stream));
     GE_CHK_STATUS_RET_NOLOG(UpdateShapeAndDataByResultSummary(output_desc, output_buffers, stream));
   }
 
@@ -689,7 +692,7 @@ Status AiCpuCCTask::LaunchKernel(const std::vector<GeTensorDesc> &input_desc,
                          "AiCpuCCTask unknown type[%d] is depend compute, it's not supported now.",
                          unknown_type_);
 
-  GE_CHK_STATUS_RET_NOLOG(UpdateExtInfo(input_desc, output_desc));
+  GE_CHK_STATUS_RET_NOLOG(UpdateExtInfo(input_desc, output_desc, stream));
 
   size_t arg_index = 0;
   auto *task_io_addr = reinterpret_cast<uintptr_t *>(io_addr_);
@@ -702,9 +705,9 @@ Status AiCpuCCTask::LaunchKernel(const std::vector<GeTensorDesc> &input_desc,
   }
 
   GE_CHK_STATUS_RET_NOLOG(LaunchKernel(stream));
-  GE_CHK_RT_RET(rtStreamSynchronize(stream));
 
   if (unknown_type_ == DEPEND_SHAPE_RANGE) {
+    GE_CHK_RT_RET(rtStreamSynchronize(stream));
     GE_CHK_STATUS_RET_NOLOG(UpdateOutputShape(output_desc));
   }
 
diff --git a/ge/single_op/task/op_task.h b/ge/single_op/task/op_task.h
index 4325da33..b9165641 100644
--- a/ge/single_op/task/op_task.h
+++ b/ge/single_op/task/op_task.h
@@ -131,7 +131,8 @@ class AiCpuBaseTask : public OpTask {
   Status SetExtInfoAndType(const std::string &kernel_ext_info);
 
   Status UpdateExtInfo(const std::vector<GeTensorDesc> &input_desc,
-                       std::vector<GeTensorDesc> &output_desc);
+                       std::vector<GeTensorDesc> &output_desc,
+                       rtStream_t stream);
   Status UpdateOutputShape(vector<GeTensorDesc> &output_desc);
   Status UpdateShapeToOutputDesc(const GeShape &shape_new, GeTensorDesc &output_desc);