From 198561626291ea14865e7346bcdc30d98c6fdaa9 Mon Sep 17 00:00:00 2001 From: Edward Chen <18449977+edgchen1@users.noreply.github.com> Date: Wed, 1 Sep 2021 18:18:32 -0700 Subject: [PATCH] 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. --- onnxruntime/core/framework/session_state.cc | 50 ++++------- .../session_state_flatbuffers_utils.cc | 77 +++++++++++++++++ .../session_state_flatbuffers_utils.h | 64 +++++++++++++- .../selector_action_transformer.cc | 2 +- onnxruntime/core/session/inference_session.cc | 83 ++++++++----------- 5 files changed, 189 insertions(+), 87 deletions(-) create mode 100644 onnxruntime/core/framework/session_state_flatbuffers_utils.cc diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc index a9975cc461..503e53d239 100644 --- a/onnxruntime/core/framework/session_state.cc +++ b/onnxruntime/core/framework/session_state.cc @@ -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 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)); } } } diff --git a/onnxruntime/core/framework/session_state_flatbuffers_utils.cc b/onnxruntime/core/framework/session_state_flatbuffers_utils.cc new file mode 100644 index 0000000000..ea194cbcaa --- /dev/null +++ b/onnxruntime/core/framework/session_state_flatbuffers_utils.cc @@ -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 diff --git a/onnxruntime/core/framework/session_state_flatbuffers_utils.h b/onnxruntime/core/framework/session_state_flatbuffers_utils.h index 832dcd560a..06c4acad29 100644 --- a/onnxruntime/core/framework/session_state_flatbuffers_utils.h +++ b/onnxruntime/core/framework/session_state_flatbuffers_utils.h @@ -5,6 +5,8 @@ #include +#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 diff --git a/onnxruntime/core/optimizer/selectors_actions/selector_action_transformer.cc b/onnxruntime/core/optimizer/selectors_actions/selector_action_transformer.cc index bed076160d..21fc13fc20 100644 --- a/onnxruntime/core/optimizer/selectors_actions/selector_action_transformer.cc +++ b/onnxruntime/core/optimizer/selectors_actions/selector_action_transformer.cc @@ -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()); diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 8e6eb4569d..a511dc0a4b 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -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()); }