From d2cbae3a04dc23a39c008122c57577fed776f849 Mon Sep 17 00:00:00 2001 From: Yi Zhang Date: Fri, 17 Jun 2022 17:07:21 +0800 Subject: [PATCH] =?UTF-8?q?Revert=20"Refactor=20ExecutionFrame=20and=20Ses?= =?UTF-8?q?sionState=20to=20reduce=20memory=20all=E2=80=A6=20(#11888)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Revert "Refactor ExecutionFrame and SessionState to reduce memory allocations and improve data locality (#11804)" This reverts commit 2ecba6fd25cce35c03f6b3891753e1bb1ee6275c. --- .../core/framework/ortmemoryinfo.h | 24 ---- include/onnxruntime/core/graph/graph_viewer.h | 3 +- .../core/framework/allocation_planner.cc | 23 ++-- .../core/framework/allocation_planner.h | 6 +- onnxruntime/core/framework/execution_frame.cc | 61 ++++---- onnxruntime/core/framework/execution_frame.h | 38 +++-- .../core/framework/execution_plan_base.h | 4 +- onnxruntime/core/framework/mem_pattern.h | 5 +- .../core/framework/mem_pattern_planner.h | 3 +- onnxruntime/core/framework/node_index_info.h | 5 +- .../core/framework/ort_value_name_idx_map.h | 22 ++- .../framework/ort_value_pattern_planner.cc | 18 ++- .../framework/ort_value_pattern_planner.h | 6 +- .../framework/orttraining_partial_executor.cc | 4 +- .../core/framework/parallel_executor.cc | 4 +- .../framework/prepacked_weights_container.h | 7 +- .../framework/sequential_execution_plan.h | 12 +- .../core/framework/sequential_executor.cc | 4 +- onnxruntime/core/framework/session_state.cc | 130 ++++++++---------- onnxruntime/core/framework/session_state.h | 64 ++++----- .../core/framework/session_state_utils.cc | 14 +- .../core/framework/session_state_utils.h | 2 +- .../core/framework/simple_tensor_allocator.cc | 2 +- .../core/framework/simple_tensor_allocator.h | 9 +- .../core/framework/tensor_allocator.cc | 2 +- onnxruntime/core/framework/tensor_allocator.h | 8 +- .../tensor_allocator_with_mem_pattern.h | 20 ++- onnxruntime/core/framework/utils.cc | 2 +- onnxruntime/core/graph/graph_viewer.cc | 4 +- .../optimizer/optimizer_execution_frame.cc | 27 +--- .../optimizer/optimizer_execution_frame.h | 12 +- .../test/framework/allocation_planner_test.cc | 2 +- .../test/framework/execution_frame_test.cc | 9 +- winml/adapter/winml_adapter_session.cpp | 2 +- 34 files changed, 232 insertions(+), 326 deletions(-) diff --git a/include/onnxruntime/core/framework/ortmemoryinfo.h b/include/onnxruntime/core/framework/ortmemoryinfo.h index ff0c90ac84..0ea9fe3a04 100644 --- a/include/onnxruntime/core/framework/ortmemoryinfo.h +++ b/include/onnxruntime/core/framework/ortmemoryinfo.h @@ -38,20 +38,6 @@ struct OrtMemoryInfo { return strcmp(name, other.name) < 0; } - static void HashCombine(size_t h, size_t& seed) { - seed ^= h + 0x9e3779b9 + (seed << 6) + (seed >> 2); - } - - // This is to make OrtMemoryInfo a valid key in hash tables - // we ignore device id - size_t Hash() const { - auto h = std::hash()(alloc_type); - HashCombine(std::hash()(mem_type), h); - HashCombine(std::hash()(id), h); - HashCombine(std::hash()(name), h); - return h; - } - std::string ToString() const { std::ostringstream ostr; ostr << "OrtMemoryInfo:[" @@ -65,7 +51,6 @@ struct OrtMemoryInfo { } }; -// Required by hash tables inline bool operator==(const OrtMemoryInfo& left, const OrtMemoryInfo& other) { return left.mem_type == other.mem_type && left.alloc_type == other.alloc_type && @@ -76,12 +61,3 @@ inline bool operator==(const OrtMemoryInfo& left, const OrtMemoryInfo& other) { inline bool operator!=(const OrtMemoryInfo& lhs, const OrtMemoryInfo& rhs) { return !(lhs == rhs); } std::ostream& operator<<(std::ostream& out, const OrtMemoryInfo& info); - -namespace std { -template<> -struct hash { - size_t operator()(const OrtMemoryInfo& i) const { - return i.Hash(); - } -}; -} diff --git a/include/onnxruntime/core/graph/graph_viewer.h b/include/onnxruntime/core/graph/graph_viewer.h index 4e6b190c35..90d3a3d42b 100644 --- a/include/onnxruntime/core/graph/graph_viewer.h +++ b/include/onnxruntime/core/graph/graph_viewer.h @@ -208,8 +208,7 @@ class GraphViewer { // if we're limiting the view to an IndexedSubGraph we need to create a few pieces of infrastructure that would // usually come from the full graph const IndexedSubGraph* filter_info_{nullptr}; - using FilteredNodeSet = InlinedHashSet; - FilteredNodeSet filtered_node_indices_; + std::unordered_set filtered_node_indices_; std::vector filtered_node_inputs_; std::vector filtered_node_inputs_including_initializers_; std::vector filtered_node_outputs_; diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 7b2e395268..842a5fef8d 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -77,16 +77,13 @@ std::ostream& operator<<(std::ostream& out, std::pair index_to_name; - index_to_name.reserve(name_idx_map.Size()); + std::unordered_map index_to_name; out << "Allocation Plan:\n"; out << "(ort_value_idx) output_name : \n"; auto plan_size = plan.allocation_plan.size(); - for (auto& name_index : name_idx_map) { + for (auto& name_index : session_state.GetOrtValueNameIdxMap()) { auto index = name_index.second; index_to_name[index] = name_index.first; out << "(" << index << ") " << name_index.first << " : "; @@ -144,10 +141,10 @@ static const KernelCreateInfo& GetKernelCreateInfo( class PlannerImpl { public: PlannerImpl(const Node* parent_node, const onnxruntime::GraphViewer& graph_viewer, - gsl::span outer_scope_node_args, const ExecutionProviders& providers, + const std::vector& outer_scope_node_args, const ExecutionProviders& providers, const KernelCreateInfoMap& kernel_create_info_map, const SubgraphsKernelCreateInfoMaps& subgraphs_kernel_create_info_maps, - const InlinedHashMap& outer_scope_node_arg_to_location_map, + const std::unordered_map& outer_scope_node_arg_to_location_map, const OrtValueNameIdxMap& ort_value_name_idx_map, const ISequentialPlannerContext& context, SequentialExecutionPlan& plan) : context_(context), @@ -169,13 +166,13 @@ class PlannerImpl { const Node* parent_node_; const onnxruntime::GraphViewer& graph_viewer_; - gsl::span outer_scope_node_args_; + const std::vector& outer_scope_node_args_; const ExecutionProviders& execution_providers_; const KernelCreateInfoMap& kernel_create_info_map_; const SubgraphsKernelCreateInfoMaps& subgraphs_kernel_create_info_maps_; - const InlinedHashMap& outer_scope_node_arg_to_location_map_; + const std::unordered_map& outer_scope_node_arg_to_location_map_; const OrtValueNameIdxMap& ort_value_name_idx_map_; @@ -1334,16 +1331,16 @@ Status PlannerImpl::CreatePlan() { Status SequentialPlanner::CreatePlan( const Node* parent_node, const onnxruntime::GraphViewer& graph_viewer, - gsl::span outer_scope_node_args, + const std::vector& outer_scope_node_args, const ExecutionProviders& providers, const KernelCreateInfoMap& kernel_create_info_map, const SubgraphsKernelCreateInfoMaps& subgraphs_kernel_create_info_maps, - const InlinedHashMap& outer_scope_node_arg_to_location_map, + const std::unordered_map& outer_scope_node_arg_to_location_map, const OrtValueNameIdxMap& ort_value_name_idx_map, const ISequentialPlannerContext& context, - std::optional& plan) { + std::unique_ptr& plan) { // allocate/reset here so we know it's clean - plan.emplace(); + plan = std::make_unique(); PlannerImpl planner(parent_node, graph_viewer, outer_scope_node_args, providers, kernel_create_info_map, subgraphs_kernel_create_info_maps, diff --git a/onnxruntime/core/framework/allocation_planner.h b/onnxruntime/core/framework/allocation_planner.h index 153e1a2025..a0d489254a 100644 --- a/onnxruntime/core/framework/allocation_planner.h +++ b/onnxruntime/core/framework/allocation_planner.h @@ -76,14 +76,14 @@ class SequentialPlanner { // This API allows user to provide a custom planner context. static Status CreatePlan( const Node* parent_node, const onnxruntime::GraphViewer& graph, - gsl::span outer_scope_node_args, + const std::vector& outer_scope_node_args, const ExecutionProviders& providers, const KernelCreateInfoMap& kernel_create_info_map, const SubgraphsKernelCreateInfoMaps& subgraphs_kernel_create_info_maps, - const InlinedHashMap& outer_scope_arg_to_location_map, + const std::unordered_map& outer_scope_arg_to_location_map, const OrtValueNameIdxMap& ort_value_name_idx_map, const ISequentialPlannerContext& context, - std::optional& plan); + std::unique_ptr& plan); }; } // namespace onnxruntime diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index d7626398b8..ae50f19352 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/onnxruntime/core/framework/execution_frame.cc @@ -27,10 +27,10 @@ namespace onnxruntime { IExecutionFrame::IExecutionFrame(const OrtValueNameIdxMap& ort_value_idx_map, const NodeIndexInfo& node_index_info, - gsl::span fetch_mlvalue_idxs) + const std::vector& fetch_mlvalue_idxs) : node_index_info_(node_index_info), all_values_size_(static_cast(ort_value_idx_map.MaxIdx()) + 1), - fetch_mlvalue_idxs_(fetch_mlvalue_idxs.begin(), fetch_mlvalue_idxs.end()), + fetch_mlvalue_idxs_(fetch_mlvalue_idxs), ort_value_idx_map_(ort_value_idx_map) { ORT_ENFORCE(node_index_info_.GetMaxMLValueIdx() == ort_value_idx_map.MaxIdx(), "node_index_info and ort_value_idx_map are out of sync and cannot be used"); @@ -55,7 +55,7 @@ Status IExecutionFrame::SetOutputMLValue(int index, const OrtValue& ort_value) { #endif #ifdef ENABLE_TRAINING -void IExecutionFrame::UpdateFeeds(gsl::span feed_mlvalue_idxs, gsl::span feeds) { +void IExecutionFrame::UpdateFeeds(const std::vector& feed_mlvalue_idxs, const std::vector& feeds) { ORT_ENFORCE(feed_mlvalue_idxs.size() == feeds.size()); for (size_t idx = 0, end = feed_mlvalue_idxs.size(); idx < end; ++idx) { @@ -68,12 +68,11 @@ void IExecutionFrame::UpdateFeeds(gsl::span feed_mlvalue_idxs, gsl::s } } -void IExecutionFrame::UpdateFetches(gsl::span fetch_mlvalue_idxs, - gsl::span fetches, const std::unordered_map& initializers) { +void IExecutionFrame::UpdateFetches(const std::vector& fetch_mlvalue_idxs, const std::vector& fetches, const std::unordered_map& initializers) { ORT_ENFORCE(fetch_mlvalue_idxs.size() == fetches.size()); if (!fetches.empty()) { - fetch_mlvalue_idxs_.assign(fetch_mlvalue_idxs.begin(), fetch_mlvalue_idxs.end()); + fetch_mlvalue_idxs_ = fetch_mlvalue_idxs; auto num_fetches = fetch_mlvalue_idxs_.size(); @@ -103,7 +102,7 @@ void IExecutionFrame::UpdateFetches(gsl::span fetch_mlvalue_idxs, } } -Status IExecutionFrame::GetOutputs(gsl::span fetch_mlvalue_idxs, std::vector& fetches) { +Status IExecutionFrame::GetOutputs(const std::vector& fetch_mlvalue_idxs, std::vector& fetches) { auto num_fetches = fetch_mlvalue_idxs.size(); if (fetches.empty()) { @@ -214,10 +213,10 @@ int IExecutionFrame::GetNodeIdxToMLValueIdx(int index) const { return ort_value_idx; } -void IExecutionFrame::Init(gsl::span feed_mlvalue_idxs, gsl::span feeds, +void IExecutionFrame::Init(const std::vector& feed_mlvalue_idxs, const std::vector& feeds, const std::unordered_map& initializers, const std::function& is_initializer_sparse_func, - gsl::span fetches) { + const std::vector& fetches) { ORT_ENFORCE(feeds.size() == feed_mlvalue_idxs.size()); ORT_ENFORCE(fetches.empty() || fetches.size() == fetch_mlvalue_idxs_.size()); @@ -281,7 +280,7 @@ void IExecutionFrame::Init(gsl::span feed_mlvalue_idxs, gsl::span())); } else { #else - ORT_UNUSED_PARAMETER(is_initializer_sparse_func); + ORT_UNUSED_PARAMETER(is_initializer_sparse_func); #endif // !defined(DISABLE_SPARSE_TENSORS) if (!dest.IsAllocated()) { // NOTE: This doesn't need to support ExecutionFrame custom allocators as they only come into play @@ -332,13 +331,14 @@ bool IExecutionFrame::IsOutput(int ort_value_idx) const { return std::find(fetch_mlvalue_idxs_.begin(), fetch_mlvalue_idxs_.end(), ort_value_idx) != fetch_mlvalue_idxs_.end(); } -ExecutionFrame::ExecutionFrame(gsl::span feed_mlvalue_idxs, gsl::span feeds, - gsl::span fetch_mlvalue_idxs, gsl::span fetches, +ExecutionFrame::ExecutionFrame(const std::vector& feed_mlvalue_idxs, const std::vector& feeds, + const std::vector& fetch_mlvalue_idxs, const std::vector& fetches, const std::unordered_map& fetch_allocators, const SessionState& session_state) : IExecutionFrame(session_state.GetOrtValueNameIdxMap(), session_state.GetNodeIndexInfo(), fetch_mlvalue_idxs), session_state_(session_state), - mem_patterns_(nullptr) { + mem_patterns_(nullptr), + planner_(nullptr) { Init( feed_mlvalue_idxs, feeds, session_state.GetInitializedTensors(), #if !defined(DISABLE_SPARSE_TENSORS) @@ -362,12 +362,12 @@ ExecutionFrame::ExecutionFrame(gsl::span feed_mlvalue_idxs, gsl::span // map the custom allocators to ort_value_idx entries if (!fetch_allocators.empty()) { - custom_allocators_.reserve(fetch_allocators.size()); - const auto idx_size = fetch_mlvalue_idxs.size(); - for (const auto& e : fetch_allocators) { - if (e.first < idx_size) { - int ort_value_idx = fetch_mlvalue_idxs[e.first]; - custom_allocators_.insert_or_assign(ort_value_idx, e.second); + for (size_t idx = 0, end = fetch_mlvalue_idxs.size(); idx < end; ++idx) { + int ort_value_idx = fetch_mlvalue_idxs[idx]; + + auto custom_alloc_entry = fetch_allocators.find(idx); + if (custom_alloc_entry != fetch_allocators.cend()) { + custom_allocators_[ort_value_idx] = custom_alloc_entry->second; } } } @@ -388,13 +388,12 @@ ExecutionFrame::ExecutionFrame(gsl::span feed_mlvalue_idxs, gsl::span //if there are some traditional ml value type in inputs disable the memory pattern optimization. if (all_tensors) { mem_patterns_ = session_state.GetMemoryPatternGroup(feeds, feed_mlvalue_idxs, inferred_shapes_); - // if no existing patterns, generate one in this execution frame + // if no existing patterns, generate one in this executionframe if (!mem_patterns_) { - planner_.emplace(*session_state.GetExecutionPlan()); + planner_ = std::make_unique(*session_state.GetExecutionPlan()); } else { // pre-allocate the big chunk requested in memory pattern. // all the internal kernel's input/output tensors will be allocated on these buffer. - buffers_.reserve(mem_patterns_->locations.size()); for (size_t i = 0; i < mem_patterns_->locations.size(); i++) { const auto& location = mem_patterns_->locations[i]; ORT_ENFORCE(buffers_.find(location) == buffers_.end()); @@ -829,7 +828,7 @@ const AllocPlanPerValue& ExecutionFrame::GetAllocationPlan(int ort_value_idx) { } void ExecutionFrame::TraceAllocate(int ort_value_idx, size_t size) { - if (planner_.has_value()) { + if (planner_) { // don't trace the output tensors or external outputs. auto& allocation_plan = GetAllocationPlan(ort_value_idx); if (allocation_plan.alloc_kind == AllocKind::kAllocateOutput || @@ -846,7 +845,7 @@ void ExecutionFrame::TraceAllocate(int ort_value_idx, size_t size) { void ExecutionFrame::TraceFree(int ort_value_idx) { // don't trace free on output tensors. - if (planner_.has_value() && !IsOutput(ort_value_idx)) { + if (planner_ && !IsOutput(ort_value_idx)) { const SequentialExecutionPlan* p_seq_exec_plan = session_state_.GetExecutionPlan(); const auto& alloc_plan = p_seq_exec_plan->allocation_plan; ORT_ENFORCE(ort_value_idx >= 0 && static_cast(ort_value_idx) < alloc_plan.size()); @@ -871,8 +870,8 @@ void ExecutionFrame::TraceFree(int ort_value_idx) { // generate memory pattern based on the tracing of memory allocation/free in current execution // return error if the planner is not setup. -Status ExecutionFrame::GeneratePatterns(MemoryPatternGroup& out) { - if (!planner_.has_value()) { +Status ExecutionFrame::GeneratePatterns(MemoryPatternGroup* out) const { + if (!planner_) { return Status(ONNXRUNTIME, FAIL, "Memory pattern planner is not enabled on this execution framework."); } @@ -890,12 +889,10 @@ bool ExecutionFrame::TryGetInferredShape(int index, TensorShape& shape) const { // Search for inferred shape. // If inferred shape is found, it's assigned to "shape" so that caller can use it. - if (inferred_shapes_ != nullptr) { - auto it = inferred_shapes_->find(ort_value_idx); - if (it != inferred_shapes_->end()) { - shape = it->second; - return true; - } + auto it = inferred_shapes_.find(ort_value_idx); + if (it != inferred_shapes_.end()) { + shape = it->second; + return true; } // Tell the caller if the search is successful or not. diff --git a/onnxruntime/core/framework/execution_frame.h b/onnxruntime/core/framework/execution_frame.h index e9097fb4be..58f094f426 100644 --- a/onnxruntime/core/framework/execution_frame.h +++ b/onnxruntime/core/framework/execution_frame.h @@ -13,7 +13,6 @@ #include "core/framework/iexecutor.h" #include "core/framework/ort_value.h" #include "core/framework/node_index_info.h" -#include "core/framework/ort_value_pattern_planner.h" #include "core/framework/sequential_execution_plan.h" #include "core/framework/tensor.h" #include "core/graph/graph_viewer.h" @@ -23,6 +22,7 @@ namespace onnxruntime { class DataTransferManager; class SessionState; class OrtValueNameIdxMap; +class OrtValuePatternPlanner; struct MemoryPatternGroup; class NodeIndexInfo; @@ -32,12 +32,12 @@ class IExecutionFrame { // initialized until the derived class is constructed. IExecutionFrame(const OrtValueNameIdxMap& ort_value_idx_map, const NodeIndexInfo& node_index_info, - gsl::span fetch_mlvalue_idxs); + const std::vector& fetch_mlvalue_idxs); - void Init(gsl::span feed_mlvalue_idxs, gsl::span feeds, + void Init(const std::vector& feed_mlvalue_idxs, const std::vector& feeds, const std::unordered_map& initializers, const std::function& is_initializer_sparse_func, - gsl::span fetches); + const std::vector& fetches); public: virtual ~IExecutionFrame(); @@ -55,13 +55,12 @@ class IExecutionFrame { // Override the index-th output with ort_value Status SetOutputMLValue(int index, const OrtValue& ort_value); #endif - -#ifdef ENABLE_TRAINING - void UpdateFeeds(gsl::span feed_mlvalue_idxs, gsl::span feeds); - void UpdateFetches(gsl::span fetch_mlvalue_idxs, gsl::span fetches, +#ifdef ENABLE_TRAINING + void UpdateFeeds(const std::vector& feed_mlvalue_idxs, const std::vector& feeds); + void UpdateFetches(const std::vector& fetch_mlvalue_idxs, const std::vector& fetches, const std::unordered_map& initializers); - Status GetOutputs(gsl::span fetch_mlvalue_idxs, std::vector& fetches); + Status GetOutputs(const std::vector& fetch_mlvalue_idxs, std::vector& fetches); #endif // TO DO: make it thread safe @@ -120,20 +119,20 @@ class IExecutionFrame { // All the intermediate values for the entire graph. // Input and Output values are passed in by executors - InlinedVector all_values_; + std::vector all_values_; // perf optimization to avoid calling all_values_.size() repeatedly as the size is fixed once constructed const size_t all_values_size_; - InlinedVector fetch_mlvalue_idxs_; + std::vector fetch_mlvalue_idxs_; const OrtValueNameIdxMap& ort_value_idx_map_; }; class ExecutionFrame final : public IExecutionFrame { public: - ExecutionFrame(gsl::span feed_mlvalue_idxs, gsl::span feeds, - gsl::span fetch_mlvalue_idxs, gsl::span fetches, + ExecutionFrame(const std::vector& feed_mlvalue_idxs, const std::vector& feeds, + const std::vector& fetch_mlvalue_idxs, const std::vector& fetches, // optional custom allocators. key is index in fetches const std::unordered_map& fetch_allocators, const SessionState& session_state); @@ -152,10 +151,10 @@ class ExecutionFrame final : public IExecutionFrame { bool create_fence = false); // thread-safe - Status GeneratePatterns(MemoryPatternGroup& out); + Status GeneratePatterns(MemoryPatternGroup* out) const; bool HasMemoryPatternPlanner() const { - return planner_.has_value(); + return planner_ != nullptr; } // This function try retrieve the inferred shapes for the given NodeArg index. @@ -213,7 +212,7 @@ class ExecutionFrame final : public IExecutionFrame { const SessionState& session_state_; // map of index to custom allocator - InlinedHashMap custom_allocators_; + std::unordered_map custom_allocators_; // If we already have cached memory pattern on these input shapes // Use this mem pattern that create a big chunk for all the internal @@ -222,17 +221,16 @@ class ExecutionFrame final : public IExecutionFrame { // If no cached memory pattern, and we enable the memory pattern optimization // use this planner_ to trace the memory allocation in current executor. - std::optional planner_; + std::unique_ptr planner_; // Big chunks on different locations that will be used by mem_pattern. - InlinedHashMap buffers_; + std::map buffers_; // Given the input shapes of the executed graph, ExecutionFrame tries inferring // all symbolic shapes. inferred_shapes_[i] is the shape of OrtValue indexed // by i, if the key i exists. // inferred_shapes_ is generated together with mem_patterns_. - // It is never updated after creation - const InlinedHashMap* inferred_shapes_{nullptr}; + std::unordered_map inferred_shapes_; #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) // Size of virtual memory allocated before any kernel execution. diff --git a/onnxruntime/core/framework/execution_plan_base.h b/onnxruntime/core/framework/execution_plan_base.h index 226fedf410..d5ccf75a80 100644 --- a/onnxruntime/core/framework/execution_plan_base.h +++ b/onnxruntime/core/framework/execution_plan_base.h @@ -3,7 +3,7 @@ #pragma once -#include "core/common/inlined_containers_fwd.h" +#include #include "core/framework/allocator.h" namespace onnxruntime { @@ -18,7 +18,7 @@ class ExecutionPlanBase { virtual const struct OrtMemoryInfo& GetLocation(size_t ort_value_index) const = 0; virtual void SetLocation(size_t ort_value_index, const struct OrtMemoryInfo&) = 0; // return all memory locations for all the MLValues - virtual InlinedHashSet GetAllLocations() const = 0; + virtual std::set GetAllLocations() const = 0; virtual ~ExecutionPlanBase() = default; }; diff --git a/onnxruntime/core/framework/mem_pattern.h b/onnxruntime/core/framework/mem_pattern.h index ae31806789..63be390d81 100644 --- a/onnxruntime/core/framework/mem_pattern.h +++ b/onnxruntime/core/framework/mem_pattern.h @@ -3,7 +3,6 @@ #pragma once #include "core/common/common.h" -#include "core/common/inlined_containers.h" #include "core/framework/allocation_planner.h" namespace onnxruntime { @@ -46,7 +45,7 @@ class MemoryPattern { return &it->second; } - const InlinedHashMap& GetPatternsMap() const { + const std::unordered_map& GetPatternsMap() const { return patterns_; } @@ -54,7 +53,7 @@ class MemoryPattern { // allow move ORT_DISALLOW_COPY_AND_ASSIGNMENT(MemoryPattern); - InlinedHashMap patterns_; + std::unordered_map patterns_; size_t peak_size_{0}; }; diff --git a/onnxruntime/core/framework/mem_pattern_planner.h b/onnxruntime/core/framework/mem_pattern_planner.h index 788c1b7d79..0c7de29919 100644 --- a/onnxruntime/core/framework/mem_pattern_planner.h +++ b/onnxruntime/core/framework/mem_pattern_planner.h @@ -235,9 +235,8 @@ class MemPatternPlanner { MemoryPattern pattern; pattern.peak_size_ = buffer_size_; - pattern.patterns_.reserve(allocs_.size()); for (auto& alloc : allocs_) { - pattern.patterns_.insert_or_assign(alloc.index_, alloc.block_); + pattern.patterns_[alloc.index_] = alloc.block_; } return pattern; diff --git a/onnxruntime/core/framework/node_index_info.h b/onnxruntime/core/framework/node_index_info.h index ca0558726f..1f8cd4d46c 100644 --- a/onnxruntime/core/framework/node_index_info.h +++ b/onnxruntime/core/framework/node_index_info.h @@ -6,7 +6,6 @@ #include #include "core/common/common.h" -#include "core/common/inlined_containers_fwd.h" #include "core/framework/ort_value.h" #include "core/graph/basic_types.h" @@ -54,7 +53,7 @@ class NodeIndexInfo final { // This vector contains the indices from the OrtValueNameIdxMap in the SessionState for each Node's input/outputs. // Order is node inputs, implicit inputs, outputs. - InlinedVector node_values_; + std::vector node_values_; // the minimum NodeIndex. we use this to minimize the size of node_offsets_. NodeIndex min_node_index_ = 0; @@ -62,7 +61,7 @@ class NodeIndexInfo final { // The entry at node_offsets_[GetNodeOffsetsIndex(Node::Index())] contains the index in node_values_ // where the information for the Node begins. size_t GetNodeOffsetsIndex(NodeIndex node_index) const { return node_index - min_node_index_; } - InlinedVector node_offsets_; + std::vector node_offsets_; const int max_mlvalue_idx_; diff --git a/onnxruntime/core/framework/ort_value_name_idx_map.h b/onnxruntime/core/framework/ort_value_name_idx_map.h index 45a7fabf91..ed0e94402c 100644 --- a/onnxruntime/core/framework/ort_value_name_idx_map.h +++ b/onnxruntime/core/framework/ort_value_name_idx_map.h @@ -4,30 +4,30 @@ #pragma once #include +#include #include "core/common/common.h" -#include "core/common/inlined_containers.h" // This class is not thread-safe // TODO: this is a static hash lookup, it's easy to do it better namespace onnxruntime { class OrtValueNameIdxMap { public: - using const_iterator = typename InlinedHashMap::const_iterator; + using const_iterator = typename std::unordered_map::const_iterator; OrtValueNameIdxMap() = default; // Add OrtValue name to map and return index associated with it. // If entry already existed the existing index value is returned. int Add(const std::string& name) { - const int idx = next_idx_; - auto p = map_.emplace(name, idx); - if (p.second) { + auto it = map_.find(name); + if (it == map_.end()) { + int idx = next_idx_++; + map_.insert(it, {name, idx}); idx_name_map_[idx] = name; - next_idx_++; return idx; } - return p.first->second; + return it->second; } common::Status GetIdx(const std::string& name, int& idx) const { @@ -54,10 +54,6 @@ class OrtValueNameIdxMap { size_t Size() const { return map_.size(); }; int MaxIdx() const { return next_idx_ - 1; } - void Reserve(size_t size) { - map_.reserve(size); - idx_name_map_.reserve(size); - } const_iterator begin() const noexcept { return map_.cbegin(); } const_iterator end() const noexcept { return map_.cend(); } @@ -66,7 +62,7 @@ class OrtValueNameIdxMap { ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OrtValueNameIdxMap); int next_idx_ = 0; - InlinedHashMap map_; - InlinedHashMap idx_name_map_; + std::unordered_map map_; + std::unordered_map idx_name_map_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/framework/ort_value_pattern_planner.cc b/onnxruntime/core/framework/ort_value_pattern_planner.cc index 06f2c9c5bd..ed90acb683 100644 --- a/onnxruntime/core/framework/ort_value_pattern_planner.cc +++ b/onnxruntime/core/framework/ort_value_pattern_planner.cc @@ -8,9 +8,8 @@ namespace onnxruntime { OrtValuePatternPlanner::OrtValuePatternPlanner(const ExecutionPlanBase& execution_plan, bool trace_using_counters) : execution_planner_(execution_plan) { - planner_map_.reserve(execution_plan.GetAllLocations().size()); for (auto& location : execution_plan.GetAllLocations()) { - planner_map_.emplace(location, trace_using_counters); + planner_map_.emplace(location, std::make_unique(trace_using_counters)); } } @@ -25,7 +24,7 @@ common::Status OrtValuePatternPlanner::TraceAllocation(int ort_value_idx, return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); } - it->second.TraceAllocation(ort_value_idx, counter, size); + it->second->TraceAllocation(ort_value_idx, counter, size); return common::Status::OK(); } #endif @@ -37,7 +36,7 @@ common::Status OrtValuePatternPlanner::TraceAllocation(int ort_value_idx, size_t return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); } - it->second.TraceAllocation(ort_value_idx, size); + it->second->TraceAllocation(ort_value_idx, size); return common::Status::OK(); } @@ -48,17 +47,16 @@ common::Status OrtValuePatternPlanner::TraceFree(int ort_value_index) { return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); } - it->second.TraceFree(ort_value_index); + it->second->TraceFree(ort_value_index); return common::Status::OK(); } -common::Status OrtValuePatternPlanner::GeneratePatterns(MemoryPatternGroup& out) { +common::Status OrtValuePatternPlanner::GeneratePatterns(MemoryPatternGroup* out) { + if (!out) return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT); - out.locations.reserve(planner_map_.size()); - out.patterns.reserve(planner_map_.size()); for (auto& it : planner_map_) { - out.locations.push_back(it.first); - out.patterns.push_back(it.second.GenerateMemPattern()); + out->locations.push_back(it.first); + out->patterns.push_back(it.second->GenerateMemPattern()); } return common::Status::OK(); diff --git a/onnxruntime/core/framework/ort_value_pattern_planner.h b/onnxruntime/core/framework/ort_value_pattern_planner.h index 2db9c265bb..77fcb41728 100644 --- a/onnxruntime/core/framework/ort_value_pattern_planner.h +++ b/onnxruntime/core/framework/ort_value_pattern_planner.h @@ -6,7 +6,6 @@ #include #include "core/common/common.h" -#include "core/common/inlined_containers.h" #include "core/framework/mem_pattern_planner.h" #include "core/framework/execution_plan_base.h" #include "core/framework/allocation_planner.h" @@ -27,13 +26,12 @@ class OrtValuePatternPlanner { #endif common::Status TraceAllocation(int ort_value_idx, size_t size); common::Status TraceFree(int ort_value_index); - common::Status GeneratePatterns(MemoryPatternGroup& out); + common::Status GeneratePatterns(MemoryPatternGroup* out); ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OrtValuePatternPlanner); private: // This map itself is const after the construction - // MemPatternPlanner has copying disabled to using node map - NodeHashMap planner_map_; + std::map> planner_map_; const ExecutionPlanBase& execution_planner_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/framework/orttraining_partial_executor.cc b/onnxruntime/core/framework/orttraining_partial_executor.cc index 40d5e2a15f..aee060e2c4 100644 --- a/onnxruntime/core/framework/orttraining_partial_executor.cc +++ b/onnxruntime/core/framework/orttraining_partial_executor.cc @@ -493,8 +493,8 @@ Status PartialExecutor::Execute(const SessionState& session_state, const std::ve } if (all_tensors) { - MemoryPatternGroup mem_patterns; - ORT_RETURN_IF_ERROR(frame.GeneratePatterns(mem_patterns)); + auto mem_patterns = std::make_unique(); + ORT_RETURN_IF_ERROR(frame.GeneratePatterns(mem_patterns.get())); ORT_RETURN_IF_ERROR(session_state.UpdateMemoryPatternGroupCache(feeds, std::move(mem_patterns))); } } diff --git a/onnxruntime/core/framework/parallel_executor.cc b/onnxruntime/core/framework/parallel_executor.cc index 1d7cf192c8..f0576b3143 100644 --- a/onnxruntime/core/framework/parallel_executor.cc +++ b/onnxruntime/core/framework/parallel_executor.cc @@ -91,8 +91,8 @@ Status ParallelExecutor::Execute(const SessionState& session_state, const std::v } if (all_tensors) { - MemoryPatternGroup mem_patterns; - ORT_RETURN_IF_ERROR(root_frame_->GeneratePatterns(mem_patterns)); + auto mem_patterns = std::make_unique(); + ORT_RETURN_IF_ERROR(root_frame_->GeneratePatterns(mem_patterns.get())); ORT_RETURN_IF_ERROR(session_state.UpdateMemoryPatternGroupCache(feeds, std::move(mem_patterns))); } } diff --git a/onnxruntime/core/framework/prepacked_weights_container.h b/onnxruntime/core/framework/prepacked_weights_container.h index bc75b8ddec..7fe317b6c4 100644 --- a/onnxruntime/core/framework/prepacked_weights_container.h +++ b/onnxruntime/core/framework/prepacked_weights_container.h @@ -3,10 +3,11 @@ #pragma once +#include +#include #include #include -#include "core/common/inlined_containers.h" #include "core/framework/buffer_deleter.h" #include "core/framework/allocator.h" @@ -57,12 +58,12 @@ class PrepackedWeightsContainer final { // Define allocators ahead of the container containing tensors because the allocators // needs to destructed after the container containing the pre-packed cached tensors // because the Tensor buffers will be de-allocated using these allocators - InlinedHashMap allocators_; + std::unordered_map allocators_; // This is an unordered map that holds a mapping between a composite key // to PrePackedWeights instances. // The key is : op_type + "+" + hash_of_prepacked_buffers_in_the_PrepackedWeights_instance. - InlinedHashMap prepacked_weights_map_; + std::unordered_map prepacked_weights_map_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/framework/sequential_execution_plan.h b/onnxruntime/core/framework/sequential_execution_plan.h index 211f3b9c4c..743670b4e1 100644 --- a/onnxruntime/core/framework/sequential_execution_plan.h +++ b/onnxruntime/core/framework/sequential_execution_plan.h @@ -4,7 +4,6 @@ #pragma once #include "core/graph/basic_types.h" -#include "core/common/inlined_containers.h" #include "core/framework/alloc_kind.h" #include "core/framework/data_types.h" #include "core/framework/execution_plan_base.h" @@ -114,10 +113,10 @@ struct SequentialExecutionPlan : public ExecutionPlanBase { std::vector execution_plan; // Records whether a given node has fence on its input or output, key is node index. - InlinedVector node_has_fence; + std::vector node_has_fence; // to_be_freed: vector elements represent indices of ml-values to be freed (as described above) - InlinedVector to_be_freed; + std::vector to_be_freed; const OrtMemoryInfo& GetLocation(size_t ort_value_index) const override { return allocation_plan[ort_value_index].location; @@ -127,11 +126,10 @@ struct SequentialExecutionPlan : public ExecutionPlanBase { allocation_plan[ort_value_index].location = info; } - InlinedHashSet GetAllLocations() const override { - InlinedHashSet locations; - locations.reserve(allocation_plan.size()); + std::set GetAllLocations() const override { + std::set locations; for (auto& alloc_plan : allocation_plan) { - locations.insert(alloc_plan.location); + if (locations.find(alloc_plan.location) == locations.end()) locations.insert(alloc_plan.location); } return locations; } diff --git a/onnxruntime/core/framework/sequential_executor.cc b/onnxruntime/core/framework/sequential_executor.cc index ed3e3fbaf1..31abfab761 100644 --- a/onnxruntime/core/framework/sequential_executor.cc +++ b/onnxruntime/core/framework/sequential_executor.cc @@ -495,8 +495,8 @@ Status SequentialExecutor::Execute(const SessionState& session_state, const std: } if (all_tensors) { - MemoryPatternGroup mem_patterns; - ORT_RETURN_IF_ERROR(frame.GeneratePatterns(mem_patterns)); + auto mem_patterns = std::make_unique(); + ORT_RETURN_IF_ERROR(frame.GeneratePatterns(mem_patterns.get())); ORT_RETURN_IF_ERROR(session_state.UpdateMemoryPatternGroupCache(feeds, std::move(mem_patterns))); } } diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc index 07c43609cd..335138916d 100644 --- a/onnxruntime/core/framework/session_state.cc +++ b/onnxruntime/core/framework/session_state.cc @@ -76,7 +76,7 @@ AllocatorPtr SessionState::GetAllocator(OrtDevice device) const noexcept { } void SessionState::CreateGraphInfo() { - graph_viewer_.emplace(graph_); + graph_viewer_ = std::make_unique(graph_); // use graph_viewer_ to initialize ort_value_name_idx_map_ LOGS(logger_, VERBOSE) << "SaveMLValueNameIndexMapping"; int idx = 0; @@ -191,16 +191,11 @@ Status SessionState::CreateKernels(const KernelRegistryManager& kernel_registry_ ORT_RETURN_IF_ERROR(kernel_registry_manager.CreateKernel(node, exec_provider, *this, kci, session_kernels_[node.Index()])); } } - node_index_info_.emplace(*graph_viewer_, ort_value_name_idx_map_); + node_index_info_ = std::make_unique(*graph_viewer_, ort_value_name_idx_map_); return Status::OK(); } -const SequentialExecutionPlan* SessionState::GetExecutionPlan() const { - if (!p_seq_exec_plan_.has_value()) { - return nullptr; - } - return &p_seq_exec_plan_.value(); -} +const SequentialExecutionPlan* SessionState::GetExecutionPlan() const { return p_seq_exec_plan_.get(); } Status SessionState::AddInitializedTensor(int ort_value_index, const OrtValue& ort_value, const OrtCallback* d, bool constant, bool sparse) { @@ -210,7 +205,7 @@ Status SessionState::AddInitializedTensor(int ort_value_index, const OrtValue& o ". Do you have duplicated calls to SessionState::AddInitializedTensor function?"); if (d != nullptr && d->f != nullptr) { - deleter_for_initialized_tensors_.insert_or_assign(ort_value_index, *d); + deleter_for_initialized_tensors_[ort_value_index] = *d; } if (constant) { @@ -314,7 +309,7 @@ static std::string GenerateKeyForPrepackedWeightsMap(const std::string& op_type, return ss_1.str(); } -Status SessionState::PrepackConstantInitializedTensors(InlinedHashMap& constant_initializers_use_count, +Status SessionState::PrepackConstantInitializedTensors(std::unordered_map& constant_initializers_use_count, const std::unordered_map& initializers_to_share_map) { auto prepacked_constant_weights = [this, &constant_initializers_use_count, &initializers_to_share_map]( bool should_cache_prepacked_weights_for_shared_initializers) -> Status { @@ -454,9 +449,8 @@ static int64_t CalculateMemoryPatternsKey(const gsl::span& tenso #ifdef ENABLE_TRAINING namespace { Status ResolveDimParams(const GraphViewer& graph, - const InlinedHashMap& feeds, - InlinedHashMap& out) { - out.reserve(graph.GetInputs().size()); + const std::map& feeds, + std::unordered_map& out) { for (const auto* input : graph.GetInputs()) { auto* shape = input->Shape(); auto it = feeds.find(input->Name()); @@ -485,7 +479,7 @@ Status ResolveDimParams(const GraphViewer& graph, Status TryResolveShape( const NodeArg* arg, - const InlinedHashMap& symbolic_dimensions, + const std::unordered_map& symbolic_dimensions, size_t& is_resolved, // indicate whether resolve successfully or not. TensorShapeVector& resolved_shape) { if (!arg->Shape()) { @@ -522,7 +516,7 @@ Status TryResolveShape( return Status::OK(); } -void TryCalculateSizeFromResolvedShape(int ml_value_idx, const InlinedHashMap& resolved_shapes, size_t& size) { +void TryCalculateSizeFromResolvedShape(int ml_value_idx, std::unordered_map& resolved_shapes, size_t& size) { size = 0; auto shape = resolved_shapes.find(ml_value_idx); if (shape != resolved_shapes.end()) { @@ -535,18 +529,17 @@ void TryCalculateSizeFromResolvedShape(int ml_value_idx, const InlinedHashMap tensor_inputs, - gsl::span feed_mlvalue_idxs, - MemoryPatternGroup& output, - InlinedHashMap& resolved_shapes) const { - InlinedHashMap feeds; - feeds.reserve(feed_mlvalue_idxs.size()); +Status SessionState::GeneratePatternGroupCache(const gsl::span& tensor_inputs, + const std::vector& feed_mlvalue_idxs, + MemoryPatternGroup* output, + std::unordered_map& resolved_shapes) const { + std::map feeds; for (size_t i = 0, end = feed_mlvalue_idxs.size(); i < end; ++i) { std::string name; ORT_RETURN_IF_ERROR(this->ort_value_name_idx_map_.GetName(feed_mlvalue_idxs[i], name)); - feeds.emplace(std::move(name), tensor_inputs[i].Get().Shape()); + feeds.insert({name, tensor_inputs[i].Get().Shape()}); } - InlinedHashMap map; + std::unordered_map map; ORT_RETURN_IF_ERROR(ResolveDimParams(*graph_viewer_, feeds, map)); auto* exe_plan = GetExecutionPlan(); ORT_ENFORCE(exe_plan); @@ -670,45 +663,36 @@ Status SessionState::GeneratePatternGroupCache(gsl::span tensor_ } return Status::OK(); } - #endif -// MemoryPatternGroup pointer is cached. It only inserted upon creation -// and is not updated if already present. -const MemoryPatternGroup* SessionState::GetMemoryPatternGroup( - gsl::span tensor_inputs, - gsl::span feed_mlvalue_idxs, - const InlinedHashMap*& out_inferred_shapes) const { - - out_inferred_shapes = nullptr; +const MemoryPatternGroup* SessionState::GetMemoryPatternGroup(const gsl::span& tensor_inputs, + const std::vector& feed_mlvalue_idxs, + std::unordered_map& inferred_shapes) const { int64_t key = CalculateMemoryPatternsKey(tensor_inputs); + std::lock_guard lock(mem_patterns_lock_); auto it = mem_patterns_.find(key); if (it == mem_patterns_.end()) { #ifdef ENABLE_TRAINING - MemoryPatternGroup mem_patterns; - InlinedHashMap inferred_shapes; - if (GeneratePatternGroupCache(tensor_inputs, feed_mlvalue_idxs, mem_patterns, inferred_shapes).IsOK()) { - auto patt_insert = mem_patterns_.insert_or_assign(key, std::move(mem_patterns)); - auto ptr = &patt_insert.first->second; - auto shape_insert = shape_patterns_.insert_or_assign(key, std::move(inferred_shapes)); - out_inferred_shapes = &shape_insert.first->second; + auto mem_patterns = std::make_unique(); + if (GeneratePatternGroupCache(tensor_inputs, feed_mlvalue_idxs, mem_patterns.get(), inferred_shapes).IsOK()) { + key = CalculateMemoryPatternsKey(tensor_inputs); + auto ptr = mem_patterns.get(); + mem_patterns_[key] = std::move(mem_patterns); + shape_patterns_[key] = inferred_shapes; return ptr; } + return nullptr; #else ORT_UNUSED_PARAMETER(feed_mlvalue_idxs); -#endif return nullptr; +#endif } - auto patt_hit = shape_patterns_.find(key); - if (patt_hit != shape_patterns_.cend()) { - out_inferred_shapes = &patt_hit->second; - } - return &it->second; + inferred_shapes = shape_patterns_[key]; + return it->second.get(); } - void SessionState::ResolveMemoryPatternFlag() { if (enable_mem_pattern_) { for (auto* input : graph_viewer_->GetInputs()) { @@ -733,13 +717,16 @@ void SessionState::ResolveMemoryPatternFlag() { } } -Status SessionState::UpdateMemoryPatternGroupCache(gsl::span tensor_inputs, - MemoryPatternGroup mem_patterns) const { +Status SessionState::UpdateMemoryPatternGroupCache(const gsl::span& tensor_inputs, + std::unique_ptr mem_patterns) const { int64_t key = CalculateMemoryPatternsKey(tensor_inputs); std::lock_guard lock(mem_patterns_lock_); - // Do not update if present, as the pointer to the existing one is cached - mem_patterns_.emplace(key, std::move(mem_patterns)); + auto it = mem_patterns_.find(key); + if (it == mem_patterns_.end()) { + mem_patterns_[key] = std::move(mem_patterns); + } + return Status::OK(); } @@ -788,7 +775,7 @@ common::Status SessionState::AddInputNameToNodeInfoMapping(const std::string& in } common::Status SessionState::GetInputNodeInfo(const std::string& input_name, - InlinedVector& node_info_vec) const { + std::vector& node_info_vec) const { auto entry = input_names_to_nodeinfo_mapping_.find(input_name); if (entry == input_names_to_nodeinfo_mapping_.cend()) { return Status(ONNXRUNTIME, FAIL, "Failed to find input name in the mapping: " + input_name); @@ -811,7 +798,7 @@ void SessionState::AddOutputNameToNodeInfoMapping(const std::string& output_name } common::Status SessionState::GetOutputNodeInfo(const std::string& output_name, - InlinedVector& node_info_vec) const { + std::vector& node_info_vec) const { auto entry = output_names_to_nodeinfo_mapping_.find(output_name); if (entry == output_names_to_nodeinfo_mapping_.cend()) { return Status(ONNXRUNTIME, FAIL, "Failed to find output name in the mapping: " + output_name); @@ -868,7 +855,7 @@ const SessionState* SessionState::GetSubgraphSessionState(onnxruntime::NodeIndex } const NodeIndexInfo& SessionState::GetNodeIndexInfo() const { - ORT_ENFORCE(node_index_info_.has_value(), "SetGraphAndCreateKernels must be called prior to GetExecutionInfo."); + ORT_ENFORCE(node_index_info_, "SetGraphAndCreateKernels must be called prior to GetExecutionInfo."); return *node_index_info_; } @@ -913,15 +900,12 @@ const InlinedHashSet* SessionState::GetToBeExecutedNodes( static Status GetSubGraphSessionStatesOrtFormat( flatbuffers::FlatBufferBuilder& builder, - const SubgraphSessionStateMap& subgraph_session_states, + const std::unordered_map>>& subgraph_session_states, std::vector>& fbs_subgraph_session_states) { - size_t number_of_states = 0; - for (const auto& pair : subgraph_session_states) { - number_of_states += pair.second.size(); - } fbs_subgraph_session_states.clear(); - fbs_subgraph_session_states.reserve(number_of_states); - for (const auto& [node_idx, session_states] : subgraph_session_states) { + 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; @@ -1116,7 +1100,7 @@ Status SessionState::LoadFromOrtFormat(const fbs::SessionState& fbs_session_stat // The main graph has a constant initializer called X, and the subgraph also has a constant initializer called X, which overrides the X from main graph. // For case like this, the current implementation will calculate the use count as 2, but they could contain completely different values so each should have a use count of 1. // This is a very rare case. If it happens and X is prepacked, the consequence is that X won't be released and memory usage of X won't be saved. This will be fine. -static void ComputeConstantInitializerUseCount(const Graph& graph, InlinedHashMap& constant_initializers_use_count) { +static void ComputeConstantInitializerUseCount(const Graph& graph, std::unordered_map& constant_initializers_use_count) { for (const auto& node : graph.Nodes()) { for (const auto* arg : node.InputDefs()) { if (arg->Exists() && graph.GetConstantInitializer(arg->Name(), true /*check_outer_scope*/)) { @@ -1230,7 +1214,7 @@ Status SessionState::FinalizeSessionState(const std::basic_string constant_initializers_use_count; + std::unordered_map constant_initializers_use_count; ComputeConstantInitializerUseCount(graph_, constant_initializers_use_count); return FinalizeSessionStateImpl(graph_location, kernel_registry_manager, nullptr, session_options, remove_initializers, constant_initializers_use_count); @@ -1271,9 +1255,8 @@ static Status OuterScopeNodeArgLocationAccumulator(const SequentialExecutionPlan const OrtValueNameIdxMap& ort_value_name_to_idx_map, const Node& parent_node, const GraphViewer& subgraph, - /*out*/ InlinedHashMap& outer_scope_arg_to_location_map) { + /*out*/ std::unordered_map& outer_scope_arg_to_location_map) { // Process implicit inputs to the node - outer_scope_arg_to_location_map.reserve(parent_node.ImplicitInputDefs().size() + parent_node.InputDefs().size()); auto process_implicit_input = [&plan, &ort_value_name_to_idx_map, &outer_scope_arg_to_location_map](const NodeArg& input, size_t /*arg_idx*/) { const auto& name = input.Name(); @@ -1366,8 +1349,8 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string& constant_initializers_use_count, - const InlinedHashMap& outer_scope_node_arg_to_location_map, + std::unordered_map& constant_initializers_use_count, + const std::unordered_map& outer_scope_node_arg_to_location_map, bool graph_info_already_created) { if (!graph_info_already_created) { CreateGraphInfo(); @@ -1379,8 +1362,7 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string> unused_initializer_names; - unused_initializer_names.reserve(graph_.GetAllInitializedTensors().size()); + std::vector unused_initializer_names; for (const auto& [name, tensor_proto] : graph_.GetAllInitializedTensors()) { ORT_UNUSED_PARAMETER(tensor_proto); int idx; @@ -1396,7 +1378,7 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string valid_outer_scope_node_args; + std::vector valid_outer_scope_node_args; if (parent_node) { auto outer_scope_node_args = parent_node->ImplicitInputDefs(); valid_outer_scope_node_args.reserve(outer_scope_node_args.size()); @@ -1427,10 +1409,10 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string Status { return AddInitializedTensor(idx, value, &d, constant, sparse); }, - logger_, data_transfer_mgr_, *p_seq_exec_plan_, session_options)); + logger_, data_transfer_mgr_, *p_seq_exec_plan_.get(), session_options)); #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) // Record Weight allocation info on device MemoryInfo::RecordInitializerAllocInfo(GetInitializedTensors()); @@ -1514,7 +1496,7 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string subgraph_outer_scope_node_arg_to_location_map; + std::unordered_map subgraph_outer_scope_node_arg_to_location_map; ORT_RETURN_IF_ERROR(OuterScopeNodeArgLocationAccumulator(*p_seq_exec_plan_, GetOrtValueNameIdxMap(), node, subgraph_session_state.GetGraphViewer(), diff --git a/onnxruntime/core/framework/session_state.h b/onnxruntime/core/framework/session_state.h index 547f76b642..9798c1abb8 100644 --- a/onnxruntime/core/framework/session_state.h +++ b/onnxruntime/core/framework/session_state.h @@ -116,7 +116,7 @@ class SessionState { } // Graph viewer. CreateGraphInfo must have been called previously. - const GraphViewer& GetGraphViewer() const noexcept { return *graph_viewer_; }; + const GraphViewer& GetGraphViewer() const noexcept { return *graph_viewer_.get(); }; // kernels // Get kernel for specified node. @@ -206,24 +206,19 @@ class SessionState { /** Get cached memory pattern based on input shapes Must be called only when all values contain tensors - In training scenarios, the cache may be updated so - the callers would receive a copy of inferred shapes - made under mutex being held. In inference scenarios, - it is not mutable, we do not obtain a lock and simply get a pointer - w/o copying a hashtable */ const MemoryPatternGroup* GetMemoryPatternGroup( - gsl::span tensor_inputs, - gsl::span feed_mlvalue_idxs, - const InlinedHashMap*& inferred_shapes) const; + const gsl::span& tensor_inputs, + const std::vector& feed_mlvalue_idxs, + std::unordered_map& inferred_shapes) const; /** Set generated memory pattern with a given input shapes. Const as it's an internal cache update only. All inputs must represent Tensors */ - Status UpdateMemoryPatternGroupCache(gsl::span tensor_inputs, - MemoryPatternGroup mem_patterns) const; + Status UpdateMemoryPatternGroupCache(const gsl::span& tensor_inputs, + std::unique_ptr mem_patterns) const; bool GetUseDeterministicCompute() const { return use_deterministic_compute_; } @@ -262,14 +257,14 @@ class SessionState { const OrtDevice* device = nullptr; }; - using NameNodeInfoMapType = InlinedHashMap>; + using NameNodeInfoMapType = std::unordered_map>; common::Status AddInputNameToNodeInfoMapping(const std::string& input_name, const NodeInfo& node_info); - common::Status GetInputNodeInfo(const std::string& input_name, InlinedVector& node_info_vec) const; + common::Status GetInputNodeInfo(const std::string& input_name, std::vector& node_info_vec) const; const NameNodeInfoMapType& GetInputNodeInfoMap() const; void AddOutputNameToNodeInfoMapping(const std::string& output_name, const NodeInfo& node_info); - common::Status GetOutputNodeInfo(const std::string& output_name, InlinedVector& node_info_vec) const; + common::Status GetOutputNodeInfo(const std::string& output_name, std::vector& node_info_vec) const; const NameNodeInfoMapType& GetOutputNodeInfoMap() const; // Get the KernelCreateInfo entry for a node. SessionState must be finalized before calling. @@ -286,7 +281,7 @@ class SessionState { const DataTransferManager& GetDataTransferMgr() const noexcept { return data_transfer_mgr_; } - InlinedVector& GetMutableWeightsBuffers() noexcept { return weights_buffers_; } + std::vector& GetMutableWeightsBuffers() noexcept { return weights_buffers_; } const NodeIndexInfo& GetNodeIndexInfo() const; @@ -360,7 +355,7 @@ class SessionState { * Prepack the constant initialized tensors for better performance. * The original constant initialized tensors will be removed to save memory. */ - Status PrepackConstantInitializedTensors(InlinedHashMap& constant_initializers_use_count, + Status PrepackConstantInitializedTensors(std::unordered_map& constant_initializers_use_count, const std::unordered_map& initializers_to_share_map); SessionState* GetMutableSubgraphSessionState(onnxruntime::NodeIndex index, const std::string& attribute_name); @@ -379,16 +374,16 @@ class SessionState { _In_opt_ const Node* parent_node, const SessionOptions& session_options, bool remove_initializers, - InlinedHashMap& constant_initializers_use_count, - const InlinedHashMap& outer_scope_node_arg_to_location_map = {}, + std::unordered_map& constant_initializers_use_count, + const std::unordered_map& outer_scope_node_arg_to_location_map = {}, bool graph_info_already_created = false); #ifdef ENABLE_TRAINING Status GeneratePatternGroupCache( - gsl::span inputs, - gsl::span feed_mlvalue_idxs, - MemoryPatternGroup& output, - InlinedHashMap& inferred_shapes) const; + const gsl::span& inputs, + const std::vector& feed_mlvalue_idxs, + MemoryPatternGroup* output, + std::unordered_map& inferred_shapes) const; #endif // the SessionState for the main Graph contains the compiled kernel hashes for the entire model @@ -409,7 +404,7 @@ class SessionState { // cache of the constructed kernels to avoid spending construction time per executor std::vector> session_kernels_; Graph& graph_; - std::optional graph_viewer_; // GraphViewer for const access to Graph + std::unique_ptr graph_viewer_; // GraphViewer for const access to Graph const ExecutionProviders& execution_providers_; @@ -468,14 +463,14 @@ class SessionState { // this is needed because we currently convert all sparse initializer into dense Tensors // if and when we actually place SparseTensor instances (we should) into OrtValues, we // will not need this structure. - InlinedHashSet sparse_initialized_tensors_; + std::unordered_set sparse_initialized_tensors_; #endif // This data structure is for uninitializing string tensors and // munmap memory region and close file descriptor - InlinedHashMap deleter_for_initialized_tensors_; - InlinedVector weights_buffers_; - std::optional p_seq_exec_plan_; + std::unordered_map deleter_for_initialized_tensors_; + std::vector weights_buffers_; + std::unique_ptr p_seq_exec_plan_ = nullptr; const logging::Logger& logger_; profiling::Profiler& profiler_; @@ -485,16 +480,10 @@ class SessionState { // lock for the mem_patterns_ mutable OrtMutex mem_patterns_lock_; + // cache for the generated mem_patterns. key is calculated based on input shapes. - // must be a node based container as a pointer is cached. - mutable NodeHashMap mem_patterns_; - // This is mutable under mutex in training scenarios so execution frame would make a copy - // of the value when created. -#ifdef ENABLE_TRAINING - mutable NodeHashMap> shape_patterns_; -#else - NodeHashMap> shape_patterns_; -#endif + mutable std::map> mem_patterns_; + mutable std::map> shape_patterns_; NameNodeInfoMapType input_names_to_nodeinfo_mapping_; NameNodeInfoMapType output_names_to_nodeinfo_mapping_; @@ -509,7 +498,8 @@ class SessionState { bool use_deterministic_compute_; bool enable_mem_reuse_; - std::optional node_index_info_; + std::unique_ptr node_index_info_; + std::multimap> cached_feeds_fetches_managers_; // Container to store pre-packed weights to share between sessions. // The life-cycle of the cache itself is maintained by the user and the user will ensure diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index 7cca8d9557..ed66874035 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -171,10 +171,8 @@ common::Status SaveInitializedTensors( // 1. first plan the memory const onnxruntime::InitializedTensorSet& initialized_tensor_set = graph.GetAllInitializedTensors(); - InlinedHashMap id_to_initialized_tensor; - id_to_initialized_tensor.reserve(initialized_tensor_set.size()); - InlinedHashSet user_supplied_initializer_ids; // set containing the ort value ids of all user supplied initializers - user_supplied_initializer_ids.reserve(initialized_tensor_set.size()); + std::unordered_map id_to_initialized_tensor; + std::set user_supplied_initializer_ids; // set containing the ort value ids of all user supplied initializers for (const auto& entry : initialized_tensor_set) { int ort_value_index; ORT_RETURN_IF_ERROR(ort_value_name_idx_map.GetIdx(entry.first, ort_value_index)); @@ -208,7 +206,7 @@ common::Status SaveInitializedTensors( // 2. allocate weight buffer on different locations // planned_initializers_memory_size_in_byte is not actual physical size. // It's the virtual size computed by planner. - InlinedHashMap planned_initializers_memory_sizes_in_byte; + std::unordered_map planned_initializers_memory_sizes_in_byte; ORT_RETURN_IF_ERROR( planner.FinalizePlan(planned_initializers_memory_sizes_in_byte)); #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) @@ -236,14 +234,14 @@ common::Status SaveInitializedTensors( } else { const ONNX_NAMESPACE::TensorProto& tensor_proto = *(entry.second); - std::optional m; + std::unique_ptr m; AllocatorPtr alloc; // TODO: if the tensor need be copied, does it have enough room? ORT_RETURN_IF_ERROR(planner.GetPreallocatedBuffer(ort_value_index, name, m, alloc)); bool use_device_allocator_for_initializers = session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsUseDeviceAllocatorForInitializers, "0") == "1"; - Status st = DeserializeTensorProto(env, graph_loc, tensor_proto, (m.has_value()) ? &*m : nullptr, alloc, default_cpu_alloc, ort_value, + Status st = DeserializeTensorProto(env, graph_loc, tensor_proto, m.get(), alloc, default_cpu_alloc, ort_value, data_transfer_mgr, use_device_allocator_for_initializers); if (!st.IsOK()) { std::ostringstream oss; @@ -281,7 +279,7 @@ static bool IsArgNameInInputsOutputs(const std::string& name, common::Status SaveInputOutputNamesToNodeMapping(const onnxruntime::GraphViewer& graph, SessionState& session_state, - gsl::span implicit_inputs) { + const std::vector& implicit_inputs) { auto& graph_inputs = graph.GetInputsIncludingInitializers(); auto& graph_outputs = graph.GetOutputs(); diff --git a/onnxruntime/core/framework/session_state_utils.h b/onnxruntime/core/framework/session_state_utils.h index fde53dc682..3d47697cc7 100644 --- a/onnxruntime/core/framework/session_state_utils.h +++ b/onnxruntime/core/framework/session_state_utils.h @@ -44,6 +44,6 @@ common::Status SaveInitializedTensors( const SessionOptions& session_options); common::Status SaveInputOutputNamesToNodeMapping(const GraphViewer& graph, SessionState& session_state, - gsl::span implicit_inputs); + const std::vector& implicit_inputs); } // namespace session_state_utils } // namespace onnxruntime diff --git a/onnxruntime/core/framework/simple_tensor_allocator.cc b/onnxruntime/core/framework/simple_tensor_allocator.cc index f76d4aaec2..08ebb67284 100644 --- a/onnxruntime/core/framework/simple_tensor_allocator.cc +++ b/onnxruntime/core/framework/simple_tensor_allocator.cc @@ -10,7 +10,7 @@ common::Status SimpleTensorAllocator::Trace(int /*id*/, const ONNX_NAMESPACE::Te } common::Status SimpleTensorAllocator::GetPreallocatedBuffer(int ort_value_index, const char* /*name*/, - std::optional& /*buf_out*/, + std::unique_ptr& /*buf_out*/, AllocatorPtr& alloc_out) { const struct OrtMemoryInfo& location = seq_plan_.GetLocation(ort_value_index); // just return allocator and let others handle it. diff --git a/onnxruntime/core/framework/simple_tensor_allocator.h b/onnxruntime/core/framework/simple_tensor_allocator.h index d2d61fa4b0..ac90d90c42 100644 --- a/onnxruntime/core/framework/simple_tensor_allocator.h +++ b/onnxruntime/core/framework/simple_tensor_allocator.h @@ -4,7 +4,6 @@ #pragma once #include -#include #include "tensor_allocator.h" #include "mem_pattern.h" #include "ort_value_pattern_planner.h" @@ -20,17 +19,17 @@ class SimpleTensorAllocator : public ITensorAllocator { public: SimpleTensorAllocator(const ExecutionPlanBase& execution_plan, const SessionState& session_state, - InlinedVector& /*weights_buffers*/) + std::vector& /*weights_buffers*/) : ITensorAllocator(session_state), seq_plan_(execution_plan) {} - common::Status FinalizePlan(InlinedHashMap& planned_memory_sizes_in_byte) override { + common::Status FinalizePlan(std::unordered_map& planned_memory_sizes_in_byte) override { // There is no memory plan to allocate a big block of memory, so // planned memory sizes in different locations are all empty. - planned_memory_sizes_in_byte.clear(); + planned_memory_sizes_in_byte = std::unordered_map(); return Status::OK(); } - common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::optional& buf_out, AllocatorPtr& alloc_out) override; + common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::unique_ptr& buf_out, AllocatorPtr& alloc_out) override; common::Status Trace(int id, const ONNX_NAMESPACE::TensorProto* value) override; const MemoryPatternGroup& GetMemPatterns() override { return mem_patterns_; diff --git a/onnxruntime/core/framework/tensor_allocator.cc b/onnxruntime/core/framework/tensor_allocator.cc index 868d3bfe24..beaf8fbd39 100644 --- a/onnxruntime/core/framework/tensor_allocator.cc +++ b/onnxruntime/core/framework/tensor_allocator.cc @@ -13,7 +13,7 @@ AllocatorPtr ITensorAllocator::GetAllocator(const OrtMemoryInfo& memory_info) { std::unique_ptr ITensorAllocator::Create(bool enable_mem_pattern, const ExecutionPlanBase& execution_plan, const SessionState& session_state, - InlinedVector& weights_buffers) { + std::vector& weights_buffers) { if (enable_mem_pattern) { return std::make_unique(execution_plan, session_state, weights_buffers); } else { diff --git a/onnxruntime/core/framework/tensor_allocator.h b/onnxruntime/core/framework/tensor_allocator.h index cd5610b2a2..081592f4e5 100644 --- a/onnxruntime/core/framework/tensor_allocator.h +++ b/onnxruntime/core/framework/tensor_allocator.h @@ -3,10 +3,8 @@ #pragma once -#include #include #include -#include #include #include #include @@ -23,7 +21,7 @@ class ITensorAllocator { // Create an ITensorAllocator instance based on enable_mem_pattern static std::unique_ptr Create(bool enable_mem_pattern, const ExecutionPlanBase& execution_plan, const SessionState& session_state, - InlinedVector& weights_buffers); + std::vector& weights_buffers); AllocatorPtr GetAllocator(const OrtMemoryInfo& memory_info); @@ -34,7 +32,7 @@ class ITensorAllocator { * When there is no more tensor to trace, call this function to finalize the * allocation. */ - virtual common::Status FinalizePlan(InlinedHashMap& planned_memory_sizes_in_byte) = 0; + virtual common::Status FinalizePlan(std::unordered_map& planned_memory_sizes_in_byte) = 0; /** * Handing out buffers reserved in @see #Trace() via parameter buf_out, @@ -49,7 +47,7 @@ class ITensorAllocator { * @return */ virtual common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, - std::optional& buf_out, + std::unique_ptr& buf_out, AllocatorPtr& alloc_out) = 0; virtual const MemoryPatternGroup& GetMemPatterns() = 0; diff --git a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h index 1bc2a1469d..4a354bc609 100644 --- a/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h +++ b/onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h @@ -3,7 +3,6 @@ #pragma once -#include #include "tensor_allocator.h" #include "mem_pattern.h" #include "ort_value_pattern_planner.h" @@ -17,15 +16,14 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { private: OrtValuePatternPlanner planner_; MemoryPatternGroup mem_patterns_; - InlinedVector& weights_buffers_; - InlinedHashMap buffers_; + std::vector& weights_buffers_; + std::map buffers_; bool is_sealed_ = false; const ExecutionPlanBase& seq_plan_; common::Status AllocatePlannedBuffersAndReportTotalSize( - InlinedHashMap& planned_memory_sizes_in_byte) { + std::unordered_map& planned_memory_sizes_in_byte) { const size_t location_len = mem_patterns_.locations.size(); - planned_memory_sizes_in_byte.reserve(location_len); for (size_t i = 0; i < location_len; ++i) { auto& location = mem_patterns_.locations[i]; auto alloc = GetAllocator(location); @@ -61,21 +59,21 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { public: TensorAllocatorWithMemPattern(const ExecutionPlanBase& execution_plan, const SessionState& session_state, - InlinedVector& weights_buffers) + std::vector& weights_buffers) : ITensorAllocator(session_state), planner_(execution_plan, /*using counters*/ false), weights_buffers_(weights_buffers), seq_plan_(execution_plan) {} - common::Status FinalizePlan(InlinedHashMap& planned_memory_sizes_in_byte) override { - ORT_RETURN_IF_ERROR(planner_.GeneratePatterns(mem_patterns_)); + common::Status FinalizePlan(std::unordered_map& planned_memory_sizes_in_byte) override { + ORT_RETURN_IF_ERROR(planner_.GeneratePatterns(&mem_patterns_)); ORT_RETURN_IF_ERROR(AllocatePlannedBuffersAndReportTotalSize(planned_memory_sizes_in_byte)); is_sealed_ = true; return Status::OK(); } common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, - std::optional& buf_out, AllocatorPtr& alloc_out) override { + std::unique_ptr& buf_out, AllocatorPtr& alloc_out) override { if (!is_sealed_) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Internal error."); } @@ -97,7 +95,7 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { if (it == buffers_.end()) { if (block != nullptr && block->size_ == 0) { // Because the size is 0, this miss find is expected. we won't allocate a buffer with size of zero. - buf_out.emplace(nullptr, 0, location); + buf_out = std::make_unique(nullptr, 0, location); return Status::OK(); } return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Weight buffer for initializer '", name, "' is not found"); @@ -107,7 +105,7 @@ class TensorAllocatorWithMemPattern : public ITensorAllocator { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Get preallocated buffer for initializer '", name, "' failed"); } - buf_out.emplace(reinterpret_cast(it->second) + block->offset_, block->size_, location); + buf_out = std::make_unique(reinterpret_cast(it->second) + block->offset_, block->size_, location); return Status::OK(); } common::Status Trace(int id, const ONNX_NAMESPACE::TensorProto* value) override { diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index 1287e879ce..58dc3d0510 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -273,7 +273,7 @@ const OrtMemoryInfo& FindMemoryInfoForValue(const SessionState& session_state, static common::Status CalculateStaticCopyInfoForFeed(const SessionState& session_state, const std::string& input_name, MLValueCopyInfo& copy_info) { - InlinedVector node_info_vec; + std::vector node_info_vec; #ifdef ENABLE_TRAINING if (session_state.GetInputNodeInfo(input_name, node_info_vec) == Status::OK()) { #else diff --git a/onnxruntime/core/graph/graph_viewer.cc b/onnxruntime/core/graph/graph_viewer.cc index 5482a8e286..4fb6e7de82 100644 --- a/onnxruntime/core/graph/graph_viewer.cc +++ b/onnxruntime/core/graph/graph_viewer.cc @@ -91,8 +91,8 @@ GraphViewer::GraphViewer(const Graph& graph, const IndexedSubGraph* filter_info) } // create set of node indexes as we need quick lookups and don't care about the order - filtered_node_indices_ = FilteredNodeSet(filter_info->nodes.cbegin(), - filter_info->nodes.cend()); + filtered_node_indices_ = std::unordered_set(filter_info->nodes.cbegin(), + filter_info->nodes.cend()); const auto& metadef = filter_info->GetMetaDef(); diff --git a/onnxruntime/core/optimizer/optimizer_execution_frame.cc b/onnxruntime/core/optimizer/optimizer_execution_frame.cc index b858e79bfb..829648a5a2 100644 --- a/onnxruntime/core/optimizer/optimizer_execution_frame.cc +++ b/onnxruntime/core/optimizer/optimizer_execution_frame.cc @@ -17,14 +17,6 @@ namespace onnxruntime { -static size_t EstimateInputsOutputs(gsl::span nodes) { - size_t num = 0; - for (auto n : nodes) { - num += n->InputDefs().size() + n->OutputDefs().size(); - } - return num; -} - OptimizerExecutionFrame::Info::Info(const std::vector& nodes, const InitializedTensorSet& initialized_tensor_set, const Path& model_path, @@ -40,7 +32,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, // Create MLValues related maps auto initialize_maps = [this, &initialized_tensor_set, &model_path](const NodeArg& arg, size_t /*index*/) -> Status { int idx = ort_value_name_idx_map_.Add(arg.Name()); - ort_value_idx_nodearg_map_.insert_or_assign(idx, &arg); + ort_value_idx_nodearg_map_[idx] = &arg; // Only create OrtValue instances for initializers used by an array of nodes. InitializedTensorSet::const_iterator it = initialized_tensor_set.find(arg.Name()); @@ -65,12 +57,6 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, }; // TODO: node->ImplicitInputDefs() need to be added here for control flow nodes. - auto num_inputs_outputs = EstimateInputsOutputs(nodes); - ort_value_name_idx_map_.Reserve(num_inputs_outputs); - ort_value_idx_nodearg_map_.reserve(num_inputs_outputs); - initializers_.reserve(initialized_tensor_set.size()); - buffer_for_initialized_tensors_.reserve(initialized_tensor_set.size()); - for (auto* node : nodes) { ORT_THROW_IF_ERROR(onnxruntime::Node::ForEachWithIndex(node->InputDefs(), initialize_maps)); ORT_THROW_IF_ERROR(onnxruntime::Node::ForEachWithIndex(node->OutputDefs(), initialize_maps)); @@ -95,10 +81,10 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, auto initialize_maps = [this, &initialized_tensor_set, &model_path](const NodeArg& arg, size_t /*index*/) -> Status { (void)model_path; int idx = ort_value_name_idx_map_.Add(arg.Name()); - ort_value_idx_nodearg_map_.insert_or_assign(idx, &arg); + ort_value_idx_nodearg_map_[idx] = &arg; // Only create OrtValue instances for initializers used by an array of nodes. - auto it = initialized_tensor_set.find(arg.Name()); + std::unordered_map::const_iterator it = initialized_tensor_set.find(arg.Name()); if (it != initialized_tensor_set.cend()) { initializers_[idx] = it->second; } @@ -106,11 +92,6 @@ OptimizerExecutionFrame::Info::Info(const std::vector& nodes, }; // TODO: node->ImplicitInputDefs() need to be added here for control flow nodes. - auto num_inputs_outputs = EstimateInputsOutputs(nodes); - ort_value_name_idx_map_.Reserve(num_inputs_outputs); - ort_value_idx_nodearg_map_.reserve(num_inputs_outputs); - initializers_.reserve(initialized_tensor_set.size()); - for (auto* node : nodes) { ORT_THROW_IF_ERROR(onnxruntime::Node::ForEachWithIndex(node->InputDefs(), initialize_maps)); ORT_THROW_IF_ERROR(onnxruntime::Node::ForEachWithIndex(node->OutputDefs(), initialize_maps)); @@ -147,7 +128,7 @@ OptimizerExecutionFrame::OptimizerExecutionFrame(const Info& info, const std::vector& fetches) : IExecutionFrame(info.GetMLValueNameIdxMap(), info.GetNodeIndexInfo(), fetch_mlvalue_idxs), info_(info) { - Init(gsl::span(), gsl::span(), info.GetInitializers(), info.GetSparseInitializerLookupFunc(), fetches); + Init(std::vector(), std::vector(), info.GetInitializers(), info.GetSparseInitializerLookupFunc(), fetches); } AllocatorPtr OptimizerExecutionFrame::GetAllocatorImpl(const OrtMemoryInfo& info) const { diff --git a/onnxruntime/core/optimizer/optimizer_execution_frame.h b/onnxruntime/core/optimizer/optimizer_execution_frame.h index c9aed24991..88e404ce74 100644 --- a/onnxruntime/core/optimizer/optimizer_execution_frame.h +++ b/onnxruntime/core/optimizer/optimizer_execution_frame.h @@ -5,7 +5,6 @@ #include -#include "core/common/inlined_containers.h" #include "core/graph/graph.h" #include "core/providers/cpu/cpu_execution_provider.h" #include "core/framework/data_transfer_manager.h" @@ -32,7 +31,11 @@ class OptimizerExecutionFrame final : public IExecutionFrame { const Path& model_path, const IExecutionProvider& execution_provider, const std::function& is_sparse_initializer_func); - ~Info() = default; + ~Info() { + for (auto& kvp : deleter_for_initialized_tensors_) { + kvp.second.f(kvp.second.param); + } + } AllocatorPtr GetAllocator(const OrtMemoryInfo& info) const { return execution_provider_.GetAllocator(info.id, info.mem_type); @@ -77,7 +80,10 @@ class OptimizerExecutionFrame final : public IExecutionFrame { OrtValueNameIdxMap ort_value_name_idx_map_; std::unordered_map ort_value_idx_nodearg_map_; std::unordered_map initializers_; - InlinedHashMap> buffer_for_initialized_tensors_; + std::unordered_map> buffer_for_initialized_tensors_; + // This data structure is for uninitializing string tensors and + // munmap memory region and close file descriptor + std::unordered_map deleter_for_initialized_tensors_; std::unique_ptr node_index_info_; const IExecutionProvider& execution_provider_; const std::function& is_sparse_initializer_func_; diff --git a/onnxruntime/test/framework/allocation_planner_test.cc b/onnxruntime/test/framework/allocation_planner_test.cc index 3004c61b6c..6be0557d67 100644 --- a/onnxruntime/test/framework/allocation_planner_test.cc +++ b/onnxruntime/test/framework/allocation_planner_test.cc @@ -181,7 +181,7 @@ class PlannerTest : public ::testing::Test { profiling::Profiler profiler_; std::unique_ptr state_; ShapeMap shape_map_; - std::optional plan_; + std::unique_ptr plan_; public: PlannerTest() diff --git a/onnxruntime/test/framework/execution_frame_test.cc b/onnxruntime/test/framework/execution_frame_test.cc index d3625c3f6b..a9c2f49447 100644 --- a/onnxruntime/test/framework/execution_frame_test.cc +++ b/onnxruntime/test/framework/execution_frame_test.cc @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/common/span_utils.h" #include "core/framework/execution_frame.h" #include "core/framework/op_kernel.h" #include "core/framework/session_state.h" @@ -270,8 +269,8 @@ TEST_F(ExecutionFrameTest, MemPatternTest) { std::vector{2, 3}, std::vector(6, 1.0f), &v3); - std::vector outputs; - ExecutionFrame frame(AsSpan({x1_idx, x2_idx, x3_idx}), AsSpan({v1, v2, v3}), {t3_idx}, outputs, {}, state); + vector outputs; + ExecutionFrame frame({x1_idx, x2_idx, x3_idx}, {v1, v2, v3}, {t3_idx}, outputs, {}, state); OrtValue& mlvalue3 = *frame.GetMutableNodeInputOrOutputMLValue(3); OrtValue& mlvalue4 = *frame.GetMutableNodeInputOrOutputMLValue(4); @@ -292,7 +291,7 @@ TEST_F(ExecutionFrameTest, MemPatternTest) { cpu_allocator->Info(), TensorShape(std::vector{2, 3}))); MemoryPatternGroup pattern; - ASSERT_STATUS_OK(frame.GeneratePatterns(pattern)); + ASSERT_STATUS_OK(frame.GeneratePatterns(&pattern)); ASSERT_EQ(pattern.patterns.size(), pattern.locations.size()); ASSERT_EQ(pattern.patterns.size(), 1u); @@ -369,7 +368,7 @@ TEST_F(ExecutionFrameTest, MemPatternWithExternalOutputsTest) { y_value, y_idx, DataTypeImpl::GetType(), cpu_allocator->Info(), TensorShape(std::vector{2, 2}))); MemoryPatternGroup pattern; - ASSERT_STATUS_OK(frame.GeneratePatterns(pattern)); + ASSERT_STATUS_OK(frame.GeneratePatterns(&pattern)); ASSERT_EQ(pattern.patterns.size(), pattern.locations.size()); ASSERT_EQ(pattern.patterns.size(), 1u); diff --git a/winml/adapter/winml_adapter_session.cpp b/winml/adapter/winml_adapter_session.cpp index 94238b2659..523758fe3d 100644 --- a/winml/adapter/winml_adapter_session.cpp +++ b/winml/adapter/winml_adapter_session.cpp @@ -214,7 +214,7 @@ static OrtDevice GetSessionGetInputDevice(_In_ OrtSession* session, _In_ const c static_cast(inference_session); const onnxruntime::SessionState& session_state = session_protected_load_accessor->GetSessionState(); - onnxruntime::InlinedVectornode_info_vec; + std::vector node_info_vec; ORT_THROW_IF_ERROR(session_state.GetInputNodeInfo(input_name, node_info_vec)); const auto& node_info = node_info_vec.front(); // all consumers of a feed have the same device so first entry is fine return *node_info.device;