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
16 #include "tensorflow/core/profiler/convert/hlo_proto_to_graph_view.h"
17
18 #include <string>
19 #include <variant>
20
21 #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h"
22 #include "tensorflow/core/platform/status_matchers.h"
23 #include "tensorflow/core/platform/test.h"
24 #include "tensorflow/core/profiler/convert/tool_options.h"
25 #include "tensorflow/core/protobuf/error_codes.pb.h"
26
27 namespace tensorflow {
28 namespace profiler {
29 namespace {
30
31 using ::tensorflow::testing::StatusIs;
32 using ::testing::HasSubstr;
33
TEST(GraphViewerParamsTest,GraphType)34 TEST(GraphViewerParamsTest, GraphType) {
35 // Default for graph type.
36 ToolOptions options1;
37 options1["type"] = "graph";
38 TF_ASSERT_OK_AND_ASSIGN(GraphViewerParams params1,
39 ParseGraphViewerParams(options1));
40 EXPECT_EQ(params1.type, "graph");
41 EXPECT_EQ(params1.node_name, "");
42 EXPECT_EQ(params1.graph_width, 3);
43 EXPECT_EQ(params1.render_options.show_backend_config, false);
44 EXPECT_EQ(params1.render_options.show_fusion_subcomputations, true);
45 EXPECT_EQ(params1.format, xla::RenderedGraphFormat::kUrl);
46
47 // User defined options for graph type.
48 ToolOptions options2;
49 options2["type"] = "graph";
50 options2["node_name"] = "fusion.111";
51 options2["graph_width"] = 10;
52 options2["show_metadata"] = 1;
53 options2["merge_fusion"] = 1;
54 options2["format"] = "html";
55 TF_ASSERT_OK_AND_ASSIGN(GraphViewerParams params2,
56 ParseGraphViewerParams(options2));
57 EXPECT_EQ(params2.type, "graph");
58 EXPECT_EQ(params2.node_name, "fusion.111");
59 EXPECT_EQ(params2.graph_width, 10);
60 EXPECT_EQ(params2.render_options.show_backend_config, true);
61 EXPECT_EQ(params2.render_options.show_fusion_subcomputations, false);
62 EXPECT_EQ(params2.format, xla::RenderedGraphFormat::kHtml);
63 }
64
TEST(GraphViewerParamsTest,ShortTxtType)65 TEST(GraphViewerParamsTest, ShortTxtType) {
66 // Default for short txt type.
67 ToolOptions options1;
68 options1["type"] = "short_txt";
69 TF_ASSERT_OK_AND_ASSIGN(GraphViewerParams params1,
70 ParseGraphViewerParams(options1));
71 EXPECT_EQ(params1.type, "short_txt");
72 EXPECT_EQ(params1.verbose, false);
73 EXPECT_EQ(params1.show_metadata, false);
74
75 // User defined options for short txt type.
76 ToolOptions options2;
77 options2["type"] = "short_txt";
78 options2["show_metadata"] = 1;
79 TF_ASSERT_OK_AND_ASSIGN(GraphViewerParams params2,
80 ParseGraphViewerParams(options2));
81 EXPECT_EQ(params2.type, "short_txt");
82 EXPECT_EQ(params2.verbose, false);
83 EXPECT_EQ(params2.show_metadata, true);
84 }
85
TEST(GraphViewerParamsTest,LongTxtType)86 TEST(GraphViewerParamsTest, LongTxtType) {
87 // Default for long txt type.
88 ToolOptions options1;
89 options1["type"] = "long_txt";
90 TF_ASSERT_OK_AND_ASSIGN(GraphViewerParams params1,
91 ParseGraphViewerParams(options1));
92 EXPECT_EQ(params1.type, "long_txt");
93 EXPECT_EQ(params1.verbose, true);
94 EXPECT_EQ(params1.show_metadata, false);
95
96 // User defined options for long txt type.
97 ToolOptions options2;
98 options2["type"] = "long_txt";
99 options2["show_metadata"] = 1;
100 TF_ASSERT_OK_AND_ASSIGN(GraphViewerParams params2,
101 ParseGraphViewerParams(options2));
102 EXPECT_EQ(params2.type, "long_txt");
103 EXPECT_EQ(params2.verbose, true);
104 EXPECT_EQ(params2.show_metadata, true);
105 }
106
TEST(GraphViewerParamsTest,OtherTypes)107 TEST(GraphViewerParamsTest, OtherTypes) {
108 ToolOptions options1;
109 EXPECT_THAT(ParseGraphViewerParams(options1),
110 StatusIs(error::INVALID_ARGUMENT,
111 HasSubstr("Graph viewer must provide a type option")));
112
113 ToolOptions options2;
114 options2["type"] = "abcd";
115 EXPECT_THAT(ParseGraphViewerParams(options2),
116 StatusIs(error::INVALID_ARGUMENT,
117 HasSubstr("Unknown graph viewer type option: abcd")));
118 }
119
120 } // namespace
121 } // namespace profiler
122 } // namespace tensorflow
123