Skip to content

Commit

Permalink
fix memory profile build (#16177)
Browse files Browse the repository at this point in the history
### Description
<!-- Describe your changes. -->
This PR is to fix the build break when onnxruntime_ENABLE_MEMORY_PROFILE
is on


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This PR is to fix the build break when onnxruntime_ENABLE_MEMORY_PROFILE
is on.
It fixes this issue
#16124

Co-authored-by: Lei Cao <leca@microsoft.com>
  • Loading branch information
2 people authored and snnn committed Jun 15, 2023
1 parent 81b9de0 commit 7d59ef2
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 16 deletions.
8 changes: 4 additions & 4 deletions onnxruntime/core/framework/execution_frame.cc
Original file line number Diff line number Diff line change
Expand Up @@ -421,14 +421,14 @@ ExecutionFrame::ExecutionFrame(gsl::span<const int> feed_mlvalue_idxs, gsl::span
// Planning of one memory type should only happen once.
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
ORT_ENFORCE(
static_activation_memory_sizes_in_byte_.find(location.name) ==
static_activation_memory_sizes_in_byte_.find(location.ToString()) ==
static_activation_memory_sizes_in_byte_.end(),
"Memory type ",
location.name,
location.ToString(),
" should only appear once.");
// static_activation_memory_in_bytes_ is max virtual memory size the planner computes.
// Memory dynamically allocated when executing kernels is not recorded using this field.
static_activation_memory_sizes_in_byte_[location.name] = peak_size;
static_activation_memory_sizes_in_byte_[location.ToString()] = peak_size;
#endif
// the memory pattern buffer will leave in the whole execution.
#ifdef ORT_ENABLE_STREAM
Expand Down Expand Up @@ -609,7 +609,7 @@ Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(OrtValue& ort_va
// Dynamic activation size would be accessed by multiple threads
// if parallel executor is used.
std::unique_lock<std::mutex> lock(mtx_);
dynamic_activation_memory_sizes_in_byte_[location.name] += size;
dynamic_activation_memory_sizes_in_byte_[location.ToString()] += size;
session_state_.GetMemoryProfiler()->GetMemoryInfo().SetDynamicAllocation(ort_value_index);
#endif
}
Expand Down
20 changes: 10 additions & 10 deletions onnxruntime/core/framework/memory_info.cc
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ void PrintInforPerTensor(const MemoryInfo::AllocInfoPerTensor& alloc_info, const
std::cout << "Index: " << alloc_info.mlvalue_index << ", ";
std::cout << "Reuse inplace: " << alloc_info.inplace_reuse << ", ";
std::cout << "Alloc type: " << alloc_info.alloc_kind << ", ";
std::cout << "Location: " << alloc_info.location.name << ", ";
std::cout << "Location: " << alloc_info.location.ToString() << ", ";
std::cout << "lifetime: (" << alloc_info.lifetime_interval.first << ", " << alloc_info.lifetime_interval.second << "), ";
std::cout << "planned block: (" << mem_info.planned_block.offset_ << ", "
<< (mem_info.planned_block.offset_ + mem_info.planned_block.size_) << "), ";
Expand All @@ -142,24 +142,24 @@ void PrintInforPerTensor(const MemoryInfo::AllocInfoPerTensor& alloc_info, const

void MemoryInfo::PrintMemoryInfoForLocation(const OrtDevice::DeviceType location) {
for (const auto& map : tensors_memory_info_map_) {
if (map.first.device.Type() != location) continue;
std::cout << "Initializer in " << map.first.name << "\n";
if (map.first.Type() != location) continue;
std::cout << "Initializer in " << map.first.ToString() << "\n";
const auto& initailizer_map = map.second.at(MapType::Initializer);
for (const auto& item : initailizer_map) {
if (AllocPlan(item.first)->location.device.Type() != location) continue;
if (AllocPlan(item.first)->location.Type() != location) continue;
PrintInforPerTensor(*AllocPlan(item.first), item.second, initailizer_map.GetAllocAddress(item.first));
}
std::cout << "\nStatic Activation in " << map.first.name << "\n";
std::cout << "\nStatic Activation in " << map.first.ToString() << "\n";
const auto& static_activation_map = map.second.at(MapType::StaticActivation);
for (const auto& item : static_activation_map) {
if (AllocPlan(item.first)->location.device.Type() != location) continue;
if (AllocPlan(item.first)->location.Type() != location) continue;
PrintInforPerTensor(*AllocPlan(item.first), item.second, static_activation_map.GetAllocAddress(item.first));
}

std::cout << "\nDynamic Activation in " << map.first.name << "\n";
std::cout << "\nDynamic Activation in " << map.first.ToString() << "\n";
const auto& dynamic_activation_map = map.second.at(MapType::DynamicActivation);
for (const auto& item : dynamic_activation_map) {
if (AllocPlan(item.first)->location.device.Type() != location) continue;
if (AllocPlan(item.first)->location.Type() != location) continue;
PrintInforPerTensor(*AllocPlan(item.first), item.second, dynamic_activation_map.GetAllocAddress(item.first));
}
}
Expand Down Expand Up @@ -273,10 +273,10 @@ void MemoryProfiler::CreateEvents(const std::string& p_name,
// Create Event for each tensor
auto& time_step_trace = GetMemoryInfo().time_step_trace_;
for (const auto& location_map : GetMemoryInfo().tensors_memory_info_map_) {
const OrtMemoryInfo& memory_info = location_map.first;
const OrtDevice& memory_info = location_map.first;
const auto& maptype_to_map_mapping = location_map.second;

if (memory_info.device.Type() != device_type) continue;
if (memory_info.Type() != device_type) continue;

// If there is no tensor of a map_type, we skip creating event for that map_type
if (maptype_to_map_mapping.find(map_type) == maptype_to_map_mapping.end()) continue;
Expand Down
4 changes: 2 additions & 2 deletions onnxruntime/core/framework/memory_info.h
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ class MemoryInfo {
bool inplace_reuse{false};
OrtValueIndex reused_buffer{0}; // The index of the reused tensor, if no reuse, it is its own tensor.
AllocKind alloc_kind{AllocKind::kAllocate};
OrtMemoryInfo location;
OrtDevice location;
};

struct AllocationSummary {
Expand Down Expand Up @@ -185,7 +185,7 @@ class MemoryInfo {
}

// Key: The map type. E.g., initializer, activation. Value: A map from the tensor index to its memory information
std::map<OrtMemoryInfo, std::map<MapType, MemoryInfoMap> > tensors_memory_info_map_;
std::map<OrtDevice, std::map<MapType, MemoryInfoMap> > tensors_memory_info_map_;

// Key: The tensor index. Value: The Allocation information per tensor
std::map<OrtValueIndex, AllocInfoPerTensor> tensor_alloc_info_map_;
Expand Down

0 comments on commit 7d59ef2

Please sign in to comment.