diff --git a/plugin/xprof/protobuf/BUILD b/plugin/xprof/protobuf/BUILD --- a/plugin/xprof/protobuf/BUILD +++ b/plugin/xprof/protobuf/BUILD @@ -1,7 +1,6 @@ # Description: # Protobufs used by TensorBoard profile plugin. -load("@com_google_protobuf//:protobuf.bzl", "py_proto_library") load("@xla//xla/tsl/platform:build_config.bzl", xprof_proto_library = "tf_proto_library") package( diff --git a/xprof/convert/hlo_to_tools_data.cc b/xprof/convert/hlo_to_tools_data.cc --- a/xprof/convert/hlo_to_tools_data.cc +++ b/xprof/convert/hlo_to_tools_data.cc @@ -62,7 +62,7 @@ absl::StatusOr ConvertHloProtoToMemoryViewer( std::string json_output; tsl::protobuf::util::JsonPrintOptions options; - options.always_print_primitive_fields = true; + options.always_print_fields_with_no_presence = true; auto encoded_status = tsl::protobuf::util::MessageToJsonString( result_or.value(), &json_output, options); if (!encoded_status.ok()) { diff --git a/xprof/convert/xplane_to_memory_profile.cc b/xprof/convert/xplane_to_memory_profile.cc --- a/xprof/convert/xplane_to_memory_profile.cc +++ b/xprof/convert/xplane_to_memory_profile.cc @@ -533,7 +533,7 @@ template absl::Status ConvertProtoToJson(const Proto& proto_output, std::string* json_output) { tsl::protobuf::util::JsonPrintOptions json_options; - json_options.always_print_primitive_fields = true; + json_options.always_print_fields_with_no_presence = true; auto status = tsl::protobuf::util::MessageToJsonString( proto_output, json_output, json_options); if (!status.ok()) { diff --git a/xprof/convert/xplane_to_tools_data.cc b/xprof/convert/xplane_to_tools_data.cc --- a/xprof/convert/xplane_to_tools_data.cc +++ b/xprof/convert/xplane_to_tools_data.cc @@ -248,7 +248,7 @@ absl::StatusOr ConvertMultiXSpacesToPodViewer( std::string json_output; tsl::protobuf::util::JsonPrintOptions opts; - opts.always_print_primitive_fields = true; + opts.always_print_fields_with_no_presence = true; auto encode_status = tsl::protobuf::util::MessageToJsonString( ConvertOpStatsToPodViewer(combined_op_stats), &json_output, opts); if (!encode_status.ok()) { @@ -325,7 +325,7 @@ absl::StatusOr ConvertMultiXSpacesToOpProfileViewer( profile); std::string json_output; tsl::protobuf::util::JsonPrintOptions opts; - opts.always_print_primitive_fields = true; + opts.always_print_fields_with_no_presence = true; auto encode_status = tsl::protobuf::util::MessageToJsonString(profile, &json_output, opts); diff --git a/xprof/convert/trace_viewer/trace_events.cc b/xprof/convert/trace_viewer/trace_events.cc --- a/xprof/convert/trace_viewer/trace_events.cc +++ b/xprof/convert/trace_viewer/trace_events.cc @@ -38,9 +38,7 @@ limitations under the License. #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/file_system.h" -#include "xla/tsl/platform/macros.h" #include "xla/tsl/profiler/utils/timespan.h" -#include "xla/tsl/platform/types.h" #include "xprof/convert/trace_viewer/trace_events_filter_interface.h" #include "xprof/convert/trace_viewer/trace_events_util.h" #include "xprof/convert/trace_viewer/trace_viewer_visibility.h" @@ -49,7 +47,6 @@ limitations under the License. namespace tensorflow { namespace profiler { -using tsl::kint64max; namespace { @@ -137,7 +134,7 @@ std::pair GetLevelBoundsForDuration(uint64_t duration_ps) { for (; i < NumLevels(); ++i) { if (duration_ps > kLayerResolutions[i]) { if (i == 0) { - return std::make_pair(kLayerResolutions[i], kint64max); + return std::make_pair(kLayerResolutions[i], std::numeric_limits::max()); } else { return std::make_pair(kLayerResolutions[i], kLayerResolutions[i - 1]); } diff --git a/xprof/utils/hlo_module_map.h b/xprof/utils/hlo_module_map.h --- a/xprof/utils/hlo_module_map.h 2025-06-10 16:15:37.000000000 +0000 +++ b/xprof/utils/hlo_module_map.h 2026-02-17 17:36:03.348162869 +0000 @@ -29,6 +29,7 @@ #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_module_metadata.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/service/hlo.pb.h" #include "xla/service/hlo_cost_analysis.h" @@ -61,7 +62,7 @@ virtual void ProcessXlaCostAnalysis( const xla::HloCostAnalysis* cost_analysis) = 0; - virtual std::string OpLocationStack(int32_t frame_id) const = 0; + virtual std::string OpLocationStack(xla::StackFrameId frame_id) const = 0; virtual tsl::profiler::OpSourceInfo SourceInfo() const = 0; virtual const ::tensorflow::profiler::PerformanceInfoWrapper* GetPerformanceInfoWrapper() const = 0; @@ -127,7 +128,7 @@ return fused_children_; } - std::string OpLocationStack(int32_t frame_id) const override { + std::string OpLocationStack(xla::StackFrameId frame_id) const override { return GetOpLocationStack(frame_id, *instr_); } diff --git a/xprof/utils/hlo_module_utils.h b/xprof/utils/hlo_module_utils.h --- a/xprof/utils/hlo_module_utils.h 2025-06-10 16:15:37.000000000 +0000 +++ b/xprof/utils/hlo_module_utils.h 2026-02-17 17:34:31.356254501 +0000 @@ -26,9 +26,18 @@ #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_module_metadata.h" #include "xla/hlo/ir/hlo_print_options.h" #include "xla/tsl/profiler/convert/xla_op_utils.h" +// This is a temporary hack to work around version skew between XLA and xprof. +// This may be removed after XLA is updated past Feb 16, 2026. +#ifndef XLA_HAVE_STACK_FRAME_ID +namespace xla { +using StackFrameId = int; +} // namespace xla +#endif + namespace tensorflow { namespace profiler { @@ -84,11 +93,15 @@ return expression; } -inline std::string GetOpLocationStack(int32_t frame_id, +inline std::string GetOpLocationStack(xla::StackFrameId frame_id, const xla::HloInstruction& instr) { std::string stack_lines; xla::HloModule* hlo_module = instr.GetModule(); +#ifdef XLA_HAVE_STACK_FRAME_ID + while (frame_id.valid()) { +#else while (frame_id != 0) { +#endif xla::HloModule::StackFrame frame = hlo_module->get_stack_frame(frame_id); if (frame.empty()) { break; @@ -107,6 +120,6 @@ return {.source_file = instr.metadata().source_file(), .source_line = instr.metadata().source_line(), .stack_frame = stack_frame_id != 0 - ? GetOpLocationStack(stack_frame_id, instr) + ? GetOpLocationStack(xla::StackFrameId{static_cast(stack_frame_id)}, instr) : ""}; } diff --git a/xprof/utils/hlo_module_utils_test.cc b/xprof/utils/hlo_module_utils_test.cc --- a/xprof/utils/hlo_module_utils_test.cc 2025-06-10 16:15:37.000000000 +0000 +++ b/xprof/utils/hlo_module_utils_test.cc 2026-02-17 17:32:56.336350473 +0000 @@ -22,6 +22,7 @@ #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_module_metadata.h" #include "xla/tests/hlo_test_base.h" #include "xla/tsl/platform/statusor.h" @@ -83,7 +84,7 @@ GetModuleWithStackFrames()); const auto* root_instruction = module_with_stack_frames->entry_computation()->root_instruction(); - EXPECT_EQ(GetOpLocationStack(1, *root_instruction), "main.py:10:5\n"); + EXPECT_EQ(GetOpLocationStack(xla::StackFrameId{1}, *root_instruction), "main.py:10:5\n"); } TEST_F(HloModuleUtilsTest, TestGetSourceInfo) {