You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
graphengine/ge/common/helper/om_file_helper.cc

430 lines
17 KiB

/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "framework/common/helper/om_file_helper.h"
#include <string>
#include <vector>
#include "common/math/math_util.h"
#include "common/auth/file_saver.h"
#include "framework/common/debug/log.h"
#include "framework/common/debug/ge_log.h"
#include "framework/common/ge_inner_error_codes.h"
#include "framework/common/util.h"
using std::string;
namespace {
const int32_t kOptionalNum = 2;
}
namespace ge {
// For Load
FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileLoadHelper::Init(const ge::ModelData &model) {
if (CheckModelValid(model) != SUCCESS) {
return FAILED;
}
uint32_t model_data_size = model.model_len - sizeof(ModelFileHeader);
uint8_t *model_data = static_cast<uint8_t *>(model.model_data) + sizeof(ModelFileHeader);
Status ret = Init(model_data, model_data_size);
return ret;
}
FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status OmFileLoadHelper::Init(uint8_t *model_data,
const uint32_t model_data_size) {
Status status = LoadModelPartitionTable(model_data, model_data_size);
if (status != SUCCESS) {
return status;
}
is_inited_ = true;
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) {
if (!is_inited_) {
GELOGE(PARAM_INVALID, "OmFileLoadHelper has not been initialized!");
return PARAM_INVALID;
}
bool found = false;
for (ModelPartition &part : context_.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;
}
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) {
GELOGE(PARAM_INVALID, "Model_data must not be null!");
return PARAM_INVALID;
}
// Model length too small
if (model.model_len < (sizeof(ModelFileHeader) + sizeof(ModelPartitionTable))) {
GELOGE(PARAM_INVALID,
"Invalid model. length[%u] < sizeof(ModelFileHeader)[%zu] + sizeof(ModelPartitionTable)[%zu].",
model.model_len, sizeof(ModelFileHeader), sizeof(ModelPartitionTable));
return PARAM_INVALID;
}
// Get file header
auto model_header = reinterpret_cast<ModelFileHeader *>(model.model_data);
// Determine whether the file length and magic number match
if ((model_header->length != model.model_len - sizeof(ModelFileHeader)) ||
(MODEL_FILE_MAGIC_NUM != model_header->magic)) {
GELOGE(PARAM_INVALID,
"Invalid model. file_header->length[%u] + sizeof(ModelFileHeader)[%zu] != model->model_len[%u] || "
"MODEL_FILE_MAGIC_NUM[%u] != file_header->magic[%u]",
model_header->length, sizeof(ModelFileHeader), model.model_len, MODEL_FILE_MAGIC_NUM, model_header->magic);
return PARAM_INVALID;
}
return SUCCESS;
}
Status OmFileLoadHelper::LoadModelPartitionTable(uint8_t *model_data, const uint32_t model_data_size) {
if (model_data == nullptr) {
GELOGE(ACL_ERROR_GE_EXEC_MODEL_ADDR_INVALID, "Param model_data must not be null!");
return ACL_ERROR_GE_EXEC_MODEL_ADDR_INVALID;
}
// Init partition table
auto partition_table = reinterpret_cast<ModelPartitionTable *>(model_data);
// Davinici model partition include graph-info weight-info task-info tbe-kernel :
// Original model partition include graph-info
if ((partition_table->num != PARTITION_SIZE) && (partition_table->num != (PARTITION_SIZE - 1)) &&
(partition_table->num != (PARTITION_SIZE - kOptionalNum)) && (partition_table->num != 1)) {
GELOGE(ACL_ERROR_GE_PARAM_INVALID, "Invalid partition_table->num:%u", partition_table->num);
return ACL_ERROR_GE_PARAM_INVALID;
}
size_t mem_offset = SIZE_OF_MODEL_PARTITION_TABLE(*partition_table);
GELOGD("ModelPartitionTable num:%u, ModelFileHeader length:%zu, ModelPartitionTable length:%zu",
partition_table->num, sizeof(ModelFileHeader), mem_offset);
if (model_data_size <= mem_offset) {
GELOGE(ACL_ERROR_GE_EXEC_MODEL_DATA_SIZE_INVALID, "invalid model data, partition_table->num:%u, model data size %u",
partition_table->num, model_data_size);
return ACL_ERROR_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 + mem_offset;
partition.type = partition_table->partition[i].type;
context_.partition_datas_.push_back(partition);
if (partition.size > model_data_size || mem_offset > model_data_size - partition.size) {
GELOGE(ACL_ERROR_GE_EXEC_MODEL_DATA_SIZE_INVALID,
"The partition size %zu is greater than the model data size %u.",
partition.size + mem_offset, model_data_size);
return ACL_ERROR_GE_EXEC_MODEL_DATA_SIZE_INVALID;
}
mem_offset += partition.size;
GELOGD("Partition, type:%d, size:%u", static_cast<int>(partition.type), partition.size);
}
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 %u: 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(ACL_ERROR_GE_EXEC_MODEL_DATA_SIZE_INVALID,
"invalid model data, partition_table->num:%u, model data size %u",
partition_table->num, model_data_size);
return ACL_ERROR_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 %u 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(ACL_ERROR_GE_EXEC_MODEL_DATA_SIZE_INVALID,
"The partition size %u is greater than the model data size %u.",
partition.size + cur_offset, model_data_size);
return ACL_ERROR_GE_EXEC_MODEL_DATA_SIZE_INVALID;
}
cur_offset += partition.size;
GELOGD("Partition, type:%d, size:%u, model_index:%u", 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:%u, all size:%u", 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_;
}
FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY ModelPartitionTable *OmFileSaveHelper::GetPartitionTable() {
auto partition_size = static_cast<uint32_t>(context_.partition_datas_.size());
// Build ModelPartitionTable, flex array
context_.partition_table_.clear();
context_.partition_table_.resize(sizeof(ModelPartitionTable) + sizeof(ModelPartitionMemInfo) * partition_size, 0);
auto partition_table = reinterpret_cast<ModelPartitionTable *>(context_.partition_table_.data());
partition_table->num = partition_size;
uint32_t mem_offset = 0;
for (uint32_t i = 0; i < partition_size; i++) {
ModelPartition partition = context_.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 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);
return FAILED;
}
context_.partition_datas_.push_back(partition);
context_.model_data_len_ += partition.size;
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;
(void)save_param.ek_file;
(void)save_param.encode_mode;
(void)save_param.hw_key_file;
(void)save_param.pri_key_file;
Status ret = SaveModelToFile(output_file, model, is_offline);
if (ret == SUCCESS) {
GELOGD("Generate model with encrypt.");
}
return ret;
}
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");
return domi::PARAM_INVALID;
}
ModelPartitionTable *partition_table = GetPartitionTable();
if (partition_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(*partition_table);
FMK_UINT32_ADDCHECK(size_of_table, model_data_len)
model_header_.length = size_of_table + model_data_len;
GELOGD("Sizeof(ModelFileHeader):%zu,sizeof(ModelPartitionTable):%u, model_data_len:%u, model_total_len:%zu",
sizeof(ModelFileHeader), size_of_table, model_data_len, model_header_.length + sizeof(ModelFileHeader));
std::vector<ModelPartition> partition_datas = context_.partition_datas_;
Status ret;
if (is_offline) {
ret = FileSaver::SaveToFile(output_file, model_header_, *partition_table, partition_datas);
} else {
ret = FileSaver::SaveToBuffWithFileHeader(model_header_, *partition_table, partition_datas, model);
}
if (ret == SUCCESS) {
GELOGD("Save model success without encrypt.");
}
return ret;
#else
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 {
ret = FileSaver::SaveToBuffWithFileHeader(model_header_, model_partition_tabels, all_model_partitions, model);
}
if (ret == SUCCESS) {
GELOGD("Save model success without encrypt.");
}
return ret;
#else
return SUCCESS;
#endif
}
} // namespace ge