Skip to content

Commit ade0861

Browse files
Profiler Teamcopybara-github
Profiler Team
authored and
copybara-github
committed
Refactor: Move profiler/convert/google, profiler/convert/oss from tensorflow to xprof
Moves the third_party/tensorflow/core/profiler/convert/google/ and third_party/tensorflow/core/profiler/convert/oss/ directories to third_party/xprof/convert/google/ and third_party/xprof/convert/oss/ respectively. Visibility rules have been updated to reflect the change. PiperOrigin-RevId: 731247147
1 parent 1989332 commit ade0861

7 files changed

+742
-0
lines changed

‎convert/google/BUILD

+57
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
load("//third_party/tensorflow:tensorflow.bzl", "tf_cc_test")
2+
3+
tf_cc_test(
4+
name = "hlo_proto_to_graph_view_test",
5+
size = "small",
6+
srcs = ["hlo_proto_to_graph_view_test.cc"],
7+
deps = [
8+
"//net/proto2/contrib/fixtures:proto-fixture-repository",
9+
"//testing/base/public:gunit_main",
10+
"@com_google_absl//absl/status:statusor",
11+
"@com_google_absl//absl/strings",
12+
"@local_xla//xla/service:hlo_graph_dumper",
13+
"@local_xla//xla/service:hlo_proto_cc",
14+
"@local_xla//xla/tsl/platform:statusor",
15+
"@org_tensorflow//tensorflow/core:test",
16+
"@org_tensorflow//tensorflow/core:test_main",
17+
"@org_tensorflow//tensorflow/core/profiler/convert:hlo_proto_to_graph_view",
18+
],
19+
)
20+
21+
tf_cc_test(
22+
name = "xspace_to_dcn_slack_analysis_test",
23+
srcs = ["xspace_to_dcn_slack_analysis_test.cc"],
24+
deps = [
25+
"//testing/base/public:gunit_main",
26+
"@com_google_absl//absl/container:flat_hash_map",
27+
"@com_google_absl//absl/strings",
28+
"@local_tsl//tsl/profiler/protobuf:xplane_proto_cc",
29+
"@local_xla//xla/hlo/ir:hlo",
30+
"@local_xla//xla/tsl/profiler/utils:math_utils",
31+
"@local_xla//xla/tsl/profiler/utils:tf_xplane_visitor",
32+
"@local_xla//xla/tsl/profiler/utils:timespan",
33+
"@local_xla//xla/tsl/profiler/utils:xplane_builder",
34+
"@local_xla//xla/tsl/profiler/utils:xplane_schema",
35+
"@local_xla//xla/tsl/profiler/utils:xplane_visitor",
36+
"@org_tensorflow//tensorflow/core/profiler/convert:xspace_to_dcn_slack_analysis",
37+
"@org_tensorflow//tensorflow/core/profiler/protobuf:dcn_collective_info_proto_cc",
38+
"@org_tensorflow//tensorflow/core/profiler/protobuf:dcn_slack_analysis_proto_cc",
39+
"@org_tensorflow//tensorflow/core/profiler/utils:hlo_proto_map",
40+
],
41+
)
42+
43+
tf_cc_test(
44+
name = "dcn_slack_analysis_combiner_test",
45+
srcs = ["dcn_slack_analysis_combiner_test.cc"],
46+
deps = [
47+
"//net/proto2/contrib/parse_proto:parse_text_proto",
48+
"//testing/base/public:gunit_main",
49+
"@org_tensorflow//tensorflow/core/profiler/convert:dcn_slack_analysis_combiner",
50+
"@org_tensorflow//tensorflow/core/profiler/protobuf:dcn_slack_analysis_proto_cc",
51+
],
52+
)
53+
54+
exports_files(
55+
["tpu_input_pipeline_analysis_constants.cc"],
56+
visibility = ["@org_tensorflow//tensorflow/core/profiler/convert:__pkg__"],
57+
)

‎convert/google/dcn_slack_analysis_combiner_test.cc

+76
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
#include "tensorflow/core/profiler/convert/dcn_slack_analysis_combiner.h"
16+
17+
#include "net/proto2/contrib/parse_proto/parse_text_proto.h"
18+
#include "testing/base/public/gmock.h"
19+
#include "testing/base/public/gunit.h"
20+
#include "tensorflow/core/profiler/protobuf/dcn_slack_analysis.pb.h"
21+
22+
namespace tensorflow {
23+
namespace profiler {
24+
namespace {
25+
using proto2::contrib::parse_proto::ParseTextProtoOrDie;
26+
using tensorflow::profiler::DcnSlackAnalysis;
27+
using testing::EqualsProto;
28+
using testing::proto::IgnoringRepeatedFieldOrdering;
29+
30+
TEST(DcnSlackAnalysisTest, VerifyCombinedSlackSummary) {
31+
DcnSlackAnalysis analysis1 =
32+
ParseTextProtoOrDie(R"pb(dcn_slack_summary {
33+
rendezvous: "collective-1"
34+
slack_us: 2
35+
occurrences: 1
36+
}
37+
dcn_slack_summary {
38+
rendezvous: "collective-2"
39+
slack_us: 2
40+
occurrences: 1
41+
})pb");
42+
DcnSlackAnalysis analysis2 =
43+
ParseTextProtoOrDie(R"pb(dcn_slack_summary {
44+
rendezvous: "collective-1"
45+
slack_us: 2
46+
occurrences: 1
47+
}
48+
dcn_slack_summary {
49+
rendezvous: "collective-2"
50+
slack_us: 2
51+
occurrences: 1
52+
})pb");
53+
54+
DcnSlackAnalysisCombiner combiner;
55+
combiner.Combine(analysis1);
56+
combiner.Combine(analysis2);
57+
58+
DcnSlackAnalysis combined = combiner.Finalize();
59+
60+
EXPECT_THAT(combined, IgnoringRepeatedFieldOrdering(EqualsProto(
61+
R"pb(dcn_slack_summary {
62+
rendezvous: "collective-1"
63+
slack_us: 2
64+
occurrences: 2
65+
}
66+
dcn_slack_summary {
67+
rendezvous: "collective-2"
68+
slack_us: 2
69+
occurrences: 2
70+
})pb")));
71+
}
72+
73+
} // namespace
74+
75+
} // namespace profiler
76+
} // namespace tensorflow

