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();
}
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<int64> 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<int64> group_id,
} // namespace
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) {
@ -79,8 +72,8 @@ void ProcessTfOpEvent(const XEventVisitor& event,
std::vector<XEvent> 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<DerivedXLineBuilder*> 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<absl::string_view> 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<XEvent> 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;

View File

@ -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<absl::string_view(
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
// an NVIDIA_GPU device trace from data passed as ScopedAnnotations and stored
// 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 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;

View File

@ -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.