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

Revert "Refactor ExecutionFrame and SessionState to reduce memory all… (#11888)

Browse source 
Revert "Refactor ExecutionFrame and SessionState to reduce memory allocations and improve data locality (#11804)"

This reverts commit 2ecba6fd25.
This commit is contained in:
Yi Zhang 2022-06-17 17:07:21 +08:00 committed by GitHub
parent bd65acd08d
commit d2cbae3a04
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
34 changed files with 232 additions and 326 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
include/onnxruntime/core/framework/ortmemoryinfo.h
View file

@ -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<int>()(alloc_type);
HashCombine(std::hash<int>()(mem_type), h);
HashCombine(std::hash<int>()(id), h);
HashCombine(std::hash<const char*>()(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<OrtMemoryInfo> {
size_t operator()(const OrtMemoryInfo& i) const {
return i.Hash();
}
};
}

3
include/onnxruntime/core/graph/graph_viewer.h
View file

@ -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<NodeIndex>;
FilteredNodeSet filtered_node_indices_;
std::unordered_set<NodeIndex> filtered_node_indices_;
std::vector<const NodeArg*> filtered_node_inputs_;
std::vector<const NodeArg*> filtered_node_inputs_including_initializers_;
std::vector<const NodeArg*> filtered_node_outputs_;

23
onnxruntime/core/framework/allocation_planner.cc
View file

@ -77,16 +77,13 @@ std::ostream& operator<<(std::ostream& out, std::pair<const SequentialExecutionP
const SequentialExecutionPlan& plan = *planinfo.first;
const SessionState& session_state = *planinfo.second;
auto& graph = session_state.GetGraphViewer();
const auto& name_idx_map = session_state.GetOrtValueNameIdxMap();
InlinedHashMap<int, std::string_view> index_to_name;
index_to_name.reserve(name_idx_map.Size());
std::unordered_map<int, std::string> index_to_name;
out << "Allocation Plan:\n";
out << "(ort_value_idx) output_name : <allocation plan>\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<const NodeArg* const> outer_scope_node_args, const ExecutionProviders& providers,
const std::vector<const NodeArg*>& outer_scope_node_args, const ExecutionProviders& providers,
const KernelCreateInfoMap& kernel_create_info_map,
const SubgraphsKernelCreateInfoMaps& subgraphs_kernel_create_info_maps,
const InlinedHashMap<OrtValueName, OrtMemoryInfo>& outer_scope_node_arg_to_location_map,
const std::unordered_map<OrtValueName, OrtMemoryInfo>& 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<const NodeArg* const> outer_scope_node_args_;
const std::vector<const NodeArg*>& outer_scope_node_args_;
const ExecutionProviders& execution_providers_;
const KernelCreateInfoMap& kernel_create_info_map_;
const SubgraphsKernelCreateInfoMaps& subgraphs_kernel_create_info_maps_;
const InlinedHashMap<OrtValueName, OrtMemoryInfo>& outer_scope_node_arg_to_location_map_;
const std::unordered_map<OrtValueName, OrtMemoryInfo>& 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<const NodeArg* const> outer_scope_node_args,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelCreateInfoMap& kernel_create_info_map,
const SubgraphsKernelCreateInfoMaps& subgraphs_kernel_create_info_maps,
const InlinedHashMap<OrtValueName, OrtMemoryInfo>& outer_scope_node_arg_to_location_map,
const std::unordered_map<OrtValueName, OrtMemoryInfo>& outer_scope_node_arg_to_location_map,
const OrtValueNameIdxMap& ort_value_name_idx_map,
const ISequentialPlannerContext& context,
std::optional<SequentialExecutionPlan>& plan) {
std::unique_ptr<SequentialExecutionPlan>& plan) {
// allocate/reset here so we know it's clean
plan.emplace();
plan = std::make_unique<SequentialExecutionPlan>();
PlannerImpl planner(parent_node, graph_viewer, outer_scope_node_args, providers,
kernel_create_info_map, subgraphs_kernel_create_info_maps,

6
onnxruntime/core/framework/allocation_planner.h
View file

@ -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<const NodeArg* const> outer_scope_node_args,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelCreateInfoMap& kernel_create_info_map,
const SubgraphsKernelCreateInfoMaps& subgraphs_kernel_create_info_maps,
const InlinedHashMap<OrtValueName, OrtMemoryInfo>& outer_scope_arg_to_location_map,
const std::unordered_map<OrtValueName, OrtMemoryInfo>& outer_scope_arg_to_location_map,
const OrtValueNameIdxMap& ort_value_name_idx_map,
const ISequentialPlannerContext& context,
std::optional<SequentialExecutionPlan>& plan);
std::unique_ptr<SequentialExecutionPlan>& plan);
};
} // namespace onnxruntime

