Refactor the derived_timeline files.
PiperOrigin-RevId: 308348454 Change-Id: Iad035e375d87ded82fcd11eb5f7f38c17b797381
This commit is contained in:
parent
ceee848cea
commit
a4e22aa1c4
@ -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;
|
||||
|
@ -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
|
||||
|
@ -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;
|
||||
|
@ -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.
|
||||
|
Loading…
Reference in New Issue
Block a user