saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix the ctx_gen tool to make sure all generated ctx.onnx have max_size (#23097)

Browse source 
### Description
Fix the qnn_ctx_gen tool to make sure all generated ctx.onnx have max_size
This commit is contained in:
Hector Li 2024-12-12 21:12:02 -08:00 committed by GitHub
parent f43f40facf
commit 2a36fd4f6e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 12 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
onnxruntime/test/qnn_ctx_gen/main.cc
View file

@ -33,8 +33,11 @@ static void CheckStatus(const Status& status) {
// from the last context cache Onnx model, find the EPContext node with main_context=1,
// and get the QNN context binary file name, this context binary contains all graphs from all Onnx models
// get the max spill fill buffer size
static void GetLastContextBinaryFileName(const std::basic_string<ORTCHAR_T> last_onnx_ctx_file,
std::string& last_ctx_bin_file) {
std::string& last_ctx_bin_file,
int64_t& max_size) {
max_size = 0;
std::shared_ptr<Model> ctx_model;
CheckStatus(Model::Load(ToPathString(last_onnx_ctx_file), ctx_model, nullptr,
(*((OrtEnv*)*ort_env.get())->GetEnvironment().GetLoggingManager()).DefaultLogger()));
@ -43,6 +46,7 @@ static void GetLastContextBinaryFileName(const std::basic_string<ORTCHAR_T> last
if (node.OpType() == "EPContext") {
NodeAttrHelper node_helper(node);
int64_t is_main_context = node_helper.Get("main_context", static_cast<int64_t>(0));
max_size = node_helper.Get("max_size", static_cast<int64_t>(0));
if (1 == is_main_context) {
last_ctx_bin_file = node_helper.Get("ep_cache_context", "");
return;
@ -55,7 +59,8 @@ static void GetLastContextBinaryFileName(const std::basic_string<ORTCHAR_T> last
// the last QNN context binary file
// Remove not used QNN context binary file, only keep the last one which contains all graphs
static void UpdateEpContextModel(const std::vector<std::basic_string<ORTCHAR_T>>& ep_ctx_files,
const std::string& last_qnn_ctx_binary_file_name) {
const std::string& last_qnn_ctx_binary_file_name,
int64_t max_size) {
for (auto ep_ctx_file : ep_ctx_files) {
std::shared_ptr<Model> ctx_model;
auto path_str = ToPathString(ep_ctx_file);
@ -75,6 +80,8 @@ static void UpdateEpContextModel(const std::vector<std::basic_string<ORTCHAR_T>>
std::remove(file_path.string().c_str());
node.ClearAttribute("ep_cache_context");
node.AddAttribute("ep_cache_context", last_qnn_ctx_binary_file_name);
node.ClearAttribute("max_size");
node.AddAttribute("max_size", max_size);
}
}
}
@ -181,7 +188,8 @@ int real_main(int argc, char* argv[]) {
// Get the last context binary file name
std::string last_qnn_ctx_binary_file_name;
GetLastContextBinaryFileName(ep_ctx_files.back(), last_qnn_ctx_binary_file_name);
int64_t max_size = 0;
GetLastContextBinaryFileName(ep_ctx_files.back(), last_qnn_ctx_binary_file_name, max_size);
std::cout << "The last context binary file: " << last_qnn_ctx_binary_file_name << std::endl;
if (last_qnn_ctx_binary_file_name.empty()) {
throw Ort::Exception("Can't find QNN context binary file from the Onnx model.", OrtErrorCode::ORT_FAIL);
@ -191,7 +199,7 @@ int real_main(int argc, char* argv[]) {
// Update generated context cache Onnx model to make the main EPContext node point to
// the last QNN context binary file
// Remove not used QNN context binary file, only keep the last one which contains all graphs
UpdateEpContextModel(ep_ctx_files, last_qnn_ctx_binary_file_name);
UpdateEpContextModel(ep_ctx_files, last_qnn_ctx_binary_file_name, max_size);
}
ORT_CATCH(const Ort::Exception& e) {
fprintf(stderr, "Failed to generate context cache file: %s \n", e.what());