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/inc/framework/common/helper/om_file_helper.h

110 lines
3.2 KiB

/**
* Copyright 2019-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.
*/
#ifndef INC_FRAMEWORK_COMMON_HELPER_OM_FILE_HELPER_H_
#define INC_FRAMEWORK_COMMON_HELPER_OM_FILE_HELPER_H_
#include <string>
#include <vector>
#include "external/ge/ge_ir_build.h"
#include "framework/common/fmk_types.h"
#include "framework/common/types.h"
#include "framework/common/ge_types.h"
using ProcParam = struct PROC_PARAM;
using std::string;
using std::vector;
namespace ge {
struct ModelPartition {
ModelPartitionType type;
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_ = 0;
};
struct SaveParam {
int32_t encode_mode;
std::string ek_file;
std::string cert_file;
std::string hw_key_file;
std::string pri_key_file;
std::string model_name;
};
class OmFileLoadHelper {
public:
Status Init(const ge::ModelData &model);
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};
};
class OmFileSaveHelper {
public:
ModelFileHeader &GetModelFileHeader() { return model_header_; }
uint32_t GetModelDataSize() const { return context_.model_data_len_; }
ModelPartitionTable *GetPartitionTable();
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 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_