@@ -34,6 +34,7 @@ limitations under the License.
34
34
#include " xla/tsl/profiler/utils/xplane_visitor.h"
35
35
#include " tensorflow/core/platform/test.h"
36
36
#include " tsl/profiler/protobuf/xplane.pb.h"
37
+ #include " xprof/utils/hlo_module_map.h"
37
38
38
39
namespace tensorflow {
39
40
namespace profiler {
@@ -53,8 +54,9 @@ using ::tsl::profiler::XStatVisitor;
53
54
54
55
TEST (DerivedTimelineTest, EmptySpaceTest) {
55
56
XSpace space;
57
+ HloModuleMap hlo_module_map;
56
58
tsl::profiler::GroupMetadataMap group_metadata_map;
57
- GenerateDerivedTimeLines (group_metadata_map, &space);
59
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
58
60
EXPECT_EQ (space.planes_size (), 0 );
59
61
}
60
62
@@ -63,6 +65,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
63
65
const absl::string_view kHloModuleName = " hlo_module" ;
64
66
const absl::string_view kKernelDetails = " kernel_details" ;
65
67
XSpace space;
68
+ HloModuleMap hlo_module_map;
66
69
tsl::profiler::GroupMetadataMap group_metadata_map;
67
70
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
68
71
XPlaneBuilder plane_builder (plane);
@@ -73,7 +76,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
73
76
CreateXEvent (&plane_builder, &line_builder, " op2" , 200 , 300 ,
74
77
{{StatType::kHloModule , kHloModuleName },
75
78
{StatType::kKernelDetails , kKernelDetails }});
76
- GenerateDerivedTimeLines (group_metadata_map, &space);
79
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
77
80
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
78
81
// Only the hlo module line is added and other empty lines are removed at the
79
82
// end.
@@ -95,6 +98,7 @@ TEST(DerivedTimelineTest, HloModuleNameSameScopeRangeIdTest) {
95
98
const absl::string_view kHloModuleName = " hlo_module" ;
96
99
const absl::string_view kKernelDetails = " kernel_details" ;
97
100
XSpace space;
101
+ HloModuleMap hlo_module_map;
98
102
tsl::profiler::GroupMetadataMap group_metadata_map;
99
103
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
100
104
XPlaneBuilder plane_builder (plane);
@@ -107,7 +111,7 @@ TEST(DerivedTimelineTest, HloModuleNameSameScopeRangeIdTest) {
107
111
{{StatType::kHloModule , XStatValue{kHloModuleName }},
108
112
{StatType::kKernelDetails , XStatValue{kKernelDetails }},
109
113
{StatType::kScopeRangeId , XStatValue{int64_t {10 }}}});
110
- GenerateDerivedTimeLines (group_metadata_map, &space);
114
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
111
115
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
112
116
// Only the hlo module line is added and other empty lines are removed at the
113
117
// end.
@@ -128,6 +132,7 @@ TEST(DerivedTimelineTest, HloModuleNameDifferentScopeRangeIdTest) {
128
132
const absl::string_view kHloModuleName = " hlo_module" ;
129
133
const absl::string_view kKernelDetails = " kernel_details" ;
130
134
XSpace space;
135
+ HloModuleMap hlo_module_map;
131
136
tsl::profiler::GroupMetadataMap group_metadata_map;
132
137
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
133
138
XPlaneBuilder plane_builder (plane);
@@ -140,7 +145,7 @@ TEST(DerivedTimelineTest, HloModuleNameDifferentScopeRangeIdTest) {
140
145
{{StatType::kHloModule , XStatValue{kHloModuleName }},
141
146
{StatType::kKernelDetails , XStatValue{kKernelDetails }},
142
147
{StatType::kScopeRangeId , XStatValue{int64_t {20 }}}});
143
- GenerateDerivedTimeLines (group_metadata_map, &space);
148
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
144
149
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
145
150
// Only the hlo module line is added and other empty lines are removed at the
146
151
// end.
@@ -160,6 +165,7 @@ TEST(DerivedTimelineTest, NoHloModuleNameTest) {
160
165
const absl::string_view kKernelDetails = " kernel_details" ;
161
166
const uint64_t kCudaGraphExecId = 1 ;
162
167
XSpace space;
168
+ HloModuleMap hlo_module_map;
163
169
tsl::profiler::GroupMetadataMap group_metadata_map;
164
170
XPlane& plane = *GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
165
171
XPlaneBuilder plane_builder (&plane);
@@ -171,7 +177,7 @@ TEST(DerivedTimelineTest, NoHloModuleNameTest) {
171
177
// Also add a CudaGraph Execution event.
172
178
CreateXEvent (&plane_builder, &line_builder, " op3" , 500 , 100 ,
173
179
{{StatType::kCudaGraphExecId , kCudaGraphExecId }});
174
- GenerateDerivedTimeLines (group_metadata_map, &space);
180
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
175
181
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (&plane);
176
182
// Only the hlo module line is added and other empty lines are removed at the
177
183
// end.
@@ -189,6 +195,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
189
195
const absl::string_view kKernelDetails = " kernel_details" ;
190
196
const uint64_t kCudaGraphExecId = 1 ;
191
197
XSpace space;
198
+ HloModuleMap hlo_module_map;
192
199
tsl::profiler::GroupMetadataMap group_metadata_map;
193
200
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
194
201
XPlaneBuilder plane_builder (plane);
@@ -203,7 +210,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
203
210
CreateXEvent (&plane_builder, &line_builder, " op3" , 500 , 100 ,
204
211
{{StatType::kTfOp , kTfOpName },
205
212
{StatType::kCudaGraphExecId , kCudaGraphExecId }});
206
- GenerateDerivedTimeLines (group_metadata_map, &space);
213
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
207
214
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
208
215
// Only the tf op line is added and other empty lines are removed at the end.
209
216
EXPECT_EQ (plane_visitor.NumLines (), 2 );
@@ -228,6 +235,7 @@ TEST(DerivedTimelineTest, DependencyTest) {
228
235
const absl::string_view kTfOpName = " mul:Mul" ;
229
236
const absl::string_view kKernelDetails = " kernel_details" ;
230
237
XSpace space;
238
+ HloModuleMap hlo_module_map;
231
239
tsl::profiler::GroupMetadataMap group_metadata_map (
232
240
{{0 , {" train 0" }}, {1 , {" train 1" }}});
233
241
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
@@ -241,7 +249,7 @@ TEST(DerivedTimelineTest, DependencyTest) {
241
249
{{StatType::kGroupId , kSecondGroupId },
242
250
{StatType::kTfOp , kTfOpName },
243
251
{StatType::kKernelDetails , kKernelDetails }});
244
- GenerateDerivedTimeLines (group_metadata_map, &space);
252
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
245
253
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
246
254
// The step line and the TF op line are added.
247
255
EXPECT_EQ (plane_visitor.NumLines (), 3 );
@@ -258,6 +266,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeTest) {
258
266
const absl::string_view kTfOpName = " scope1/scope2/mul:Mul" ;
259
267
const absl::string_view kKernelDetails = " kernel_details" ;
260
268
XSpace space;
269
+ HloModuleMap hlo_module_map;
261
270
tsl::profiler::GroupMetadataMap group_metadata_map;
262
271
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
263
272
XPlaneBuilder plane_builder (plane);
@@ -268,7 +277,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeTest) {
268
277
CreateXEvent (&plane_builder, &line_builder, " op2" , 200 , 300 ,
269
278
{{StatType::kTfOp , kTfOpName },
270
279
{StatType::kKernelDetails , kKernelDetails }});
271
- GenerateDerivedTimeLines (group_metadata_map, &space);
280
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
272
281
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
273
282
// The TF name scope line and the TF op line are added.
274
283
EXPECT_EQ (plane_visitor.NumLines (), 3 );
@@ -298,6 +307,7 @@ TEST(DerivedTimelineTest, TfNameScopeMaintainsOrder) {
298
307
const absl::string_view kTfOpName = " scope1/scope2/mul:Mul" ;
299
308
const absl::string_view kKernelDetails = " kernel_details" ;
300
309
XSpace space;
310
+ HloModuleMap hlo_module_map;
301
311
tsl::profiler::GroupMetadataMap group_metadata_map;
302
312
XPlane* plane = tsl::profiler::GetOrCreateTpuXPlane (
303
313
&space, /* device_ordinal=*/ 0 , " TPU V4" , 0 , 0 );
@@ -306,7 +316,7 @@ TEST(DerivedTimelineTest, TfNameScopeMaintainsOrder) {
306
316
CreateXEvent (&plane_builder, &line_builder, " op1" , 0 , 10000 ,
307
317
{{StatType::kTfOp , kTfOpName },
308
318
{StatType::kKernelDetails , kKernelDetails }});
309
- GenerateDerivedTimeLines (group_metadata_map, &space);
319
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
310
320
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
311
321
// The TF name scope line and the TF op line are added.
312
322
EXPECT_EQ (plane_visitor.NumLines (), 3 );
@@ -329,6 +339,7 @@ TEST(DerivedTimelineTest, OnlyDerivedEventsFromLineWithMostEvents) {
329
339
const absl::string_view kTfOpName = " scope1/scope2/mul:Mul" ;
330
340
const absl::string_view kKernelDetails = " kernel_details" ;
331
341
XSpace space;
342
+ HloModuleMap hlo_module_map;
332
343
tsl::profiler::GroupMetadataMap group_metadata_map;
333
344
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
334
345
XPlaneBuilder plane_builder (plane);
@@ -346,7 +357,7 @@ TEST(DerivedTimelineTest, OnlyDerivedEventsFromLineWithMostEvents) {
346
357
{{StatType::kTfOp , kTfOpName },
347
358
{StatType::kKernelDetails , kKernelDetails }});
348
359
// Derive lines for the plane.
349
- GenerateDerivedTimeLines (group_metadata_map, &space);
360
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
350
361
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
351
362
// The TF name scope line and the TF op line are added.
352
363
EXPECT_EQ (plane_visitor.NumLines (), 4 );
@@ -381,6 +392,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
381
392
{
382
393
// Case 1: shirnk is possible.
383
394
XSpace space;
395
+ HloModuleMap hlo_module_map;
384
396
tsl::profiler::GroupMetadataMap group_metadata_map;
385
397
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
386
398
XPlaneBuilder plane_builder (plane);
@@ -391,7 +403,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
391
403
CreateXEvent (
392
404
&plane_builder, &line_builder, " op2" , 20000 , 30000 ,
393
405
{{StatType::kTfOp , " a/d/Mul:Mul" }, {StatType::kKernelDetails , " blah" }});
394
- GenerateDerivedTimeLines (group_metadata_map, &space);
406
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
395
407
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
396
408
// The TF name scope line and the TF op line are added.
397
409
EXPECT_EQ (plane_visitor.NumLines (), 3 );
@@ -415,6 +427,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
415
427
{
416
428
// Case 2: shirnk is impossible due to top event is too small.
417
429
XSpace space;
430
+ HloModuleMap hlo_module_map;
418
431
tsl::profiler::GroupMetadataMap group_metadata_map;
419
432
XPlane* plane = GetOrCreateGpuXPlane (&space, /* device_ordinal=*/ 0 );
420
433
XPlaneBuilder plane_builder (plane);
@@ -428,7 +441,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
428
441
CreateXEvent (
429
442
&plane_builder, &line_builder, " op3" , 20000 , 30000 ,
430
443
{{StatType::kTfOp , " a/g/Mul:Mul" }, {StatType::kKernelDetails , " blah" }});
431
- GenerateDerivedTimeLines (group_metadata_map, &space);
444
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
432
445
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
433
446
// The TF name scope line and the TF op line are added.
434
447
EXPECT_EQ (plane_visitor.NumLines (), 3 );
@@ -466,6 +479,7 @@ TEST(DerivedTimelineTest, XloOpHasCudaGraphStats) {
466
479
constexpr int64_t kCorrelationIdValue = 10000 ;
467
480
const uint64_t kCudaGraphIdValue = 20 ;
468
481
XSpace space;
482
+ HloModuleMap hlo_module_map;
469
483
tsl::profiler::GroupMetadataMap group_metadata_map;
470
484
471
485
// Build Input Plane/Line/Events and derive events from them.
@@ -486,7 +500,7 @@ TEST(DerivedTimelineTest, XloOpHasCudaGraphStats) {
486
500
{StatType::kHloOp , kHloOpName },
487
501
{StatType::kCorrelationId , kCorrelationIdValue },
488
502
{StatType::kCudaGraphId , kCudaGraphIdValue }});
489
- GenerateDerivedTimeLines (group_metadata_map, &space);
503
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
490
504
491
505
// Check that the HLO op line is added and has the extra stats for the first
492
506
// derived event.
@@ -552,6 +566,7 @@ TEST(DerivedTimelineTest, MergeAndNoMerge) {
552
566
constexpr absl::string_view kHloModuleName = " Framework Ops" ;
553
567
static constexpr absl::string_view kTfOpName = " abc:model/layer/MatMul_1" ;
554
568
XSpace space;
569
+ HloModuleMap hlo_module_map;
555
570
tsl::profiler::GroupMetadataMap group_metadata_map;
556
571
XPlane* plane = tsl::profiler::GetOrCreateTpuXPlane (
557
572
&space, /* device_ordinal=*/ 0 , " DummyTPU" , 1.0 , 1.0 );
@@ -568,7 +583,7 @@ TEST(DerivedTimelineTest, MergeAndNoMerge) {
568
583
CreateXEvent (
569
584
&plane_builder, &line_builder, " op3" , 1501 , 300 ,
570
585
{{StatType::kHloModule , kHloModuleName }, {StatType::kTfOp , kTfOpName }});
571
- GenerateDerivedTimeLines (group_metadata_map, &space);
586
+ GenerateDerivedTimeLines (group_metadata_map, &space, hlo_module_map );
572
587
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor (plane);
573
588
// Only the hlo module line is added and other empty lines are removed at the
574
589
// end.
0 commit comments