Read producer/consumer ids as uint64. Also, add uint64 to XStatValue to support test cases with uint64 stats and add IntOrUintValue() to XStatVisitor.

PiperOrigin-RevId: 315414571
Change-Id: I61f98161488c39df8557aa24c7f4f94e50a708f7
This commit is contained in:
Jiho Choi 2020-06-08 21:25:10 -07:00 committed by TensorFlower Gardener
parent 1a90749db9
commit 590a27e7db
8 changed files with 236 additions and 94 deletions

View File

@ -40,15 +40,15 @@ TEST(ConvertXPlaneToMemoryProfile, OneAllocatorMultiActivitiesTest) {
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &tf_executor_thread, "MemoryAllocation",
40000, 1000,
{{StatType::kBytesReserved, 2000},
{StatType::kBytesAllocated, 3000},
{StatType::kBytesAvailable, 5000},
{StatType::kPeakBytesInUse, 8500},
{StatType::kRequestedBytes, 200},
{StatType::kAllocationBytes, 256},
{StatType::kAddress, 222333},
{StatType::kStepId, -93746},
{StatType::kDataType, 1},
{{StatType::kBytesReserved, 2000LL},
{StatType::kBytesAllocated, 3000LL},
{StatType::kBytesAvailable, 5000LL},
{StatType::kPeakBytesInUse, 8500LL},
{StatType::kRequestedBytes, 200LL},
{StatType::kAllocationBytes, 256LL},
{StatType::kAddress, 222333LL},
{StatType::kStepId, -93746LL},
{StatType::kDataType, 1LL},
{StatType::kAllocatorName, "GPU_0_bfc"},
{StatType::kTfOp, "foo/bar"},
{StatType::kRegionType, "output"},
@ -56,30 +56,30 @@ TEST(ConvertXPlaneToMemoryProfile, OneAllocatorMultiActivitiesTest) {
CreateXEvent(&host_plane_builder, &tf_executor_thread, "MemoryDeallocation",
50000, 1000,
{{StatType::kBytesReserved, 2000},
{StatType::kBytesAllocated, 2744},
{StatType::kBytesAvailable, 5256},
{StatType::kPeakBytesInUse, 8500},
{StatType::kRequestedBytes, 200},
{StatType::kAllocationBytes, 256},
{StatType::kAddress, 222333},
{StatType::kStepId, 0},
{StatType::kDataType, 0},
{{StatType::kBytesReserved, 2000LL},
{StatType::kBytesAllocated, 2744LL},
{StatType::kBytesAvailable, 5256LL},
{StatType::kPeakBytesInUse, 8500LL},
{StatType::kRequestedBytes, 200LL},
{StatType::kAllocationBytes, 256LL},
{StatType::kAddress, 222333LL},
{StatType::kStepId, 0LL},
{StatType::kDataType, 0LL},
{StatType::kAllocatorName, "GPU_0_bfc"},
{StatType::kRegionType, ""},
{StatType::kTensorShapes, ""}});
CreateXEvent(&host_plane_builder, &tf_executor_thread, "MemoryAllocation",
70000, 1000,
{{StatType::kBytesReserved, 2000},
{StatType::kBytesAllocated, 5000},
{StatType::kBytesAvailable, 3000},
{StatType::kPeakBytesInUse, 9500},
{StatType::kRequestedBytes, 300},
{StatType::kAllocationBytes, 300},
{StatType::kAddress, 345678},
{StatType::kStepId, -93746},
{StatType::kDataType, 9},
{{StatType::kBytesReserved, 2000LL},
{StatType::kBytesAllocated, 5000LL},
{StatType::kBytesAvailable, 3000LL},
{StatType::kPeakBytesInUse, 9500LL},
{StatType::kRequestedBytes, 300LL},
{StatType::kAllocationBytes, 300LL},
{StatType::kAddress, 345678LL},
{StatType::kStepId, -93746LL},
{StatType::kDataType, 9LL},
{StatType::kAllocatorName, "GPU_0_bfc"},
{StatType::kTfOp, "mul_grad/Sum"},
{StatType::kRegionType, "temp"},

View File

@ -86,6 +86,9 @@ TEST(ConvertXPlaneToOpStats, RunEnvironment) {
}
TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
constexpr int64 kStepNum = 123;
constexpr int64 kStepId = 0;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
@ -93,14 +96,14 @@ TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
auto main_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kTraceContext,
0, 100, {{StatType::kStepNum, 123}});
0, 100, {{StatType::kStepNum, kStepNum}});
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kFunctionRun,
10, 90, {{StatType::kStepId, 0}});
10, 90, {{StatType::kStepId, kStepId}});
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(1);
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 20, 80,
{{StatType::kStepId, 0}});
{{StatType::kStepId, kStepId}});
CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 30, 70);
GroupTfEvents(&space, /*event_group_name_map=*/nullptr);
@ -111,6 +114,10 @@ TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
}
TEST(ConvertXPlaneToOpStats, GpuStepDbTest) {
constexpr int64 kStepNum = 123;
constexpr int64 kStepId = 0;
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
@ -118,16 +125,16 @@ TEST(ConvertXPlaneToOpStats, GpuStepDbTest) {
auto main_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kTraceContext,
0, 100, {{StatType::kStepNum, 123}});
0, 100, {{StatType::kStepNum, kStepNum}});
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kFunctionRun,
10, 90, {{StatType::kStepId, 0}});
10, 90, {{StatType::kStepId, kStepId}});
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(1);
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 20, 20,
{{StatType::kStepId, 0}});
{{StatType::kStepId, kStepId}});
CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 30, 10,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
XPlaneBuilder device_plane_builder(space.add_planes());
device_plane_builder.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
@ -135,7 +142,7 @@ TEST(ConvertXPlaneToOpStats, GpuStepDbTest) {
auto stream = device_plane_builder.GetOrCreateLine(0);
CreateXEvent(&device_plane_builder, &stream, "matmul", 50, 40,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
GroupTfEvents(&space, /*event_group_name_map=*/nullptr);
OpStats op_stats = ConvertXSpaceToOpStats(space);

View File

@ -37,6 +37,13 @@ namespace {
// consists of matmul. The host's step db should be created only for the step
// observed on the host.
TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
constexpr int64 kFirstStepNum = 123;
constexpr int64 kSecondStepNum = 456;
constexpr int64 kFirstStepId = 0;
constexpr int64 kSecondStepId = 1;
constexpr int64 kFirstCorrelationId = 100;
constexpr int64 kSecondCorrelationId = 200;
XSpace space;
XPlane* host_plane = space.add_planes();
XPlaneBuilder host_plane_builder(host_plane);
@ -45,25 +52,25 @@ TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
auto main_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kTraceContext,
0, 100, {{StatType::kStepNum, 123}});
0, 100, {{StatType::kStepNum, kFirstStepNum}});
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kFunctionRun,
10, 90, {{StatType::kStepId, 0}});
10, 90, {{StatType::kStepId, kFirstStepId}});
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kTraceContext,
300, 100, {{StatType::kStepNum, 456}});
300, 100, {{StatType::kStepNum, kSecondStepNum}});
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kFunctionRun,
310, 90, {{StatType::kStepId, 1}});
310, 90, {{StatType::kStepId, kSecondStepId}});
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(1);
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 20, 20,
{{StatType::kStepId, 0}});
{{StatType::kStepId, kFirstStepId}});
CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 30, 10,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kFirstCorrelationId}});
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 320, 20,
{{StatType::kStepId, 1}});
{{StatType::kStepId, kSecondStepId}});
CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 330, 10,
{{StatType::kCorrelationId, 200}});
{{StatType::kCorrelationId, kSecondCorrelationId}});
XPlane* device_plane = space.add_planes();
XPlaneBuilder device_plane_builder(device_plane);
@ -71,7 +78,7 @@ TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
auto stream = device_plane_builder.GetOrCreateLine(0);
CreateXEvent(&device_plane_builder, &stream, "matmul", 50, 40,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kFirstCorrelationId}});
GroupTfEvents(&space, nullptr);
StepEvents device_step_events =

