From e7df70fd4e86a3c4ca895fb41be8d8cc4e397702 Mon Sep 17 00:00:00 2001
From: zhaozhixuan <zhaozhixuan2@hisilicon.com>
Date: Sun, 28 Mar 2021 09:35:35 +0800
Subject: [PATCH] Add ut.

---
 ge/graph/load/model_manager/davinci_model.cc | 6 +++---
 ge/hybrid/executor/hybrid_model_executor.cc  | 2 +-
 tests/depends/runtime/src/runtime_stub.cc    | 2 ++
 tests/ut/ge/CMakeLists.txt                   | 1 +
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/ge/graph/load/model_manager/davinci_model.cc b/ge/graph/load/model_manager/davinci_model.cc
index dfe6f390..74e2be9d 100755
--- a/ge/graph/load/model_manager/davinci_model.cc
+++ b/ge/graph/load/model_manager/davinci_model.cc
@@ -3700,7 +3700,7 @@ Status DavinciModel::NnExecute(rtStream_t stream, bool async_mode, const InputDa
     if (profiling_model_execute_on) {
       GE_CHK_STATUS_RET_NOLOG(
         ProfilingManager::Instance().ProfileStepInfo(index_id, model_id, 0, rt_model_stream_, device_id));
-    }       
+    }
     GELOGD("rtModelExecute do");
     GE_IF_BOOL_EXEC(profiling_model_execute_on, SetProfileTime(MODEL_INFER_START));
     rtError_t rt_ret = rtModelExecute(rt_model_handle_, rt_model_stream_, 0);
@@ -3709,8 +3709,8 @@ Status DavinciModel::NnExecute(rtStream_t stream, bool async_mode, const InputDa
     GELOGD("rtModelExecute end");
     if (profiling_model_execute_on) {
       GE_CHK_STATUS_RET_NOLOG(
-        ProfilingManager::Instance().ProfileStepInfo(index_id, model_id, 0, rt_model_stream_, device_id));
-    }   
+        ProfilingManager::Instance().ProfileStepInfo(index_id, model_id, 1, rt_model_stream_, device_id));
+    }
     iterator_count_++;
   }
 
diff --git a/ge/hybrid/executor/hybrid_model_executor.cc b/ge/hybrid/executor/hybrid_model_executor.cc
index 80dc4184..f1251586 100755
--- a/ge/hybrid/executor/hybrid_model_executor.cc
+++ b/ge/hybrid/executor/hybrid_model_executor.cc
@@ -78,11 +78,11 @@ Status HybridModelExecutor::ExecuteGraphInternal(SubgraphExecutor &executor,
   GE_CHK_STATUS_RET_NOLOG(ResetExecutionContext(context_));
   RECORD_MODEL_EXECUTION_EVENT(&context_, "[InitContext] End");
 
-  // tag_id 0 means step begin, 1 meas step end.
   uint64_t index_id = context_.iteration + 1;
   uint64_t model_id = static_cast<uint64_t>(model_->GetModelId());
   int32_t device_id = static_cast<int32_t>(device_id_);
   auto &prof_mgr = ProfilingManager::Instance();
+  // tag_id 0 means step begin, 1 meas step end.
   if (prof_mgr.ProfilingModelExecuteOn()) {
     GE_CHK_STATUS_RET_NOLOG(prof_mgr.ProfileStepInfo(index_id, model_id, 0, stream_, device_id));
   }
diff --git a/tests/depends/runtime/src/runtime_stub.cc b/tests/depends/runtime/src/runtime_stub.cc
index f8eedcbc..b062ec80 100644
--- a/tests/depends/runtime/src/runtime_stub.cc
+++ b/tests/depends/runtime/src/runtime_stub.cc
@@ -313,6 +313,8 @@ rtError_t rtFlushCache(uint64_t base, uint32_t len) { return RT_ERROR_NONE; }
 
 rtError_t rtProfilerTrace(uint64_t id, bool notify, uint32_t flags, rtStream_t stream_) { return RT_ERROR_NONE; }
 
+rtError_t rtProfilerTraceEx(uint64_t id, uint64_t modelId, uint16_t tagId, rtStream_t stream) { return RT_ERROR_NONE; }
+
 rtError_t rtMemSetRC(const void *dev_ptr, uint64_t size, uint32_t read_count) { return RT_ERROR_NONE; }
 
 rtError_t rtStreamSwitch(void *ptr, rtCondition_t condition, int64_t value, rtStream_t true_stream, rtStream_t stream) {
diff --git a/tests/ut/ge/CMakeLists.txt b/tests/ut/ge/CMakeLists.txt
index 02ff7fcc..90b8b0ed 100755
--- a/tests/ut/ge/CMakeLists.txt
+++ b/tests/ut/ge/CMakeLists.txt
@@ -854,6 +854,7 @@ target_link_libraries(ge_ut_common PRIVATE
     ascend_protobuf
     json
     ge_ut_graph
+    runtime_stub
 )
 
 # build common format