From b2570a1dc548ed278ca11301887af8939a003040 Mon Sep 17 00:00:00 2001
From: lianghao <lianghao24@hisilicon.com>
Date: Thu, 5 Nov 2020 15:50:16 +0800
Subject: [PATCH] SetMemoryMallocSize in inference

---
 ge/offline/main.cc | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/ge/offline/main.cc b/ge/offline/main.cc
index 2af8a832..b10671ce 100755
--- a/ge/offline/main.cc
+++ b/ge/offline/main.cc
@@ -32,6 +32,7 @@
 #include "graph/anchor.h"
 #include "graph/debug/ge_attr_define.h"
 #include "graph/graph.h"
+#include "graph/manager/graph_var_manager.h"
 #include "graph/op_desc.h"
 #include "graph/utils/graph_utils.h"
 #include "graph/utils/type_utils.h"
@@ -253,7 +254,7 @@ class GFlagUtils {
         "  --out_nodes         Output nodes designated by users. Separate multiple nodes with semicolons (;)."
         "Use double quotation marks (\") to enclose each argument.\n"
         "                      E.g.: \"node_name1:0;node_name1:1;node_name2:0\"\n"
-        "  --input_fp16_nodes  Input node datatype is fp16. Separate multiple nodes with semicolons (;)."
+        "  --input_fp16_nodes  Input node datatype is fp16. Separate multiple nodes with semicolons (;). "
         "Use double quotation marks (\") to enclose each argument. "
         "E.g.: \"node_name1;node_name2\"\n"
         "  --insert_op_conf    Config file to insert new op\n"
@@ -894,6 +895,13 @@ domi::Status GenerateModel(std::map<string, string> &options, std::string output
     return domi::FAILED;
   }
 
+  geRet = ge::VarManager::Instance(0)->SetMemoryMallocSize(options);
+  if (geRet != ge::SUCCESS) {
+    GELOGE(ge::FAILED, "SetMemoryMallocSize failed.");
+    (void)ge::GELib::GetInstance()->Finalize();
+    return domi::FAILED;
+  }
+
   ge::Graph graph;
   std::vector<ge::GeTensor> inputs;
   if (FLAGS_framework == domi::MINDSPORE) {
@@ -1027,6 +1035,13 @@ domi::Status GenerateSingleOp(const std::string& json_file_path) {
     return domi::FAILED;
   }
 
+  ret = ge::VarManager::Instance(0)->SetMemoryMallocSize(options);
+  if (ret != ge::SUCCESS) {
+    GELOGE(ge::FAILED, "SetMemoryMallocSize failed.");
+    (void)ge::GELib::GetInstance()->Finalize();
+    return domi::FAILED;
+  }
+
   vector<ge::SingleOpBuildParam> build_params;
   if (ge::SingleOpParser::ParseSingleOpList(json_file_path, build_params) != ge::SUCCESS) {
     DOMI_LOGE("parse single op json file failed");