61
onnxruntime/core/framework/execution_frame.cc
View file

@ -27,10 +27,10 @@ namespace onnxruntime {
IExecutionFrame::IExecutionFrame(const OrtValueNameIdxMap& ort_value_idx_map,
const NodeIndexInfo& node_index_info,
gsl::span<const int> fetch_mlvalue_idxs)
const std::vector<int>& fetch_mlvalue_idxs)
: node_index_info_(node_index_info),
all_values_size_(static_cast<size_t>(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<const int> feed_mlvalue_idxs, gsl::span<const OrtValue> feeds) {
void IExecutionFrame::UpdateFeeds(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& 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<const int> feed_mlvalue_idxs, gsl::s
}
}
void IExecutionFrame::UpdateFetches(gsl::span<const int> fetch_mlvalue_idxs,
gsl::span<const OrtValue> fetches, const std::unordered_map<int, OrtValue>& initializers) {
void IExecutionFrame::UpdateFetches(const std::vector<int>& fetch_mlvalue_idxs, const std::vector<OrtValue>& fetches, const std::unordered_map<int, OrtValue>& 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<const int> fetch_mlvalue_idxs,
}
}
Status IExecutionFrame::GetOutputs(gsl::span<const int> fetch_mlvalue_idxs, std::vector<OrtValue>& fetches) {
Status IExecutionFrame::GetOutputs(const std::vector<int>& fetch_mlvalue_idxs, std::vector<OrtValue>& 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<const int> feed_mlvalue_idxs, gsl::span<const OrtValue> feeds,
void IExecutionFrame::Init(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::unordered_map<int, OrtValue>& initializers,
const std::function<bool(const std::string& name)>& is_initializer_sparse_func,
gsl::span<const OrtValue> fetches) {
const std::vector<OrtValue>& 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<const int> feed_mlvalue_idxs, gsl::span<con
*dest.GetMutable<SparseTensor>()));
} 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<const int> feed_mlvalue_idxs, gsl::span<const OrtValue> feeds,
gsl::span<const int> fetch_mlvalue_idxs, gsl::span<const OrtValue> fetches,
ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::vector<int>& fetch_mlvalue_idxs, const std::vector<OrtValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>& 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<const int> 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<const int> 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<OrtValuePatternPlanner>(*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<size_t>(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.

38
onnxruntime/core/framework/execution_frame.h
View file

@ -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<const int> fetch_mlvalue_idxs);
const std::vector<int>& fetch_mlvalue_idxs);
void Init(gsl::span<const int> feed_mlvalue_idxs, gsl::span<const OrtValue> feeds,
void Init(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::unordered_map<int, OrtValue>& initializers,
const std::function<bool(const std::string& name)>& is_initializer_sparse_func,
gsl::span<const OrtValue> fetches);
const std::vector<OrtValue>& 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<const int> feed_mlvalue_idxs, gsl::span<const OrtValue> feeds);
void UpdateFetches(gsl::span<const int> fetch_mlvalue_idxs, gsl::span<const OrtValue> fetches,
#ifdef ENABLE_TRAINING
void UpdateFeeds(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds);
void UpdateFetches(const std::vector<int>& fetch_mlvalue_idxs, const std::vector<OrtValue>& fetches,
const std::unordered_map<int, OrtValue>& initializers);
Status GetOutputs(gsl::span<const int> fetch_mlvalue_idxs, std::vector<OrtValue>& fetches);
Status GetOutputs(const std::vector<int>& fetch_mlvalue_idxs, std::vector<OrtValue>& 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<OrtValue> all_values_;
std::vector<OrtValue> 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<int> fetch_mlvalue_idxs_;
std::vector<int> fetch_mlvalue_idxs_;
const OrtValueNameIdxMap& ort_value_idx_map_;
};
class ExecutionFrame final : public IExecutionFrame {
public:
ExecutionFrame(gsl::span<const int> feed_mlvalue_idxs, gsl::span<const OrtValue> feeds,
gsl::span<const int> fetch_mlvalue_idxs, gsl::span<const OrtValue> fetches,
ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::vector<int>& fetch_mlvalue_idxs, const std::vector<OrtValue>& fetches,
// optional custom allocators. key is index in fetches
const std::unordered_map<size_t, IExecutor::CustomAllocator>& 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<int, IExecutor::CustomAllocator> custom_allocators_;
std::unordered_map<int, IExecutor::CustomAllocator> 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<OrtValuePatternPlanner> planner_;
std::unique_ptr<OrtValuePatternPlanner> planner_;
// Big chunks on different locations that will be used by mem_pattern.
InlinedHashMap<OrtMemoryInfo, BufferUniquePtr> buffers_;
std::map<OrtMemoryInfo, BufferUniquePtr> 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<int, TensorShape>* inferred_shapes_{nullptr};
std::unordered_map<int, TensorShape> inferred_shapes_;
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
// Size of virtual memory allocated before any kernel execution.

4
onnxruntime/core/framework/execution_plan_base.h
View file

@ -3,7 +3,7 @@
#pragma once
#include "core/common/inlined_containers_fwd.h"
#include <set>
#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<struct OrtMemoryInfo> GetAllLocations() const = 0;
virtual std::set<struct OrtMemoryInfo> GetAllLocations() const = 0;
virtual ~ExecutionPlanBase() = default;
};

5
onnxruntime/core/framework/mem_pattern.h
View file

@ -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<int, MemoryBlock>& GetPatternsMap() const {
const std::unordered_map<int, MemoryBlock>& GetPatternsMap() const {
return patterns_;
}
@ -54,7 +53,7 @@ class MemoryPattern {
// allow move
ORT_DISALLOW_COPY_AND_ASSIGNMENT(MemoryPattern);
InlinedHashMap<int, MemoryBlock> patterns_;
std::unordered_map<int, MemoryBlock> patterns_;
size_t peak_size_{0};
};

3
onnxruntime/core/framework/mem_pattern_planner.h
View file

@ -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;

5
onnxruntime/core/framework/node_index_info.h
View file

@ -6,7 +6,6 @@
#include <vector>
#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<int> node_values_;
std::vector<int> 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<int> node_offsets_;
std::vector<int> node_offsets_;
const int max_mlvalue_idx_;

22
onnxruntime/core/framework/ort_value_name_idx_map.h
View file

@ -4,30 +4,30 @@
#pragma once
#include <atomic>
#include <unordered_map>
#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<std::string, int>::const_iterator;
using const_iterator = typename std::unordered_map<std::string, int>::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<std::string, int> map_;
InlinedHashMap<int, std::string> idx_name_map_;
std::unordered_map<std::string, int> map_;
std::unordered_map<int, std::string> idx_name_map_;
};
} // namespace onnxruntime

18
onnxruntime/core/framework/ort_value_pattern_planner.cc
View file

@ -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<MemPatternPlanner>(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();

6
onnxruntime/core/framework/ort_value_pattern_planner.h
View file

@ -6,7 +6,6 @@
#include <memory>
#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<OrtMemoryInfo, MemPatternPlanner> planner_map_;
std::map<OrtMemoryInfo, std::unique_ptr<MemPatternPlanner>> planner_map_;
const ExecutionPlanBase& execution_planner_;
};
} // namespace onnxruntime

4
onnxruntime/core/framework/orttraining_partial_executor.cc
View file

@ -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<MemoryPatternGroup>();
ORT_RETURN_IF_ERROR(frame.GeneratePatterns(mem_patterns.get()));
ORT_RETURN_IF_ERROR(session_state.UpdateMemoryPatternGroupCache(feeds, std::move(mem_patterns)));
}
}

