From a4e22aa1c416cb5f1924a0bda95603fbb725d776 Mon Sep 17 00:00:00 2001 From: Jiho Choi Date: Fri, 24 Apr 2020 16:26:20 -0700 Subject: [PATCH] Refactor the derived_timeline files. PiperOrigin-RevId: 308348454 Change-Id: Iad035e375d87ded82fcd11eb5f7f38c17b797381 --- .../core/profiler/utils/derived_timeline.cc | 91 +++++++++---------- .../core/profiler/utils/derived_timeline.h | 15 +-- .../core/profiler/utils/xplane_schema.cc | 7 ++ .../core/profiler/utils/xplane_schema.h | 8 ++ 4 files changed, 69 insertions(+), 52 deletions(-) diff --git a/tensorflow/core/profiler/utils/derived_timeline.cc b/tensorflow/core/profiler/utils/derived_timeline.cc index c682c1515a3..c99d8e82cb7 100644 --- a/tensorflow/core/profiler/utils/derived_timeline.cc +++ b/tensorflow/core/profiler/utils/derived_timeline.cc @@ -38,24 +38,16 @@ absl::string_view DummySymbolResolver(absl::string_view hlo_module, return absl::string_view(); } -const absl::string_view kDerivedLineSteps = "Steps"; -const absl::string_view kDerivedLineTensorFlowNameScope = - "TensorFlow Name Scope"; -const absl::string_view kDerivedLineTensorFlowOps = "TensorFlow Ops"; -const absl::string_view kDerivedLineXlaModules = "XLA Modules"; -const absl::string_view kDerivedLineXlaOps = "XLA Ops"; -const absl::string_view kDerivedLineKernelLaunch = "Launch Stats"; const absl::string_view kAnnotationDelimiter = "::"; -XEvent CreateXEvent(const XEventVisitor& src_event_visitor, - const XEventMetadata& metadata, - int64 group_id_stat_metadata_id, +XEvent CreateXEvent(const XEventMetadata& metadata, int64 offset_ps, + int64 duration_ps, int64 group_id_stat_metadata_id, absl::optional group_id) { XEvent event; event.set_metadata_id(metadata.id()); // TODO(b/150498419): Normalize with the line start time. - event.set_offset_ps(src_event_visitor.OffsetPs()); - event.set_duration_ps(src_event_visitor.DurationPs()); + event.set_offset_ps(offset_ps); + event.set_duration_ps(duration_ps); if (group_id) { XStat* stat = event.add_stats(); stat->set_metadata_id(group_id_stat_metadata_id); @@ -64,9 +56,10 @@ XEvent CreateXEvent(const XEventVisitor& src_event_visitor, return event; } -void ProcessTfOpEvent(const XEventVisitor& event, - absl::string_view tf_op_full_name, - absl::optional group_id, +} // namespace + +void ProcessTfOpEvent(absl::string_view tf_op_full_name, int64 offset_ps, + int64 duration_ps, absl::optional group_id, XPlaneBuilder* plane_builder, DerivedXLineBuilder* tf_name_scope_line_builder, DerivedXLineBuilder* tf_op_line_builder) { @@ -79,8 +72,8 @@ void ProcessTfOpEvent(const XEventVisitor& event, std::vector name_scope_event_per_level; for (const auto& tf_name_scope : ParseTfNameScopes(tf_op)) { name_scope_event_per_level.push_back(CreateXEvent( - event, *plane_builder->GetOrCreateEventMetadata(tf_name_scope), - group_id_stat_metadata_id, group_id)); + *plane_builder->GetOrCreateEventMetadata(tf_name_scope), offset_ps, + duration_ps, group_id_stat_metadata_id, group_id)); } tf_name_scope_line_builder->ExpandOrAddEvents(name_scope_event_per_level); } @@ -89,12 +82,11 @@ void ProcessTfOpEvent(const XEventVisitor& event, // Set the display name to op_type so that the events of the same op_type have // the same color in the trace viewer. tf_op_event_metadata->set_display_name(TfOpEventName(tf_op)); - tf_op_line_builder->ExpandOrAddEvent(CreateXEvent( - event, *tf_op_event_metadata, group_id_stat_metadata_id, group_id)); + tf_op_line_builder->ExpandOrAddEvent( + CreateXEvent(*tf_op_event_metadata, offset_ps, duration_ps, + group_id_stat_metadata_id, group_id)); } -} // namespace - DerivedXLineBuilder::DerivedXLineBuilder( XPlaneBuilder* plane, int64 line_id, absl::string_view name, int64 timestamp_ns, std::vector dependent_lines) @@ -118,14 +110,20 @@ void DerivedXLineBuilder::ExpandOrAddLevelEvent(const XEvent& event, last_event->SetDurationPs((offset_ps + duration_ps) - last_event->OffsetPs()); } else { - // Otherwise, create a new event for the given level. - last_event = line_.AddEvent(event); - // Reset last events lower than the given level. + // Otherwise, reset the last events lower than or equal to the given level. ResetLastEvents(level); - if (level == 0) ResetDependentLines(); + // And create a new event for the given level. + last_event = line_.AddEvent(event); } } +void DerivedXLineBuilder::ResetLastEvents(int level) { + for (int i = level; i < last_event_by_level_.size(); ++i) { + last_event_by_level_[i] = absl::nullopt; + } + if (level == 0) ResetDependentLines(); +} + void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver, const EventGroupNameMap& event_group_name_map, XPlane* device_trace, bool step_info_only) { @@ -142,19 +140,18 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver, absl::c_sort(events); XPlaneBuilder plane(device_trace); - DerivedXLineBuilder tf_ops(&plane, kThreadIdTfOp, kDerivedLineTensorFlowOps, + DerivedXLineBuilder tf_ops(&plane, kThreadIdTfOp, kTensorFlowOpLineName, start_timestamp_ns, {}); DerivedXLineBuilder tf_name_scope(&plane, kThreadIdTfNameScope, - kDerivedLineTensorFlowNameScope, + kTensorFlowNameScopeLineName, start_timestamp_ns, {&tf_ops}); - DerivedXLineBuilder hlo_ops(&plane, kThreadIdHloOp, kDerivedLineXlaOps, + DerivedXLineBuilder hlo_ops(&plane, kThreadIdHloOp, kXlaOpLineName, start_timestamp_ns, {}); DerivedXLineBuilder hlo_modules(&plane, kThreadIdHloModule, - kDerivedLineXlaModules, start_timestamp_ns, - {&tf_ops, &tf_name_scope, &hlo_ops}); - DerivedXLineBuilder steps(&plane, kThreadIdStepInfo, kDerivedLineSteps, - start_timestamp_ns, - {&tf_ops, &tf_name_scope, &hlo_ops, &hlo_modules}); + kXlaModuleLineName, start_timestamp_ns, + {&tf_name_scope, &hlo_ops}); + DerivedXLineBuilder steps(&plane, kThreadIdStepInfo, kStepLineName, + start_timestamp_ns, {&hlo_modules}); int64 group_id_stat_metadata_id = plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kGroupId))->id(); int64 step_name_stat_metadata_id = @@ -162,6 +159,8 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver, // Process events in order by start time. for (const XEventVisitor& event : events) { + int64 offset_ps = event.OffsetPs(); + int64 duration_ps = event.DurationPs(); absl::string_view tf_op_full_name; absl::string_view hlo_module_name; std::vector hlo_op_names; @@ -184,8 +183,8 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver, if (group_id) { XEvent step_event = CreateXEvent( - event, *plane.GetOrCreateEventMetadata(absl::StrCat(*group_id)), - group_id_stat_metadata_id, group_id); + *plane.GetOrCreateEventMetadata(absl::StrCat(*group_id)), offset_ps, + duration_ps, group_id_stat_metadata_id, group_id); if (auto group_name = gtl::FindOrNull(event_group_name_map, *group_id)) { XStat* stat = step_event.add_stats(); stat->set_metadata_id(step_name_stat_metadata_id); @@ -201,9 +200,9 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver, if (!is_kernel) continue; if (!hlo_module_name.empty()) { - hlo_modules.ExpandOrAddEvent( - CreateXEvent(event, *plane.GetOrCreateEventMetadata(hlo_module_name), - group_id_stat_metadata_id, group_id)); + hlo_modules.ExpandOrAddEvent(CreateXEvent( + *plane.GetOrCreateEventMetadata(hlo_module_name), offset_ps, + duration_ps, group_id_stat_metadata_id, group_id)); } if (!hlo_op_names.empty()) { // GPU kernel compiled by XLA @@ -211,19 +210,19 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver, std::vector hlo_op_event_per_level; for (absl::string_view hlo_op_name : hlo_op_names) { DCHECK(!hlo_op_name.empty()); - hlo_op_event_per_level.push_back( - CreateXEvent(event, *plane.GetOrCreateEventMetadata(hlo_op_name), - group_id_stat_metadata_id, group_id)); + hlo_op_event_per_level.push_back(CreateXEvent( + *plane.GetOrCreateEventMetadata(hlo_op_name), offset_ps, + duration_ps, group_id_stat_metadata_id, group_id)); } hlo_ops.ExpandOrAddEvents(hlo_op_event_per_level); auto tf_op_name = symbol_resolver(hlo_module_name, hlo_op_names.back()); if (!tf_op_name.empty()) { - ProcessTfOpEvent(event, tf_op_name, group_id, &plane, &tf_name_scope, - &tf_ops); + ProcessTfOpEvent(tf_op_name, offset_ps, duration_ps, group_id, &plane, + &tf_name_scope, &tf_ops); } } else if (!tf_op_full_name.empty()) { // GPU kernel not compiled by XLA - ProcessTfOpEvent(event, tf_op_full_name, group_id, &plane, &tf_name_scope, - &tf_ops); + ProcessTfOpEvent(tf_op_full_name, offset_ps, duration_ps, group_id, + &plane, &tf_name_scope, &tf_ops); } } RemoveEmptyLines(device_trace); @@ -294,7 +293,7 @@ void DeriveEventsFromHostTrace(const XPlane* host_trace, XPlaneBuilder device_plane(device_traces[i]); XLineBuilder launch_line = device_plane.GetOrCreateLine(kThreadIdKernelLaunch); - launch_line.SetName(kDerivedLineKernelLaunch); + launch_line.SetName(kKernelLaunchLineName); launch_line.SetTimestampNs(std::min(device_plane_start, host_plane_start)); for (const auto& it : per_device_launch_info[i]) { uint64 group_id = it.first; diff --git a/tensorflow/core/profiler/utils/derived_timeline.h b/tensorflow/core/profiler/utils/derived_timeline.h index 64f6cad88d5..61b62bdc8da 100644 --- a/tensorflow/core/profiler/utils/derived_timeline.h +++ b/tensorflow/core/profiler/utils/derived_timeline.h @@ -40,12 +40,8 @@ class DerivedXLineBuilder { ExpandOrAddLevelEvent(event, /*level=*/0); } - // Reset last events lower than the given level. - void ResetLastEvents(int level = -1) { - for (int i = level + 1; i < last_event_by_level_.size(); ++i) { - last_event_by_level_[i] = absl::nullopt; - } - } + // Reset the last events lower than or equal to the given level. + void ResetLastEvents(int level = 0); private: // If the last event of the given level has the same metadata, expands it to @@ -70,6 +66,13 @@ class DerivedXLineBuilder { using SymbolResolver = std::function; +// Derives TF name scope and op events from the TF op's fully qualified name. +void ProcessTfOpEvent(absl::string_view tf_op_full_name, int64 offset_ps, + int64 duration_ps, absl::optional group_id, + XPlaneBuilder* plane_builder, + DerivedXLineBuilder* tf_name_scope_line_builder, + DerivedXLineBuilder* tf_op_line_builder); + // Derives "Step Info", "Tensorflow Ops", "XLA Ops" and "XLA Module" lines in // an NVIDIA_GPU device trace from data passed as ScopedAnnotations and stored // as XStats in XEvents corresponding to GPU Kernels. Consecutive annotations diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc index 402b755329c..2851b572593 100644 --- a/tensorflow/core/profiler/utils/xplane_schema.cc +++ b/tensorflow/core/profiler/utils/xplane_schema.cc @@ -28,6 +28,13 @@ const absl::string_view kCuptiDriverApiPlaneName = "/host:CUPTI"; const absl::string_view kMetadataPlane = "/host:metadata"; const absl::string_view kTFStreamzPlane = "/host:tfstreamz"; +const absl::string_view kStepLineName = "Steps"; +const absl::string_view kTensorFlowNameScopeLineName = "TensorFlow Name Scope"; +const absl::string_view kTensorFlowOpLineName = "TensorFlow Ops"; +const absl::string_view kXlaModuleLineName = "XLA Modules"; +const absl::string_view kXlaOpLineName = "XLA Ops"; +const absl::string_view kKernelLaunchLineName = "Launch Stats"; + const int32 kHostPlaneId = 49; const int32 kGpuPlaneBaseId = 0; const int32 kCuptiDriverApiPlaneId = 50; diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h index 6ffb4f90f7c..427a3a2968d 100644 --- a/tensorflow/core/profiler/utils/xplane_schema.h +++ b/tensorflow/core/profiler/utils/xplane_schema.h @@ -36,6 +36,14 @@ ABSL_CONST_INIT extern const absl::string_view kMetadataPlane; // Name of XPlane that contains kpi related metrics. ABSL_CONST_INIT extern const absl::string_view kTFStreamzPlane; +// Names of XLines that contain ML-level events. +ABSL_CONST_INIT extern const absl::string_view kStepLineName; +ABSL_CONST_INIT extern const absl::string_view kTensorFlowNameScopeLineName; +ABSL_CONST_INIT extern const absl::string_view kTensorFlowOpLineName; +ABSL_CONST_INIT extern const absl::string_view kXlaModuleLineName; +ABSL_CONST_INIT extern const absl::string_view kXlaOpLineName; +ABSL_CONST_INIT extern const absl::string_view kKernelLaunchLineName; + // Id of XPlane that contains TraceMe events. ABSL_CONST_INIT extern const int32 kHostPlaneId; // Ids prefix of XPlane that contains GPU events.