Skip to content

Update Single Host XPlane Processing and OpStat conversion to share HloModuleMap. #1310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Update Single Host XPlane Processing and OpStat conversion to share HloModuleMap. #1310

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions xprof/utils/derived_timeline.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -619,9 +619,9 @@ void DeriveEventsFromHostTrace(
}

void GenerateDerivedTimeLines(
const tsl::profiler::GroupMetadataMap& group_metadata_map, XSpace* space) {
HloModuleMap hlo_module_map;
{
const tsl::profiler::GroupMetadataMap& group_metadata_map, XSpace* space,
HloModuleMap& hlo_module_map) {
if (hlo_module_map.empty()) {
HloProtoMap hlo_proto_map;
hlo_proto_map.AddHloProtosFromXSpace(*space);
for (const auto& [program_id, hlo_proto] : hlo_proto_map) {
Expand Down
4 changes: 3 additions & 1 deletion xprof/utils/derived_timeline.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ limitations under the License.
#include "xla/tsl/profiler/utils/timespan.h"
#include "xla/tsl/profiler/utils/xplane_builder.h"
#include "tsl/profiler/protobuf/xplane.pb.h"
#include "xprof/utils/hlo_module_map.h"

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

// Derives `Tensorflow Ops`, `Tensorflow Name Scope` and `Source Code` lines
// from device_trace.
Expand Down
43 changes: 29 additions & 14 deletions xprof/utils/derived_timeline_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ limitations under the License.
#include "xla/tsl/profiler/utils/xplane_visitor.h"
#include "tensorflow/core/platform/test.h"
#include "tsl/profiler/protobuf/xplane.pb.h"
#include "xprof/utils/hlo_module_map.h"

namespace tensorflow {
namespace profiler {
Expand All @@ -53,8 +54,9 @@ using ::tsl::profiler::XStatVisitor;

TEST(DerivedTimelineTest, EmptySpaceTest) {
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
EXPECT_EQ(space.planes_size(), 0);
}

Expand All @@ -63,6 +65,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
const absl::string_view kHloModuleName = "hlo_module";
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
Expand All @@ -73,7 +76,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
CreateXEvent(&plane_builder, &line_builder, "op2", 200, 300,
{{StatType::kHloModule, kHloModuleName},
{StatType::kKernelDetails, kKernelDetails}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// Only the hlo module line is added and other empty lines are removed at the
// end.
Expand All @@ -95,6 +98,7 @@ TEST(DerivedTimelineTest, HloModuleNameSameScopeRangeIdTest) {
const absl::string_view kHloModuleName = "hlo_module";
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
Expand All @@ -107,7 +111,7 @@ TEST(DerivedTimelineTest, HloModuleNameSameScopeRangeIdTest) {
{{StatType::kHloModule, XStatValue{kHloModuleName}},
{StatType::kKernelDetails, XStatValue{kKernelDetails}},
{StatType::kScopeRangeId, XStatValue{int64_t{10}}}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// Only the hlo module line is added and other empty lines are removed at the
// end.
Expand All @@ -128,6 +132,7 @@ TEST(DerivedTimelineTest, HloModuleNameDifferentScopeRangeIdTest) {
const absl::string_view kHloModuleName = "hlo_module";
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
Expand All @@ -140,7 +145,7 @@ TEST(DerivedTimelineTest, HloModuleNameDifferentScopeRangeIdTest) {
{{StatType::kHloModule, XStatValue{kHloModuleName}},
{StatType::kKernelDetails, XStatValue{kKernelDetails}},
{StatType::kScopeRangeId, XStatValue{int64_t{20}}}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// Only the hlo module line is added and other empty lines are removed at the
// end.
Expand All @@ -160,6 +165,7 @@ TEST(DerivedTimelineTest, NoHloModuleNameTest) {
const absl::string_view kKernelDetails = "kernel_details";
const uint64_t kCudaGraphExecId = 1;
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane& plane = *GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(&plane);
Expand All @@ -171,7 +177,7 @@ TEST(DerivedTimelineTest, NoHloModuleNameTest) {
// Also add a CudaGraph Execution event.
CreateXEvent(&plane_builder, &line_builder, "op3", 500, 100,
{{StatType::kCudaGraphExecId, kCudaGraphExecId}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(&plane);
// Only the hlo module line is added and other empty lines are removed at the
// end.
Expand All @@ -189,6 +195,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
const absl::string_view kKernelDetails = "kernel_details";
const uint64_t kCudaGraphExecId = 1;
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
Expand All @@ -203,7 +210,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
CreateXEvent(&plane_builder, &line_builder, "op3", 500, 100,
{{StatType::kTfOp, kTfOpName},
{StatType::kCudaGraphExecId, kCudaGraphExecId}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// Only the tf op line is added and other empty lines are removed at the end.
EXPECT_EQ(plane_visitor.NumLines(), 2);
Expand All @@ -228,6 +235,7 @@ TEST(DerivedTimelineTest, DependencyTest) {
const absl::string_view kTfOpName = "mul:Mul";
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map(
{{0, {"train 0"}}, {1, {"train 1"}}});
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
Expand All @@ -241,7 +249,7 @@ TEST(DerivedTimelineTest, DependencyTest) {
{{StatType::kGroupId, kSecondGroupId},
{StatType::kTfOp, kTfOpName},
{StatType::kKernelDetails, kKernelDetails}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// The step line and the TF op line are added.
EXPECT_EQ(plane_visitor.NumLines(), 3);
Expand All @@ -258,6 +266,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeTest) {
const absl::string_view kTfOpName = "scope1/scope2/mul:Mul";
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
Expand All @@ -268,7 +277,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeTest) {
CreateXEvent(&plane_builder, &line_builder, "op2", 200, 300,
{{StatType::kTfOp, kTfOpName},
{StatType::kKernelDetails, kKernelDetails}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// The TF name scope line and the TF op line are added.
EXPECT_EQ(plane_visitor.NumLines(), 3);
Expand Down Expand Up @@ -298,6 +307,7 @@ TEST(DerivedTimelineTest, TfNameScopeMaintainsOrder) {
const absl::string_view kTfOpName = "scope1/scope2/mul:Mul";
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = tsl::profiler::GetOrCreateTpuXPlane(
&space, /*device_ordinal=*/0, "TPU V4", 0, 0);
Expand All @@ -306,7 +316,7 @@ TEST(DerivedTimelineTest, TfNameScopeMaintainsOrder) {
CreateXEvent(&plane_builder, &line_builder, "op1", 0, 10000,
{{StatType::kTfOp, kTfOpName},
{StatType::kKernelDetails, kKernelDetails}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// The TF name scope line and the TF op line are added.
EXPECT_EQ(plane_visitor.NumLines(), 3);
Expand All @@ -329,6 +339,7 @@ TEST(DerivedTimelineTest, OnlyDerivedEventsFromLineWithMostEvents) {
const absl::string_view kTfOpName = "scope1/scope2/mul:Mul";
const absl::string_view kKernelDetails = "kernel_details";
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
Expand All @@ -346,7 +357,7 @@ TEST(DerivedTimelineTest, OnlyDerivedEventsFromLineWithMostEvents) {
{{StatType::kTfOp, kTfOpName},
{StatType::kKernelDetails, kKernelDetails}});
// Derive lines for the plane.
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// The TF name scope line and the TF op line are added.
EXPECT_EQ(plane_visitor.NumLines(), 4);
Expand Down Expand Up @@ -381,6 +392,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
{
// Case 1: shirnk is possible.
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
Expand All @@ -391,7 +403,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
CreateXEvent(
&plane_builder, &line_builder, "op2", 20000, 30000,
{{StatType::kTfOp, "a/d/Mul:Mul"}, {StatType::kKernelDetails, "blah"}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// The TF name scope line and the TF op line are added.
EXPECT_EQ(plane_visitor.NumLines(), 3);
Expand All @@ -415,6 +427,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
{
// Case 2: shirnk is impossible due to top event is too small.
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;
XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
XPlaneBuilder plane_builder(plane);
Expand All @@ -428,7 +441,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeShrinkTest) {
CreateXEvent(
&plane_builder, &line_builder, "op3", 20000, 30000,
{{StatType::kTfOp, "a/g/Mul:Mul"}, {StatType::kKernelDetails, "blah"}});
GenerateDerivedTimeLines(group_metadata_map, &space);
GenerateDerivedTimeLines(group_metadata_map, &space, hlo_module_map);
XPlaneVisitor plane_visitor = tsl::profiler::CreateTfXPlaneVisitor(plane);
// The TF name scope line and the TF op line are added.
EXPECT_EQ(plane_visitor.NumLines(), 3);
Expand Down Expand Up @@ -466,6 +479,7 @@ TEST(DerivedTimelineTest, XloOpHasCudaGraphStats) {
constexpr int64_t kCorrelationIdValue = 10000;
const uint64_t kCudaGraphIdValue = 20;
XSpace space;
HloModuleMap hlo_module_map;
tsl::profiler::GroupMetadataMap group_metadata_map;

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

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