4
onnxruntime/core/framework/parallel_executor.cc
View file

@ -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<MemoryPatternGroup>();
ORT_RETURN_IF_ERROR(root_frame_->GeneratePatterns(mem_patterns.get()));
ORT_RETURN_IF_ERROR(session_state.UpdateMemoryPatternGroupCache(feeds, std::move(mem_patterns)));
}
}

7
onnxruntime/core/framework/prepacked_weights_container.h
View file

@ -3,10 +3,11 @@
#pragma once
#include <unordered_map>
#include <unordered_set>
#include <string>
#include <cstdint>
#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<std::string, AllocatorPtr> allocators_;
std::unordered_map<std::string, AllocatorPtr> 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<std::string, PrePackedWeights> prepacked_weights_map_;
std::unordered_map<std::string, PrePackedWeights> prepacked_weights_map_;
};
} // namespace onnxruntime

12
onnxruntime/core/framework/sequential_execution_plan.h
View file

@ -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<NodeExecutionPlan> execution_plan;
// Records whether a given node has fence on its input or output, key is node index.
InlinedVector<bool> node_has_fence;
std::vector<bool> node_has_fence;
// to_be_freed: vector elements represent indices of ml-values to be freed (as described above)
InlinedVector<OrtValueIndex> to_be_freed;
std::vector<OrtValueIndex> 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<OrtMemoryInfo> GetAllLocations() const override {
InlinedHashSet<OrtMemoryInfo> locations;
locations.reserve(allocation_plan.size());
std::set<OrtMemoryInfo> GetAllLocations() const override {
std::set<OrtMemoryInfo> 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;
}

4
onnxruntime/core/framework/sequential_executor.cc
View file

@ -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<MemoryPatternGroup>();
ORT_RETURN_IF_ERROR(frame.GeneratePatterns(mem_patterns.get()));
ORT_RETURN_IF_ERROR(session_state.UpdateMemoryPatternGroupCache(feeds, std::move(mem_patterns)));
}
}

