Profiling AR version 5

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

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

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);