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_to_tools_data.h"
17
18 #include <optional>
19 #include <string>
20 #include <utility>
21 #include <vector>
22
23 #include "absl/status/statusor.h"
24 #include "absl/strings/string_view.h"
25 #include "tensorflow/compiler/xla/service/hlo.pb.h"
26 #include "tensorflow/core/platform/env.h"
27 #include "tensorflow/core/platform/errors.h"
28 #include "tensorflow/core/platform/logging.h"
29 #include "tensorflow/core/platform/path.h"
30 #include "tensorflow/core/platform/statusor.h"
31 #include "tensorflow/core/profiler/convert/hlo_proto_to_graph_view.h"
32 #include "tensorflow/core/profiler/convert/hlo_proto_to_memory_visualization_utils.h"
33 #include "tensorflow/core/profiler/convert/tool_options.h"
34 #include "tensorflow/core/profiler/convert/xplane_to_hlo.h"
35
36 namespace tensorflow {
37 namespace profiler {
38
39 namespace {
40
ConvertHloProtoToMemoryViewer(const xla::HloProto & hlo_proto)41 StatusOr<std::string> ConvertHloProtoToMemoryViewer(
42 const xla::HloProto& hlo_proto) {
43 static constexpr int kSmallBufferSize = 16 * 1024; // 16KB
44 static constexpr int kMemorySpaceColor = 0; // HBM
45
46 auto result_or = ConvertHloProtoToPreprocessResult(
47 hlo_proto, kSmallBufferSize,
48 GetHeapSimulatorTraceId(hlo_proto, kMemorySpaceColor), kMemorySpaceColor);
49 if (!result_or.ok()) {
50 return errors::Internal(
51 "Failed to convert HLO proto to memory viewer result: ",
52 result_or.status().message());
53 }
54
55 std::string json_output;
56 tensorflow::protobuf::util::JsonPrintOptions options;
57 options.always_print_primitive_fields = true;
58 auto encoded_status = tensorflow::protobuf::util::MessageToJsonString(
59 result_or.value(), &json_output, options);
60 if (!encoded_status.ok()) {
61 const auto& error_message = encoded_status.message();
62 return errors::InvalidArgument(
63 "Failed to convert memory viewer result to JSON format: ",
64 absl::string_view(error_message.data(), error_message.length()));
65 }
66
67 return json_output;
68 }
69
ConvertHloProtoToGraphViewer(const xla::HloProto & hlo_proto,const ToolOptions & options)70 StatusOr<std::string> ConvertHloProtoToGraphViewer(
71 const xla::HloProto& hlo_proto, const ToolOptions& options) {
72 TF_ASSIGN_OR_RETURN(GraphViewerParams params,
73 ParseGraphViewerParams(options));
74 if (params.type == "graph") {
75 return ConvertHloProtoToGraph(hlo_proto, params.node_name,
76 params.graph_width, params.render_options,
77 params.format);
78 } else {
79 return ConvertHloProtoToStringView(hlo_proto, params.verbose,
80 params.show_metadata);
81 }
82 }
83
84 } // namespace
85
ConvertHloProtoToToolData(const std::vector<std::string> & xspace_paths,const absl::string_view tool_name,const ToolOptions & options)86 StatusOr<std::string> ConvertHloProtoToToolData(
87 const std::vector<std::string>& xspace_paths,
88 const absl::string_view tool_name, const ToolOptions& options) {
89 if (xspace_paths.empty()) {
90 return std::string();
91 }
92
93 // <options> must provide a hlo_module_name field to identify the HLO module.
94 std::optional<std::string> hlo_module_name =
95 GetParam<std::string>(options, "hlo_module_name");
96 if (!hlo_module_name.has_value() || hlo_module_name->empty()) {
97 return errors::InvalidArgument(
98 "Can not find HLO module name from options.");
99 }
100
101 // Load HLO module from file.
102 absl::string_view base_dir = tensorflow::io::Dirname(xspace_paths[0]);
103 std::string hlo_proto_file_name =
104 GetHloProtoFileName(base_dir, *hlo_module_name);
105 xla::HloProto hlo_proto;
106 TF_RETURN_IF_ERROR(tensorflow::ReadBinaryProto(
107 tensorflow::Env::Default(), hlo_proto_file_name, &hlo_proto));
108
109 // Convert from HLO proto to tools data.
110 if (tool_name == "memory_viewer") {
111 return ConvertHloProtoToMemoryViewer(hlo_proto);
112 } else if (tool_name == "graph_viewer") {
113 return ConvertHloProtoToGraphViewer(hlo_proto, options);
114 } else {
115 return errors::InvalidArgument(
116 "Can not find tool: ", tool_name,
117 ". Please update to the latest version of Tensorflow.");
118 }
119 }
120
121 } // namespace profiler
122 } // namespace tensorflow
123