130
onnxruntime/core/framework/session_state.cc
View file

@ -76,7 +76,7 @@ AllocatorPtr SessionState::GetAllocator(OrtDevice device) const noexcept {
}
void SessionState::CreateGraphInfo() {
graph_viewer_.emplace(graph_);
graph_viewer_ = std::make_unique<onnxruntime::GraphViewer>(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<NodeIndexInfo>(*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<std::string, size_t>& constant_initializers_use_count,
Status SessionState::PrepackConstantInitializedTensors(std::unordered_map<std::string, size_t>& constant_initializers_use_count,
const std::unordered_map<std::string, const OrtValue*>& 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<const OrtValue>& tenso
#ifdef ENABLE_TRAINING
namespace {
Status ResolveDimParams(const GraphViewer& graph,
const InlinedHashMap<std::string, TensorShape>& feeds,
InlinedHashMap<std::string, int64_t>& out) {
out.reserve(graph.GetInputs().size());
const std::map<std::string, TensorShape>& feeds,
std::unordered_map<std::string, int64_t>& 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<std::string, int64_t>& symbolic_dimensions,
const std::unordered_map<std::string, int64_t>& 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<int, TensorShape>& resolved_shapes, size_t& size) {
void TryCalculateSizeFromResolvedShape(int ml_value_idx, std::unordered_map<int, TensorShape>& 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<in
} // namespace
// If this function fails NO memory planning will take place, hence lets ONLY FAIL and stop training where warranted, example SIZE overflow.
Status SessionState::GeneratePatternGroupCache(gsl::span<const OrtValue> tensor_inputs,
gsl::span<const int> feed_mlvalue_idxs,
MemoryPatternGroup& output,
InlinedHashMap<int, TensorShape>& resolved_shapes) const {
InlinedHashMap<std::string, TensorShape> feeds;
feeds.reserve(feed_mlvalue_idxs.size());
Status SessionState::GeneratePatternGroupCache(const gsl::span<const OrtValue>& tensor_inputs,
const std::vector<int>& feed_mlvalue_idxs,
MemoryPatternGroup* output,
std::unordered_map<int, TensorShape>& resolved_shapes) const {
std::map<std::string, TensorShape> 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<Tensor>().Shape());
feeds.insert({name, tensor_inputs[i].Get<Tensor>().Shape()});
}
InlinedHashMap<std::string, int64_t> map;
std::unordered_map<std::string, int64_t> 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<const OrtValue> 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<const OrtValue> tensor_inputs,
gsl::span<const int> feed_mlvalue_idxs,
const InlinedHashMap<int, TensorShape>*& out_inferred_shapes) const {
out_inferred_shapes = nullptr;
const MemoryPatternGroup* SessionState::GetMemoryPatternGroup(const gsl::span<const OrtValue>& tensor_inputs,
const std::vector<int>& feed_mlvalue_idxs,
std::unordered_map<int, TensorShape>& inferred_shapes) const {
int64_t key = CalculateMemoryPatternsKey(tensor_inputs);
std::lock_guard<OrtMutex> lock(mem_patterns_lock_);
auto it = mem_patterns_.find(key);
if (it == mem_patterns_.end()) {
#ifdef ENABLE_TRAINING
MemoryPatternGroup mem_patterns;
InlinedHashMap<int, TensorShape> 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<MemoryPatternGroup>();
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<const OrtValue> tensor_inputs,
MemoryPatternGroup mem_patterns) const {
Status SessionState::UpdateMemoryPatternGroupCache(const gsl::span<const OrtValue>& tensor_inputs,
std::unique_ptr<MemoryPatternGroup> mem_patterns) const {
int64_t key = CalculateMemoryPatternsKey(tensor_inputs);
std::lock_guard<OrtMutex> 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<NodeInfo>& node_info_vec) const {
std::vector<NodeInfo>& 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<NodeInfo>& node_info_vec) const {
std::vector<NodeInfo>& 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<NodeIndex>* SessionState::GetToBeExecutedNodes(
static Status GetSubGraphSessionStatesOrtFormat(
flatbuffers::FlatBufferBuilder& builder,
const SubgraphSessionStateMap& subgraph_session_states,
const std::unordered_map<NodeIndex, std::unordered_map<std::string, std::unique_ptr<SessionState>>>& subgraph_session_states,
std::vector<flatbuffers::Offset<fbs::SubGraphSessionState>>& 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<std::string, size_t>& constant_initializers_use_count) {
static void ComputeConstantInitializerUseCount(const Graph& graph, std::unordered_map<std::string, size_t>& 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<PATH_CHAR_TYPE
#endif
}
InlinedHashMap<std::string, size_t> constant_initializers_use_count;
std::unordered_map<std::string, size_t> 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<OrtValueName, OrtMemoryInfo>& outer_scope_arg_to_location_map) {
/*out*/ std::unordered_map<OrtValueName, OrtMemoryInfo>& 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<PATH_CHAR_
_In_opt_ const Node* parent_node,
const SessionOptions& session_options,
bool remove_initializers,
InlinedHashMap<std::string, size_t>& constant_initializers_use_count,
const InlinedHashMap<OrtValueName, OrtMemoryInfo>& outer_scope_node_arg_to_location_map,
std::unordered_map<std::string, size_t>& constant_initializers_use_count,
const std::unordered_map<OrtValueName, OrtMemoryInfo>& 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<PATH_CHAR_
// Not needed in a basic minimal build because only runtime optimizations are expected to possibly result in unused
// initializers and they are only enabled in an extended minimal build.
{
InlinedVector<std::reference_wrapper<const std::string>> unused_initializer_names;
unused_initializer_names.reserve(graph_.GetAllInitializedTensors().size());
std::vector<std::string> 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<PATH_CHAR_
#endif // defined(ORT_EXTENDED_MINIMAL_BUILD)
// ignore any outer scope args we don't know about. this can happen if a node contains multiple subgraphs.
InlinedVector<const NodeArg*> valid_outer_scope_node_args;
std::vector<const NodeArg*> 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<PATH_CHAR_
MemoryInfo::GenerateTensorMap(GetExecutionPlan(), GetOrtValueNameIdxMap());
#endif
// Memory pattern tracer allocates all initializers on a single contiguous
// buffer. This has the effect of reducing memory fragmentation.
// Memory pattern tracer allocates all initializers on a single continous
// buffer. This has the effect of reducing memory fragementation.
// Further more, NCCL kernels require initializers to be allocated
// contiguously.
// continously.
//
// In inferencing scenarios, however, we often want to pre-process and then
// release some initializers. See OpKernel::PrePack(). Letting all initializers
@ -1460,7 +1442,7 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
[this](int idx, const OrtValue& value, const OrtCallback& d, bool constant, bool sparse) -> 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<PATH_CHAR_
// is used in OuterScopeNodeArgLocationAccumulator()
subgraph_session_state.CreateGraphInfo();
InlinedHashMap<OrtValueName, OrtMemoryInfo> subgraph_outer_scope_node_arg_to_location_map;
std::unordered_map<OrtValueName, OrtMemoryInfo> subgraph_outer_scope_node_arg_to_location_map;
ORT_RETURN_IF_ERROR(OuterScopeNodeArgLocationAccumulator(*p_seq_exec_plan_, GetOrtValueNameIdxMap(),
node,
subgraph_session_state.GetGraphViewer(),

64
onnxruntime/core/framework/session_state.h
View file

@ -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<const OrtValue> tensor_inputs,
gsl::span<const int> feed_mlvalue_idxs,
const InlinedHashMap<int, TensorShape>*& inferred_shapes) const;
const gsl::span<const OrtValue>& tensor_inputs,
const std::vector<int>& feed_mlvalue_idxs,
std::unordered_map<int, TensorShape>& 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<const OrtValue> tensor_inputs,
MemoryPatternGroup mem_patterns) const;
Status UpdateMemoryPatternGroupCache(const gsl::span<const OrtValue>& tensor_inputs,
std::unique_ptr<MemoryPatternGroup> mem_patterns) const;
bool GetUseDeterministicCompute() const { return use_deterministic_compute_; }
@ -262,14 +257,14 @@ class SessionState {
const OrtDevice* device = nullptr;
};
using NameNodeInfoMapType = InlinedHashMap<std::string, InlinedVector<NodeInfo>>;
using NameNodeInfoMapType = std::unordered_map<std::string, std::vector<NodeInfo>>;
common::Status AddInputNameToNodeInfoMapping(const std::string& input_name, const NodeInfo& node_info);
common::Status GetInputNodeInfo(const std::string& input_name, InlinedVector<NodeInfo>& node_info_vec) const;
common::Status GetInputNodeInfo(const std::string& input_name, std::vector<NodeInfo>& 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<NodeInfo>& node_info_vec) const;
common::Status GetOutputNodeInfo(const std::string& output_name, std::vector<NodeInfo>& 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<BufferUniquePtr>& GetMutableWeightsBuffers() noexcept { return weights_buffers_; }
std::vector<BufferUniquePtr>& 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<std::string, size_t>& constant_initializers_use_count,
Status PrepackConstantInitializedTensors(std::unordered_map<std::string, size_t>& constant_initializers_use_count,
const std::unordered_map<std::string, const OrtValue*>& 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<std::string, size_t>& constant_initializers_use_count,
const InlinedHashMap<OrtValueName, OrtMemoryInfo>& outer_scope_node_arg_to_location_map = {},
std::unordered_map<std::string, size_t>& constant_initializers_use_count,
const std::unordered_map<OrtValueName, OrtMemoryInfo>& outer_scope_node_arg_to_location_map = {},
bool graph_info_already_created = false);
#ifdef ENABLE_TRAINING
Status GeneratePatternGroupCache(
gsl::span<const OrtValue> inputs,
gsl::span<const int > feed_mlvalue_idxs,
MemoryPatternGroup& output,
InlinedHashMap<int, TensorShape>& inferred_shapes) const;
const gsl::span<const OrtValue>& inputs,
const std::vector<int>& feed_mlvalue_idxs,
MemoryPatternGroup* output,
std::unordered_map<int, TensorShape>& 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<std::unique_ptr<OpKernel>> session_kernels_;
Graph& graph_;
std::optional<GraphViewer> graph_viewer_; // GraphViewer for const access to Graph
std::unique_ptr<GraphViewer> 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<int> sparse_initialized_tensors_;
std::unordered_set<int> sparse_initialized_tensors_;
#endif
// This data structure is for uninitializing string tensors and
// munmap memory region and close file descriptor
InlinedHashMap<int, OrtCallback> deleter_for_initialized_tensors_;
InlinedVector<BufferUniquePtr> weights_buffers_;
std::optional<SequentialExecutionPlan> p_seq_exec_plan_;
std::unordered_map<int, OrtCallback> deleter_for_initialized_tensors_;
std::vector<BufferUniquePtr> weights_buffers_;
std::unique_ptr<SequentialExecutionPlan> 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<int64_t, MemoryPatternGroup> 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<int64_t, InlinedHashMap<int, TensorShape>> shape_patterns_;
#else
NodeHashMap<int64_t, InlinedHashMap<int, TensorShape>> shape_patterns_;
#endif
mutable std::map<int64_t, std::unique_ptr<MemoryPatternGroup>> mem_patterns_;
mutable std::map<int64_t, std::unordered_map<int, TensorShape>> 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<NodeIndexInfo> node_index_info_;
std::unique_ptr<NodeIndexInfo> node_index_info_;
std::multimap<int, std::unique_ptr<FeedsFetchesManager>> 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

14
onnxruntime/core/framework/session_state_utils.cc
View file

@ -171,10 +171,8 @@ common::Status SaveInitializedTensors(
// 1. first plan the memory
const onnxruntime::InitializedTensorSet& initialized_tensor_set = graph.GetAllInitializedTensors();
InlinedHashMap<int, const ONNX_NAMESPACE::TensorProto*> id_to_initialized_tensor;
id_to_initialized_tensor.reserve(initialized_tensor_set.size());
InlinedHashSet<int> 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<int, const ONNX_NAMESPACE::TensorProto*> id_to_initialized_tensor;
std::set<int> 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<std::string, size_t> planned_initializers_memory_sizes_in_byte;
std::unordered_map<std::string, size_t> 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<MemBuffer> m;
std::unique_ptr<MemBuffer> 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<const NodeArg* const> implicit_inputs) {
const std::vector<const NodeArg*>& implicit_inputs) {
auto& graph_inputs = graph.GetInputsIncludingInitializers();
auto& graph_outputs = graph.GetOutputs();

2
onnxruntime/core/framework/session_state_utils.h
View file

@ -44,6 +44,6 @@ common::Status SaveInitializedTensors(
const SessionOptions& session_options);
common::Status SaveInputOutputNamesToNodeMapping(const GraphViewer& graph,
SessionState& session_state,
gsl::span<const NodeArg* const> implicit_inputs);
const std::vector<const NodeArg*>& implicit_inputs);
} // namespace session_state_utils
} // namespace onnxruntime

2
onnxruntime/core/framework/simple_tensor_allocator.cc
View file

@ -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<MemBuffer>& /*buf_out*/,
std::unique_ptr<MemBuffer>& /*buf_out*/,
AllocatorPtr& alloc_out) {
const struct OrtMemoryInfo& location = seq_plan_.GetLocation(ort_value_index);
// just return allocator and let others handle it.

9
onnxruntime/core/framework/simple_tensor_allocator.h
View file

@ -4,7 +4,6 @@
#pragma once
#include <unordered_map>
#include <core/common/inlined_containers.h>
#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<BufferUniquePtr>& /*weights_buffers*/)
std::vector<BufferUniquePtr>& /*weights_buffers*/)
: ITensorAllocator(session_state),
seq_plan_(execution_plan) {}
common::Status FinalizePlan(InlinedHashMap<std::string, size_t>& planned_memory_sizes_in_byte) override {
common::Status FinalizePlan(std::unordered_map<std::string, size_t>& 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<std::string, size_t>();
return Status::OK();
}
common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::optional<MemBuffer>& buf_out, AllocatorPtr& alloc_out) override;
common::Status GetPreallocatedBuffer(int ort_value_index, const char* name, std::unique_ptr<MemBuffer>& buf_out, AllocatorPtr& alloc_out) override;
common::Status Trace(int id, const ONNX_NAMESPACE::TensorProto* value) override;
const MemoryPatternGroup& GetMemPatterns() override {
return mem_patterns_;

2
onnxruntime/core/framework/tensor_allocator.cc
View file

@ -13,7 +13,7 @@ AllocatorPtr ITensorAllocator::GetAllocator(const OrtMemoryInfo& memory_info) {
std::unique_ptr<ITensorAllocator> ITensorAllocator::Create(bool enable_mem_pattern,
const ExecutionPlanBase& execution_plan,
const SessionState& session_state,
InlinedVector<BufferUniquePtr>& weights_buffers) {
std::vector<BufferUniquePtr>& weights_buffers) {
if (enable_mem_pattern) {
return std::make_unique<TensorAllocatorWithMemPattern>(execution_plan, session_state, weights_buffers);
} else {

8
onnxruntime/core/framework/tensor_allocator.h
View file

@ -3,10 +3,8 @@
#pragma once
#include <optional>
#include <core/common/status.h>
#include <core/common/common.h>
#include <core/common/inlined_containers_fwd.h>
#include <core/graph/onnx_protobuf.h>
#include <core/framework/allocator.h>
#include <core/framework/tensor.h>
@ -23,7 +21,7 @@ class ITensorAllocator {
// Create an ITensorAllocator instance based on enable_mem_pattern
static std::unique_ptr<ITensorAllocator> Create(bool enable_mem_pattern, const ExecutionPlanBase& execution_plan,
const SessionState& session_state,
InlinedVector<BufferUniquePtr>& weights_buffers);
std::vector<BufferUniquePtr>& 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<std::string, size_t>& planned_memory_sizes_in_byte) = 0;
virtual common::Status FinalizePlan(std::unordered_map<std::string, size_t>& 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<MemBuffer>& buf_out,
std::unique_ptr<MemBuffer>& buf_out,
AllocatorPtr& alloc_out) = 0;
virtual const MemoryPatternGroup& GetMemPatterns() = 0;

20
onnxruntime/core/framework/tensor_allocator_with_mem_pattern.h
View file

@ -3,7 +3,6 @@
#pragma once
#include <core/common/inlined_containers.h>
#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<BufferUniquePtr>& weights_buffers_;
InlinedHashMap<OrtMemoryInfo, void*> buffers_;
std::vector<BufferUniquePtr>& weights_buffers_;
std::map<OrtMemoryInfo, void*> buffers_;
bool is_sealed_ = false;
const ExecutionPlanBase& seq_plan_;
common::Status AllocatePlannedBuffersAndReportTotalSize(
InlinedHashMap<std::string, size_t>& planned_memory_sizes_in_byte) {
std::unordered_map<std::string, size_t>& 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<BufferUniquePtr>& weights_buffers)
std::vector<BufferUniquePtr>& weights_buffers)
: ITensorAllocator(session_state),
planner_(execution_plan, /*using counters*/ false),
weights_buffers_(weights_buffers),
seq_plan_(execution_plan) {}
common::Status FinalizePlan(InlinedHashMap<std::string, size_t>& planned_memory_sizes_in_byte) override {
ORT_RETURN_IF_ERROR(planner_.GeneratePatterns(mem_patterns_));
common::Status FinalizePlan(std::unordered_map<std::string, size_t>& 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<MemBuffer>& buf_out, AllocatorPtr& alloc_out) override {
std::unique_ptr<MemBuffer>& 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<MemBuffer>(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<char*>(it->second) + block->offset_, block->size_, location);
buf_out = std::make_unique<MemBuffer>(reinterpret_cast<char*>(it->second) + block->offset_, block->size_, location);
return Status::OK();
}
common::Status Trace(int id, const ONNX_NAMESPACE::TensorProto* value) override {

2
onnxruntime/core/framework/utils.cc
View file

@ -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<SessionState::NodeInfo> node_info_vec;
std::vector<SessionState::NodeInfo> node_info_vec;
#ifdef ENABLE_TRAINING
if (session_state.GetInputNodeInfo(input_name, node_info_vec) == Status::OK()) {
#else

4
onnxruntime/core/graph/graph_viewer.cc
View file

@ -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<NodeIndex>(filter_info->nodes.cbegin(),
filter_info->nodes.cend());
const auto& metadef = filter_info->GetMetaDef();

27
onnxruntime/core/optimizer/optimizer_execution_frame.cc
View file

@ -17,14 +17,6 @@
namespace onnxruntime {
static size_t EstimateInputsOutputs(gsl::span<const Node* const> nodes) {
size_t num = 0;
for (auto n : nodes) {
num += n->InputDefs().size() + n->OutputDefs().size();
}
return num;
}
OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
const InitializedTensorSet& initialized_tensor_set,
const Path& model_path,
@ -40,7 +32,7 @@ OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& 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<const Node*>& 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<const Node*>& 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<std::string, OrtValue>::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<const Node*>& 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<OrtValue>& fetches)
: IExecutionFrame(info.GetMLValueNameIdxMap(), info.GetNodeIndexInfo(), fetch_mlvalue_idxs),
info_(info) {
Init(gsl::span<const int>(), gsl::span<const OrtValue>(), info.GetInitializers(), info.GetSparseInitializerLookupFunc(), fetches);
Init(std::vector<int>(), std::vector<OrtValue>(), info.GetInitializers(), info.GetSparseInitializerLookupFunc(), fetches);
}
AllocatorPtr OptimizerExecutionFrame::GetAllocatorImpl(const OrtMemoryInfo& info) const {

12
onnxruntime/core/optimizer/optimizer_execution_frame.h
View file

@ -5,7 +5,6 @@
#include <unordered_map>
#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<bool(const std::string&)>& 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<int, const NodeArg*> ort_value_idx_nodearg_map_;
std::unordered_map<int, OrtValue> initializers_;
InlinedHashMap<int, std::unique_ptr<char[]>> buffer_for_initialized_tensors_;
std::unordered_map<int, std::unique_ptr<char[]>> buffer_for_initialized_tensors_;
// This data structure is for uninitializing string tensors and
// munmap memory region and close file descriptor
std::unordered_map<int, OrtCallback> deleter_for_initialized_tensors_;
std::unique_ptr<NodeIndexInfo> node_index_info_;
const IExecutionProvider& execution_provider_;
const std::function<bool(const std::string&)>& is_sparse_initializer_func_;

2
onnxruntime/test/framework/allocation_planner_test.cc
View file

@ -181,7 +181,7 @@ class PlannerTest : public ::testing::Test {
profiling::Profiler profiler_;
std::unique_ptr<SessionState> state_;
ShapeMap shape_map_;
std::optional<SequentialExecutionPlan> plan_;
std::unique_ptr<SequentialExecutionPlan> plan_;
public:
PlannerTest()

9
onnxruntime/test/framework/execution_frame_test.cc
View file

@ -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<int64_t>{2, 3},
std::vector<float>(6, 1.0f), &v3);
std::vector<OrtValue> outputs;
ExecutionFrame frame(AsSpan({x1_idx, x2_idx, x3_idx}), AsSpan({v1, v2, v3}), {t3_idx}, outputs, {}, state);
vector<OrtValue> 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<int64_t>{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<float>(), cpu_allocator->Info(), TensorShape(std::vector<int64_t>{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);

2
winml/adapter/winml_adapter_session.cpp
View file

@ -214,7 +214,7 @@ static OrtDevice GetSessionGetInputDevice(_In_ OrtSession* session, _In_ const c
static_cast<InferenceSessionProtectedLoadAccessor*>(inference_session);
const onnxruntime::SessionState& session_state = session_protected_load_accessor->GetSessionState();
onnxruntime::InlinedVector<onnxruntime::SessionState::NodeInfo>node_info_vec;
std::vector<onnxruntime::SessionState::NodeInfo> 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;