Remove intermediate data type AggregationStats and ActivityMetadata.

PiperOrigin-RevId: 357103126
Change-Id: Ib9dba2b55325d04311f33f4907adea7992309b4c
This commit is contained in:
Tianrun Li 2021-02-11 18:40:08 -08:00 committed by TensorFlower Gardener
parent b26252b383
commit a0f8ebc1cb

View File

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