Refactor the derived_timeline files.

PiperOrigin-RevId: 308348454
Change-Id: Iad035e375d87ded82fcd11eb5f7f38c17b797381
This commit is contained in:
Jiho Choi 2020-04-24 16:26:20 -07:00 committed by TensorFlower Gardener
parent ceee848cea
commit a4e22aa1c4
4 changed files with 69 additions and 52 deletions

View File

@ -38,24 +38,16 @@ absl::string_view DummySymbolResolver(absl::string_view hlo_module,
return absl::string_view(); 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 = "::"; const absl::string_view kAnnotationDelimiter = "::";
XEvent CreateXEvent(const XEventVisitor& src_event_visitor, XEvent CreateXEvent(const XEventMetadata& metadata, int64 offset_ps,
const XEventMetadata& metadata, int64 duration_ps, int64 group_id_stat_metadata_id,
int64 group_id_stat_metadata_id,
absl::optional<int64> group_id) { absl::optional<int64> group_id) {
XEvent event; XEvent event;
event.set_metadata_id(metadata.id()); event.set_metadata_id(metadata.id());
// TODO(b/150498419): Normalize with the line start time. // TODO(b/150498419): Normalize with the line start time.
event.set_offset_ps(src_event_visitor.OffsetPs()); event.set_offset_ps(offset_ps);
event.set_duration_ps(src_event_visitor.DurationPs()); event.set_duration_ps(duration_ps);
if (group_id) { if (group_id) {
XStat* stat = event.add_stats(); XStat* stat = event.add_stats();
stat->set_metadata_id(group_id_stat_metadata_id); stat->set_metadata_id(group_id_stat_metadata_id);
@ -64,9 +56,10 @@ XEvent CreateXEvent(const XEventVisitor& src_event_visitor,
return event; return event;
} }
void ProcessTfOpEvent(const XEventVisitor& event, } // namespace
absl::string_view tf_op_full_name,
absl::optional<int64> group_id, void ProcessTfOpEvent(absl::string_view tf_op_full_name, int64 offset_ps,
int64 duration_ps, absl::optional<int64> group_id,
XPlaneBuilder* plane_builder, XPlaneBuilder* plane_builder,
DerivedXLineBuilder* tf_name_scope_line_builder, DerivedXLineBuilder* tf_name_scope_line_builder,
DerivedXLineBuilder* tf_op_line_builder) { DerivedXLineBuilder* tf_op_line_builder) {
@ -79,8 +72,8 @@ void ProcessTfOpEvent(const XEventVisitor& event,
std::vector<XEvent> name_scope_event_per_level; std::vector<XEvent> name_scope_event_per_level;
for (const auto& tf_name_scope : ParseTfNameScopes(tf_op)) { for (const auto& tf_name_scope : ParseTfNameScopes(tf_op)) {
name_scope_event_per_level.push_back(CreateXEvent( name_scope_event_per_level.push_back(CreateXEvent(
event, *plane_builder->GetOrCreateEventMetadata(tf_name_scope), *plane_builder->GetOrCreateEventMetadata(tf_name_scope), offset_ps,
group_id_stat_metadata_id, group_id)); duration_ps, group_id_stat_metadata_id, group_id));
} }
tf_name_scope_line_builder->ExpandOrAddEvents(name_scope_event_per_level); 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 // Set the display name to op_type so that the events of the same op_type have
// the same color in the trace viewer. // the same color in the trace viewer.
tf_op_event_metadata->set_display_name(TfOpEventName(tf_op)); tf_op_event_metadata->set_display_name(TfOpEventName(tf_op));
tf_op_line_builder->ExpandOrAddEvent(CreateXEvent( tf_op_line_builder->ExpandOrAddEvent(
event, *tf_op_event_metadata, group_id_stat_metadata_id, group_id)); CreateXEvent(*tf_op_event_metadata, offset_ps, duration_ps,
group_id_stat_metadata_id, group_id));
} }
} // namespace
DerivedXLineBuilder::DerivedXLineBuilder( DerivedXLineBuilder::DerivedXLineBuilder(
XPlaneBuilder* plane, int64 line_id, absl::string_view name, XPlaneBuilder* plane, int64 line_id, absl::string_view name,
int64 timestamp_ns, std::vector<DerivedXLineBuilder*> dependent_lines) int64 timestamp_ns, std::vector<DerivedXLineBuilder*> dependent_lines)
@ -118,14 +110,20 @@ void DerivedXLineBuilder::ExpandOrAddLevelEvent(const XEvent& event,
last_event->SetDurationPs((offset_ps + duration_ps) - last_event->SetDurationPs((offset_ps + duration_ps) -
last_event->OffsetPs()); last_event->OffsetPs());
} else { } else {
// Otherwise, create a new event for the given level. // Otherwise, reset the last events lower than or equal to the given level.
last_event = line_.AddEvent(event);
// Reset last events lower than the given level.
ResetLastEvents(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, void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
const EventGroupNameMap& event_group_name_map, const EventGroupNameMap& event_group_name_map,
XPlane* device_trace, bool step_info_only) { XPlane* device_trace, bool step_info_only) {
@ -142,19 +140,18 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
absl::c_sort(events); absl::c_sort(events);
XPlaneBuilder plane(device_trace); XPlaneBuilder plane(device_trace);
DerivedXLineBuilder tf_ops(&plane, kThreadIdTfOp, kDerivedLineTensorFlowOps, DerivedXLineBuilder tf_ops(&plane, kThreadIdTfOp, kTensorFlowOpLineName,
start_timestamp_ns, {}); start_timestamp_ns, {});
DerivedXLineBuilder tf_name_scope(&plane, kThreadIdTfNameScope, DerivedXLineBuilder tf_name_scope(&plane, kThreadIdTfNameScope,
kDerivedLineTensorFlowNameScope, kTensorFlowNameScopeLineName,
start_timestamp_ns, {&tf_ops}); start_timestamp_ns, {&tf_ops});
DerivedXLineBuilder hlo_ops(&plane, kThreadIdHloOp, kDerivedLineXlaOps, DerivedXLineBuilder hlo_ops(&plane, kThreadIdHloOp, kXlaOpLineName,
start_timestamp_ns, {}); start_timestamp_ns, {});
DerivedXLineBuilder hlo_modules(&plane, kThreadIdHloModule, DerivedXLineBuilder hlo_modules(&plane, kThreadIdHloModule,
kDerivedLineXlaModules, start_timestamp_ns, kXlaModuleLineName, start_timestamp_ns,
{&tf_ops, &tf_name_scope, &hlo_ops}); {&tf_name_scope, &hlo_ops});
DerivedXLineBuilder steps(&plane, kThreadIdStepInfo, kDerivedLineSteps, DerivedXLineBuilder steps(&plane, kThreadIdStepInfo, kStepLineName,
start_timestamp_ns, start_timestamp_ns, {&hlo_modules});
{&tf_ops, &tf_name_scope, &hlo_ops, &hlo_modules});
int64 group_id_stat_metadata_id = int64 group_id_stat_metadata_id =
plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kGroupId))->id(); plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kGroupId))->id();
int64 step_name_stat_metadata_id = int64 step_name_stat_metadata_id =
@ -162,6 +159,8 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
// Process events in order by start time. // Process events in order by start time.
for (const XEventVisitor& event : events) { 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 tf_op_full_name;
absl::string_view hlo_module_name; absl::string_view hlo_module_name;
std::vector<absl::string_view> hlo_op_names; std::vector<absl::string_view> hlo_op_names;
@ -184,8 +183,8 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
if (group_id) { if (group_id) {
XEvent step_event = CreateXEvent( XEvent step_event = CreateXEvent(
event, *plane.GetOrCreateEventMetadata(absl::StrCat(*group_id)), *plane.GetOrCreateEventMetadata(absl::StrCat(*group_id)), offset_ps,
group_id_stat_metadata_id, group_id); duration_ps, group_id_stat_metadata_id, group_id);
if (auto group_name = gtl::FindOrNull(event_group_name_map, *group_id)) { if (auto group_name = gtl::FindOrNull(event_group_name_map, *group_id)) {
XStat* stat = step_event.add_stats(); XStat* stat = step_event.add_stats();
stat->set_metadata_id(step_name_stat_metadata_id); stat->set_metadata_id(step_name_stat_metadata_id);
@ -201,9 +200,9 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
if (!is_kernel) continue; if (!is_kernel) continue;
if (!hlo_module_name.empty()) { if (!hlo_module_name.empty()) {
hlo_modules.ExpandOrAddEvent( hlo_modules.ExpandOrAddEvent(CreateXEvent(
CreateXEvent(event, *plane.GetOrCreateEventMetadata(hlo_module_name), *plane.GetOrCreateEventMetadata(hlo_module_name), offset_ps,
group_id_stat_metadata_id, group_id)); duration_ps, group_id_stat_metadata_id, group_id));
} }
if (!hlo_op_names.empty()) { // GPU kernel compiled by XLA if (!hlo_op_names.empty()) { // GPU kernel compiled by XLA
@ -211,19 +210,19 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
std::vector<XEvent> hlo_op_event_per_level; std::vector<XEvent> hlo_op_event_per_level;
for (absl::string_view hlo_op_name : hlo_op_names) { for (absl::string_view hlo_op_name : hlo_op_names) {
DCHECK(!hlo_op_name.empty()); DCHECK(!hlo_op_name.empty());
hlo_op_event_per_level.push_back( hlo_op_event_per_level.push_back(CreateXEvent(
CreateXEvent(event, *plane.GetOrCreateEventMetadata(hlo_op_name), *plane.GetOrCreateEventMetadata(hlo_op_name), offset_ps,
group_id_stat_metadata_id, group_id)); duration_ps, group_id_stat_metadata_id, group_id));
} }
hlo_ops.ExpandOrAddEvents(hlo_op_event_per_level); hlo_ops.ExpandOrAddEvents(hlo_op_event_per_level);
auto tf_op_name = symbol_resolver(hlo_module_name, hlo_op_names.back()); auto tf_op_name = symbol_resolver(hlo_module_name, hlo_op_names.back());
if (!tf_op_name.empty()) { if (!tf_op_name.empty()) {
ProcessTfOpEvent(event, tf_op_name, group_id, &plane, &tf_name_scope, ProcessTfOpEvent(tf_op_name, offset_ps, duration_ps, group_id, &plane,
&tf_ops); &tf_name_scope, &tf_ops);
} }
} else if (!tf_op_full_name.empty()) { // GPU kernel not compiled by XLA } 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, ProcessTfOpEvent(tf_op_full_name, offset_ps, duration_ps, group_id,
&tf_ops); &plane, &tf_name_scope, &tf_ops);
} }
} }
RemoveEmptyLines(device_trace); RemoveEmptyLines(device_trace);
@ -294,7 +293,7 @@ void DeriveEventsFromHostTrace(const XPlane* host_trace,
XPlaneBuilder device_plane(device_traces[i]); XPlaneBuilder device_plane(device_traces[i]);
XLineBuilder launch_line = XLineBuilder launch_line =
device_plane.GetOrCreateLine(kThreadIdKernelLaunch); device_plane.GetOrCreateLine(kThreadIdKernelLaunch);
launch_line.SetName(kDerivedLineKernelLaunch); launch_line.SetName(kKernelLaunchLineName);
launch_line.SetTimestampNs(std::min(device_plane_start, host_plane_start)); launch_line.SetTimestampNs(std::min(device_plane_start, host_plane_start));
for (const auto& it : per_device_launch_info[i]) { for (const auto& it : per_device_launch_info[i]) {
uint64 group_id = it.first; uint64 group_id = it.first;

View File

@ -40,12 +40,8 @@ class DerivedXLineBuilder {
ExpandOrAddLevelEvent(event, /*level=*/0); ExpandOrAddLevelEvent(event, /*level=*/0);
} }
// Reset last events lower than the given level. // Reset the last events lower than or equal to the given level.
void ResetLastEvents(int level = -1) { void ResetLastEvents(int level = 0);
for (int i = level + 1; i < last_event_by_level_.size(); ++i) {
last_event_by_level_[i] = absl::nullopt;
}
}
private: private:
// If the last event of the given level has the same metadata, expands it to // 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<absl::string_view( using SymbolResolver = std::function<absl::string_view(
absl::string_view hlo_module_name, absl::string_view hlo_op)>; absl::string_view hlo_module_name, absl::string_view hlo_op)>;
// 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<int64> 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 // Derives "Step Info", "Tensorflow Ops", "XLA Ops" and "XLA Module" lines in
// an NVIDIA_GPU device trace from data passed as ScopedAnnotations and stored // an NVIDIA_GPU device trace from data passed as ScopedAnnotations and stored
// as XStats in XEvents corresponding to GPU Kernels. Consecutive annotations // as XStats in XEvents corresponding to GPU Kernels. Consecutive annotations

View File

@ -28,6 +28,13 @@ const absl::string_view kCuptiDriverApiPlaneName = "/host:CUPTI";
const absl::string_view kMetadataPlane = "/host:metadata"; const absl::string_view kMetadataPlane = "/host:metadata";
const absl::string_view kTFStreamzPlane = "/host:tfstreamz"; 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 kHostPlaneId = 49;
const int32 kGpuPlaneBaseId = 0; const int32 kGpuPlaneBaseId = 0;
const int32 kCuptiDriverApiPlaneId = 50; const int32 kCuptiDriverApiPlaneId = 50;

View File

@ -36,6 +36,14 @@ ABSL_CONST_INIT extern const absl::string_view kMetadataPlane;
// Name of XPlane that contains kpi related metrics. // Name of XPlane that contains kpi related metrics.
ABSL_CONST_INIT extern const absl::string_view kTFStreamzPlane; 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. // Id of XPlane that contains TraceMe events.
ABSL_CONST_INIT extern const int32 kHostPlaneId; ABSL_CONST_INIT extern const int32 kHostPlaneId;
// Ids prefix of XPlane that contains GPU events. // Ids prefix of XPlane that contains GPU events.