Remove intermediate data type AggregationStats and ActivityMetadata.
PiperOrigin-RevId: 357103126 Change-Id: Ib9dba2b55325d04311f33f4907adea7992309b4c
This commit is contained in:
parent
b26252b383
commit
a0f8ebc1cb
@ -51,29 +51,6 @@ constexpr int64 kInvalidStepId = -1;
|
||||
// MemoryActivityMetadata proto it contains.
|
||||
using IndexMetaPair = std::pair<int64 /*index*/, const MemoryActivityMetadata*>;
|
||||
|
||||
// Aggregated memory stats from an allocator. Temporary container to fill
|
||||
// MemoryAggregationStats.
|
||||
struct AggregationStats {
|
||||
int64 bytes_reserved = 0;
|
||||
int64 bytes_allocated = 0;
|
||||
int64 bytes_available = 0;
|
||||
double fragmentation = 0;
|
||||
int64 peak_bytes_in_use = 0;
|
||||
};
|
||||
|
||||
// Metadata associated with each memory allocation/deallocation activity.
|
||||
// Temporary container to fill MemoryActivityMetadata.
|
||||
struct ActivityMetadata {
|
||||
int64 requested_bytes = 0;
|
||||
int64 allocation_bytes = 0;
|
||||
uint64 address = 0;
|
||||
absl::string_view tf_op_name;
|
||||
int64 step_id = kInvalidStepId;
|
||||
absl::string_view region_type;
|
||||
int64 data_type = 0;
|
||||
absl::string_view tensor_shape;
|
||||
};
|
||||
|
||||
bool IsMemoryAllocation(int64 event_type) {
|
||||
return event_type == HostEventType::kMemoryAllocation;
|
||||
}
|
||||
@ -82,51 +59,22 @@ bool IsMemoryDeallocation(int64 event_type) {
|
||||
return event_type == HostEventType::kMemoryDeallocation;
|
||||
}
|
||||
|
||||
void FillAggregationStats(const AggregationStats& src,
|
||||
MemoryAggregationStats* dst) {
|
||||
dst->set_stack_reserved_bytes(src.bytes_reserved);
|
||||
dst->set_heap_allocated_bytes(src.bytes_allocated);
|
||||
dst->set_free_memory_bytes(src.bytes_available);
|
||||
dst->set_fragmentation(src.fragmentation);
|
||||
dst->set_peak_bytes_in_use(src.peak_bytes_in_use);
|
||||
}
|
||||
|
||||
void FillActivityMetadata(int64 event_type, const ActivityMetadata& src,
|
||||
MemoryActivityMetadata* dst) {
|
||||
if (IsMemoryAllocation(event_type)) {
|
||||
dst->set_memory_activity(ALLOCATION);
|
||||
} else if (IsMemoryDeallocation(event_type)) {
|
||||
dst->set_memory_activity(DEALLOCATION);
|
||||
}
|
||||
dst->set_requested_bytes(src.requested_bytes);
|
||||
dst->set_allocation_bytes(src.allocation_bytes);
|
||||
dst->set_address(src.address);
|
||||
dst->set_tf_op_name(std::string(src.tf_op_name));
|
||||
dst->set_step_id(src.step_id);
|
||||
dst->set_region_type(std::string(src.region_type));
|
||||
dst->set_data_type(tensorflow::DataTypeString(
|
||||
static_cast<tensorflow::DataType>(src.data_type)));
|
||||
dst->set_tensor_shape(std::string(src.tensor_shape));
|
||||
}
|
||||
|
||||
void UpdateProfileSummary(const AggregationStats& stats, int64 time_offset_ps,
|
||||
MemoryProfileSummary* summary) {
|
||||
void UpdateProfileSummary(const MemoryAggregationStats& stats,
|
||||
int64 time_offset_ps, MemoryProfileSummary* summary) {
|
||||
// Update the peak memory usage over allocator's lifetime.
|
||||
summary->set_peak_bytes_usage_lifetime(stats.peak_bytes_in_use);
|
||||
summary->set_peak_bytes_usage_lifetime(stats.peak_bytes_in_use());
|
||||
MemoryAggregationStats* peak_stats = summary->mutable_peak_stats();
|
||||
// If we reach (or stay at) peak memory usage within the profiling window,
|
||||
// update memory profile summary.
|
||||
if (stats.bytes_reserved + stats.bytes_allocated >=
|
||||
if (stats.stack_reserved_bytes() + stats.heap_allocated_bytes() >=
|
||||
peak_stats->peak_bytes_in_use()) {
|
||||
peak_stats->set_peak_bytes_in_use(stats.bytes_reserved +
|
||||
stats.bytes_allocated);
|
||||
peak_stats->set_stack_reserved_bytes(stats.bytes_reserved);
|
||||
peak_stats->set_heap_allocated_bytes(stats.bytes_allocated);
|
||||
peak_stats->set_free_memory_bytes(stats.bytes_available);
|
||||
peak_stats->set_fragmentation(stats.fragmentation);
|
||||
*peak_stats = stats;
|
||||
peak_stats->set_peak_bytes_in_use(stats.stack_reserved_bytes() +
|
||||
stats.heap_allocated_bytes());
|
||||
summary->set_peak_stats_time_ps(time_offset_ps);
|
||||
summary->set_memory_capacity(stats.bytes_reserved + stats.bytes_allocated +
|
||||
stats.bytes_available);
|
||||
summary->set_memory_capacity(stats.stack_reserved_bytes() +
|
||||
stats.heap_allocated_bytes() +
|
||||
stats.free_memory_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
@ -145,8 +93,15 @@ MemoryProfile GenerateMemoryProfile(const XPlane* host_trace) {
|
||||
return;
|
||||
}
|
||||
|
||||
AggregationStats stats;
|
||||
ActivityMetadata metadata;
|
||||
MemoryAggregationStats stats;
|
||||
MemoryActivityMetadata metadata;
|
||||
if (IsMemoryAllocation(event_type)) {
|
||||
metadata.set_memory_activity(ALLOCATION);
|
||||
} else if (IsMemoryDeallocation(event_type)) {
|
||||
metadata.set_memory_activity(DEALLOCATION);
|
||||
}
|
||||
metadata.set_step_id(kInvalidStepId);
|
||||
|
||||
std::string memory_id;
|
||||
event.ForEachStat([&](const XStatVisitor& stat) {
|
||||
if (!stat.Type().has_value()) return;
|
||||
@ -159,59 +114,59 @@ MemoryProfile GenerateMemoryProfile(const XPlane* host_trace) {
|
||||
memory_id = std::string(stat.StrOrRefValue());
|
||||
break;
|
||||
case StatType::kBytesReserved:
|
||||
stats.bytes_reserved = stat.IntValue();
|
||||
stats.set_stack_reserved_bytes(stat.IntValue());
|
||||
break;
|
||||
case StatType::kBytesAllocated:
|
||||
stats.bytes_allocated = stat.IntValue();
|
||||
stats.set_heap_allocated_bytes(stat.IntValue());
|
||||
break;
|
||||
case StatType::kBytesAvailable:
|
||||
stats.bytes_available = stat.IntValue();
|
||||
stats.set_free_memory_bytes(stat.IntValue());
|
||||
break;
|
||||
case StatType::kFragmentation:
|
||||
stats.fragmentation = stat.DoubleValue();
|
||||
stats.set_fragmentation(stat.DoubleValue());
|
||||
break;
|
||||
case StatType::kPeakBytesInUse:
|
||||
stats.peak_bytes_in_use = stat.IntValue();
|
||||
stats.set_peak_bytes_in_use(stat.IntValue());
|
||||
break;
|
||||
case StatType::kRequestedBytes:
|
||||
metadata.requested_bytes = stat.IntValue();
|
||||
metadata.set_requested_bytes(stat.IntValue());
|
||||
break;
|
||||
case StatType::kAllocationBytes:
|
||||
metadata.allocation_bytes = stat.IntValue();
|
||||
metadata.set_allocation_bytes(stat.IntValue());
|
||||
break;
|
||||
case StatType::kAddress:
|
||||
metadata.address = stat.IntValue();
|
||||
metadata.set_address(stat.IntValue());
|
||||
break;
|
||||
case StatType::kTfOp:
|
||||
metadata.tf_op_name = stat.StrOrRefValue();
|
||||
metadata.set_tf_op_name(std::string(stat.StrOrRefValue()));
|
||||
break;
|
||||
case StatType::kGroupId:
|
||||
metadata.step_id = stat.IntValue();
|
||||
metadata.set_step_id(stat.IntValue());
|
||||
break;
|
||||
case StatType::kRegionType:
|
||||
metadata.region_type = stat.StrOrRefValue();
|
||||
metadata.set_region_type(std::string(stat.StrOrRefValue()));
|
||||
break;
|
||||
case StatType::kDataType:
|
||||
metadata.data_type = stat.IntValue();
|
||||
metadata.set_data_type(tensorflow::DataTypeString(
|
||||
static_cast<tensorflow::DataType>(stat.IntValue())));
|
||||
break;
|
||||
case StatType::kTensorShapes:
|
||||
metadata.tensor_shape = stat.StrOrRefValue();
|
||||
metadata.set_tensor_shape(std::string(stat.StrOrRefValue()));
|
||||
break;
|
||||
}
|
||||
});
|
||||
|
||||
MemoryProfileSnapshot* snapshot =
|
||||
(*memory_profile.mutable_memory_profile_per_allocator())[memory_id]
|
||||
.add_memory_profile_snapshots();
|
||||
snapshot->set_time_offset_ps(event.OffsetPs());
|
||||
FillAggregationStats(stats, snapshot->mutable_aggregation_stats());
|
||||
FillActivityMetadata(event_type, metadata,
|
||||
snapshot->mutable_activity_metadata());
|
||||
|
||||
MemoryProfileSummary* summary =
|
||||
(*memory_profile.mutable_memory_profile_per_allocator())[memory_id]
|
||||
.mutable_profile_summary();
|
||||
UpdateProfileSummary(stats, event.OffsetPs(), summary);
|
||||
|
||||
MemoryProfileSnapshot* snapshot =
|
||||
(*memory_profile.mutable_memory_profile_per_allocator())[memory_id]
|
||||
.add_memory_profile_snapshots();
|
||||
snapshot->set_time_offset_ps(event.OffsetPs());
|
||||
*snapshot->mutable_aggregation_stats() = std::move(stats);
|
||||
*snapshot->mutable_activity_metadata() = std::move(metadata);
|
||||
});
|
||||
});
|
||||
return memory_profile;
|
||||
@ -320,11 +275,16 @@ void InsertSpecialAllocations(int64 unmapped_allocation_bytes, int64 step_id,
|
||||
if (unmapped_allocation_bytes > 0) {
|
||||
MemoryActivityMetadata* special_allocation =
|
||||
memory_profile->add_special_allocations();
|
||||
FillActivityMetadata(HostEventType::kMemoryAllocation,
|
||||
{unmapped_allocation_bytes, unmapped_allocation_bytes,
|
||||
0, "unused preallocated device memory", step_id,
|
||||
"persist/dynamic", 0, "unknown"},
|
||||
special_allocation);
|
||||
special_allocation->set_memory_activity(ALLOCATION);
|
||||
special_allocation->set_requested_bytes(unmapped_allocation_bytes);
|
||||
special_allocation->set_allocation_bytes(unmapped_allocation_bytes);
|
||||
special_allocation->set_address(0);
|
||||
special_allocation->set_tf_op_name("unused preallocated device memory");
|
||||
special_allocation->set_step_id(step_id);
|
||||
special_allocation->set_region_type("persist/dynamic");
|
||||
special_allocation->set_data_type(
|
||||
tensorflow::DataTypeString(static_cast<tensorflow::DataType>(0)));
|
||||
special_allocation->set_tensor_shape("unknown");
|
||||
active_allocs->push_back({--index, special_allocation});
|
||||
}
|
||||
int64 stack_bytes =
|
||||
@ -332,10 +292,16 @@ void InsertSpecialAllocations(int64 unmapped_allocation_bytes, int64 step_id,
|
||||
if (stack_bytes > 0) {
|
||||
MemoryActivityMetadata* special_allocation =
|
||||
memory_profile->add_special_allocations();
|
||||
FillActivityMetadata(
|
||||
HostEventType::kMemoryAllocation,
|
||||
{stack_bytes, stack_bytes, 0, "stack", step_id, "stack", 0, "unknown"},
|
||||
special_allocation);
|
||||
special_allocation->set_memory_activity(ALLOCATION);
|
||||
special_allocation->set_requested_bytes(stack_bytes);
|
||||
special_allocation->set_allocation_bytes(stack_bytes);
|
||||
special_allocation->set_address(0);
|
||||
special_allocation->set_tf_op_name("stack");
|
||||
special_allocation->set_step_id(step_id);
|
||||
special_allocation->set_region_type("stack");
|
||||
special_allocation->set_data_type(
|
||||
tensorflow::DataTypeString(static_cast<tensorflow::DataType>(0)));
|
||||
special_allocation->set_tensor_shape("unknown");
|
||||
active_allocs->push_back({--index, special_allocation});
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user