Use TraceMe::InstantActivity with TraceMeEncode in memory profiling

PiperOrigin-RevId: 316103229
Change-Id: Iab8743b81720c0db4adc7c3b908f9def5e450777
This commit is contained in:
Jose Baiocchi 2020-06-12 07:47:35 -07:00 committed by TensorFlower Gardener
parent dfaeeac538
commit 777ac03c0e
2 changed files with 35 additions and 21 deletions

View File

@ -465,30 +465,32 @@ void BFCAllocator::AddTraceMe(absl::string_view traceme_name, const void* ptr) {
void BFCAllocator::AddTraceMe(absl::string_view traceme_name, void BFCAllocator::AddTraceMe(absl::string_view traceme_name,
const void* chunk_ptr, int64 req_bytes, const void* chunk_ptr, int64 req_bytes,
int64 alloc_bytes) { int64 alloc_bytes) {
tensorflow::profiler::TraceMe trace_me( tensorflow::profiler::TraceMe::InstantActivity(
[&]() TF_EXCLUSIVE_LOCKS_REQUIRED(lock_) { [this, traceme_name, chunk_ptr, req_bytes,
AllocatorStats stats = stats_; alloc_bytes]() TF_NO_THREAD_SAFETY_ANALYSIS {
int64 bytes_available = int64 bytes_available =
memory_limit_ - stats.bytes_reserved - stats.bytes_in_use; memory_limit_ - stats_.bytes_reserved - stats_.bytes_in_use;
const auto& annotation = const auto& annotation =
ScopedMemoryDebugAnnotation::CurrentAnnotation(); ScopedMemoryDebugAnnotation::CurrentAnnotation();
std::string tensor_shape = annotation.pending_shape std::string tensor_shape;
? annotation.pending_shape->DebugString() if (annotation.pending_shape) {
: ""; tensor_shape = annotation.pending_shape->DebugString();
return absl::StrCat(traceme_name, "#allocator_name=", name_, }
",bytes_reserved=", stats.bytes_reserved, return tensorflow::profiler::TraceMeEncode(
",bytes_allocated=", stats.bytes_in_use, traceme_name, {{"allocator_name", name_},
",bytes_available=", bytes_available, {"bytes_reserved", stats_.bytes_reserved},
",fragmentation=", GetFragmentation(), {"bytes_allocated", stats_.bytes_in_use},
",peak_bytes_in_use=", stats.peak_bytes_in_use, {"bytes_available", bytes_available},
",requested_bytes=", req_bytes, {"fragmentation", GetFragmentation()},
",allocation_bytes=", alloc_bytes, {"peak_bytes_in_use", stats_.peak_bytes_in_use},
",addr=", reinterpret_cast<uint64>(chunk_ptr), {"requested_bytes", req_bytes},
",tf_op=", annotation.pending_op_name, {"allocation_bytes", alloc_bytes},
",id=", annotation.pending_step_id, {"addr", reinterpret_cast<uint64>(chunk_ptr)},
",region_type=", annotation.pending_region_type, {"tf_op", annotation.pending_op_name},
",data_type=", annotation.pending_data_type, {"id", annotation.pending_step_id},
",shape=", tensor_shape, "#"); {"region_type", annotation.pending_region_type},
{"data_type", annotation.pending_data_type},
{"shape", tensor_shape}});
}, },
/*level=*/profiler::TraceMeLevel::kInfo); /*level=*/profiler::TraceMeLevel::kInfo);
} }

View File

@ -227,6 +227,18 @@ class TraceMe {
#endif #endif
} }
// Records the time of an instant activity.
template <typename NameGeneratorT>
static void InstantActivity(NameGeneratorT name_generator, int level = 1) {
#if !defined(IS_MOBILE_PLATFORM)
if (TF_PREDICT_FALSE(TraceMeRecorder::Active(level))) {
uint64 now = EnvTime::NowNanos();
TraceMeRecorder::Record({kCompleteActivity, name_generator(),
/*start_time=*/now, /*end_time=*/now});
}
#endif
}
static bool Active(int level = 1) { static bool Active(int level = 1) {
#if !defined(IS_MOBILE_PLATFORM) #if !defined(IS_MOBILE_PLATFORM)
return TraceMeRecorder::Active(level); return TraceMeRecorder::Active(level);