Skip to content

Commit fea9da8

Browse files
cliveverghesecopybara-github
cliveverghese
authored and
copybara-github
committed
Update Single Host XPlane Processing and OpStat conversion to share HloModuleMap.
PiperOrigin-RevId: 743706393
1 parent 1f64f58 commit fea9da8

File tree

3 files changed

+35
-18
lines changed

3 files changed

+35
-18
lines changed

‎xprof/utils/derived_timeline.cc

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -619,9 +619,9 @@ void DeriveEventsFromHostTrace(
619619
}
620620

621621
void GenerateDerivedTimeLines(
622-
const tsl::profiler::GroupMetadataMap& group_metadata_map, XSpace* space) {
623-
HloModuleMap hlo_module_map;
624-
{
622+
const tsl::profiler::GroupMetadataMap& group_metadata_map, XSpace* space,
623+
HloModuleMap& hlo_module_map) {
624+
if (hlo_module_map.empty()) {
625625
HloProtoMap hlo_proto_map;
626626
hlo_proto_map.AddHloProtosFromXSpace(*space);
627627
for (const auto& [program_id, hlo_proto] : hlo_proto_map) {

‎xprof/utils/derived_timeline.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ limitations under the License.
2828
#include "xla/tsl/profiler/utils/timespan.h"
2929
#include "xla/tsl/profiler/utils/xplane_builder.h"
3030
#include "tsl/profiler/protobuf/xplane.pb.h"
31+
#include "xprof/utils/hlo_module_map.h"
3132

3233
namespace tensorflow {
3334
namespace profiler {
@@ -189,7 +190,8 @@ void DeriveEventsFromHostTrace(
189190
// Loops through XPlanes of input XSpace, if it is "device" XPlane, generating
190191
// derived timelines for the plane by calling DeriveEventsFromAnnotations.
191192
void GenerateDerivedTimeLines(
192-
const tsl::profiler::GroupMetadataMap& group_metadata_map, XSpace* space);
193+
const tsl::profiler::GroupMetadataMap& group_metadata_map, XSpace* space,
194+
HloModuleMap& hlo_module_map);
193195

194196
// Derives `Tensorflow Ops`, `Tensorflow Name Scope` and `Source Code` lines
195197
// from device_trace.

‎xprof/utils/derived_timeline_test.cc

Lines changed: 29 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ limitations under the License.
3434
#include "xla/tsl/profiler/utils/xplane_visitor.h"
3535
#include "tensorflow/core/platform/test.h"
3636
#include "tsl/profiler/protobuf/xplane.pb.h"
37+
#include "xprof/utils/hlo_module_map.h"
3738

3839
namespace tensorflow {
3940
namespace profiler {
@@ -53,8 +54,9 @@ using ::tsl::profiler::XStatVisitor;
5354

5455
TEST(DerivedTimelineTest, EmptySpaceTest) {
5556
XSpace space;
57+
HloModuleMap hlo_module_map;
5658
tsl::profiler::GroupMetadataMap group_metadata_map;
57-
GenerateDerivedTimeLines(group_metadata_map, &space);
59+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
5860
EXPECT_EQ(space.planes_size(), 0);
5961
}
6062

@@ -63,6 +65,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
6365
const absl::string_view kHloModuleName = "hlo_module";
6466
const absl::string_view kKernelDetails = "kernel_details";
6567
XSpace space;
68+
HloModuleMap hlo_module_map;
6669
tsl::profiler::GroupMetadataMap group_metadata_map;
6770
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
6871
XPlaneBuilder plane_builder(plane);
@@ -73,7 +76,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
7376
CreateXEvent(&plane_builder, &line_builder, "op2", 200, 300,
7477
{{StatType::kHloModule, kHloModuleName},
7578
{StatType::kKernelDetails, kKernelDetails}});
76-
GenerateDerivedTimeLines(group_metadata_map, &space);
79+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
7780
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
7881
// Only the hlo module line is added and other empty lines are removed at the
7982
// end.
@@ -95,6 +98,7 @@ TEST(DerivedTimelineTest, HloModuleNameSameScopeRangeIdTest) {
9598
const absl::string_view kHloModuleName = "hlo_module";
9699
const absl::string_view kKernelDetails = "kernel_details";
97100
XSpace space;
101+
HloModuleMap hlo_module_map;
98102
tsl::profiler::GroupMetadataMap group_metadata_map;
99103
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
100104
XPlaneBuilder plane_builder(plane);
@@ -107,7 +111,7 @@ TEST(DerivedTimelineTest, HloModuleNameSameScopeRangeIdTest) {
107111
{{StatType::kHloModule, XStatValue{kHloModuleName}},
108112
{StatType::kKernelDetails, XStatValue{kKernelDetails}},
109113
{StatType::kScopeRangeId, XStatValue{int64_t{10}}}});
110-
GenerateDerivedTimeLines(group_metadata_map, &space);
114+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
111115
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
112116
// Only the hlo module line is added and other empty lines are removed at the
113117
// end.
@@ -128,6 +132,7 @@ TEST(DerivedTimelineTest, HloModuleNameDifferentScopeRangeIdTest) {
128132
const absl::string_view kHloModuleName = "hlo_module";
129133
const absl::string_view kKernelDetails = "kernel_details";
130134
XSpace space;
135+
HloModuleMap hlo_module_map;
131136
tsl::profiler::GroupMetadataMap group_metadata_map;
132137
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
133138
XPlaneBuilder plane_builder(plane);
@@ -140,7 +145,7 @@ TEST(DerivedTimelineTest, HloModuleNameDifferentScopeRangeIdTest) {
140145
{{StatType::kHloModule, XStatValue{kHloModuleName}},
141146
{StatType::kKernelDetails, XStatValue{kKernelDetails}},
142147
{StatType::kScopeRangeId, XStatValue{int64_t{20}}}});
143-
GenerateDerivedTimeLines(group_metadata_map, &space);
148+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
144149
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
145150
// Only the hlo module line is added and other empty lines are removed at the
146151
// end.
@@ -160,6 +165,7 @@ TEST(DerivedTimelineTest, NoHloModuleNameTest) {
160165
const absl::string_view kKernelDetails = "kernel_details";
161166
const uint64_t kCudaGraphExecId = 1;
162167
XSpace space;
168+
HloModuleMap hlo_module_map;
163169
tsl::profiler::GroupMetadataMap group_metadata_map;
164170
XPlane& plane = *GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
165171
XPlaneBuilder plane_builder(&plane);
@@ -171,7 +177,7 @@ TEST(DerivedTimelineTest, NoHloModuleNameTest) {
171177
// Also add a CudaGraph Execution event.
172178
CreateXEvent(&plane_builder, &line_builder, "op3", 500, 100,
173179
{{StatType::kCudaGraphExecId, kCudaGraphExecId}});
174-
GenerateDerivedTimeLines(group_metadata_map, &space);
180+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
175181
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(&plane);
176182
// Only the hlo module line is added and other empty lines are removed at the
177183
// end.
@@ -189,6 +195,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
189195
const absl::string_view kKernelDetails = "kernel_details";
190196
const uint64_t kCudaGraphExecId = 1;
191197
XSpace space;
198+
HloModuleMap hlo_module_map;
192199
tsl::profiler::GroupMetadataMap group_metadata_map;
193200
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
194201
XPlaneBuilder plane_builder(plane);
@@ -203,7 +210,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
203210
CreateXEvent(&plane_builder, &line_builder, "op3", 500, 100,
204211
{{StatType::kTfOp, kTfOpName},
205212
{StatType::kCudaGraphExecId, kCudaGraphExecId}});
206-
GenerateDerivedTimeLines(group_metadata_map, &space);
213+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
207214
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
208215
// Only the tf op line is added and other empty lines are removed at the end.
209216
EXPECT_EQ(plane_visitor.NumLines(), 2);
@@ -228,6 +235,7 @@ TEST(DerivedTimelineTest, DependencyTest) {
228235
const absl::string_view kTfOpName = "mul:Mul";
229236
const absl::string_view kKernelDetails = "kernel_details";
230237
XSpace space;
238+
HloModuleMap hlo_module_map;
231239
tsl::profiler::GroupMetadataMap group_metadata_map(
232240
{{0, {"train 0"}}, {1, {"train 1"}}});
233241
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
@@ -241,7 +249,7 @@ TEST(DerivedTimelineTest, DependencyTest) {
241249
{{StatType::kGroupId, kSecondGroupId},
242250
{StatType::kTfOp, kTfOpName},
243251
{StatType::kKernelDetails, kKernelDetails}});
244-
GenerateDerivedTimeLines(group_metadata_map, &space);
252+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
245253
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
246254
// The step line and the TF op line are added.
247255
EXPECT_EQ(plane_visitor.NumLines(), 3);
@@ -258,6 +266,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeTest) {
258266
const absl::string_view kTfOpName = "scope1/scope2/mul:Mul";
259267
const absl::string_view kKernelDetails = "kernel_details";
260268
XSpace space;
269+
HloModuleMap hlo_module_map;
261270
tsl::profiler::GroupMetadataMap group_metadata_map;
262271
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
263272
XPlaneBuilder plane_builder(plane);
@@ -268,7 +277,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeTest) {
268277
CreateXEvent(&plane_builder, &line_builder, "op2", 200, 300,
269278
{{StatType::kTfOp, kTfOpName},
270279
{StatType::kKernelDetails, kKernelDetails}});
271-
GenerateDerivedTimeLines(group_metadata_map, &space);
280+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
272281
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
273282
// The TF name scope line and the TF op line are added.
274283
EXPECT_EQ(plane_visitor.NumLines(), 3);
@@ -298,6 +307,7 @@ TEST(DerivedTimelineTest, TfNameScopeMaintainsOrder) {
298307
const absl::string_view kTfOpName = "scope1/scope2/mul:Mul";
299308
const absl::string_view kKernelDetails = "kernel_details";
300309
XSpace space;
310+
HloModuleMap hlo_module_map;
301311
tsl::profiler::GroupMetadataMap group_metadata_map;
302312
XPlane* plane = tsl::profiler::GetOrCreateTpuXPlane(
303313
&space, /*device_ordinal=*/0, "TPU V4", 0, 0);
@@ -306,7 +316,7 @@ TEST(DerivedTimelineTest, TfNameScopeMaintainsOrder) {
306316
CreateXEvent(&plane_builder, &line_builder, "op1", 0, 10000,
307317
{{StatType::kTfOp, kTfOpName},
308318
{StatType::kKernelDetails, kKernelDetails}});
309-
GenerateDerivedTimeLines(group_metadata_map, &space);
319+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
310320
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
311321
// The TF name scope line and the TF op line are added.
312322
EXPECT_EQ(plane_visitor.NumLines(), 3);
@@ -329,6 +339,7 @@ TEST(DerivedTimelineTest, OnlyDerivedEventsFromLineWithMostEvents) {
329339
const absl::string_view kTfOpName = "scope1/scope2/mul:Mul";
330340
const absl::string_view kKernelDetails = "kernel_details";
331341
XSpace space;
342+
HloModuleMap hlo_module_map;
332343
tsl::profiler::GroupMetadataMap group_metadata_map;
333344
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
334345
XPlaneBuilder plane_builder(plane);
@@ -346,7 +357,7 @@ TEST(DerivedTimelineTest, OnlyDerivedEventsFromLineWithMostEvents) {
346357
{{StatType::kTfOp, kTfOpName},
347358
{StatType::kKernelDetails, kKernelDetails}});
348359
// Derive lines for the plane.
349-
GenerateDerivedTimeLines(group_metadata_map, &space);
360+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
350361
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
351362
// The TF name scope line and the TF op line are added.
352363
EXPECT_EQ(plane_visitor.NumLines(), 4);
@@ -381,6 +392,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
381392
{
382393
// Case 1: shirnk is possible.
383394
XSpace space;
395+
HloModuleMap hlo_module_map;
384396
tsl::profiler::GroupMetadataMap group_metadata_map;
385397
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
386398
XPlaneBuilder plane_builder(plane);
@@ -391,7 +403,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
391403
CreateXEvent(
392404
&plane_builder, &line_builder, "op2", 20000, 30000,
393405
{{StatType::kTfOp, "a/d/Mul:Mul"}, {StatType::kKernelDetails, "blah"}});
394-
GenerateDerivedTimeLines(group_metadata_map, &space);
406+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
395407
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
396408
// The TF name scope line and the TF op line are added.
397409
EXPECT_EQ(plane_visitor.NumLines(), 3);
@@ -415,6 +427,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
415427
{
416428
// Case 2: shirnk is impossible due to top event is too small.
417429
XSpace space;
430+
HloModuleMap hlo_module_map;
418431
tsl::profiler::GroupMetadataMap group_metadata_map;
419432
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
420433
XPlaneBuilder plane_builder(plane);
@@ -428,7 +441,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
428441
CreateXEvent(
429442
&plane_builder, &line_builder, "op3", 20000, 30000,
430443
{{StatType::kTfOp, "a/g/Mul:Mul"}, {StatType::kKernelDetails, "blah"}});
431-
GenerateDerivedTimeLines(group_metadata_map, &space);
444+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
432445
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
433446
// The TF name scope line and the TF op line are added.
434447
EXPECT_EQ(plane_visitor.NumLines(), 3);
@@ -466,6 +479,7 @@ TEST(DerivedTimelineTest, XloOpHasCudaGraphStats) {
466479
constexpr int64_t kCorrelationIdValue = 10000;
467480
const uint64_t kCudaGraphIdValue = 20;
468481
XSpace space;
482+
HloModuleMap hlo_module_map;
469483
tsl::profiler::GroupMetadataMap group_metadata_map;
470484

471485
// Build Input Plane/Line/Events and derive events from them.
@@ -486,7 +500,7 @@ TEST(DerivedTimelineTest, XloOpHasCudaGraphStats) {
486500
{StatType::kHloOp, kHloOpName},
487501
{StatType::kCorrelationId, kCorrelationIdValue},
488502
{StatType::kCudaGraphId, kCudaGraphIdValue}});
489-
GenerateDerivedTimeLines(group_metadata_map, &space);
503+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
490504

491505
// Check that the HLO op line is added and has the extra stats for the first
492506
// derived event.
@@ -552,6 +566,7 @@ TEST(DerivedTimelineTest, MergeAndNoMerge) {
552566
constexpr absl::string_view kHloModuleName = "Framework Ops";
553567
static constexpr absl::string_view kTfOpName = "abc:model/layer/MatMul_1";
554568
XSpace space;
569+
HloModuleMap hlo_module_map;
555570
tsl::profiler::GroupMetadataMap group_metadata_map;
556571
XPlane* plane = tsl::profiler::GetOrCreateTpuXPlane(
557572
&space, /*device_ordinal=*/0, "DummyTPU", 1.0, 1.0);
@@ -568,7 +583,7 @@ TEST(DerivedTimelineTest, MergeAndNoMerge) {
568583
CreateXEvent(
569584
&plane_builder, &line_builder, "op3", 1501, 300,
570585
{{StatType::kHloModule, kHloModuleName}, {StatType::kTfOp, kTfOpName}});
571-
GenerateDerivedTimeLines(group_metadata_map, &space);
586+
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
572587
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
573588
// Only the hlo module line is added and other empty lines are removed at the
574589
// end.

0 commit comments

Comments
 (0)