From 2e87815ba5382d6b3af45a5f39215a431aac39f4 Mon Sep 17 00:00:00 2001
From: taoxiangdong <taoxiangdong1@huawei.com>
Date: Thu, 10 Dec 2020 19:56:58 +0800
Subject: [PATCH] Profiling AR version 5

---
 ge/init/gelib.cc                              |  6 +++-
 .../ops_kernel_builder_manager.cc             | 28 +++++++++++++++++++
 .../ops_kernel_builder_manager.h              |  6 ++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/ge/init/gelib.cc b/ge/init/gelib.cc
index 92700179..20859c07 100755
--- a/ge/init/gelib.cc
+++ b/ge/init/gelib.cc
@@ -145,6 +145,10 @@ Status GELib::InnerInitialize(const map<string, string> &options) {
     return initOpsBuilderStatus;
   }
 
+  if (is_train_mode_ || (options_.device_id != kDefaultDeviceIdForInfer)) {
+    GE_CHK_RT_RET(rtSetDevice(options.device_id));
+  }
+
   GELOGI("sessionManager initial.");
   GE_TIMESTAMP_START(SessionManagerInitialize);
   Status initSmStatus = sessionManager_.Initialize(options);
@@ -338,7 +342,7 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status GELib::InitSystemWithOpt
   // set device id
   GELOGI("set logical device id:%u", options.device_id);
   GetContext().SetCtxDeviceId(static_cast<uint32_t>(options.device_id));
-  GE_CHK_RT_RET(rtSetDevice(options.device_id));
+  // GE_CHK_RT_RET(rtSetDevice(options.device_id));
 
   // In the scenario that the automatic add fusion is set, but there is no cleanaddr operator,
   // maybe need to check it
diff --git a/ge/opskernel_manager/ops_kernel_builder_manager.cc b/ge/opskernel_manager/ops_kernel_builder_manager.cc
index e0001fcd..30c3ffcc 100644
--- a/ge/opskernel_manager/ops_kernel_builder_manager.cc
+++ b/ge/opskernel_manager/ops_kernel_builder_manager.cc
@@ -167,4 +167,32 @@ Status OpsKernelBuilderManager::GenerateTask(const Node &node,
   GELOGD("Done invoking GenerateTask successfully");
   return SUCCESS;
 }
+
+Status OpsKernelBuilderManager::RegProfReporterCallBack(void *func) const {
+  GE_CHECK_NOTNULL(func);
+  for (auto it = ops_kernel_builders_.begin(); it != ops_kernel_builders_.end(); it++) {
+    GE_CHK_STATUS_RET(it->second->RegProfReporterCallBack(func),
+                      "Failed to invoke RegProfReporterCallBack.");
+  }
+  GELOGD("Done invoking RegProfReporterCallBack successfully");
+  return SUCCESS;
+}
+
+Status OpsKernelBuilderManager::ProfStart(uint64_t prof_config) const {
+  for (auto it = ops_kernel_builders_.begin(); it != ops_kernel_builders_.end(); it++) {
+    GE_CHK_STATUS_RET(it->second->ProfStart(prof_config),
+                      "Failed to invoke ProfStart.");
+  }
+  GELOGD("Done invoking ProfStart successfully");
+  return SUCCESS;
+}
+
+Status OpsKernelBuilderManager::ProfStop((uint64_t prof_config) const {
+  for (auto it = ops_kernel_builders_.begin(); it != ops_kernel_builders_.end(); it++) {
+    GE_CHK_STATUS_RET(it->second->ProfStop(prof_config),
+                      "Failed to invoke ProfStop.");
+  }
+  GELOGD("Done invoking ProfStop successfully");
+  return SUCCESS;
+}
 }  // namespace ge
\ No newline at end of file
diff --git a/ge/opskernel_manager/ops_kernel_builder_manager.h b/ge/opskernel_manager/ops_kernel_builder_manager.h
index 7a95ddfa..bcd05fe0 100644
--- a/ge/opskernel_manager/ops_kernel_builder_manager.h
+++ b/ge/opskernel_manager/ops_kernel_builder_manager.h
@@ -46,6 +46,12 @@ class OpsKernelBuilderManager {
   Status GenerateTask(const Node &node, RunContext &context,
                       std::vector<domi::TaskDef> &tasks) const;
 
+  Status RegProfReporterCallBack(void *func) const;
+
+  Status ProfStart(uint64_t prof_config) const;
+
+  Status ProfStop((uint64_t prof_config) const;
+
  private:
   OpsKernelBuilderManager() = default;
   static Status GetLibPaths(const std::map<std::string, std::string> &options, std::string &lib_paths);