mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-07 17:15:29 +00:00
Trim InferenceSession binary size. (#8917)
- Move flatbuffers SessionState access code into helper functions instead of duplicating them between InferenceSession and SessionState. - Trim VerifyEachNodeIsAssignedToAnEp(), e.g., disable verbose log output in a minimal build.
This commit is contained in:
Edward Chen
GitHub
parent
332c2ba4f4
commit
1985616262
5 changed files with 189 additions and 87 deletions
|
|
@ -878,7 +878,7 @@ static Status GetSubGraphSessionStatesOrtFormat(
|
|||
for (const auto& name_to_subgraph_session_state : session_states) {
|
||||
const std::string& attr_name = name_to_subgraph_session_state.first;
|
||||
SessionState& subgraph_session_state = *name_to_subgraph_session_state.second;
|
||||
auto graph_id = builder.CreateString(experimental::utils::GetSubGraphId(node_idx, attr_name));
|
||||
auto graph_id = builder.CreateString(experimental::utils::GetSubgraphId(node_idx, attr_name));
|
||||
flatbuffers::Offset<fbs::SessionState> session_state;
|
||||
ORT_RETURN_IF_ERROR(
|
||||
subgraph_session_state.SaveToOrtFormat(builder, session_state));
|
||||
|
|
@ -952,15 +952,9 @@ Status SessionState::CreateSubgraphSessionState() {
|
|||
#if defined(ENABLE_ORT_FORMAT_LOAD)
|
||||
Status SessionState::LoadFromOrtFormat(const fbs::SessionState& fbs_session_state,
|
||||
const KernelRegistryManager& kernel_registry_manager) {
|
||||
const auto* const fbs_kcis = fbs_session_state.kernels();
|
||||
ORT_RETURN_IF(nullptr == fbs_kcis, "Kernel create info is null. Invalid ORT format model.");
|
||||
const auto* const node_indices = fbs_kcis->node_indices();
|
||||
const auto* const kernel_def_hashes = fbs_kcis->kernel_def_hashes();
|
||||
ORT_RETURN_IF(nullptr == node_indices, "Kernel create info node indices are null. Invalid ORT format model.");
|
||||
ORT_RETURN_IF(nullptr == kernel_def_hashes, "Kernel create info hashes are null. Invalid ORT format model.");
|
||||
ORT_RETURN_IF_NOT(node_indices->size() == kernel_def_hashes->size(),
|
||||
"Size mismatch for kernel create info node indexes and hashes. Invalid ORT format model.",
|
||||
node_indices->size(), " != ", kernel_def_hashes->size());
|
||||
using experimental::utils::FbsSessionStateViewer;
|
||||
const FbsSessionStateViewer fbs_session_state_viewer{fbs_session_state};
|
||||
ORT_RETURN_IF_ERROR(fbs_session_state_viewer.Validate());
|
||||
|
||||
auto add_kernel_by_hash =
|
||||
[&kernel_registry_manager, this](const Node& node, uint64_t hash) {
|
||||
|
|
@ -976,19 +970,18 @@ Status SessionState::LoadFromOrtFormat(const fbs::SessionState& fbs_session_stat
|
|||
const auto& compiled_kernel_hashes = GetCompiledKernelHashes();
|
||||
|
||||
// process the nodes that existed when the model was created
|
||||
for (flatbuffers::uoffset_t i = 0; i < node_indices->size(); i++) {
|
||||
const auto node_idx = node_indices->Get(i);
|
||||
const auto kernel_hash = kernel_def_hashes->Get(i);
|
||||
for (FbsSessionStateViewer::Index i = 0, end = fbs_session_state_viewer.GetNumNodeKernelInfos(); i < end; ++i) {
|
||||
const auto node_kernel_info = fbs_session_state_viewer.GetNodeKernelInfo(i);
|
||||
|
||||
Node* const node = graph_.GetNode(node_idx);
|
||||
Node* const node = graph_.GetNode(node_kernel_info.node_index);
|
||||
if (node == nullptr) {
|
||||
// this is OK if we have compiled kernels and the original node was replaced. if not the model is invalid.
|
||||
ORT_RETURN_IF(compiled_kernel_hashes.empty(),
|
||||
"Can't find node with index ", node_idx, ". Invalid ORT format model.");
|
||||
"Can't find node with index ", node_kernel_info.node_index, ". Invalid ORT format model.");
|
||||
continue;
|
||||
}
|
||||
|
||||
ORT_RETURN_IF_ERROR(add_kernel_by_hash(*node, kernel_hash));
|
||||
ORT_RETURN_IF_ERROR(add_kernel_by_hash(*node, node_kernel_info.kernel_def_hash));
|
||||
}
|
||||
|
||||
// lookup the hashes for any nodes we compiled. the nodes indexes for compiled nodes are not in node_indices
|
||||
|
|
@ -1006,27 +999,12 @@ Status SessionState::LoadFromOrtFormat(const fbs::SessionState& fbs_session_stat
|
|||
}
|
||||
|
||||
if (!subgraph_session_states_.empty()) {
|
||||
const auto* const fbs_sub_graph_session_states = fbs_session_state.sub_graph_session_states();
|
||||
ORT_RETURN_IF(nullptr == fbs_sub_graph_session_states,
|
||||
"SessionState for subgraphs is null. Invalid ORT format model.");
|
||||
for (const auto& [node_idx, session_states] : subgraph_session_states_) {
|
||||
for (const auto& [attr_name, subgraph_session_state] : session_states) {
|
||||
const fbs::SessionState* fbs_subgraph_session_state;
|
||||
ORT_RETURN_IF_ERROR(fbs_session_state_viewer.GetSubgraphSessionState(node_idx, attr_name, fbs_subgraph_session_state));
|
||||
|
||||
for (const auto& pair : subgraph_session_states_) {
|
||||
const auto node_idx = pair.first;
|
||||
const auto& session_states = pair.second;
|
||||
for (const auto& name_to_subgraph_session_state : session_states) {
|
||||
const std::string& attr_name = name_to_subgraph_session_state.first;
|
||||
SessionState& subgraph_session_state = *name_to_subgraph_session_state.second;
|
||||
|
||||
// Use the graphid as the key to search the for the fbs::SubGraphSessionState
|
||||
const std::string key = experimental::utils::GetSubGraphId(node_idx, attr_name);
|
||||
const auto* const fbs_sub_graph_ss = fbs_sub_graph_session_states->LookupByKey(key.c_str());
|
||||
ORT_RETURN_IF(nullptr == fbs_sub_graph_ss,
|
||||
"Subgraph SessionState entry for ", key, " is missing. Invalid ORT format model.");
|
||||
|
||||
const auto* const fbs_sub_session_state = fbs_sub_graph_ss->session_state();
|
||||
ORT_RETURN_IF(nullptr == fbs_sub_session_state,
|
||||
"Subgraph SessionState for ", key, " is null. Invalid ORT format model.");
|
||||
ORT_RETURN_IF_ERROR(subgraph_session_state.LoadFromOrtFormat(*fbs_sub_session_state, kernel_registry_manager));
|
||||
ORT_RETURN_IF_ERROR(subgraph_session_state->LoadFromOrtFormat(*fbs_subgraph_session_state, kernel_registry_manager));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,77 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/framework/session_state_flatbuffers_utils.h"
|
||||
|
||||
namespace onnxruntime::experimental::utils {
|
||||
|
||||
std::string GetSubgraphId(const NodeIndex node_idx, const std::string& attr_name) {
|
||||
return std::to_string(node_idx) + "_" + attr_name;
|
||||
}
|
||||
|
||||
FbsSessionStateViewer::FbsSessionStateViewer(const fbs::SessionState& fbs_session_state)
|
||||
: fbs_session_state_{fbs_session_state} {
|
||||
}
|
||||
|
||||
Status FbsSessionStateViewer::Validate() const {
|
||||
if (fbs_session_state_.sub_graph_session_states() == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "SessionState for subgraphs is null. Invalid ORT format model.");
|
||||
}
|
||||
|
||||
const auto* const fbs_kcis = fbs_session_state_.kernels();
|
||||
if (fbs_kcis == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Kernel create info is null. Invalid ORT format model.");
|
||||
}
|
||||
|
||||
const auto* const fbs_node_indices = fbs_kcis->node_indices();
|
||||
if (fbs_node_indices == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Kernel create info node indices are null. Invalid ORT format model.");
|
||||
}
|
||||
|
||||
const auto* const fbs_kernel_def_hashes = fbs_kcis->kernel_def_hashes();
|
||||
if (fbs_kernel_def_hashes == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Kernel create info hashes are null. Invalid ORT format model.");
|
||||
}
|
||||
|
||||
if (fbs_node_indices->size() != fbs_kernel_def_hashes->size()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
|
||||
"Size mismatch for kernel create info node indexes and hashes. Invalid ORT format model.",
|
||||
fbs_node_indices->size(), " != ", fbs_kernel_def_hashes->size());
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
FbsSessionStateViewer::NodeKernelInfo FbsSessionStateViewer::GetNodeKernelInfo(Index idx) const {
|
||||
const auto* const fbs_kcis = fbs_session_state_.kernels();
|
||||
const auto* const fbs_node_indices = fbs_kcis->node_indices();
|
||||
const auto* const fbs_kernel_def_hashes = fbs_kcis->kernel_def_hashes();
|
||||
|
||||
return {fbs_node_indices->Get(idx), fbs_kernel_def_hashes->Get(idx)};
|
||||
}
|
||||
|
||||
FbsSessionStateViewer::Index FbsSessionStateViewer::GetNumNodeKernelInfos() const {
|
||||
return fbs_session_state_.kernels()->node_indices()->size();
|
||||
}
|
||||
|
||||
Status FbsSessionStateViewer::GetSubgraphSessionState(NodeIndex node_idx, const std::string& attr_name,
|
||||
const fbs::SessionState*& fbs_subgraph_session_state_out) const {
|
||||
const auto key = GetSubgraphId(node_idx, attr_name);
|
||||
const auto* const fbs_subgraph_session_state_entry =
|
||||
fbs_session_state_.sub_graph_session_states()->LookupByKey(key.c_str());
|
||||
if (fbs_subgraph_session_state_entry == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
|
||||
"Subgraph SessionState entry for ", key, " is missing. Invalid ORT format model.");
|
||||
}
|
||||
|
||||
const auto* const fbs_subgraph_session_state = fbs_subgraph_session_state_entry->session_state();
|
||||
if (fbs_subgraph_session_state == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
|
||||
"Subgraph SessionState for ", key, " is null. Invalid ORT format model.");
|
||||
}
|
||||
|
||||
fbs_subgraph_session_state_out = fbs_subgraph_session_state;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace onnxruntime::experimental::utils
|
||||
|
|
@ -5,6 +5,8 @@
|
|||
|
||||
#include <string>
|
||||
|
||||
#include "core/common/common.h"
|
||||
#include "core/flatbuffers/schema/ort.fbs.h"
|
||||
#include "core/graph/basic_types.h"
|
||||
|
||||
namespace onnxruntime::experimental::utils {
|
||||
|
|
@ -16,8 +18,64 @@ namespace onnxruntime::experimental::utils {
|
|||
* @param attr_name The name of the node attribute that contains the subgraph.
|
||||
* @return The subgraph key.
|
||||
*/
|
||||
inline std::string GetSubGraphId(const NodeIndex node_idx, const std::string& attr_name) {
|
||||
return std::to_string(node_idx) + "_" + attr_name;
|
||||
}
|
||||
std::string GetSubgraphId(const NodeIndex node_idx, const std::string& attr_name);
|
||||
|
||||
/**
|
||||
* Provides read-only helper functions for a fbs::SessionState instance.
|
||||
*/
|
||||
class FbsSessionStateViewer {
|
||||
public:
|
||||
/**
|
||||
* Creates an instance.
|
||||
* Validation is not performed here, but in Validate().
|
||||
*
|
||||
* @param fbs_session_state The fbs::SessionState instance.
|
||||
*/
|
||||
FbsSessionStateViewer(const fbs::SessionState& fbs_session_state);
|
||||
|
||||
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(FbsSessionStateViewer);
|
||||
|
||||
/**
|
||||
* Validates the underlying fbs::SessionState instance.
|
||||
* WARNING: Other methods assume that the fbs::SessionState is valid!
|
||||
*
|
||||
* @return Whether the fbs::SessionState instance is valid.
|
||||
*/
|
||||
Status Validate() const;
|
||||
|
||||
using Index = flatbuffers::uoffset_t;
|
||||
|
||||
struct NodeKernelInfo {
|
||||
NodeIndex node_index;
|
||||
uint64_t kernel_def_hash;
|
||||
};
|
||||
|
||||
/**
|
||||
* Retrieves the node kernel info element.
|
||||
*
|
||||
* @param index The index of the node kernel info element.
|
||||
* @return The node kernel info element.
|
||||
*/
|
||||
NodeKernelInfo GetNodeKernelInfo(Index idx) const;
|
||||
|
||||
/**
|
||||
* Gets the number of node kernel info elements.
|
||||
*/
|
||||
Index GetNumNodeKernelInfos() const;
|
||||
|
||||
/**
|
||||
* Retrieves the subgraph session state from the fbs::SessionState instance.
|
||||
*
|
||||
* @param node_idx The index of the node containing the subgraph.
|
||||
* @param attr_name The name of the attribute containing the subgraph.
|
||||
* @param[out] fbs_subgraph_session_state The subgraph session state. Non-null if successful.
|
||||
* @return Whether the retrieval was successful.
|
||||
*/
|
||||
Status GetSubgraphSessionState(NodeIndex node_idx, const std::string& attr_name,
|
||||
const fbs::SessionState*& fbs_subgraph_session_state) const;
|
||||
|
||||
private:
|
||||
const fbs::SessionState& fbs_session_state_;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime::experimental::utils
|
||||
|
|
|
|||
|
|
@ -131,7 +131,7 @@ Status SelectorActionTransformer::ApplySaved(Graph& graph, bool& modified, const
|
|||
const auto& actions_map = selectors_and_actions_.ActionsMap();
|
||||
const auto actions_map_end = actions_map.cend();
|
||||
|
||||
for (const auto entry : saved_actions) {
|
||||
for (const auto& entry : saved_actions) {
|
||||
auto action_iter = actions_map.find(entry.action_name);
|
||||
if (action_iter == actions_map_end) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Missing action ", entry.action_name, " for transformer ", Name());
|
||||
|
|
|
|||
|
|
@ -110,28 +110,25 @@ Status VerifyEachNodeIsAssignedToAnEpImpl(const Graph& graph, bool is_verbose,
|
|||
for (const auto& node : graph.Nodes()) {
|
||||
const auto& node_provider = node.GetExecutionProviderType();
|
||||
if (node_provider.empty()) {
|
||||
std::ostringstream oss;
|
||||
oss << "Could not find an implementation for the node ";
|
||||
if (!node.Name().empty()) {
|
||||
oss << node.Name() << ":";
|
||||
}
|
||||
oss << node.OpType() << "(" << node.SinceVersion() << ")";
|
||||
|
||||
return Status(common::ONNXRUNTIME, common::NOT_IMPLEMENTED, oss.str());
|
||||
} else {
|
||||
if (is_verbose) { // TODO: should we disable this if the number of nodes are above a certain threshold?
|
||||
std::string node_str = node.OpType();
|
||||
node_str += " (";
|
||||
node_str += node.Name();
|
||||
node_str += ")";
|
||||
node_placements[node_provider].push_back(node_str);
|
||||
}
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED,
|
||||
"Could not find an implementation for ",
|
||||
node.OpType(), "(", node.SinceVersion(), ") node with name '", node.Name(), "'");
|
||||
}
|
||||
|
||||
#if !defined(ORT_MINIMAL_BUILD)
|
||||
if (is_verbose) { // TODO: should we disable this if the number of nodes is above a certain threshold?
|
||||
const std::string node_str = node.OpType() + " (" + node.Name() + ")";
|
||||
node_placements[node_provider].push_back(node_str);
|
||||
}
|
||||
#endif // !defined(ORT_MINIMAL_BUILD)
|
||||
|
||||
// recurse into subgraphs
|
||||
const auto subgraphs = node.GetSubgraphs();
|
||||
for (const auto& subgraph : subgraphs) {
|
||||
ORT_RETURN_IF_ERROR(VerifyEachNodeIsAssignedToAnEpImpl(*subgraph, is_verbose, node_placements));
|
||||
const auto status = VerifyEachNodeIsAssignedToAnEpImpl(*subgraph, is_verbose, node_placements);
|
||||
if (!status.IsOK()) {
|
||||
return status;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -140,10 +137,16 @@ Status VerifyEachNodeIsAssignedToAnEpImpl(const Graph& graph, bool is_verbose,
|
|||
|
||||
Status VerifyEachNodeIsAssignedToAnEp(const Graph& graph, const logging::Logger& logger) {
|
||||
NodePlacementMap node_placements{};
|
||||
#if !defined(ORT_MINIMAL_BUILD)
|
||||
const bool is_verbose_mode = logger.GetSeverity() == logging::Severity::kVERBOSE;
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(logger);
|
||||
const bool is_verbose_mode = false;
|
||||
#endif // !defined(ORT_MINIMAL_BUILD)
|
||||
|
||||
ORT_RETURN_IF_ERROR(VerifyEachNodeIsAssignedToAnEpImpl(graph, is_verbose_mode, node_placements));
|
||||
const auto status = VerifyEachNodeIsAssignedToAnEpImpl(graph, is_verbose_mode, node_placements);
|
||||
|
||||
#if !defined(ORT_MINIMAL_BUILD)
|
||||
// print placement info
|
||||
if (is_verbose_mode) {
|
||||
LOGS(logger, VERBOSE) << "Node placements";
|
||||
|
|
@ -158,8 +161,9 @@ Status VerifyEachNodeIsAssignedToAnEp(const Graph& graph, const logging::Logger&
|
|||
}
|
||||
}
|
||||
}
|
||||
#endif // !defined(ORT_MINIMAL_BUILD)
|
||||
|
||||
return Status::OK();
|
||||
return status;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
|
|
@ -1187,44 +1191,29 @@ Status TransformGraphForOrtFormatModel(onnxruntime::Graph& graph, const logging:
|
|||
|
||||
Status AssignNodesToEpsFromHashesImpl(Graph& graph, const fbs::SessionState& fbs_session_state,
|
||||
const KernelRegistryManager& kernel_registry_manager) {
|
||||
const auto* const fbs_sub_graph_session_states = fbs_session_state.sub_graph_session_states();
|
||||
ORT_RETURN_IF(nullptr == fbs_sub_graph_session_states,
|
||||
"SessionState for subgraphs is null. Invalid ORT format model.");
|
||||
using experimental::utils::FbsSessionStateViewer;
|
||||
const FbsSessionStateViewer fbs_session_state_viewer{fbs_session_state};
|
||||
ORT_RETURN_IF_ERROR(fbs_session_state_viewer.Validate());
|
||||
|
||||
for (auto& node : graph.Nodes()) {
|
||||
for (auto& [attribute, subgraph] : node.GetAttributeNameToMutableSubgraphMap()) {
|
||||
const auto key = experimental::utils::GetSubGraphId(node.Index(), attribute);
|
||||
const auto* const fbs_sub_graph_ss = fbs_sub_graph_session_states->LookupByKey(key.c_str());
|
||||
ORT_RETURN_IF(nullptr == fbs_sub_graph_ss,
|
||||
"Subgraph SessionState entry for ", key, " is missing. Invalid ORT format model.");
|
||||
const fbs::SessionState* fbs_subgraph_session_state;
|
||||
ORT_RETURN_IF_ERROR(fbs_session_state_viewer.GetSubgraphSessionState(node.Index(), attribute,
|
||||
fbs_subgraph_session_state));
|
||||
|
||||
const auto* const fbs_sub_session_state = fbs_sub_graph_ss->session_state();
|
||||
ORT_RETURN_IF(nullptr == fbs_sub_session_state,
|
||||
"Subgraph SessionState for ", key, " is null. Invalid ORT format model.");
|
||||
|
||||
ORT_RETURN_IF_ERROR(AssignNodesToEpsFromHashesImpl(*subgraph, *fbs_sub_session_state, kernel_registry_manager));
|
||||
ORT_RETURN_IF_ERROR(AssignNodesToEpsFromHashesImpl(*subgraph, *fbs_subgraph_session_state,
|
||||
kernel_registry_manager));
|
||||
}
|
||||
}
|
||||
|
||||
const auto* const fbs_kcis = fbs_session_state.kernels();
|
||||
ORT_RETURN_IF(nullptr == fbs_kcis, "Kernel create info is null. Invalid ORT format model.");
|
||||
const auto* const node_indices = fbs_kcis->node_indices();
|
||||
const auto* const kernel_def_hashes = fbs_kcis->kernel_def_hashes();
|
||||
ORT_RETURN_IF(nullptr == node_indices, "Kernel create info node indices are null. Invalid ORT format model.");
|
||||
ORT_RETURN_IF(nullptr == kernel_def_hashes, "Kernel create info hashes are null. Invalid ORT format model.");
|
||||
ORT_RETURN_IF_NOT(node_indices->size() == kernel_def_hashes->size(),
|
||||
"Size mismatch for kernel create info node indexes and hashes. Invalid ORT format model.",
|
||||
node_indices->size(), " != ", kernel_def_hashes->size());
|
||||
|
||||
for (flatbuffers::uoffset_t i = 0; i < node_indices->size(); ++i) {
|
||||
const auto node_idx = node_indices->Get(i);
|
||||
const auto kernel_hash = kernel_def_hashes->Get(i);
|
||||
Node* node = graph.GetNode(node_idx);
|
||||
for (FbsSessionStateViewer::Index i = 0, end = fbs_session_state_viewer.GetNumNodeKernelInfos(); i < end; ++i) {
|
||||
const auto node_kernel_info = fbs_session_state_viewer.GetNodeKernelInfo(i);
|
||||
Node* node = graph.GetNode(node_kernel_info.node_index);
|
||||
if (!node || !node->GetExecutionProviderType().empty()) continue;
|
||||
|
||||
const KernelCreateInfo* kci = nullptr;
|
||||
ORT_RETURN_IF_NOT(kernel_registry_manager.SearchKernelRegistriesByHash(kernel_hash, &kci),
|
||||
"Failed to find kernel def hash (", kernel_hash, ") in kernel registries for ",
|
||||
ORT_RETURN_IF_NOT(kernel_registry_manager.SearchKernelRegistriesByHash(node_kernel_info.kernel_def_hash, &kci),
|
||||
"Failed to find kernel def hash (", node_kernel_info.kernel_def_hash, ") in kernel registries for ",
|
||||
node->OpType(), "(", node->SinceVersion(), ") node with name '", node->Name(), "'.");
|
||||
node->SetExecutionProviderType(kci->kernel_def->Provider());
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue