From e77925035abf6e68eba89c2a0e18fc72def9a804 Mon Sep 17 00:00:00 2001
From: wuweikang <wuweikang@huawei.com>
Date: Tue, 26 Jan 2021 16:09:24 +0800
Subject: [PATCH] set AiCoreNum in options

---
 ge/init/gelib.cc                          | 26 +++++++++++++++++++++++
 ge/init/gelib.h                           |  1 +
 tests/depends/runtime/src/runtime_stub.cc |  5 +++++
 3 files changed, 32 insertions(+)

diff --git a/ge/init/gelib.cc b/ge/init/gelib.cc
index b81632bd..29e89c95 100755
--- a/ge/init/gelib.cc
+++ b/ge/init/gelib.cc
@@ -54,6 +54,7 @@ const int kSocVersionLen = 50;
 const int kDefaultDeviceIdForTrain = 0;
 const int kDefaultDeviceIdForInfer = -1;
 const char *const kGlobalOptionFpCeilingModeDefault = "2";
+const int32_t ACL_ERROR_RT_FEATURE_NOT_SUPPORT = 207000;
 }  // namespace
 static std::shared_ptr<GELib> instancePtr_ = nullptr;
 
@@ -76,6 +77,13 @@ Status GELib::Initialize(const map<string, string> &options) {
     GELOGE(ret, "GeLib initial failed.");
     return ret;
   }
+
+  ret = instancePtr_->SetAiCoreNum(new_options);
+  if (ret != SUCCESS) {
+    GELOGE(ret, "GeLib initial: SetAiCoreNum failed.");
+    return ret;
+  }
+
   instancePtr_->SetDefaultPrecisionMode(new_options);
 
   if (new_options.find("ge.fpCeilingMode") == new_options.end()) {
@@ -251,6 +259,24 @@ Status GELib::SetRTSocVersion(const map<string, string> &options, map<string, st
   return SUCCESS;
 }
 
+Status GELib::SetAiCoreNum(map<string, string> &options) {
+  // Already set or get AICORE_NUM from options in offline mode
+  if (options.find(AICORE_NUM) != options.end()) {
+    return SUCCESS;
+  }
+
+  uint32_t aicore_num = 0;
+  rtError_t ret = rtGetAiCoreCount(&aicore_num);
+  if (ret == ACL_ERROR_RT_FEATURE_NOT_SUPPORT) {  // offline without ATC Input of AiCoreNum
+    return SUCCESS;
+  } else if (ret == RT_ERROR_NONE) {  // online-mode
+    options.emplace(std::make_pair(AICORE_NUM, std::to_string(aicore_num)));
+    return SUCCESS;
+  }
+  GELOGE(FAILED, "rtGetAiCoreCount failed.");
+  return FAILED;
+}
+
 void GELib::InitOptions(const map<string, string> &options) {
   this->options_.session_id = 0;
   auto iter = options.find(OPTION_EXEC_SESSION_ID);
diff --git a/ge/init/gelib.h b/ge/init/gelib.h
index e52b8dd6..39dd9bda 100644
--- a/ge/init/gelib.h
+++ b/ge/init/gelib.h
@@ -81,6 +81,7 @@ class GELib {
   Status InnerInitialize(const map<string, string> &options);
   Status SystemInitialize(const map<string, string> &options);
   Status SetRTSocVersion(const map<string, string> &options, map<string, string> &new_options);
+  Status SetAiCoreNum(map<string, string> &options);
   void SetDefaultPrecisionMode(map<string, string> &new_options);
   void RollbackInit();
   void InitOptions(const map<string, string> &options);
diff --git a/tests/depends/runtime/src/runtime_stub.cc b/tests/depends/runtime/src/runtime_stub.cc
index 9b45e7e2..d4e7824d 100644
--- a/tests/depends/runtime/src/runtime_stub.cc
+++ b/tests/depends/runtime/src/runtime_stub.cc
@@ -319,6 +319,11 @@ rtError_t rtGetSocVersion(char *version, const uint32_t maxLen)
   return RT_ERROR_NONE;
 }
 
+rtError_t rtGetAiCoreCount(uint32_t *aiCoreCnt)
+{
+  return RT_ERROR_NONE;
+}
+
 rtError_t rtSetTaskFailCallback(rtTaskFailCallback callback)
 {
   return RT_ERROR_NONE;