‎convert/google/hlo_proto_to_graph_view_test.cc

+116
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
#include "tensorflow/core/profiler/convert/hlo_proto_to_graph_view.h"
16+
17+
#include <string>
18+
19+
#include "net/proto2/contrib/fixtures/proto-fixture-repository.h"
20+
#include "testing/base/public/gmock.h"
21+
#include "absl/status/statusor.h"
22+
#include "absl/strings/string_view.h"
23+
#include "xla/service/hlo.pb.h"
24+
#include "xla/service/hlo_graph_dumper.h"
25+
#include "xla/tsl/platform/statusor.h"
26+
#include "tensorflow/core/platform/test.h"
27+
28+
namespace tensorflow {
29+
namespace profiler {
30+
namespace {
31+
32+
using ::proto2::contrib::fixtures::ProtoFixtureRepository;
33+
using ::testing::Eq;
34+
using ::testing::HasSubstr;
35+
using ::testing::StartsWith;
36+
using ::tsl::StatusOr;
37+
using ::xla::HloProto;
38+
39+
HloProto CreateModule(absl::string_view name, int id,
40+
absl::string_view computation) {
41+
HloProto hlo_proto =
42+
ProtoFixtureRepository()
43+
.RegisterValue("@name", name)
44+
.RegisterValue("@id", id)
45+
.RegisterValue("@computation", computation)
46+
.ParseTextProtoOrDie(R"pb(
47+
hlo_module {
48+
name: @name
49+
id: @id
50+
entry_computation_name: @computation
51+
entry_computation_id: 0
52+
computations {
53+
id: 0
54+
name: @computation
55+
instructions: {
56+
name: @computation
57+
id: 1
58+
opcode: "constant"
59+
shape: {
60+
element_type: S32
61+
layout {}
62+
}
63+
literal: {
64+
shape: {
65+
element_type: S32
66+
layout {}
67+
}
68+
s32s: 0
69+
}
70+
}
71+
instructions: {
72+
name: "constant.7"
73+
id: 0
74+
opcode: "constant"
75+
shape: {
76+
element_type: S32
77+
layout {}
78+
}
79+
literal: {
80+
shape: {
81+
element_type: S32
82+
layout {}
83+
}
84+
s32s: 0
85+
}
86+
}
87+
root_id: 1
88+
}
89+
host_program_shape: { result: { element_type: 4 } }
90+
})pb");
91+
return hlo_proto;
92+
}
93+
94+
TEST(HloProtoToGraphViewTest, ConvertHloProtoToGraphView) {
95+
HloProto hlo_proto = CreateModule("module_name", 5, "computation_name");
96+
absl::StatusOr<std::string> response = ConvertHloProtoToGraph(
97+
hlo_proto, "computation_name", 10, xla::HloRenderOptions(),
98+
xla::RenderedGraphFormat::kDot);
99+
EXPECT_THAT(*response, StartsWith("digraph"));
100+
EXPECT_THAT(*response, HasSubstr("computation_name"));
101+
}
102+
103+
TEST(HloProtoToGraphViewTest, ConvertHloProtoToStringView) {
104+
HloProto hlo_proto = CreateModule("module_name", 5, "computation_name");
105+
absl::StatusOr<std::string> response =
106+
ConvertHloProtoToStringView(hlo_proto, false, false);
107+
EXPECT_THAT(*response,
108+
Eq("HloModule module_name, "
109+
"entry_computation_layout={()->s32[]}\n\nENTRY "
110+
"computation_name {\n constant.7 = s32[] constant(0)\n ROOT "
111+
"computation_name = s32[] constant(0)\n}\n\n"));
112+
}
113+
114+
} // namespace
115+
} // namespace profiler
116+
} // namespace tensorflow

‎convert/google/tpu_input_pipeline_analysis_constants.cc

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
#include "tensorflow/core/profiler/convert/tpu_input_pipeline_analysis_constants.h"
16+
17+
#include "absl/strings/string_view.h"
18+
19+
namespace tensorflow {
20+
namespace profiler {
21+
22+
constexpr absl::string_view kProfileAllHostsDoc =
23+
"https://goto.corp.google.com/xprof-faqs#profile-multiple-hosts";
24+
constexpr absl::string_view kSparseCoreV0Name = "BarnaCore";
25+
26+
} // namespace profiler
27+
} // namespace tensorflow

0 commit comments

Comments
 (0)