1- diff --git a/tensorflow/core/profiler/convert/xplane_to_tools_data.cc b/tensorflow/core/profiler/convert/xplane_to_tools_data.cc
2- index 6a3c045db5c..fbb8a84786e 100644
3- --- a/tensorflow/core/profiler/convert/xplane_to_tools_data.cc
4- +++ b/tensorflow/core/profiler/convert/xplane_to_tools_data.cc
5- @@ -20,9 +20,13 @@ limitations under the License.
6-
7- #include "absl/strings/str_format.h"
8- #include "absl/strings/string_view.h"
9- + #include "tensorflow/compiler/xla/service/hlo.pb.h"
10- + #include "tensorflow/core/lib/gtl/map_util.h"
11- #include "tensorflow/core/platform/env.h"
12- #include "tensorflow/core/platform/logging.h"
13- + #include "tensorflow/core/platform/path.h"
14- #include "tensorflow/core/platform/protobuf.h"
15- + #include "tensorflow/core/profiler/convert/hlo_proto_to_memory_visualization_utils.h"
16- #include "tensorflow/core/profiler/convert/hlo_to_tools_data.h"
17- #include "tensorflow/core/profiler/convert/op_stats_to_input_pipeline_analysis.h"
18- #include "tensorflow/core/profiler/convert/op_stats_to_op_profile.h"
19- @@ -232,7 +236,77 @@ std::pair<std::string, bool> ConvertMultiXSpacesToOpProfileViewer(
20-
21- return std::make_pair(profile.SerializeAsString(), true);
22- }
23- - } // namespace
24- +
25- + std::pair<std::string, bool> ConvertHloProtoToMemoryViewer(
26- + const xla::HloProto& hlo_proto) {
27- + static constexpr int kSmallBufferSize = 16 * 1024; // 16KB
28- + static constexpr int kMemorySpaceColor = 0; // HBM
29- +
30- + auto result_or = ConvertHloProtoToPreprocessResult(
31- + hlo_proto, kSmallBufferSize,
32- + GetHeapSimulatorTraceId(hlo_proto, kMemorySpaceColor), kMemorySpaceColor);
33- + if (!result_or.ok()) {
34- + LOG(ERROR) << "Failed to convert HLO proto to memory viewer result: "
35- + << result_or.status().message();
36- + return std::make_pair("", false);
37- + }
38- +
39- + std::string json_output;
40- + tensorflow::protobuf::util::JsonPrintOptions options;
41- + options.always_print_primitive_fields = true;
42- + auto encoded_status = tensorflow::protobuf::util::MessageToJsonString(
43- + result_or.value(), &json_output, options);
44- + if (!encoded_status.ok()) {
45- + LOG(ERROR) << "Failed to convert memory viewer result to JSON format: "
46- + << encoded_status.message();
47- + return std::make_pair("", false);
48- + }
49- +
50- + return std::make_pair(json_output, true);
51- + }
52- +
53- + std::pair<std::string, bool> ConvertHloProtoToToolData(
54- + const std::vector<std::string>& xspace_paths,
55- + const absl::string_view tool_name,
56- + const absl::flat_hash_map<std::string, std::variant<int, std::string>>&
57- + options) {
58- + if (xspace_paths.empty()) {
59- + return std::make_pair("", false);
60- + }
61- +
62- + // <options> must provide a hlo_module_name field to identify the HLO module.
63- + auto* result = gtl::FindOrNull(options, "hlo_module_name");
64- + if (!result) {
65- + LOG(ERROR) << "Can not find HLO module name from options.";
66- + return std::make_pair("", false);
67- + }
68- + const std::string* hlo_module_name = std::get_if<std::string>(result);
69- + if (!hlo_module_name || hlo_module_name->empty()) {
70- + LOG(ERROR) << "Can not find HLO module name from options.";
71- + return std::make_pair("", false);
72- + }
73- +
74- + // Load HLO module from file.
75- + absl::string_view base_dir = tensorflow::io::Dirname(xspace_paths[0]);
76- + std::string hlo_proto_file_name =
77- + GetHloProtoFileName(base_dir, *hlo_module_name);
78- + xla::HloProto hlo_proto;
79- + tensorflow::Status status = tensorflow::ReadBinaryProto(
80- + tensorflow::Env::Default(), hlo_proto_file_name, &hlo_proto);
81- + if (!status.ok()) {
82- + LOG(ERROR) << "Failed to read HLO proto: " << status.error_message();
83- + return std::make_pair("", false);
84- + }
85- +
86- + // Convert from HLO proto to tools data.
87- + if (tool_name == "memory_viewer") {
88- + return ConvertHloProtoToMemoryViewer(hlo_proto);
89- + } else {
90- + LOG(ERROR) << "Can not find tool: " << tool_name
91- + << ". Please update to the latest version of Tensorflow.";
92- + return std::make_pair("", false);
93- + }
94- + }
95-
96- std::pair<std::string, bool> ConvertMultiXSpacesToToolData(
97- const std::vector<XSpace>& xspaces,
98- @@ -278,5 +352,6 @@ std::pair<std::string, bool> ConvertMultiXSpacesToToolData(
99- }
100- }
101-
102- + } // namespace
103- } // namespace profiler
104- } // namespace tensorflow
105- --- a/tensorflow/core/platform/macros.h
106- +++ b/tensorflow/core/platform/macros.h
107- @@ -21,7 +21,7 @@ limitations under the License.
108- // Compiler supports GCC-style attributes
109- #define TF_ATTRIBUTE_NORETURN __attribute__((noreturn))
110- #define TF_ATTRIBUTE_ALWAYS_INLINE __attribute__((always_inline))
111- - #define TF_ATTRIBUTE_NOINLINE __attribute__((noinline))
112- + #define TF_ATTRIBUTE_NOINLINE [[noinline]]
113- #define TF_ATTRIBUTE_UNUSED __attribute__((unused))
114- #define TF_ATTRIBUTE_COLD __attribute__((cold))
115- #define TF_ATTRIBUTE_WEAK __attribute__((weak))
1+ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_fp16.h b/tensorflow/compiler/xla/service/cpu/runtime_fp16.h
2+ index 9fe020d5937..32774c2f3c0 100644
1163--- a/tensorflow/compiler/xla/service/cpu/runtime_fp16.h
1174+++ b/tensorflow/compiler/xla/service/cpu/runtime_fp16.h
1185@@ -18,12 +18,7 @@ limitations under the License.
@@ -129,4 +16,19 @@ index 6a3c045db5c..fbb8a84786e 100644
12916 // Older versions of Clang don't have _Float16. Since both float and _Float16
13017 // are passed in the same register we can use the wider type and careful casting
13118 // to conform to x86_64 psABI. This only works with the assumption that we're
132-
19+ diff --git a/tensorflow/core/profiler/convert/xplane_to_tools_data.cc b/tensorflow/core/profiler/convert/xplane_to_tools_data.cc
20+ index b70ea8af5df..b54a895e306 100644
21+ --- a/tensorflow/core/profiler/convert/xplane_to_tools_data.cc
22+ +++ b/tensorflow/core/profiler/convert/xplane_to_tools_data.cc
23+ @@ -230,8 +230,9 @@ StatusOr<std::string> ConvertMultiXSpacesToToolData(
24+ return ConvertMultiXSpacesToTfDataBottleneckAnalysis(session_snapshot);
25+ } else if (tool_name == "op_profile") {
26+ return ConvertMultiXSpacesToOpProfileViewer(session_snapshot);
27+ - } else if (tool_name == "memory_viewer" || tool_name == "graph_viewer") {
28+ - return ConvertHloProtoToToolData(session_snapshot, tool_name, options);
29+ + // this function is not being used by xla
30+ + // } else if (tool_name == "memory_viewer" || tool_name == "graph_viewer") {
31+ + // return ConvertHloProtoToToolData(session_snapshot, tool_name, options);
32+ } else if (tool_name == "tool_names") {
33+ return GetAvailableToolNames(session_snapshot);
34+ } else {
0 commit comments