View File

@ -102,6 +102,9 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
// Checks that the dependency between the step line and the TF op line prevents
// TF op events from being expanded.
TEST(DerivedTimelineTest, DependencyTest) {
constexpr int64 kFirstGroupId = 0;
constexpr int64 kSecondGroupId = 1;
const absl::string_view kTfOpName = "mul:Mul";
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
@ -110,11 +113,11 @@ TEST(DerivedTimelineTest, DependencyTest) {
XPlaneBuilder plane_builder(plane);
auto line_builder = plane_builder.GetOrCreateLine(0);
CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,
{{StatType::kGroupId, 0},
{{StatType::kGroupId, kFirstGroupId},
{StatType::kLevel0, kTfOpName},
{StatType::kKernelDetails, kKernelDetails}});
CreateXEvent(&plane_builder, &line_builder, "op2", 200, 300,
{{StatType::kGroupId, 1},
{{StatType::kGroupId, kSecondGroupId},
{StatType::kLevel0, kTfOpName},
{StatType::kKernelDetails, kKernelDetails}});
GenerateDerivedTimeLines(event_group_name_map, &space);

View File

@ -111,6 +111,8 @@ void SetGroupId(const XPlaneVisitor& visitor, int64 group_id, XEvent* event) {
void SetContextGroup(EventNode* event, ContextGroupMap* context_groups) {
auto producer = event->GetProducerContext();
if (producer.has_value()) {
DCHECK_EQ(((*context_groups)[producer->type][producer->id]).producer,
nullptr);
((*context_groups)[producer->type][producer->id]).producer = event;
}
auto consumer = event->GetConsumerContext();
@ -124,9 +126,10 @@ void ConnectContextGroups(const ContextGroupMap& context_groups) {
for (auto& type_id_group : context_groups) {
for (auto& id_group : type_id_group.second) {
const ContextGroup& group = id_group.second;
EventNode* parent = group.producer;
for (EventNode* child : group.consumers) {
parent->AddChild(child);
if (EventNode* parent = group.producer) {
for (EventNode* child : group.consumers) {
parent->AddChild(child);
}
}
}
}
@ -194,13 +197,13 @@ EventNode::EventNode(const XPlaneVisitor* plane, XLine* raw_line,
producer_type = stat.IntValue();
break;
case StatType::kProducerId:
producer_id = stat.IntValue();
producer_id = stat.UintValue();
break;
case StatType::kConsumerType:
consumer_type = stat.IntValue();
break;
case StatType::kConsumerId:
consumer_id = stat.IntValue();
consumer_id = stat.UintValue();
break;
case StatType::kIsRoot:
is_root_ = stat.IntValue();
@ -330,7 +333,7 @@ void EventForest::ConnectIntraThread(const XPlaneVisitor& visitor,
void EventForest::ConnectInterThread(
const std::vector<InterThreadConnectInfo>& connect_info_list) {
for (const auto& connect_info : connect_info_list) {
absl::flat_hash_map<std::vector<int64>, EventNode*> connect_map;
absl::flat_hash_map<std::vector<uint64>, EventNode*> connect_map;
const std::vector<int64>& parent_stat_types =
connect_info.parent_stat_types;
const std::vector<int64>* child_stat_types = &connect_info.child_stat_types;
@ -340,14 +343,12 @@ void EventForest::ConnectInterThread(
if (auto parent_event_node_list =
gtl::FindOrNull(event_node_map_, connect_info.parent_event_type)) {
for (const auto& parent_event_node : *parent_event_node_list) {
std::vector<int64> stats;
std::vector<uint64> stats;
for (auto stat_type : parent_stat_types) {
absl::optional<XStatVisitor> stat =
parent_event_node->GetContextStat(stat_type);
if (!stat) break;
stats.push_back((stat->ValueCase() == XStat::kInt64Value)
? stat->IntValue()
: stat->UintValue());
stats.push_back(stat->IntOrUintValue());
}
if (stats.size() == parent_stat_types.size()) {
connect_map[stats] = parent_event_node.get();
@ -357,14 +358,12 @@ void EventForest::ConnectInterThread(
if (auto child_event_node_list =
gtl::FindOrNull(event_node_map_, connect_info.child_event_type)) {
for (const auto& child_event_node : *child_event_node_list) {
std::vector<int64> stats;
std::vector<uint64> stats;
for (auto stat_type : *child_stat_types) {
absl::optional<XStatVisitor> stat =
child_event_node->GetContextStat(stat_type);
if (!stat) break;
stats.push_back((stat->ValueCase() == XStat::kInt64Value)
? stat->IntValue()
: stat->UintValue());
stats.push_back(stat->IntOrUintValue());
}
if (stats.size() == child_stat_types->size()) {
if (auto parent_event_node = gtl::FindPtrOrNull(connect_map, stats)) {

View File

@ -31,6 +31,10 @@ namespace profiler {
namespace {
TEST(GroupEventsTest, GroupGpuTraceTest) {
constexpr int64 kStepNum = 123;
constexpr int64 kStepId = 0;
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
@ -38,16 +42,16 @@ TEST(GroupEventsTest, GroupGpuTraceTest) {
auto main_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kTraceContext,
0, 100, {{StatType::kStepNum, 123}});
0, 100, {{StatType::kStepNum, kStepNum}});
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kFunctionRun,
10, 90, {{StatType::kStepId, 0}});
10, 90, {{StatType::kStepId, kStepId}});
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(1);
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 20, 80,
{{StatType::kStepId, 0}});
{{StatType::kStepId, kStepId}});
CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 30, 70,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
XPlane* device_plane = space.add_planes();
XPlaneBuilder device_plane_builder(device_plane);
@ -55,7 +59,7 @@ TEST(GroupEventsTest, GroupGpuTraceTest) {
auto stream = device_plane_builder.GetOrCreateLine(0);
CreateXEvent(&device_plane_builder, &stream, "matmul", 200, 300,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
EventGroupNameMap event_group_name_map;
GroupTfEvents(&space, &event_group_name_map);
@ -69,6 +73,10 @@ TEST(GroupEventsTest, GroupGpuTraceTest) {
}
TEST(GroupEventsTest, GroupTensorFlowLoopTest) {
constexpr int64 kStepId = 0;
constexpr int64 kIterNum = 10;
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
@ -77,12 +85,12 @@ TEST(GroupEventsTest, GroupTensorFlowLoopTest) {
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 5, 10,
{{StatType::kStepId, 0}, {StatType::kIterNum, 10}});
{{StatType::kStepId, kStepId}, {StatType::kIterNum, kIterNum}});
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 20, 80,
{{StatType::kStepId, 0}, {StatType::kIterNum, 10}});
{{StatType::kStepId, kStepId}, {StatType::kIterNum, kIterNum}});
CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 30, 70,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
XPlane* device_plane = space.add_planes();
XPlaneBuilder device_plane_builder(device_plane);
@ -90,7 +98,7 @@ TEST(GroupEventsTest, GroupTensorFlowLoopTest) {
auto stream = device_plane_builder.GetOrCreateLine(0);
CreateXEvent(&device_plane_builder, &stream, "matmul", 200, 300,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
EventGroupNameMap event_group_name_map;
GroupTfEvents(&space, &event_group_name_map);
@ -111,6 +119,11 @@ TEST(GroupEventsTest, GroupTensorFlowLoopTest) {
// group_id is initialized to the first TF loop's first iter_num (10) and then
// monotonically increased.
TEST(GroupEventsTest, GroupMultipleTensorFlowLoopsTest) {
constexpr int64 kFirstStepId = 0;
constexpr int64 kSecondStepId = 1;
constexpr int64 kFirstIterNumStart = 10;
constexpr int64 kSecondIterNumStart = 0;
XSpace space;
XPlaneBuilder host_plane_builder(space.add_planes());
host_plane_builder.SetName(kHostThreads);
@ -119,17 +132,21 @@ TEST(GroupEventsTest, GroupMultipleTensorFlowLoopsTest) {
auto first_tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &first_tf_executor_thread,
HostEventType::kExecutorStateProcess, 220, 80,
{{StatType::kStepId, 1}, {StatType::kIterNum, 0}});
{{StatType::kStepId, kSecondStepId},
{StatType::kIterNum, kSecondIterNumStart}});
CreateXEvent(&host_plane_builder, &first_tf_executor_thread,
HostEventType::kExecutorStateProcess, 320, 80,
{{StatType::kStepId, 1}, {StatType::kIterNum, 1}});
{{StatType::kStepId, kSecondStepId},
{StatType::kIterNum, kSecondIterNumStart + 1}});
auto second_tf_executor_thread = host_plane_builder.GetOrCreateLine(1);
CreateXEvent(&host_plane_builder, &second_tf_executor_thread,
HostEventType::kExecutorStateProcess, 20, 80,
{{StatType::kStepId, 0}, {StatType::kIterNum, 10}});
{{StatType::kStepId, kFirstStepId},
{StatType::kIterNum, kFirstIterNumStart}});
CreateXEvent(&host_plane_builder, &second_tf_executor_thread,
HostEventType::kExecutorStateProcess, 120, 80,
{{StatType::kStepId, 0}, {StatType::kIterNum, 11}});
{{StatType::kStepId, kFirstStepId},
{StatType::kIterNum, kFirstIterNumStart + 1}});
EventGroupNameMap event_group_name_map;
GroupTfEvents(&space, &event_group_name_map);
@ -141,6 +158,10 @@ TEST(GroupEventsTest, GroupMultipleTensorFlowLoopsTest) {
}
TEST(GroupEventsTest, GroupFunctionalOp) {
constexpr int64 kStepNum = 123;
constexpr int64 kStepId = 0;
constexpr int64 kFunctionStepId = 1;
XSpace space;
XPlane* host_plane = space.add_planes();
XPlaneBuilder host_plane_builder(host_plane);
@ -149,20 +170,20 @@ TEST(GroupEventsTest, GroupFunctionalOp) {
auto main_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kTraceContext,
0, 200, {{StatType::kStepNum, 123}});
0, 200, {{StatType::kStepNum, kStepNum}});
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kFunctionRun,
10, 190, {{StatType::kStepId, 0}});
10, 190, {{StatType::kStepId, kStepId}});
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 20, 80,
{{StatType::kStepId, 0}});
{{StatType::kStepId, kStepId}});
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kRemoteCallOp, 30, 70,
{{StatType::kFunctionStepId, 1}});
{{StatType::kFunctionStepId, kFunctionStepId}});
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 100, 150,
{{StatType::kStepId, 1}});
{{StatType::kStepId, kFunctionStepId}});
EventGroupNameMap event_group_name_map;
GroupTfEvents(&space, &event_group_name_map);
@ -185,6 +206,8 @@ TEST(GroupEventsTest, GroupFunctionalOp) {
}
TEST(GroupEventsTest, EagerOpTest) {
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlane* host_plane = space.add_planes();
XPlaneBuilder host_plane_builder(host_plane);
@ -194,12 +217,12 @@ TEST(GroupEventsTest, EagerOpTest) {
auto main_thread = host_plane_builder.GetOrCreateLine(0);
// Eagerly scheduled GPU kernel.
CreateXEvent(&host_plane_builder, &main_thread,
HostEventType::kEagerKernelExecute, 10, 100, {});
HostEventType::kEagerKernelExecute, 10, 100);
CreateXEvent(&host_plane_builder, &main_thread, "matmul", 10, 100,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
// Eagerly executed CPU TF op.
CreateXEvent(&host_plane_builder, &main_thread,
HostEventType::kEagerKernelExecute, 120, 80, {});
HostEventType::kEagerKernelExecute, 120, 80);
CreateXEvent(&host_plane_builder, &main_thread, "add:Add", 120, 80);
XPlane* device_plane = space.add_planes();
@ -209,7 +232,7 @@ TEST(GroupEventsTest, EagerOpTest) {
auto stream = device_plane_builder.GetOrCreateLine(0);
// Eagerly executed GPU kernel.
CreateXEvent(&device_plane_builder, &stream, "matmul", 200, 300,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
GroupTfEvents(&space, /*event_group_name_map=*/nullptr);
XPlaneVisitor host_plane_visitor = CreateTfXPlaneVisitor(host_plane);
@ -227,6 +250,10 @@ TEST(GroupEventsTest, EagerOpTest) {
}
TEST(GroupEventsTest, FunctionOpTest) {
constexpr int64 kStepNum = 123;
constexpr int64 kStepId = 0;
constexpr int64 kCorrelationId = 100;
XSpace space;
XPlane* host_plane = space.add_planes();
XPlaneBuilder host_plane_builder(host_plane);
@ -235,19 +262,19 @@ TEST(GroupEventsTest, FunctionOpTest) {
auto main_thread = host_plane_builder.GetOrCreateLine(0);
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kTraceContext,
0, 100, {{StatType::kStepNum, 123}});
0, 100, {{StatType::kStepNum, kStepNum}});
CreateXEvent(&host_plane_builder, &main_thread,
HostEventType::kEagerKernelExecute, 10, 90, {});
HostEventType::kEagerKernelExecute, 10, 90);
CreateXEvent(&host_plane_builder, &main_thread, HostEventType::kFunctionRun,
10, 90, {{StatType::kStepId, 0}});
10, 90, {{StatType::kStepId, kStepId}});
auto tf_executor_thread = host_plane_builder.GetOrCreateLine(1);
CreateXEvent(&host_plane_builder, &tf_executor_thread,
HostEventType::kExecutorStateProcess, 20, 80,
{{StatType::kStepId, 0}});
{{StatType::kStepId, kStepId}});
// GPU kernel scheduled inside tf.function.
CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 30, 30,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
// CPU TF op executed inside tf.function.
CreateXEvent(&host_plane_builder, &tf_executor_thread, "add:Add", 70, 20);
@ -258,7 +285,7 @@ TEST(GroupEventsTest, FunctionOpTest) {
auto stream = device_plane_builder.GetOrCreateLine(0);
// GPU kernel executed as part of tf.function.
CreateXEvent(&device_plane_builder, &stream, "matmul", 200, 300,
{{StatType::kCorrelationId, 100}});
{{StatType::kCorrelationId, kCorrelationId}});
GroupTfEvents(&space, /*event_group_name_map=*/nullptr);
XPlaneVisitor host_plane_visitor = CreateTfXPlaneVisitor(host_plane);
@ -276,8 +303,9 @@ TEST(GroupEventsTest, FunctionOpTest) {
}
TEST(GroupEventsTest, SemanticArgTest) {
constexpr int64 kIsRoot = 1;
constexpr int64 kStepNum = 100;
constexpr int kContextType = 123;
constexpr int64 kContextType = 123;
constexpr uint64 kContextId = 456;
XSpace raw_space;
@ -286,7 +314,7 @@ TEST(GroupEventsTest, SemanticArgTest) {
plane.ReserveLines(2);
auto root_producer = plane.GetOrCreateLine(0);
CreateXEvent(&plane, &root_producer, HostEventType::kTraceContext, 0, 100,
{{StatType::kIsRoot, 1}, {StatType::kStepNum, kStepNum}});
{{StatType::kIsRoot, kIsRoot}, {StatType::kStepNum, kStepNum}});
CreateXEvent(&plane, &root_producer, HostEventType::kFunctionRun, 10, 90,
{{StatType::kProducerType, kContextType},
{StatType::kProducerId, kContextId}});
@ -314,7 +342,99 @@ TEST(GroupEventsTest, SemanticArgTest) {
EXPECT_EQ(num_events, 3);
}
TEST(GroupEventsTest, SemanticIntArgNoMatchTest) {
constexpr int64 kIsRoot = 1;
constexpr int64 kStepNum = 100;
constexpr int64 kContextType = 123;
constexpr uint64 kProducerId = 456;
constexpr uint64 kConsumerId = 789;
XSpace raw_space;
XPlane* raw_plane = raw_space.add_planes();
XPlaneBuilder plane(raw_plane);
plane.ReserveLines(2);
auto root_producer = plane.GetOrCreateLine(0);
CreateXEvent(&plane, &root_producer, HostEventType::kTraceContext, 0, 100,
{{StatType::kIsRoot, kIsRoot}, {StatType::kStepNum, kStepNum}});
CreateXEvent(&plane, &root_producer, HostEventType::kFunctionRun, 10, 90,
{{StatType::kProducerType, kContextType},
{StatType::kProducerId, kProducerId}});
auto consumer = plane.GetOrCreateLine(1);
CreateXEvent(&plane, &consumer, HostEventType::kExecutorStateProcess, 20, 80,
{{StatType::kConsumerType, kContextType},
{StatType::kConsumerId, kConsumerId}});
GroupTfEvents(&raw_space, /*event_group_name_map=*/nullptr);
int num_events = 0;
CreateTfXPlaneVisitor(raw_plane).ForEachLine(
[&](const tensorflow::profiler::XLineVisitor& line) {
num_events += line.NumEvents();
line.ForEachEvent(
[&](const tensorflow::profiler::XEventVisitor& event) {
absl::optional<int64> group_id;
if (absl::optional<XStatVisitor> stat =
event.GetStat(StatType::kGroupId)) {
group_id = stat->IntValue();
}
if (event.Type() == HostEventType::kExecutorStateProcess) {
EXPECT_FALSE(group_id.has_value());
} else {
EXPECT_TRUE(group_id.has_value());
EXPECT_EQ(*group_id, 0);
}
});
});
EXPECT_EQ(num_events, 3);
}
TEST(GroupEventsTest, SemanticUintArgNoMatchTest) {
constexpr int64 kIsRoot = 1;
constexpr int64 kStepNum = 100;
constexpr int64 kContextType = 123;
constexpr uint64 kProducerId = UINT64_MAX;
constexpr uint64 kConsumerId = UINT64_MAX - 1;
XSpace raw_space;
XPlane* raw_plane = raw_space.add_planes();
XPlaneBuilder plane(raw_plane);
plane.ReserveLines(2);
auto root_producer = plane.GetOrCreateLine(0);
CreateXEvent(&plane, &root_producer, HostEventType::kTraceContext, 0, 100,
{{StatType::kIsRoot, kIsRoot}, {StatType::kStepNum, kStepNum}});
CreateXEvent(&plane, &root_producer, HostEventType::kFunctionRun, 10, 90,
{{StatType::kProducerType, kContextType},
{StatType::kProducerId, kProducerId}});
auto consumer = plane.GetOrCreateLine(1);
CreateXEvent(&plane, &consumer, HostEventType::kExecutorStateProcess, 20, 80,
{{StatType::kConsumerType, kContextType},
{StatType::kConsumerId, kConsumerId}});
GroupTfEvents(&raw_space, /*event_group_name_map=*/nullptr);
int num_events = 0;
CreateTfXPlaneVisitor(raw_plane).ForEachLine(
[&](const tensorflow::profiler::XLineVisitor& line) {
num_events += line.NumEvents();
line.ForEachEvent(
[&](const tensorflow::profiler::XEventVisitor& event) {
absl::optional<int64> group_id;
if (absl::optional<XStatVisitor> stat =
event.GetStat(StatType::kGroupId)) {
group_id = stat->IntValue();
}
if (event.Type() == HostEventType::kExecutorStateProcess) {
EXPECT_FALSE(group_id.has_value());
} else {
EXPECT_TRUE(group_id.has_value());
EXPECT_EQ(*group_id, 0);
}
});
});
EXPECT_EQ(num_events, 3);
}
TEST(GroupEventsTest, AsyncEventTest) {
constexpr int64 kIsRoot = 1;
constexpr int64 kIsAsync = 1;
constexpr absl::string_view kParent = "parent";
constexpr absl::string_view kAsync = "async";
constexpr absl::string_view kChild = "child";
@ -324,8 +444,9 @@ TEST(GroupEventsTest, AsyncEventTest) {
XPlaneBuilder plane(raw_plane);
plane.ReserveLines(1);
auto line = plane.GetOrCreateLine(0);
CreateXEvent(&plane, &line, kParent, 0, 100, {{StatType::kIsRoot, 1}});
CreateXEvent(&plane, &line, kAsync, 10, 200, {{StatType::kIsAsync, 1}});
CreateXEvent(&plane, &line, kParent, 0, 100, {{StatType::kIsRoot, kIsRoot}});
CreateXEvent(&plane, &line, kAsync, 10, 200,
{{StatType::kIsAsync, kIsAsync}});
CreateXEvent(&plane, &line, kChild, 20, 80);
GroupTfEvents(&raw_space, /*event_group_name_map=*/nullptr);

View File

@ -26,7 +26,7 @@ limitations under the License.
namespace tensorflow {
namespace profiler {
using XStatValue = absl::variant<int64, absl::string_view>;
using XStatValue = absl::variant<int64, uint64, absl::string_view>;
void CreateXEvent(
XPlaneBuilder* plane_builder, XLineBuilder* line_builder,
@ -36,7 +36,7 @@ void CreateXEvent(
void CreateXEvent(
XPlaneBuilder* plane_builder, XLineBuilder* line_builder,
HostEventType event_type, int64 offset_ps, int64 duration_ps,
std::initializer_list<std::pair<StatType, XStatValue>> stats);
std::initializer_list<std::pair<StatType, XStatValue>> stats = {});
void CreateTfFunctionCallEvent(XPlaneBuilder* plane_builder,
XLineBuilder* line_builder,

View File

@ -53,6 +53,11 @@ class XStatVisitor {
uint64 UintValue() const { return stat_->uint64_value(); }
uint64 IntOrUintValue() const {
return ValueCase() == XStat::kUint64Value ? UintValue()
: static_cast<uint64>(IntValue());
}
double DoubleValue() const { return stat_->double_value(); }
// Returns a string view.