!512 om file save ge_root_model

From: @yskhhh
Reviewed-by: @xchu42,@wqtshg,@ji_chen
Signed-off-by: @ji_chen

ge/common/CMakeLists.txt

@@ -24,6 +24,7 @@ set(SRC_LIST
     "helper/om_file_helper.cc"
     "helper/model_helper.cc"
     "../model/ge_model.cc"
+    "../model/ge_root_model.cc"
     "auth/file_saver.cc"
     "fp16_t.cc"
     "math/fp16_math.cc"

ge/common/auth/file_saver.cc

@@ -258,6 +258,65 @@ FileSaver::SaveToFile(const string &file_path, ModelFileHeader &file_header, Mod
   return SUCCESS;
 }
 
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status
+FileSaver::SaveToFile(const string &file_path, ModelFileHeader &file_header,
+                      vector<ModelPartitionTable *> &model_partition_tables,
+                      const vector<vector<ModelPartition>> &all_partition_datas) {
+  file_header.is_encrypt = ModelEncryptType::UNENCRYPTED;
+
+  const Status ret = SaveWithFileHeader(file_path, file_header, model_partition_tables, all_partition_datas);
+  GE_CHK_BOOL_RET_STATUS(ret == SUCCESS, FAILED, "save file failed, file_path:%s, file header len:%u.",
+                         file_path.c_str(), file_header.length);
+  return SUCCESS;
+}
+
+Status FileSaver::SaveWithFileHeader(const std::string &file_path, const ModelFileHeader &file_header,
+                                     vector<ModelPartitionTable *> &model_partition_tables,
+                                     const vector<vector<ModelPartition>> &all_partition_datas) {
+
+  GE_CHK_BOOL_EXEC(model_partition_tables.size() == all_partition_datas.size(),
+                   return PARAM_INVALID,
+                   "model table size %zu does not match partition size %zu",
+                   model_partition_tables.size(), all_partition_datas.size())
+  for (size_t index = 0; index < model_partition_tables.size(); ++index) {
+    auto &cur_partiton_data = all_partition_datas[index];
+    auto &cur_model_partition_table = *model_partition_tables[index];
+    GE_CHK_BOOL_RET_STATUS(!cur_partiton_data.empty() && cur_model_partition_table.num != 0
+                           && cur_model_partition_table.num == cur_partiton_data.size(), FAILED,
+                           "Invalid param:partition data size is (%u), model_partition_table.num is (%zu).",
+                           cur_model_partition_table.num, cur_partiton_data.size());
+  }
+
+  // Open file
+  int32_t fd = 0;
+  GE_CHK_BOOL_TRUE_EXEC_WITH_LOG(OpenFile(fd, file_path) != SUCCESS, return FAILED);
+  Status ret = SUCCESS;
+  do {
+    // Write file header
+    GE_CHK_BOOL_TRUE_EXEC_WITH_LOG(
+        WriteData(static_cast<const void *>(&file_header), sizeof(ModelFileHeader), fd) != SUCCESS, ret = FAILED;
+        break);
+    for (size_t index = 0; index < model_partition_tables.size(); ++index) {
+      // Write model partition table
+      auto &cur_tabel = *model_partition_tables[index];
+      uint32_t table_size = static_cast<uint32_t>(SIZE_OF_MODEL_PARTITION_TABLE(cur_tabel));
+      GE_CHK_BOOL_TRUE_EXEC_WITH_LOG(
+          WriteData(static_cast<const void *>(&cur_tabel), table_size, fd) != SUCCESS, ret = FAILED; break);
+      // Write partition data
+      auto &cur_partition_datas = all_partition_datas[index];
+      for (const auto &partition_data : cur_partition_datas) {
+        GELOGI("GC:size[%zu]", partition_data.size);
+        GE_CHK_BOOL_TRUE_EXEC_WITH_LOG(
+            WriteData(static_cast<const void *>(partition_data.data), partition_data.size, fd) != SUCCESS, ret = FAILED;
+            break);
+      }
+    }
+  } while (0);
+  // Close file
+  GE_CHK_BOOL_RET_STATUS(mmClose(fd) == EN_OK, FAILED, "Close file failed.");
+  return ret;
+}
+
 FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status FileSaver::SaveToFile(const string &file_path, const void *data,
                                                                               int len) {
   if (data == nullptr || len <= 0) {

ge/common/auth/file_saver.h

@@ -74,6 +74,10 @@ class FileSaver {
                            ModelPartitionTable &model_partition_table,
                            const std::vector<ModelPartition> &partition_datas);
 
+  static Status SaveToFile(const string &file_path, ModelFileHeader &file_header,
+                        vector<ModelPartitionTable *> &model_partition_tables,
+                        const vector<vector<ModelPartition>> &all_partition_datas);
+
   static Status SaveToBuffWithFileHeader(const ModelFileHeader &file_header,
                                             ModelPartitionTable &model_partition_table,
                                             const std::vector<ModelPartition> &partitionDatas,
@@ -108,6 +112,9 @@ class FileSaver {
   static Status SaveWithFileHeader(const std::string &file_path, const ModelFileHeader &file_header,
                                    ModelPartitionTable &model_partition_table,
                                    const std::vector<ModelPartition> &partition_datas);
+  static Status SaveWithFileHeader(const std::string &file_path, const ModelFileHeader &file_header,
+                                       vector<ModelPartitionTable *> &model_partition_tables,
+                                       const vector<vector<ModelPartition>> &all_partition_datas);
 };
 }  // namespace ge
 #endif  // GE_COMMON_AUTH_FILE_SAVER_H_

ge/common/ge_common.mk

@@ -7,6 +7,7 @@ GE_COMMON_LOCAL_SRC_FILES := \
     helper/om_file_helper.cc \
     helper/model_helper.cc \
     ../model/ge_model.cc \
+    ../model/ge_root_model.cc \
     auth/file_saver.cc \
     fp16_t.cc \
     math/fp16_math.cc \

ge/common/helper/model_helper.cc

@@ -32,6 +32,7 @@ using domi::ModelTaskDef;
 
 namespace {
 const int64_t kOriginalOmPartitionNum = 1;
+const uint32_t kStatiOmFileModelNum = 1;
 }
 
 
@@ -39,7 +40,7 @@ namespace ge {
 FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY ModelHelper::~ModelHelper() { (void)ReleaseLocalModelData(); }
 
 Status ModelHelper::SaveModelPartition(std::shared_ptr<OmFileSaveHelper> &om_file_save_helper, ModelPartitionType type,
-                                       const uint8_t *data, size_t size) {
+                                       const uint8_t *data, size_t size, size_t model_index) {
   if (size < 1 || size > UINT32_MAX) {
     GELOGE(PARAM_INVALID, "Add model partition failed, partition size %zu invalid", size);
     if (size > UINT32_MAX) {
@@ -68,25 +69,16 @@ Status ModelHelper::SaveModelPartition(std::shared_ptr<OmFileSaveHelper> &om_fil
   partition_model.data = const_cast<uint8_t *>(data);
   partition_model.size = static_cast<uint32_t>(size);
   partition_model.type = type;
-  if (om_file_save_helper->AddPartition(partition_model) != SUCCESS) {
+  if (om_file_save_helper->AddPartition(partition_model, model_index) != SUCCESS) {
     GELOGE(PARAM_INVALID, "Add model partition failed, partition size %zu", size);
     return PARAM_INVALID;
   }
   return SUCCESS;
 }
 
-FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmModel(const GeModelPtr &ge_model,
-                                                                                   const SaveParam &save_param,
-                                                                                   const std::string &output_file,
-                                                                                   ModelBufferData& model) {
-  if (output_file.empty()) {
-    GELOGE(FAILED, "GraphBuilder SaveModel received invalid file name prefix");
-    return FAILED;
-  }
 
-  GE_IF_BOOL_EXEC(ge_model == nullptr, GELOGE(FAILED, "Ge_model is nullptr"); return FAILED);
-  std::shared_ptr<OmFileSaveHelper> om_file_save_helper = ge::MakeShared<OmFileSaveHelper>();
-  GE_CHECK_NOTNULL(om_file_save_helper);
+Status ModelHelper::SaveModelDef(std::shared_ptr<OmFileSaveHelper> &om_file_save_helper,
+                                 const GeModelPtr &ge_model, ge::Buffer &model_buffer, size_t model_index) {
   ModelPtr model_tmp = ge::MakeShared<ge::Model>(ge_model->GetName(), ge_model->GetPlatformVersion());
   if (model_tmp == nullptr) {
     GELOGE(FAILED, "Create Model %s Ptr failed", ge_model->GetName().c_str());
@@ -96,16 +88,21 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmMod
   model_tmp->SetVersion(ge_model->GetVersion());
   model_tmp->SetAttr(ge_model->MutableAttrMap());
 
-  ge::Buffer model_buffer;
+
   (void)model_tmp->Save(model_buffer);
   GELOGD("MODEL_DEF size is %zu", model_buffer.GetSize());
   if (model_buffer.GetSize() > 0) {
     if (SaveModelPartition(om_file_save_helper, ModelPartitionType::MODEL_DEF, model_buffer.GetData(),
-                           model_buffer.GetSize()) != SUCCESS) {
+                           model_buffer.GetSize(), model_index) != SUCCESS) {
       GELOGE(PARAM_INVALID, "Add model graph partition failed");
       return PARAM_INVALID;
     }
   }
+  return SUCCESS;
+}
+
+Status ModelHelper::SaveModelWeights(std::shared_ptr<OmFileSaveHelper> &om_file_save_helper,
+                                     const GeModelPtr &ge_model, size_t model_index) {
   auto ge_model_weight = ge_model->GetWeight();
   GELOGD("WEIGHTS_DATA size is %zu, %p", ge_model_weight.GetSize(), ge_model_weight.GetData());
   // weight is not necessary
@@ -113,31 +110,44 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmMod
     GE_CHK_STATUS_RET(SaveModelPartition(om_file_save_helper,
                                          ModelPartitionType::WEIGHTS_DATA,
                                          ge_model_weight.GetData(),
-                                         ge_model_weight.GetSize()), "Add weight partition failed");
+                                         ge_model_weight.GetSize(), model_index), "Add weight partition failed");
   }
+  return SUCCESS;
+}
 
+Status ModelHelper::SaveModelTbeKernel(std::shared_ptr<OmFileSaveHelper> &om_file_save_helper,
+                                       const GeModelPtr &ge_model, size_t model_index) {
   TBEKernelStore tbe_kernel_store = ge_model->GetTBEKernelStore();
   GELOGD("TBE_KERNELS size is %zu", tbe_kernel_store.DataSize());
   if (tbe_kernel_store.DataSize() > 0) {
     GE_CHK_STATUS_RET(SaveModelPartition(om_file_save_helper,
                                          ModelPartitionType::TBE_KERNELS,
-                                         tbe_kernel_store.Data(),
-                                         tbe_kernel_store.DataSize()), "Add tbe kernel partition failed");
+                                         ge_model->GetTBEKernelStore().Data(),
+                                         ge_model->GetTBEKernelStore().DataSize(), model_index), 
+                                         "Add tbe kernel partition failed");
   }
-
   // no need to check value, DATA->NetOutput
   (void)tbe_kernel_store.Load(tbe_kernel_store.Data(), tbe_kernel_store.DataSize());
 
+  return SUCCESS;
+}
+
+Status ModelHelper::SaveModelCustAICPU(std::shared_ptr<OmFileSaveHelper> &om_file_save_helper,
+                                       const GeModelPtr &ge_model, size_t model_index) {
   CustAICPUKernelStore cust_aicpu_kernel_store = ge_model->GetCustAICPUKernelStore();
   GELOGD("cust aicpu kernels size is %zu", cust_aicpu_kernel_store.DataSize());
   if (cust_aicpu_kernel_store.DataSize() > 0) {
     GE_CHK_STATUS_RET(SaveModelPartition(om_file_save_helper,
                                          ModelPartitionType::CUST_AICPU_KERNELS,
-                                         cust_aicpu_kernel_store.Data(),
-                                         cust_aicpu_kernel_store.DataSize()),
+                                         ge_model->GetCustAICPUKernelStore().Data(),
+                                         cust_aicpu_kernel_store.DataSize(), model_index),
                       "Add cust aicpu kernel partition failed");
   }
+  return SUCCESS;
+}
 
+Status ModelHelper::SaveModelTaskDef(std::shared_ptr<OmFileSaveHelper> &om_file_save_helper,
+                                     const GeModelPtr &ge_model, ge::Buffer &task_buffer, size_t model_index) {
   std::shared_ptr<ModelTaskDef> model_task_def = ge_model->GetModelTaskDefPtr();
   if (model_task_def == nullptr) {
     GELOGE(ACL_ERROR_GE_MEMORY_ALLOCATION, "Create model task def ptr failed");
@@ -146,9 +156,9 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmMod
   size_t partition_task_size = model_task_def->ByteSizeLong();
   GE_IF_BOOL_EXEC(partition_task_size == 0 || partition_task_size > INT_MAX,
                   GELOGE(FAILED, "Model_def's byte size (%zu) is invalid!", partition_task_size);
-                  return FAILED);
+                      return FAILED);
 
-  ge::Buffer task_buffer(partition_task_size);
+  task_buffer = ge::Buffer(partition_task_size);
   if (task_buffer.GetSize() == 0) {
     GELOGE(ACL_ERROR_GE_MEMORY_ALLOCATION, "Alloc model task def buffer failed");
     return ACL_ERROR_GE_MEMORY_ALLOCATION;
@@ -159,14 +169,20 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmMod
   GELOGD("TASK_INFO size is %zu", partition_task_size);
 
   if (SaveModelPartition(om_file_save_helper, ModelPartitionType::TASK_INFO, task_buffer.GetData(),
-                         partition_task_size) != SUCCESS) {
+                         partition_task_size, model_index) != SUCCESS) {
     GELOGE(PARAM_INVALID, "Add model task def partition failed");
     return PARAM_INVALID;
   }
+  return SUCCESS;
+}
+
+Status ModelHelper::SaveModelHeader(std::shared_ptr<OmFileSaveHelper> &om_file_save_helper,
+                                    const GeModelPtr &ge_model, size_t model_num) {
   // Save target/version to model_header
   ModelFileHeader &model_header = om_file_save_helper->GetModelFileHeader();
   model_header.platform_type = ge_model->GetPlatformType();
   model_header.om_ir_version = ge_model->GetVersion();
+  model_header.model_num = model_num;
   std::string platform_version = ge_model->GetPlatformVersion();
 
   errno_t err;
@@ -188,8 +204,142 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmMod
   }
   string model_name = reinterpret_cast<char *>(model_header.name);
   GELOGD("Model name save:%s", model_name.c_str());
+  return SUCCESS;
+}
+
+Status ModelHelper::SaveAllModelPartiton(std::shared_ptr<OmFileSaveHelper>& om_file_save_helper,
+                                         const GeModelPtr &ge_model, ge::Buffer &model_buffer,
+                                         ge::Buffer &task_buffer, size_t model_index) {
+  if (SaveModelDef(om_file_save_helper, ge_model, model_buffer, model_index) != SUCCESS) {
+    GELOGE(FAILED, "save model def failed");
+    return FAILED;
+  }
+
+  if (SaveModelWeights(om_file_save_helper, ge_model, model_index) != SUCCESS) {
+    GELOGE(FAILED, "save model weights failed");
+    return FAILED;
+  }
+
+  if (SaveModelTbeKernel(om_file_save_helper, ge_model, model_index) != SUCCESS) {
+    GELOGE(FAILED, "save model tbe kernel failed");
+    return FAILED;
+  }
+
+  if (SaveModelCustAICPU(om_file_save_helper, ge_model, model_index) != SUCCESS) {
+    GELOGE(FAILED, "save model cust ai cpu failed");
+    return FAILED;
+  }
+
+
+  if (SaveModelTaskDef(om_file_save_helper, ge_model, task_buffer, model_index) != SUCCESS) {
+    GELOGE(FAILED, "save task def failed");
+    return FAILED;
+  }
+  return SUCCESS;
+}
+
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmModel(const GeModelPtr &ge_model,
+                                                                                   const SaveParam &save_param,
+                                                                                   const std::string &output_file,
+                                                                                   ModelBufferData& model) {
+  if (output_file.empty()) {
+    GELOGE(FAILED, "GraphBuilder SaveModel received invalid file name prefix");
+    return FAILED;
+  }
 
-  Status ret = om_file_save_helper->SaveModel(save_param, output_file.c_str(), model, is_offline_);
+  GE_IF_BOOL_EXEC(ge_model == nullptr, GELOGE(FAILED, "Ge_model is nullptr"); return FAILED);
+  std::shared_ptr<OmFileSaveHelper> om_file_save_helper = ge::MakeShared<OmFileSaveHelper>();
+  GE_CHECK_NOTNULL(om_file_save_helper);
+  ge::Buffer model_buffer;
+  ge::Buffer task_buffer;
+
+  auto ret = SaveAllModelPartiton(om_file_save_helper, ge_model, model_buffer, task_buffer);
+  if (ret != SUCCESS) {
+    GELOGE(ret, "save all model partition failed");
+    return ret;
+  }
+
+  ret = SaveModelHeader(om_file_save_helper, ge_model);
+  if (ret != SUCCESS) {
+    GELOGE(ret, "save model header failed");
+    return ret;
+  }
+
+  ret = om_file_save_helper->SaveModel(save_param, output_file.c_str(), model, is_offline_);
+  if (ret != SUCCESS) {
+    GELOGE(FAILED, "OmFileSaveHelper SaveModel return fail.");
+    return ret;
+  }
+  return SUCCESS;
+}
+
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmRootModel(
+    const GeRootModelPtr &ge_root_model,
+    const SaveParam &save_param,
+    const std::string &output_file,
+    ModelBufferData& model,
+    bool is_unknown_shape) {
+
+  GE_CHECK_NOTNULL(ge_root_model);
+  GE_IF_BOOL_EXEC(ge_root_model == nullptr, GELOGE(FAILED, "Ge_root_model is nullptr"); return FAILED);
+
+  auto &name_to_ge_model = ge_root_model->GetSubgraphInstanceNameToModel();
+  GE_IF_BOOL_EXEC(name_to_ge_model.empty(), GELOGE(FAILED, "Ge_root_model has no sub model"); return FAILED);
+  GE_IF_BOOL_EXEC(output_file.empty(),
+                  GELOGE(FAILED, "GraphBuilder SaveModel received invalid file name prefix");
+                  return FAILED);
+
+  if (!is_unknown_shape) {
+    auto &model_root = name_to_ge_model.begin()->second;
+    return SaveToOmModel(model_root, save_param, output_file, model);
+  }
+
+  std::shared_ptr<OmFileSaveHelper> om_file_save_helper = ge::MakeShared<OmFileSaveHelper>();
+  GE_CHECK_NOTNULL(om_file_save_helper);
+
+  auto &first_ge_model = name_to_ge_model.at(ge_root_model->GetRootGraph()->GetName());
+
+  // ge root model must be the first to be loaded
+  vector<string> model_names{ge_root_model->GetRootGraph()->GetName()};
+  for (auto &item : name_to_ge_model) {
+    if (item.first != model_names.front()) {
+      model_names.emplace_back(item.first);
+    }
+  }
+  
+  vector<ge::Buffer> model_buffers(model_names.size());
+  vector<ge::Buffer> task_buffers(model_names.size());
+
+  size_t cur_index = 0;
+
+  if (model_names.size() > 1) {
+    GELOGD("only save first model MODEL_DEF");
+    if (SaveModelDef(om_file_save_helper, first_ge_model, model_buffers[cur_index], cur_index) != SUCCESS) {
+      GELOGE(FAILED, "save model def failed");
+      return FAILED;
+    }
+    ++cur_index;
+  }
+
+  for (; cur_index < model_names.size(); ++cur_index) {
+    auto model_name = model_names[cur_index];
+    GELOGD("cur model %s index is %zu", model_name.c_str(), cur_index);
+    const GeModelPtr &ge_model = name_to_ge_model.at(model_name);
+    auto ret = SaveAllModelPartiton(om_file_save_helper, ge_model, model_buffers[cur_index],
+                                    task_buffers[cur_index], cur_index);
+    if (ret != SUCCESS) {
+      GELOGE(INTERNAL_ERROR, "Save model %s failed", model_name.c_str());
+      return INTERNAL_ERROR;
+    }
+  }
+
+  auto ret = SaveModelHeader(om_file_save_helper, first_ge_model, model_names.size());
+  if (ret != SUCCESS) {
+    GELOGE(INTERNAL_ERROR, "Save model %s header failed", first_ge_model->GetName().c_str());
+    return INTERNAL_ERROR;
+  }
+
+  ret = om_file_save_helper->SaveRootModel(save_param, output_file.c_str(), model, is_offline_);
   if (ret != SUCCESS) {
     GELOGE(FAILED, "OmFileSaveHelper SaveModel return fail.");
     return FAILED;
@@ -288,7 +438,6 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::LoadModel(c
   }
 
   file_header_ = reinterpret_cast<ModelFileHeader *>(model_data.model_data);
-
   OmFileLoadHelper om_load_helper;
   status = om_load_helper.Init(model_addr_tmp_, model_len_tmp_);
   if (status != SUCCESS) {
@@ -310,7 +459,61 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::LoadModel(c
     GELOGE(status, "GenerateGeModel failed");
     return status;
   }
+  GELOGD("in ModelHelper::LoadModel, is_assign_model_ is setted to true!");
+  is_assign_model_ = true;
+  return SUCCESS;
+}
+
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::LoadRootModel(const ge::ModelData &model_data) {
+  if (model_data.model_data == nullptr || model_data.model_len == 0) {
+    GELOGE(GE_EXEC_MODEL_DATA_SIZE_INVALID, "Model_data is nullptr, or model_data_size is 0");
+    return GE_EXEC_MODEL_DATA_SIZE_INVALID;
+  }
+
+  if (is_assign_model_) {
+    GELOGE(GE_EXEC_LOAD_MODEL_REPEATED, "Model helper has already loaded!");
+    return GE_EXEC_LOAD_MODEL_REPEATED;
+  }
+
+  if (ReleaseLocalModelData() != SUCCESS) {
+    GELOGE(INTERNAL_ERROR, "ReleaseLocalModelData failed.");
+    return INTERNAL_ERROR;
+  }
+
+  Status status = ge::DavinciModelParser::ParseModelContent(model_data, model_addr_tmp_, model_len_tmp_);
+  if (status != SUCCESS) {
+    GELOGE(status, "Parse model content failed!");
+    return status;
+  }
+
+  file_header_ = reinterpret_cast<ModelFileHeader *>(model_data.model_data);
+
+  //model verison 1.0 file header does not have model_num member
+  is_unknown_shape_model_ = file_header_->version >= ge::MODEL_VERSION &&
+                            file_header_->model_num > kStatiOmFileModelNum;
+  GELOGD("cur om model is ge root model or no %d, model version %zu", is_unknown_shape_model_, file_header_->version);
 
+  OmFileLoadHelper om_load_helper;
+  if (is_unknown_shape_model_) {
+    auto model_num = file_header_->model_num;
+    status = om_load_helper.Init(model_addr_tmp_, model_len_tmp_, model_num);
+  } else {
+    status = om_load_helper.Init(model_addr_tmp_, model_len_tmp_);
+  }
+  if (status != SUCCESS) {
+    GELOGE(status, "Om_load_helper init failed");
+    model_addr_tmp_ = nullptr;
+    return status;
+  }
+  // Encrypt model need to del temp model/no encrypt model don't need to del model
+  model_addr_tmp_ = nullptr;
+
+  status = GenerateGeRootModel(om_load_helper);
+  if (status != SUCCESS) {
+    GELOGE(status, "GenerateGeRootModel failed");
+    return status;
+  }
+  GELOGD("in ModelHelper::LoadRootModel, is_assign_model_ is setted to true!");
   is_assign_model_ = true;
   return SUCCESS;
 }
@@ -341,6 +544,61 @@ Status ModelHelper::GenerateGeModel(OmFileLoadHelper &om_load_helper) {
   return SUCCESS;
 }
 
+Status ModelHelper::GenerateGeRootModel(OmFileLoadHelper &om_load_helper) {
+  GELOGD("Begin to generate ge root model");
+  root_model_ = ge::MakeShared<ge::GeRootModel>();
+  GE_CHECK_NOTNULL(root_model_);
+  if (!is_unknown_shape_model_) {
+    if (GenerateGeModel(om_load_helper) != SUCCESS) {
+      GELOGE(FAILED, "GenerateGeModel failed");
+      return FAILED;
+    }
+    GE_CHECK_NOTNULL(model_);
+    root_model_->SetRootGraph(GraphUtils::GetComputeGraph(model_->GetGraph()));
+    return SUCCESS;
+  }
+
+  bool is_first_model = true;
+  for (size_t mode_index = 0;  mode_index < file_header_->model_num; ++mode_index) {
+    GeModelPtr cur_model = ge::MakeShared<ge::GeModel>();
+    Status ret = LoadModelData(om_load_helper, cur_model, mode_index);
+    if (ret != SUCCESS) {
+      return GE_EXEC_LOAD_MODEL_PARTITION_FAILED;
+    }
+
+    if (is_first_model) {
+      is_first_model = false;
+      root_model_->SetRootGraph(GraphUtils::GetComputeGraph(cur_model->GetGraph()));
+      root_model_->SetModelId(cur_model->GetModelId());
+      model_ = cur_model;
+      continue;
+    }
+
+    ret = LoadWeights(om_load_helper, cur_model, mode_index);
+    if (ret != SUCCESS) {
+      return GE_EXEC_LOAD_WEIGHT_PARTITION_FAILED;
+    }
+
+    ret = LoadTBEKernelStore(om_load_helper, cur_model, mode_index);
+    if (ret != SUCCESS) {
+      return GE_EXEC_LOAD_KERNEL_PARTITION_FAILED;
+    }
+
+    ret = LoadCustAICPUKernelStore(om_load_helper, cur_model, mode_index);
+    if (ret != SUCCESS) {
+      return GE_EXEC_LOAD_KERNEL_PARTITION_FAILED;
+    }
+
+    ret = LoadTask(om_load_helper, cur_model, mode_index);
+    if (ret != SUCCESS) {
+      return GE_EXEC_LOAD_TASK_PARTITION_FAILED;
+    }
+    root_model_->SetSubgraphInstanceNameToModel(cur_model->GetName(), cur_model);
+  }
+
+  return SUCCESS;
+}
+
 Status ModelHelper::LoadModelData(OmFileLoadHelper &om_load_helper) {
   ModelPartition partition_model_def;
   // no need to check value, DATA->NetOutput
@@ -366,6 +624,28 @@ void ModelHelper::SetModelToGeModel(ge::Model &model) {
   model_->SetAttr(model.MutableAttrMap());
 }
 
+Status ModelHelper::LoadModelData(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index) {
+  ModelPartition partition_model_def;
+  // no need to check value, DATA->NetOutput
+  om_load_helper.GetModelPartition(ModelPartitionType::MODEL_DEF, partition_model_def, mode_index);
+  GELOGD("Model_def partition addr:%p,size:%u", partition_model_def.data, partition_model_def.size);
+
+  ge::Model model;
+  if (ge::Model::Load(partition_model_def.data, partition_model_def.size, model) != SUCCESS) {
+    GELOGE(INTERNAL_ERROR, "Load model failed.");
+    return INTERNAL_ERROR;
+  }
+
+  cur_model->SetGraph(model.GetGraph());
+  cur_model->SetName(model.GetName());
+  cur_model->SetVersion(model.GetVersion());
+  cur_model->SetPlatformVersion(model.GetPlatformVersion());
+  cur_model->SetAttr(model.MutableAttrMap());
+
+  return SUCCESS;
+}
+
+
 Status ModelHelper::LoadWeights(OmFileLoadHelper &om_load_helper) {
   ModelPartition partition;
   if (om_load_helper.GetModelPartition(ModelPartitionType::WEIGHTS_DATA, partition) != SUCCESS) {
@@ -379,6 +659,19 @@ Status ModelHelper::LoadWeights(OmFileLoadHelper &om_load_helper) {
   return SUCCESS;
 }
 
+Status ModelHelper::LoadWeights(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index) {
+  ModelPartition partition;
+  if (om_load_helper.GetModelPartition(ModelPartitionType::WEIGHTS_DATA, partition, mode_index) != SUCCESS) {
+    GELOGE(FAILED, "Get weight model partition failed.");
+    return FAILED;
+  }
+  ge::Buffer weight = ge::Buffer::CopyFrom(partition.data, partition.size);
+  cur_model->SetWeight(weight);
+
+  GELOGD("GetWeight size:%u", partition.size);
+  return SUCCESS;
+}
+
 FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::LoadTask(OmFileLoadHelper &om_load_helper) {
   ModelPartition task_partition;
   if (om_load_helper.GetModelPartition(ModelPartitionType::TASK_INFO, task_partition) != SUCCESS) {
@@ -398,6 +691,27 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::LoadTask(Om
   return SUCCESS;
 }
 
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::LoadTask(OmFileLoadHelper &om_load_helper,
+                                                                              GeModelPtr &cur_model,
+                                                                              size_t mode_index) {
+  ModelPartition task_partition;
+  if (om_load_helper.GetModelPartition(ModelPartitionType::TASK_INFO, task_partition, mode_index) != SUCCESS) {
+    GELOGE(FAILED, "Get task model partition failed.");
+    return FAILED;
+  }
+  std::shared_ptr<ModelTaskDef> task = ge::MakeShared<ModelTaskDef>();
+  GE_CHECK_NOTNULL(task);
+  if (task_partition.size != 0) {
+    if (!ReadProtoFromArray(task_partition.data, task_partition.size, task.get())) {
+      GELOGE(INTERNAL_ERROR, "ReadProtoFromArray failed.");
+      return INTERNAL_ERROR;
+    }
+    GELOGD("TASK_INFO op_size:%zu, stream_num:%u", task->op().size(), task->stream_num());
+  }
+  cur_model->SetModelTaskDef(task);
+  return SUCCESS;
+}
+
 Status ModelHelper::LoadTBEKernelStore(OmFileLoadHelper &om_load_helper) {
   // Load tbe kernels
   ModelPartition partition_kernel_def;
@@ -414,6 +728,23 @@ Status ModelHelper::LoadTBEKernelStore(OmFileLoadHelper &om_load_helper) {
   return SUCCESS;
 }
 
+Status ModelHelper::LoadTBEKernelStore(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index) {
+  // Load tbe kernels
+  ModelPartition partition_kernel_def;
+  TBEKernelStore kernel_store;
+  if (om_load_helper.GetModelPartition(ModelPartitionType::TBE_KERNELS, partition_kernel_def, mode_index) ==
+      SUCCESS) {
+    GELOGD("Kernels partition size:%u", partition_kernel_def.size);
+    if (kernel_store.Load(partition_kernel_def.data, partition_kernel_def.size)) {
+      GELOGD("Load tbe kernels success");
+    } else {
+      GELOGW("Load tbe kernels failed");
+    }
+  }
+  cur_model->SetTBEKernelStore(kernel_store);
+  return SUCCESS;
+}
+
 Status ModelHelper::LoadCustAICPUKernelStore(OmFileLoadHelper &om_load_helper) {
   // Load cust aicpu kernels
   ModelPartition partition_kernel_def;
@@ -421,19 +752,39 @@ Status ModelHelper::LoadCustAICPUKernelStore(OmFileLoadHelper &om_load_helper) {
   if (om_load_helper.GetModelPartition(ModelPartitionType::CUST_AICPU_KERNELS, partition_kernel_def) == SUCCESS) {
     GELOGD("Kernels partition size:%u", partition_kernel_def.size);
     if (kernel_store.Load(partition_kernel_def.data, partition_kernel_def.size)) {
-      GELOGI("Load cust aicpu kernels success");
+      GELOGD("Load cust aicpu kernels success");
+    } else {
+      GELOGW("Load cust aicpu kernels failed");
     }
   }
   model_->SetCustAICPUKernelStore(kernel_store);
   return SUCCESS;
 }
 
+Status ModelHelper::LoadCustAICPUKernelStore(OmFileLoadHelper &om_load_helper,
+                                             GeModelPtr &cur_model, size_t mode_index) {
+  // Load cust aicpu kernels
+  ModelPartition partition_kernel_def;
+  CustAICPUKernelStore kernel_store;
+  if (om_load_helper.GetModelPartition(ModelPartitionType::CUST_AICPU_KERNELS, partition_kernel_def, mode_index)
+      == SUCCESS) {
+    GELOGD("Kernels partition size:%u", partition_kernel_def.size);
+    if (kernel_store.Load(partition_kernel_def.data, partition_kernel_def.size)) {
+      GELOGD("Load cust aicpu kernels success");
+    } else {
+      GELOGW("Load cust aicpu kernels failed");
+    }
+  }
+  cur_model->SetCustAICPUKernelStore(kernel_store);
+  return SUCCESS;
+}
+
 FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY GeModelPtr ModelHelper::GetGeModel() {
   if (model_ != nullptr) {
     return model_;
   }
 
-  GELOGI("Model has not been loaded!");
+  GELOGD("Model has not been loaded!");
   std::shared_ptr<ge::GeModel> out_model = ge::MakeShared<ge::GeModel>();
   if (out_model == nullptr) {
     return nullptr;
@@ -441,6 +792,20 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY GeModelPtr ModelHelper::GetGeMo
   return out_model;
 }
 
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY GeRootModelPtr ModelHelper::GetGeRootModel() {
+  if (root_model_ != nullptr) {
+    return root_model_;
+  }
+
+  GELOGD("Model has not been loaded!");
+  std::shared_ptr<ge::GeRootModel> out_model = ge::MakeShared<ge::GeRootModel>();
+  if (out_model == nullptr) {
+    return nullptr;
+  }
+  return out_model;
+}
+
+
 Status ModelHelper::ReleaseLocalModelData() noexcept {
   Status result = SUCCESS;
   if (model_addr_tmp_ != nullptr) {

ge/common/helper/om_file_helper.cc

@@ -52,6 +52,17 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileLoadHelper::Init(u
   return SUCCESS;
 }
 
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileLoadHelper::Init(uint8_t *model_data,
+                                                                               uint32_t model_data_size,
+                                                                               uint32_t model_num) {
+  Status status = LoadModelPartitionTable(model_data, model_data_size, model_num);
+  if (status != SUCCESS) {
+    return status;
+  }
+  is_inited_ = true;
+  return SUCCESS;
+}
+
 // Use both
 FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileLoadHelper::GetModelPartition(ModelPartitionType type,
                                                                                             ModelPartition &partition) {
@@ -79,6 +90,37 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileLoadHelper::GetMod
   return SUCCESS;
 }
 
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileLoadHelper::GetModelPartition(ModelPartitionType type,
+                                                                                            ModelPartition &partition,
+                                                                                            size_t model_index) {
+  if (!is_inited_) {
+    GELOGE(PARAM_INVALID, "OmFileLoadHelper has not been initialized!");
+    return PARAM_INVALID;
+  }
+  if (model_index >= model_contexts_.size()) {
+    GELOGE(PARAM_INVALID, "cur index : %zu, model_contexts size:%zu", model_index, model_contexts_.size());
+    return PARAM_INVALID;
+  }
+  auto &cur_ctx = model_contexts_[model_index];
+  bool found = false;
+  for (ModelPartition &part : cur_ctx.partition_datas_) {
+    if (part.type == type) {
+      partition = part;
+      found = true;
+      break;
+    }
+  }
+
+  if (!found) {
+    if (type != ModelPartitionType::TBE_KERNELS && type != ModelPartitionType::WEIGHTS_DATA &&
+        type != ModelPartitionType::CUST_AICPU_KERNELS) {
+      GELOGE(FAILED, "GetModelPartition:type:%d is not in partition_datas!", static_cast<int>(type));
+      return FAILED;
+    }
+  }
+  return SUCCESS;
+}
+
 Status OmFileLoadHelper::CheckModelValid(const ge::ModelData &model) const {
   // Parameter validity check
   if (model.model_data == nullptr) {
@@ -148,6 +190,61 @@ Status OmFileLoadHelper::LoadModelPartitionTable(uint8_t *model_data, const uint
   return SUCCESS;
 }
 
+Status OmFileLoadHelper::LoadModelPartitionTable(uint8_t *model_data, uint32_t model_data_size, uint32_t model_num) {
+  if (model_data == nullptr) {
+    GELOGE(PARAM_INVALID, "Param model_data must not be null!");
+    return PARAM_INVALID;
+  }
+
+  uint32_t cur_offset = 0;
+  for (uint32_t index = 0; index < model_num; ++index) {
+    // Init partition table
+    auto partition_table = reinterpret_cast<ModelPartitionTable *>(model_data + cur_offset);
+    size_t partition_table_size = SIZE_OF_MODEL_PARTITION_TABLE(*partition_table);
+    cur_offset += partition_table_size;
+    GELOGD("Cur model index %zu: ModelPartitionTable num :%u, "
+           "ModelFileHeader length :%zu, ModelPartitionTable length :%zu",
+           index, partition_table->num, sizeof(ModelFileHeader), partition_table_size);
+    if (model_data_size <= cur_offset) {
+      GELOGE(GE_EXEC_MODEL_DATA_SIZE_INVALID, "invalid model data, partition_table->num:%u, model data size %u",
+             partition_table->num, model_data_size);
+      return GE_EXEC_MODEL_DATA_SIZE_INVALID;
+    }
+
+    for (uint32_t i = 0; i < partition_table->num; i++) {
+      ModelPartition partition;
+      partition.size = partition_table->partition[i].mem_size;
+      partition.data = model_data + cur_offset;
+      partition.type = partition_table->partition[i].type;
+      if (index >= model_contexts_.size()) {
+        if (index != model_contexts_.size()) {
+          GELOGE(FAILED, "cur index is %zu make model_contexts_ overflow", index);
+          return FAILED;
+        }
+
+        OmFileContext tmp_ctx;
+        tmp_ctx.partition_datas_.push_back(partition);
+        model_contexts_.push_back(tmp_ctx);
+      } else {
+        model_contexts_[index].partition_datas_.push_back(partition);
+      }
+
+      if (partition.size > model_data_size || cur_offset > model_data_size - partition.size) {
+        GELOGE(GE_EXEC_MODEL_DATA_SIZE_INVALID, "The partition size %zu is greater than the model data size %u.",
+               partition.size + cur_offset, model_data_size);
+        return GE_EXEC_MODEL_DATA_SIZE_INVALID;
+      }
+      cur_offset += partition.size;
+      GELOGD("Partition, type:%d, size:%u, model_index:%zu", static_cast<int>(partition.type), partition.size, index);
+    }
+  }
+  if (cur_offset != model_data_size) {
+    GELOGE(FAILED, "do not get the complete model, read end offset:%zu, all size:%zu", cur_offset, model_data_size);
+    return FAILED;
+  }
+  return SUCCESS;
+}
+
 FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY const std::vector<ModelPartition>
   &OmFileSaveHelper::GetModelPartitions() const {
   return context_.partition_datas_;
@@ -172,6 +269,28 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY ModelPartitionTable *OmFileSave
   return partition_table;
 }
 
+FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY ModelPartitionTable *OmFileSaveHelper::GetPartitionTable(
+    size_t cur_ctx_index) {
+  auto &cur_ctx = model_contexts_[cur_ctx_index];
+  auto partition_size = static_cast<uint32_t>(cur_ctx.partition_datas_.size());
+  // Build ModelPartitionTable, flex array
+  cur_ctx.partition_table_.clear();
+  cur_ctx.partition_table_.resize(sizeof(ModelPartitionTable) + sizeof(ModelPartitionMemInfo) * partition_size, 0);
+
+  auto partition_table = reinterpret_cast<ModelPartitionTable *>(cur_ctx.partition_table_.data());
+  partition_table->num = partition_size;
+
+  uint32_t mem_offset = 0;
+  for (uint32_t i = 0; i < partition_size; i++) {
+    ModelPartition partition = cur_ctx.partition_datas_[i];
+    partition_table->partition[i] = {partition.type, mem_offset, partition.size};
+    mem_offset += partition.size;
+    GELOGD("Partition, type:%d, size:%u", static_cast<int>(partition.type), partition.size);
+  }
+  return partition_table;
+}
+
+
 FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileSaveHelper::AddPartition(ModelPartition &partition) {
   if (ge::CheckUint32AddOverflow(context_.model_data_len_, partition.size) != SUCCESS) {
     GELOGE(FAILED, "UINT32 %u and %u addition can result in overflow!", context_.model_data_len_, partition.size);
@@ -182,6 +301,27 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileSaveHelper::AddPar
   return SUCCESS;
 }
 
+Status OmFileSaveHelper::AddPartition(ModelPartition &partition, size_t cur_index) {
+  if (ge::CheckUint32AddOverflow(context_.model_data_len_, partition.size) != SUCCESS) {
+    GELOGE(FAILED, "UINT32 %u and %u addition can result in overflow!", context_.model_data_len_, partition.size);
+    return FAILED;
+  }
+  if (cur_index >= model_contexts_.size()) {
+    if (cur_index != model_contexts_.size()) {
+      GELOGE(FAILED, "cur index is %zu make model_contexts_ overflow", cur_index);
+      return FAILED;
+    }
+    OmFileContext tmp_ctx;
+    tmp_ctx.model_data_len_ += partition.size;
+    tmp_ctx.partition_datas_.push_back(partition);
+    model_contexts_.push_back(tmp_ctx);
+  } else {
+    model_contexts_[cur_index].model_data_len_ += partition.size;
+    model_contexts_[cur_index].partition_datas_.push_back(partition);
+  }
+  return SUCCESS;
+}
+
 Status OmFileSaveHelper::SaveModel(const SaveParam &save_param, const char *output_file, ModelBufferData &model,
                                    bool is_offline) {
   (void)save_param.cert_file;
@@ -198,6 +338,10 @@ Status OmFileSaveHelper::SaveModel(const SaveParam &save_param, const char *outp
 
 Status OmFileSaveHelper::SaveModelToFile(const char *output_file, ModelBufferData &model, bool is_offline) {
 #if !defined(NONSUPPORT_SAVE_TO_FILE)
+  if (context_.partition_datas_.empty()) {
+    GE_CHK_BOOL_EXEC(!model_contexts_.empty(), return FAILED, "mode contexts empty");
+    context_ = model_contexts_.front();
+  }
   uint32_t model_data_len = context_.model_data_len_;
   if (model_data_len == 0) {
     GELOGE(domi::PARAM_INVALID, "Model data len error! should not be 0");
@@ -231,4 +375,53 @@ Status OmFileSaveHelper::SaveModelToFile(const char *output_file, ModelBufferDat
   return SUCCESS;
 #endif
 }
+
+Status OmFileSaveHelper::SaveRootModel(const SaveParam &save_param, const char *output_file,
+                                       ModelBufferData &model, bool is_offline) {
+  (void)save_param.cert_file;
+  (void)save_param.ek_file;
+  (void)save_param.encode_mode;
+  (void)save_param.hw_key_file;
+  (void)save_param.pri_key_file;
+
+#if !defined(NONSUPPORT_SAVE_TO_FILE)
+  vector<ModelPartitionTable *> model_partition_tabels;
+  vector<vector<ModelPartition>> all_model_partitions;
+  for (size_t ctx_index = 0; ctx_index < model_contexts_.size(); ++ctx_index) {
+    auto &cur_ctx = model_contexts_[ctx_index];
+    uint32_t cur_model_data_len = cur_ctx.model_data_len_;
+    if (cur_model_data_len == 0) {
+      GELOGE(domi::PARAM_INVALID, "Model data len error! should not be 0");
+      return domi::PARAM_INVALID;
+    }
+
+    auto tmp_table = GetPartitionTable(ctx_index);
+    if (tmp_table == nullptr) {
+      GELOGE(ge::GE_GRAPH_SAVE_FAILED, "SaveModelToFile execute failed: partition_table is NULL.");
+      return ge::GE_GRAPH_SAVE_FAILED;
+    }
+    uint32_t size_of_table = SIZE_OF_MODEL_PARTITION_TABLE(*tmp_table);
+    FMK_UINT32_ADDCHECK(size_of_table, cur_model_data_len)
+    FMK_UINT32_ADDCHECK(size_of_table + cur_model_data_len, model_header_.length)
+    model_header_.length += size_of_table + cur_model_data_len;
+    model_partition_tabels.push_back(tmp_table);
+    all_model_partitions.push_back(cur_ctx.partition_datas_);
+    GELOGD("sizeof(ModelPartitionTable):%u, cur_model_data_len:%u, cur_context_index:%zu",
+           size_of_table, cur_model_data_len, ctx_index);
+  }
+  Status ret;
+  if (is_offline) {
+    ret = FileSaver::SaveToFile(output_file, model_header_, model_partition_tabels, all_model_partitions);
+  } else {
+    GELOGW("do not support save ge root model to buff now");
+    return FAILED;
+  }
+  if (ret == SUCCESS) {
+    GELOGD("Save model success without encrypt.");
+  }
+  return ret;
+#else
+  return SUCCESS;
+#endif
+}
 }  // namespace ge

ge/common/types.cc

@@ -801,7 +801,7 @@ const uint32_t XRGB_CHN_NUM = 4;
 ///
 const bool DEFAULT_GLOBAL_POOLING = false;
 
-const uint32_t MODEL_VERSION = 0x10000000; ///< Model version 1.0///
+const uint32_t MODEL_VERSION = 0x20000000; ///< Model version 2.0///
 
 // Eltwise's input size
 const int ELTWISE_MIN_INPUT_SIZE = 2;

ge/generator/ge_generator.cc

@@ -240,6 +240,8 @@ class GeGenerator::Impl {
 
   Status SaveModel(const string &file_name_prefix, GeModelPtr &models, ModelBufferData &model);
 
+  Status SaveRootModel(const string &file_name_prefix, GeRootModelPtr &model, ModelBufferData &model_buff);
+
   Status SaveParams(GeModelPtr &ge_model, const string &type, const map<string, GeAttrValue> &attrs,
                     const vector<GeTensor> &inputs, const vector<GeTensor> &outputs);
 
@@ -505,19 +507,7 @@ Status GeGenerator::GenerateModel(const Graph &graph, const string &file_name_pr
 
   GE_CHECK_NOTNULL(ge_root_model);
   GE_CHECK_NOTNULL(ge_root_model->GetRootGraph());
-  ModelHelper model_helper;
-  string model_name = "";
-  Status name_ret = model_helper.GetModelNameFromMergedGraphName(ge_root_model->GetRootGraph()->GetName(), model_name);
-  if (name_ret != SUCCESS) {
-    ErrorManager::GetInstance().ATCReportErrMessage("E10000", {"parameter"}, {"output"});
-    GELOGE(FAILED, "Get model_name failed. Param --output is invalid");
-    return PARAM_INVALID;
-  }
-  map<string, GeModelPtr> name_to_ge_model = ge_root_model->GetSubgraphInstanceNameToModel();
-  GeModelPtr &ge_model = name_to_ge_model[ge_root_model->GetRootGraph()->GetName()];
-  GE_RETURN_WITH_LOG_IF_FALSE(ge_model != nullptr, "ge_model can not be null");
-  ge_model->SetName(model_name);
-  ret = impl_->SaveModel(file_name_prefix, ge_model, model);
+  ret = impl_->SaveRootModel(file_name_prefix, ge_root_model, model);
   if (ret != SUCCESS) {
     GELOGE(ret, "Save model failed");
     if (impl_->graph_manager_.Finalize() != SUCCESS) {
@@ -712,6 +702,44 @@ Status GeGenerator::Impl::SaveModel(const string &file_name_prefix, GeModelPtr &
   return SUCCESS;
 }
 
+Status GeGenerator::Impl::SaveRootModel(const string &file_name_prefix, GeRootModelPtr &ge_root_model,
+                                        ModelBufferData &model_buff) {
+  bool is_unknown_shape = false;
+  auto ret = ge_root_model->CheckIsUnknownShape(is_unknown_shape);
+  if (ret != SUCCESS) {
+    GELOGE(FAILED, "Check root model is unkonwn shape failed");
+    return FAILED;
+  }
+  GELOGD("begin save root model, cur model is unkonwn shape model ? : %d", is_unknown_shape);
+  GE_CHK_BOOL_EXEC(!ge_root_model->GetSubgraphInstanceNameToModel().empty(), return FAILED,
+                   "ge root model has no sub model")
+  GeModelPtr model_root = nullptr;
+  if (is_unknown_shape) {
+    model_root = make_shared<GeModel>();
+    model_root->SetGraph(GraphUtils::CreateGraphFromComputeGraph(ge_root_model->GetRootGraph()));
+    ge_root_model->SetSubgraphInstanceNameToModel(ge_root_model->GetRootGraph()->GetName(), model_root);
+    model_root->SetName(ge_root_model->GetRootGraph()->GetName());
+  } else {
+    model_root = ge_root_model->GetSubgraphInstanceNameToModel().begin()->second;
+  }
+  // set atc version
+  if (!SetAtcVersionInfo(*(model_root.get()))) {
+    GELOGW("SetPackageVersionInfo of atc failed!");
+  }
+  // set opp version
+  if (!SetOppVersionInfo(*(model_root.get()))) {
+    GELOGW("SetPackageVersionInfo of ops failed!");
+  }
+  ModelHelper model_helper;
+  model_helper.SetSaveMode(is_offline_);
+  ret = model_helper.SaveToOmRootModel(ge_root_model, save_param_, file_name_prefix, model_buff, is_unknown_shape);
+  if (ret != SUCCESS) {
+    GELOGE(ret, "Save to om model failed");
+    return ret;
+  }
+  return SUCCESS;
+}
+
 Status GeGenerator::Impl::BuildModel(const Graph &graph, const vector<GeTensor> &inputs,
                                      GeRootModelPtr &ge_root_model) {
   static std::atomic<GraphId> atomic_graph_id(0);

ge/model/ge_root_model.h

@@ -23,6 +23,7 @@
 namespace ge {
 class GeRootModel {
  public:
+  GeRootModel() = default;
   explicit GeRootModel(ComputeGraphPtr &root_graph) : root_graph_(root_graph), model_id_(INVALID_MODEL_ID) {};
   ~GeRootModel() = default;
 
@@ -35,11 +36,11 @@ class GeRootModel {
   void SetModelId(uint32_t model_id) { model_id_ = model_id; }
   uint32_t GetModelId() const { return model_id_; }
   Status CheckIsUnknownShape(bool &is_dynamic_shape);
-
+  void SetRootGraph(ComputeGraphPtr graph) { root_graph_ = graph; }
  private:
-  ComputeGraphPtr root_graph_;
+  ComputeGraphPtr root_graph_ = nullptr;
   std::map<std::string, GeModelPtr> subgraph_instance_name_to_model_;
-  uint32_t model_id_;
+  uint32_t model_id_ = 0;
 };
 }  // namespace ge
 using GeRootModelPtr = std::shared_ptr<ge::GeRootModel>;

inc/framework/common/helper/model_helper.h

@@ -25,6 +25,7 @@
 #include "common/types.h"
 #include "graph/model.h"
 #include "model/ge_model.h"
+#include "model/ge_root_model.h"
 
 namespace ge {
 class ModelHelper {
@@ -32,17 +33,22 @@ class ModelHelper {
   ModelHelper() = default;
   ~ModelHelper();
 
-  Status SaveToOmModel(const GeModelPtr &ge_model, const SaveParam &save_param,
-                       const std::string &output_file, ge::ModelBufferData &model);
-  Status SaveOriginalGraphToOmModel(const ge::Graph& graph, const std::string& output_file);
+  Status SaveToOmModel(const GeModelPtr &ge_model, const SaveParam &save_param, const std::string &output_file,
+                       ge::ModelBufferData &model);
+  Status SaveToOmRootModel(const GeRootModelPtr &ge_root_model, const SaveParam &save_param, const string &output_file,
+                           ModelBufferData &model, bool is_unknown_shape);
+  Status SaveOriginalGraphToOmModel(const ge::Graph &graph, const std::string &output_file);
   Status LoadModel(const ge::ModelData &model_data);
-  Status GetModelBufferData(ge::ModelBufferData& model);
+  Status LoadRootModel(const ge::ModelData &model_data);
+  Status GetModelBufferData(ge::ModelBufferData &model);
 
-  const ModelFileHeader* GetFileHeader() const { return file_header_; }
+  const ModelFileHeader *GetFileHeader() const { return file_header_; }
 
   GeModelPtr GetGeModel();
+  GeRootModelPtr GetGeRootModel();
   void SetSaveMode(bool val) { is_offline_ = val; }
   bool GetSaveMode(void) const { return is_offline_; }
+  bool GetModelType() const { return is_unknown_shape_model_; };
 
   Status GetBaseNameFromFileName(const std::string &file_name, std::string &base_name);
   Status GetModelNameFromMergedGraphName(const std::string &graph_name, std::string &model_name);
@@ -50,24 +56,46 @@ class ModelHelper {
  private:
   bool is_assign_model_ = false;
   bool is_offline_ = true;
-  ModelFileHeader* file_header_ = nullptr;
+  bool is_unknown_shape_model_ = false;
+  ModelFileHeader *file_header_ = nullptr;
   // Encrypted model need delete temp model and unencrypted model need not delete model
   uint8_t *model_addr_tmp_ = nullptr;
   uint32_t model_len_tmp_ = 0;
   GeModelPtr model_;
+  GeRootModelPtr root_model_;
 
-  ModelHelper(const ModelHelper&);
-  ModelHelper& operator=(const ModelHelper&);
-  Status GenerateGeModel(OmFileLoadHelper& om_load_helper);
-  Status LoadModelData(OmFileLoadHelper& om_load_helper);
-  void SetModelToGeModel(ge::Model& model);
-  Status LoadWeights(OmFileLoadHelper& om_load_helper);
-  Status LoadTask(OmFileLoadHelper& om_load_helper);
-  Status LoadTBEKernelStore(OmFileLoadHelper& om_load_helper);
-  Status LoadCustAICPUKernelStore(OmFileLoadHelper& om_load_helper);
+  ModelHelper(const ModelHelper &);
+  ModelHelper &operator=(const ModelHelper &);
+  Status GenerateGeModel(OmFileLoadHelper &om_load_helper);
+  Status GenerateGeRootModel(OmFileLoadHelper &om_load_helper);
+  Status LoadModelData(OmFileLoadHelper &om_load_helper);
+  void SetModelToGeModel(ge::Model &model);
+  Status LoadModelData(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index);
+  Status LoadWeights(OmFileLoadHelper &om_load_helper);
+  Status LoadWeights(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index);
+  Status LoadTask(OmFileLoadHelper &om_load_helper);
+  Status LoadTask(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index);
+  Status LoadTBEKernelStore(OmFileLoadHelper &om_load_helper);
+  Status LoadTBEKernelStore(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index);
+  Status LoadCustAICPUKernelStore(OmFileLoadHelper &om_load_helper);
+  Status LoadCustAICPUKernelStore(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index);
   Status ReleaseLocalModelData() noexcept;
-  Status SaveModelPartition(std::shared_ptr<OmFileSaveHelper>& om_file_save_helper,
-                            ModelPartitionType type, const uint8_t* data, size_t size);
+  Status SaveModelPartition(std::shared_ptr<OmFileSaveHelper> &om_file_save_helper, ModelPartitionType type,
+                            const uint8_t *data, size_t size, size_t model_index);
+  Status SaveModelDef(shared_ptr<OmFileSaveHelper> &om_file_save_helper, const GeModelPtr &ge_model,
+                      Buffer &model_buffer, size_t model_index = 0);
+  Status SaveModelWeights(shared_ptr<OmFileSaveHelper> &om_file_save_helper, const GeModelPtr &ge_model,
+                          size_t model_index = 0);
+  Status SaveModelTbeKernel(shared_ptr<OmFileSaveHelper> &om_file_save_helper, const GeModelPtr &ge_model,
+                            size_t model_index = 0);
+  Status SaveModelCustAICPU(shared_ptr<OmFileSaveHelper> &om_file_save_helper, const GeModelPtr &ge_model,
+                            size_t model_index = 0);
+  Status SaveModelTaskDef(shared_ptr<OmFileSaveHelper> &om_file_save_helper, const GeModelPtr &ge_model,
+                          Buffer &task_buffer, size_t model_index = 0);
+  Status SaveModelHeader(shared_ptr<OmFileSaveHelper> &om_file_save_helper, const GeModelPtr &ge_model,
+                         size_t model_num = 1);
+  Status SaveAllModelPartiton(shared_ptr<OmFileSaveHelper> &om_file_save_helper, const GeModelPtr &ge_model,
+                              Buffer &model_buffer, Buffer &task_buffer, size_t model_index = 0);
 };
 }  // namespace ge
 #endif  // INC_FRAMEWORK_COMMON_HELPER_MODEL_HELPER_H_

inc/framework/common/helper/om_file_helper.h

@@ -32,14 +32,14 @@ using std::vector;
 namespace ge {
 struct ModelPartition {
   ModelPartitionType type;
-  uint8_t* data = 0;
+  uint8_t *data = 0;
   uint32_t size = 0;
 };
 
 struct OmFileContext {
   std::vector<ModelPartition> partition_datas_;
   std::vector<char> partition_table_;
-  uint32_t model_data_len_;
+  uint32_t model_data_len_ = 0;
 };
 
 struct SaveParam {
@@ -57,15 +57,23 @@ class OmFileLoadHelper {
 
   Status Init(uint8_t *model_data, const uint32_t model_data_size);
 
+  Status Init(uint8_t *model_data, const uint32_t model_data_size, uint32_t model_num);
+
   Status GetModelPartition(ModelPartitionType type, ModelPartition &partition);
 
+  Status GetModelPartition(ModelPartitionType type, ModelPartition &partition, size_t model_index);
+
   OmFileContext context_;
 
+  vector<OmFileContext> model_contexts_;
+
  private:
   Status CheckModelValid(const ge::ModelData &model) const;
 
   Status LoadModelPartitionTable(uint8_t *model_data, const uint32_t model_data_size);
 
+  Status LoadModelPartitionTable(uint8_t *model_data, const uint32_t model_data_size, uint32_t model_num);
+
   bool is_inited_{false};
 };
 
@@ -79,15 +87,23 @@ class OmFileSaveHelper {
 
   Status AddPartition(ModelPartition &partition);
 
+  Status AddPartition(ModelPartition &partition, size_t cur_index);
+
   const std::vector<ModelPartition> &GetModelPartitions() const;
 
-  Status SaveModel(const SaveParam &save_param, const char *target_file,
-                   ge::ModelBufferData& model, bool is_offline = true);
+  Status SaveModel(const SaveParam &save_param, const char *target_file, ge::ModelBufferData &model,
+                   bool is_offline = true);
 
   Status SaveModelToFile(const char *output_file, ge::ModelBufferData &model, bool is_offline = true);
 
+  vector<OmFileContext> model_contexts_;
+
   ModelFileHeader model_header_;
   OmFileContext context_;
+
+  ModelPartitionTable *GetPartitionTable(size_t cur_ctx_index);
+
+  Status SaveRootModel(const SaveParam &save_param, const char *output_file, ModelBufferData &model, bool is_offline);
 };
 }  // namespace ge
 #endif  // INC_FRAMEWORK_COMMON_HELPER_OM_FILE_HELPER_H_

inc/framework/common/types.h

@@ -529,7 +529,7 @@ REGISTER_OPTYPE_DECLARE(HVDWAIT, "HorovodWait");
 // aicpu op for online_infer dynamic_dims
 REGISTER_OPTYPE_DECLARE(GETDYNAMICDIMS, "GetDynamicDims");
 
-enum InputMode { INPUT = 0, CONST_INPUT};
+enum InputMode { INPUT = 0, CONST_INPUT };
 
 // Definition of the processing status enum of the process module
 enum ModelProcessState {
@@ -605,7 +605,7 @@ static constexpr uint32_t MODEL_FILE_CHECKSUM_LENGTH = 64;
 ///
 /// @brief length of the reserved field in the model file header
 ///
-static constexpr uint32_t MODEL_FILE_RESERVED_LENGTH = 79;
+static constexpr uint32_t MODEL_FILE_RESERVED_LENGTH = 75;
 
 ///
 /// @ingroup domi_omg
@@ -843,9 +843,10 @@ struct ModelFileHeader {
   uint32_t ops = 0;                                       // Computing power (Kops)
   uint8_t userdefineinfo[USER_DEFINE_INFO_LENGTH] = {0};  // User-defined information. The value contains 32 characters
   uint32_t om_ir_version = 0;
+  uint32_t model_num = 0;
   uint8_t platform_version[PLATFORM_VERSION_LEN] = {0};
   uint8_t platform_type = {0};
-  uint8_t reserved[MODEL_FILE_RESERVED_LENGTH] = {0};  // Reserved field 79
+  uint8_t reserved[MODEL_FILE_RESERVED_LENGTH] = {0};  // Reserved field 75
 };
 
 static constexpr uint8_t TARGET_TYPE_LTTE_8BIT = 0;