experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

XPlane schema cleanup

Browse Source 
PiperOrigin-RevId: 316564434
Change-Id: Icabe3400c82cb4aff84c2bdae7ecf59424b36fa4
This commit is contained in:
Jose Baiocchi 2020-06-15 16:04:21 -07:00 committed by TensorFlower Gardener
parent d29d8af754
commit f4e20ec5ae
26 changed files with 117 additions and 112 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/profiler/convert/BUILD
View File

@ -47,6 +47,7 @@ tf_cc_test(
"//tensorflow/core/profiler/utils:time_utils",
"//tensorflow/core/profiler/utils:xplane_builder",
"//tensorflow/core/profiler/utils:xplane_schema",
"//tensorflow/core/profiler/utils:xplane_test_utils",
"@com_google_absl//absl/strings",
],
)
@ -171,6 +172,7 @@ tf_cc_test(
"//tensorflow/core/profiler/utils:time_utils",
"//tensorflow/core/profiler/utils:xplane_builder",
"//tensorflow/core/profiler/utils:xplane_schema",
"//tensorflow/core/profiler/utils:xplane_test_utils",
"@com_google_absl//absl/strings",
],
)

5
tensorflow/core/profiler/convert/op_stats_to_tf_stats_test.cc
View File

@ -26,6 +26,7 @@ limitations under the License.
#include "tensorflow/core/profiler/utils/time_utils.h"
#include "tensorflow/core/profiler/utils/xplane_builder.h"
#include "tensorflow/core/profiler/utils/xplane_schema.h"
#include "tensorflow/core/profiler/utils/xplane_test_utils.h"
namespace tensorflow {
namespace profiler {
@ -59,8 +60,8 @@ TEST(OpStatsToTfStats, GpuTfStats) {
constexpr int64 kKernel3DurationNs = 10000;
XSpace space;
XPlaneBuilder device_plane(space.add_planes());
device_plane.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
XPlaneBuilder device_plane(
GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
XLineBuilder stream1 = device_plane.GetOrCreateLine(/*line_id=*/10);
AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kKernel1StartNs,
kKernel1DurationNs, /*on_device=*/true, kKernel1,

3
tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc
View File

@ -32,9 +32,8 @@ namespace {
// activities within one memory allocator captured in host trace.
TEST(ConvertXPlaneToMemoryProfile, OneAllocatorMultiActivitiesTest) {
XSpace space;
XPlane* host_plane = space.add_planes();
XPlane* host_plane = GetOrCreateHostXPlane(&space);
XPlaneBuilder host_plane_builder(host_plane);
host_plane_builder.SetName(kHostThreads);
host_plane_builder.ReserveLines(1);
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(0);

24
tensorflow/core/profiler/convert/xplane_to_op_metrics_db_test.cc
View File

@ -25,6 +25,7 @@ limitations under the License.
#include "tensorflow/core/profiler/utils/time_utils.h"
#include "tensorflow/core/profiler/utils/xplane_builder.h"
#include "tensorflow/core/profiler/utils/xplane_schema.h"
#include "tensorflow/core/profiler/utils/xplane_test_utils.h"
namespace tensorflow {
namespace profiler {
@ -43,12 +44,6 @@ void AddTensorFlowOpEvent(absl::string_view tf_op_fullname,
tf_op_fullname);
}
void SetXPlaneNameAndId(absl::string_view name, int64 id,
XPlaneBuilder* plane) {
plane->SetName(name);
plane->SetId(id);
}
TEST(ConvertXPlaneToOpMetricsDb, HostOpMetricsDb) {
static constexpr char kTfOp1[] = "TfOp1";
static constexpr char kTfOp2[] = "TfOp2";
@ -57,9 +52,9 @@ TEST(ConvertXPlaneToOpMetricsDb, HostOpMetricsDb) {
constexpr int64 kTfOp2StartNs = 110000;
constexpr int64 kTfOp2DurationNs = 10000;
XPlane xplane;
XPlaneBuilder host_plane(&xplane);
SetXPlaneNameAndId(kHostThreads, /*id=*/0, &host_plane);
XSpace xspace;
XPlane* xplane = GetOrCreateHostXPlane(&xspace);
XPlaneBuilder host_plane(xplane);
XLineBuilder thread1 = host_plane.GetOrCreateLine(/*line_id=*/10);
AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kTfOp1StartNs,
kTfOp1DurationNs, /*on_device=*/false,
@ -72,7 +67,7 @@ TEST(ConvertXPlaneToOpMetricsDb, HostOpMetricsDb) {
kTfOp2DurationNs, /*on_device=*/false,
/*kernel_name=*/"", &host_plane, &thread2);
OpMetricsDb op_metrics = ConvertHostThreadsXPlaneToOpMetricsDb(xplane);
OpMetricsDb op_metrics = ConvertHostThreadsXPlaneToOpMetricsDb(*xplane);
// Op1, Op2, Idle.
EXPECT_EQ(3, op_metrics.metrics_db_size());
uint64 total_op_duration =
@ -115,10 +110,9 @@ TEST(ConvertXPlaneToOpMetricsDb, DeviceOpMetricsDb) {
constexpr int64 kKernel3StartNs = 120000;
constexpr int64 kKernel3DurationNs = 10000;
XPlane xplane;
XPlaneBuilder device_plane(&xplane);
SetXPlaneNameAndId(absl::StrCat(kGpuPlanePrefix, ":0"), /*id=*/1,
&device_plane);
XSpace xspace;
XPlane* xplane = GetOrCreateGpuXPlane(&xspace, /*device_ordinal=*/0);
XPlaneBuilder device_plane(xplane);
XLineBuilder stream1 = device_plane.GetOrCreateLine(/*line_id=*/10);
AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kKernel1StartNs,
kKernel1DurationNs, /*on_device=*/true, kKernel1,
@ -138,7 +132,7 @@ TEST(ConvertXPlaneToOpMetricsDb, DeviceOpMetricsDb) {
&device_plane, &stream2);
OpMetricsDb op_metrics = ConvertDeviceTraceXPlaneToOpMetricsDb(
xplane, /*peak_tera_flops_per_second=*/0,
*xplane, /*peak_tera_flops_per_second=*/0,
/*peak_hbm_bw_giga_bytes_per_second=*/0);
// kernel1, kernel2, kernel3, Idle.

2
tensorflow/core/profiler/convert/xplane_to_op_stats.cc
View File

@ -129,7 +129,7 @@ void PropagateXSpaceDiagnosticsToOpStats(const XSpace& space,
}
OpStats ConvertXSpaceToOpStats(const XSpace& space) {
const XPlane* host_plane = FindPlaneWithName(space, kHostThreads);
const XPlane* host_plane = FindPlaneWithName(space, kHostThreadsPlaneName);
std::vector<const XPlane*> device_planes =
FindPlanesWithPrefix(space, kGpuPlanePrefix);
OpStats op_stats;

22
tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc
View File

@ -43,8 +43,8 @@ TEST(ConvertXPlaneToOpStats, PerfEnv) {
constexpr int kComputeCapMajor = 7;
constexpr int kComputeCapMinor = 0;
XPlaneBuilder device_plane(space.add_planes());
device_plane.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
XPlaneBuilder device_plane(
GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
device_plane.ParseAndAddStatValue(
*device_plane.GetOrCreateStatMetadata("clock_rate"),
absl::StrCat(kClockRateKHz));
@ -71,10 +71,10 @@ TEST(ConvertXPlaneToOpStats, PerfEnv) {
TEST(ConvertXPlaneToOpStats, RunEnvironment) {
XSpace space;
XPlaneBuilder device_plane1(space.add_planes());
device_plane1.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
XPlaneBuilder device_plane2(space.add_planes());
device_plane2.SetName(absl::StrCat(kGpuPlanePrefix, ":1"));
XPlaneBuilder device_plane1(
GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
XPlaneBuilder device_plane2(
GetOrCreateGpuXPlane(&space, /*device_ordinal=*/1));
GroupTfEvents(&space, /*event_group_name_map=*/nullptr);
OpStats op_stats = ConvertXSpaceToOpStats(space);
@ -91,8 +91,7 @@ TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
constexpr int64 kStepId = 0;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
host_plane_builder.ReserveLines(2);
auto main_thread = host_plane_builder.GetOrCreateLine(0);
@ -120,8 +119,7 @@ TEST(ConvertXPlaneToOpStats, GpuStepDbTest) {
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
host_plane_builder.ReserveLines(2);
auto main_thread = host_plane_builder.GetOrCreateLine(0);
@ -137,8 +135,8 @@ TEST(ConvertXPlaneToOpStats, GpuStepDbTest) {
CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 30, 10,
{{StatType::kCorrelationId, kCorrelationId}});
XPlaneBuilder device_plane_builder(space.add_planes());
device_plane_builder.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
XPlaneBuilder device_plane_builder(
GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
device_plane_builder.ReserveLines(1);
auto stream = device_plane_builder.GetOrCreateLine(0);

3
tensorflow/core/profiler/convert/xplane_to_profile_response.cc
View File

@ -139,7 +139,8 @@ Status ConvertXSpaceToProfileResponse(const XSpace& xspace,
AddToolData(ToolName(kKernelStats), op_stats.kernel_stats_db(), response);
}
if (tools.contains(kMemoryProfile)) {
if (const XPlane* host_plane = FindPlaneWithName(xspace, kHostThreads)) {
if (const XPlane* host_plane =
FindPlaneWithName(xspace, kHostThreadsPlaneName)) {
MemoryProfile memory_profile = ConvertXPlaneToMemoryProfile(*host_plane);
std::string json_output;
TF_RETURN_IF_ERROR(ConvertProtoToJson(memory_profile, &json_output));

3
tensorflow/core/profiler/convert/xplane_to_step_events_test.cc
View File

@ -45,9 +45,8 @@ TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
constexpr int64 kSecondCorrelationId = 200;
XSpace space;
XPlane* host_plane = space.add_planes();
XPlane* host_plane = GetOrCreateHostXPlane(&space);
XPlaneBuilder host_plane_builder(host_plane);
host_plane_builder.SetName(kHostThreads);
host_plane_builder.ReserveLines(2);
auto main_thread = host_plane_builder.GetOrCreateLine(0);

12
tensorflow/core/profiler/convert/xplane_to_tf_functions_test.cc
View File

@ -43,11 +43,12 @@ constexpr double kMaxError = 0.001;
TfFunctionDb ConvertXSpaceToTfFunctionDb(const XSpace& space) {
TfFunctionDb result;
const XPlane* host_plane = FindPlaneWithName(space, kHostThreads);
const XPlane* host_plane = FindPlaneWithName(space, kHostThreadsPlaneName);
if (host_plane) {
XPlaneVisitor plane = CreateTfXPlaneVisitor(host_plane);
plane.ForEachLine([&result](const XLineVisitor& line) {
CombineTfFunctionDb(ConvertHostThreadsXLineToTfFunctionDb(line), &result);
TfFunctionDb tf_function_db = ConvertHostThreadsXLineToTfFunctionDb(line);
CombineTfFunctionDb(tf_function_db, &result);
});
}
return result;
@ -56,7 +57,7 @@ TfFunctionDb ConvertXSpaceToTfFunctionDb(const XSpace& space) {
TEST(ConvertXPlaneToTfFunctions, CombineTwoThreads) {
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
host_plane_builder.SetName(kHostThreadsPlaneName);
host_plane_builder.ReserveLines(2);
std::string kFunctionName = "decrement";
@ -100,7 +101,7 @@ TEST(ConvertXPlaneToTfFunctions, CombineTwoThreads) {
TEST(ConvertXPlaneToTfFunctions, NestedFunctions) {
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
host_plane_builder.SetName(kHostThreadsPlaneName);
host_plane_builder.ReserveLines(1);
std::string kOuterFunctionName = "outer";
std::string kInnerFunctionName = "inner";
@ -140,8 +141,7 @@ TEST(ConvertXPlaneToTfFunctions, NestedFunctions) {
TEST(ConvertXPlaneToTfFunctions, EagerPlusConcrete) {
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
host_plane_builder.ReserveLines(2);
std::string kEagerFunctionName = "i_am_eager";
std::string kConcreteFunctionName = "i_am_concrete";

2
tensorflow/core/profiler/convert/xplane_to_trace_events.cc
View File

@ -41,7 +41,7 @@ Device BuildDeviceAndResource(const XPlaneVisitor& plane) {
device.set_name(std::string(plane.Name()));
device.set_device_id(plane.Id());
bool sort_by_ordinal = plane.Name() == kHostThreads;
bool sort_by_ordinal = (plane.Name() == kHostThreadsPlaneName);
int ordinal = 0;
plane.ForEachLine([&](const XLineVisitor& line) {
Resource resource;

2
tensorflow/core/profiler/internal/cpu/host_tracer.cc
View File

@ -146,7 +146,7 @@ Status HostTracer::CollectData(XSpace* space) {
return errors::Internal("TraceMeRecorder not stopped");
}
MakeCompleteEvents(&events_);
XPlane* plane = GetOrCreatePlane(space, kHostThreads);
XPlane* plane = FindOrAddMutablePlaneWithName(space, kHostThreadsPlaneName);
plane->set_id(kHostPlaneId);
ConvertCompleteEventsToXPlane(start_timestamp_ns_, events_, plane);
events_.clear();

2
tensorflow/core/profiler/internal/cpu/host_tracer_test.cc
View File

@ -150,7 +150,7 @@ TEST(HostTracerTest, CollectsTraceMeEventsAsXSpace) {
ASSERT_EQ(space.planes_size(), 1);
const auto& plane = space.planes(0);
XPlaneVisitor xplane(&plane);
ASSERT_EQ(plane.name(), kHostThreads);
ASSERT_EQ(plane.name(), kHostThreadsPlaneName);
ASSERT_EQ(plane.lines_size(), 1);
ASSERT_EQ(plane.event_metadata_size(), 7);
ASSERT_EQ(plane.stat_metadata_size(), 2);

2
tensorflow/core/profiler/internal/cpu/metadata_collector.cc
View File

@ -65,7 +65,7 @@ class MetadataCollector : public ProfilerInterface {
Status CollectData(XSpace* space) override {
if (!debug_info_.empty()) {
XPlane* plane = GetOrCreatePlane(space, kMetadataPlane);
XPlane* plane = FindOrAddMutablePlaneWithName(space, kMetadataPlaneName);
plane->set_id(kMetadataPlaneId);
XPlaneBuilder xplane(plane);
const XStatMetadata& hlo_proto_stat =

7
tensorflow/core/profiler/internal/gpu/device_tracer.cc
View File

@ -223,11 +223,12 @@ class CuptiTraceCollectorImpl : public CuptiTraceCollector {
<< " callback api events and " << num_activity_events_
<< " activity events. " << ReportDroppedEvents();
uint64 end_gpu_ns = CuptiTracer::GetTimestamp();
XPlaneBuilder host_plane(GetOrCreatePlane(space, kCuptiDriverApiPlaneName));
XPlaneBuilder host_plane(
FindOrAddMutablePlaneWithName(space, kCuptiDriverApiPlaneName));
host_plane.SetId(kCuptiDriverApiPlaneId);
for (int device_ordinal = 0; device_ordinal < num_gpus_; ++device_ordinal) {
std::string name = absl::StrCat(kGpuPlanePrefix, device_ordinal);
XPlaneBuilder device_plane(GetOrCreatePlane(space, name));
std::string name = GpuPlaneName(device_ordinal);
XPlaneBuilder device_plane(FindOrAddMutablePlaneWithName(space, name));
device_plane.SetId(kGpuPlaneBaseId + device_ordinal);
per_device_collector_[device_ordinal].Flush(start_gpu_ns_, end_gpu_ns,
&device_plane, &host_plane);

5
tensorflow/core/profiler/lib/profiler_session.cc
View File

@ -98,11 +98,10 @@ Status ProfilerSession::CollectData(profiler::XSpace* space) {
const profiler::XPlane* cupti_driver_api_plane =
profiler::FindPlaneWithName(*space, profiler::kCuptiDriverApiPlaneName);
if (cupti_driver_api_plane) {
profiler::XPlane* host_plane =
profiler::GetOrCreatePlane(space, profiler::kHostThreads);
profiler::XPlane* host_plane = profiler::FindOrAddMutablePlaneWithName(
space, profiler::kHostThreadsPlaneName);
profiler::MergePlanes(*cupti_driver_api_plane, host_plane);
profiler::SortXLinesBy(host_plane, profiler::XLinesComparatorByName());
// This might invalidate host_plane pointer.
profiler::RemovePlaneWithName(space, profiler::kCuptiDriverApiPlaneName);
}
// 2. Normalize all timestamps by shifting timeline to profiling start time.

1
tensorflow/core/profiler/utils/BUILD
View File

@ -233,6 +233,7 @@ cc_library(
deps = [
":xplane_builder",
":xplane_schema",
":xplane_utils",
"//tensorflow/core/platform:types",
"//tensorflow/core/profiler/protobuf:xplane_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",

7
tensorflow/core/profiler/utils/derived_timeline.cc
View File

@ -339,9 +339,10 @@ void GenerateDerivedTimeLines(const EventGroupNameMap& event_group_name_map,
XSpace* space, bool step_info_only) {
for (XPlane& plane : *space->mutable_planes()) {
// Derived timelines only generated for device traces.
if (plane.id() == kHostPlaneId) continue;
DeriveEventsFromAnnotations(DummySymbolResolver, event_group_name_map,
&plane, step_info_only);
if (IsGpuPlaneName(plane.name())) {
DeriveEventsFromAnnotations(DummySymbolResolver, event_group_name_map,
&plane, step_info_only);
}
}
}

8
tensorflow/core/profiler/utils/derived_timeline_test.cc
View File

@ -44,7 +44,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
EventGroupNameMap event_group_name_map;
XPlane* plane = space.add_planes();
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
auto line_builder = plane_builder.GetOrCreateLine(0);
CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,
@ -74,7 +74,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
EventGroupNameMap event_group_name_map;
XPlane* plane = space.add_planes();
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
auto line_builder = plane_builder.GetOrCreateLine(0);
CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,
@ -109,7 +109,7 @@ TEST(DerivedTimelineTest, DependencyTest) {
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
EventGroupNameMap event_group_name_map({{0, "train 0"}, {1, "train 1"}});
XPlane* plane = space.add_planes();
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
auto line_builder = plane_builder.GetOrCreateLine(0);
CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,
@ -138,7 +138,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeTest) {
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
EventGroupNameMap event_group_name_map;
XPlane* plane = space.add_planes();
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
auto line_builder = plane_builder.GetOrCreateLine(0);
CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,

19
tensorflow/core/profiler/utils/group_events.cc
View File

@ -54,13 +54,13 @@ void CreateStatMetadata(XPlane* plane) {
}
// Returns event type if it is a KernelLaunch or KernelExecute event.
absl::optional<int64> GetKernelEventType(const XPlaneVisitor& visitor,
absl::optional<int64> GetKernelEventType(bool is_host_plane,
const XPlaneVisitor& visitor,
const XEvent& event) {
for (const auto& stat : event.stats()) {
if (visitor.GetStatType(stat) == StatType::kCorrelationId) {
// TODO(b/149095099): avoid string comparison.
return visitor.Name() == kHostThreads ? HostEventType::kKernelLaunch
: HostEventType::kKernelExecute;
return is_host_plane ? HostEventType::kKernelLaunch
: HostEventType::kKernelExecute;
}
}
return absl::nullopt;
@ -72,14 +72,15 @@ bool IsTfOpEvent(const XPlaneVisitor& visitor, const XEvent& event) {
return tf_op.category == Category::kTensorFlow;
}
int64 GetEventType(const XPlaneVisitor& visitor, const XEvent& event) {
int64 GetEventType(bool is_host_plane, const XPlaneVisitor& visitor,
const XEvent& event) {
if (absl::optional<int64> event_type = visitor.GetEventType(event)) {
return *event_type;
} else if (absl::optional<int64> kernel_event_type =
GetKernelEventType(visitor, event)) {
GetKernelEventType(is_host_plane, visitor, event)) {
// KernelLaunch and KernelExecute event types are not supported by
// XPlaneVisitor and should be checked separately.
// TODO(148346217): Make XPlaneVisitor support KernelLaunch and
// TODO(b/148346217): Make XPlaneVisitor support KernelLaunch and
// KernelExecute event types.
return *kernel_event_type;
} else if (IsTfOpEvent(visitor, event)) {
@ -396,6 +397,8 @@ bool EventNode::StartsBefore(const EventNode& other) const {
void EventForest::ConnectIntraThread(const XPlaneVisitor& visitor,
XPlane* plane,
ContextGroupMap* context_groups) {
// TODO(b/149095099): avoid string comparison.
bool is_host_plane = (visitor.Name() == kHostThreadsPlaneName);
for (auto& line : *plane->mutable_lines()) {
std::vector<EventNode*> parent_nodes;
for (auto& event : *line.mutable_events()) {
@ -418,7 +421,7 @@ void EventForest::ConnectIntraThread(const XPlaneVisitor& visitor,
}
parent_nodes.push_back(cur_node.get());
// event_node_map_ keeps cur_node alive.
event_node_map_[GetEventType(visitor, event)].push_back(
event_node_map_[GetEventType(is_host_plane, visitor, event)].push_back(
std::move(cur_node));
}
}

18
tensorflow/core/profiler/utils/group_events_test.cc
View File

@ -36,8 +36,7 @@ TEST(GroupEventsTest, GroupGpuTraceTest) {
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
host_plane_builder.ReserveLines(2);
auto main_thread = host_plane_builder.GetOrCreateLine(0);
@ -78,8 +77,7 @@ TEST(GroupEventsTest, GroupTensorFlowLoopTest) {
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
host_plane_builder.ReserveLines(1);
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
@ -125,8 +123,7 @@ TEST(GroupEventsTest, GroupMultipleTensorFlowLoopsTest) {
constexpr int64 kSecondIterNumStart = 0;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
host_plane_builder.ReserveLines(2);
auto first_tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
@ -163,9 +160,8 @@ TEST(GroupEventsTest, GroupFunctionalOp) {
constexpr int64 kFunctionStepId = 1;
XSpace space;
XPlane* host_plane = space.add_planes();
XPlane* host_plane = GetOrCreateHostXPlane(&space);
XPlaneBuilder host_plane_builder(host_plane);
host_plane_builder.SetName(kHostThreads);
host_plane_builder.ReserveLines(2);
auto main_thread = host_plane_builder.GetOrCreateLine(0);
@ -209,9 +205,8 @@ TEST(GroupEventsTest, EagerOpTest) {
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlane* host_plane = space.add_planes();
XPlane* host_plane = GetOrCreateHostXPlane(&space);
XPlaneBuilder host_plane_builder(host_plane);
host_plane_builder.SetName(kHostThreads);
host_plane_builder.ReserveLines(1);
auto main_thread = host_plane_builder.GetOrCreateLine(0);
@ -255,9 +250,8 @@ TEST(GroupEventsTest, FunctionOpTest) {
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlane* host_plane = space.add_planes();
XPlane* host_plane = GetOrCreateHostXPlane(&space);
XPlaneBuilder host_plane_builder(host_plane);
host_plane_builder.SetName(kHostThreads);
host_plane_builder.ReserveLines(2);
auto main_thread = host_plane_builder.GetOrCreateLine(0);

6
tensorflow/core/profiler/utils/xplane_schema.cc
View File

@ -26,11 +26,11 @@ limitations under the License.
namespace tensorflow {
namespace profiler {
const absl::string_view kHostThreads = "/host:CPU";
const absl::string_view kHostThreadsPlaneName = "/host:CPU";
const absl::string_view kGpuPlanePrefix = "/device:GPU:";
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 kMetadataPlaneName = "/host:metadata";
const absl::string_view kTFStreamzPlaneName = "/host:tfstreamz";
const absl::string_view kStepLineName = "Steps";
const absl::string_view kTensorFlowNameScopeLineName = "TensorFlow Name Scope";

16
tensorflow/core/profiler/utils/xplane_schema.h
View File

@ -16,6 +16,8 @@ limitations under the License.
#ifndef TENSORFLOW_CORE_PROFILER_UTILS_XPLANE_SCHEMA_H_
#define TENSORFLOW_CORE_PROFILER_UTILS_XPLANE_SCHEMA_H_
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "tensorflow/core/platform/logging.h"
@ -25,15 +27,15 @@ namespace tensorflow {
namespace profiler {
// Name of XPlane that contains TraceMe events.
ABSL_CONST_INIT extern const absl::string_view kHostThreads;
ABSL_CONST_INIT extern const absl::string_view kHostThreadsPlaneName;
// Name prefix of XPlane that contains GPU events.
ABSL_CONST_INIT extern const absl::string_view kGpuPlanePrefix;
// Name of XPlane that contains CUPTI driver API generated events.
ABSL_CONST_INIT extern const absl::string_view kCuptiDriverApiPlaneName;
// Name of XPlane that contains profile metadata such as XLA debug info.
ABSL_CONST_INIT extern const absl::string_view kMetadataPlane;
ABSL_CONST_INIT extern const absl::string_view kMetadataPlaneName;
// 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 kTFStreamzPlaneName;
// Names of XLines that contain ML-level events.
ABSL_CONST_INIT extern const absl::string_view kStepLineName;
@ -184,6 +186,14 @@ enum StatType {
kLastStatType = kDevCapComputeCapMinor,
};
inline std::string GpuPlaneName(int32 device_ordinal) {
return absl::StrCat(kGpuPlanePrefix, device_ordinal);
}
inline bool IsGpuPlaneName(absl::string_view plane_name) {
return absl::StartsWith(plane_name, kGpuPlanePrefix);
}
absl::string_view GetHostEventTypeStr(HostEventType event_type);
bool IsHostEventType(HostEventType event_type, absl::string_view event_name);

10
tensorflow/core/profiler/utils/xplane_test_utils.cc
View File

@ -22,6 +22,7 @@ limitations under the License.
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
#include "tensorflow/core/profiler/utils/xplane_builder.h"
#include "tensorflow/core/profiler/utils/xplane_schema.h"
#include "tensorflow/core/profiler/utils/xplane_utils.h"
namespace tensorflow {
namespace profiler {
@ -44,6 +45,15 @@ class XStatValueVisitor {
} // namespace
XPlane* GetOrCreateHostXPlane(XSpace* space) {
return FindOrAddMutablePlaneWithName(space, kHostThreadsPlaneName);
}
XPlane* GetOrCreateGpuXPlane(XSpace* space, int32 device_ordinal) {
std::string name = GpuPlaneName(device_ordinal);
return FindOrAddMutablePlaneWithName(space, name);
}
void CreateXEvent(
XPlaneBuilder* plane_builder, XLineBuilder* line_builder,
absl::string_view event_name, int64 offset_ps, int64 duration_ps,

4
tensorflow/core/profiler/utils/xplane_test_utils.h
View File

@ -28,6 +28,10 @@ namespace profiler {
using XStatValue = absl::variant<int64, uint64, absl::string_view>;
XPlane* GetOrCreateHostXPlane(XSpace* space);
XPlane* GetOrCreateGpuXPlane(XSpace* space, int32 device_ordinal);
void CreateXEvent(
XPlaneBuilder* plane_builder, XLineBuilder* line_builder,
absl::string_view event_name, int64 offset_ps, int64 duration_ps,

29
tensorflow/core/profiler/utils/xplane_utils.cc
View File

@ -65,12 +65,19 @@ std::vector<const XPlane*> FindPlanesWithPrefix(const XSpace& space,
return result;
}
XPlane* GetOrCreatePlane(XSpace* space, absl::string_view name) {
XPlane* FindMutablePlaneWithName(XSpace* space, absl::string_view name) {
for (XPlane& plane : *space->mutable_planes()) {
if (plane.name() == name) return &plane;
}
XPlane* plane = space->add_planes();
plane->set_name(std::string(name));
return nullptr;
}
XPlane* FindOrAddMutablePlaneWithName(XSpace* space, absl::string_view name) {
XPlane* plane = FindMutablePlaneWithName(space, name);
if (plane == nullptr) {
plane = space->add_planes();
plane->set_name(name.data(), name.size());
}
return plane;
}
@ -128,22 +135,6 @@ void RemoveEmptyLines(XPlane* plane) {
lines->end());
}
XPlane* FindMutablePlaneWithName(XSpace* space, absl::string_view name) {
for (XPlane& plane : *space->mutable_planes()) {
if (plane.name() == name) return &plane;
}
return nullptr;
}
XPlane* FindOrAddMutablePlaneWithName(XSpace* space, absl::string_view name) {
XPlane* plane = FindMutablePlaneWithName(space, name);
if (plane == nullptr) {
plane = space->add_planes();
plane->set_name(std::string(name));
}
return plane;
}
void SortXPlane(XPlane* plane) {
for (XLine& line : *plane->mutable_lines()) {
auto& events = *line.mutable_events();

15
tensorflow/core/profiler/utils/xplane_utils.h
View File

@ -32,8 +32,12 @@ const XPlane* FindPlaneWithName(const XSpace& space, absl::string_view name);
std::vector<const XPlane*> FindPlanesWithPrefix(const XSpace& space,
absl::string_view prefix);
// Returns the plane with the given name, create it if necessary.
XPlane* GetOrCreatePlane(XSpace* space, absl::string_view name);
// Returns the plane with the given name in the container or null if not found.
XPlane* FindMutablePlaneWithName(XSpace* space, absl::string_view name);
// Returns the plane with the given name in the container. If necessary, adds a
// new plane to the container.
XPlane* FindOrAddMutablePlaneWithName(XSpace* space, absl::string_view name);
// Returns true if event is nested by parent.
bool IsNested(const tensorflow::profiler::XEvent& event,
@ -49,13 +53,6 @@ void RemovePlaneWithName(XSpace* space, absl::string_view name);
void RemoveEmptyPlanes(XSpace* space);
void RemoveEmptyLines(XPlane* plane);
// Returns the plane with the given name in the container or null if not found.
XPlane* FindMutablePlaneWithName(XSpace* space, absl::string_view name);
// Returns the plane with the given name in the container. If necessary, adds a
// new plane to the container.
XPlane* FindOrAddMutablePlaneWithName(XSpace* space, absl::string_view name);
// Sort lines in plane with a provided comparator.
template <class Compare>
void SortXLinesBy(XPlane* plane, Compare comp) {