Skip to content

[XLA] Add StackFrames and StackFrameId classes, and use them on HloModule instead of exposing the StackFrameIndex proto directly. #2307

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[XLA] Add StackFrames and StackFrameId classes, and use them on HloModule instead of exposing the StackFrameIndex proto directly. #2307

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions xprof/utils/hlo_module_map.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ limitations under the License.
#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"
Expand Down Expand Up @@ -62,7 +63,7 @@ class HloInstructionInterface {

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;
Expand Down Expand Up @@ -128,7 +129,7 @@ class HloInstructionWrapper : public HloInstructionInterface {
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_);
}

Expand Down
5 changes: 4 additions & 1 deletion xprof/utils/hlo_module_utils.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ limitations under the License.
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/hlo/ir/hlo_module_metadata.h"
#include "xla/service/hlo.pb.h"
#include "xla/tsl/profiler/convert/xla_op_utils.h"

Expand Down Expand Up @@ -159,7 +160,9 @@ OpSourceInfo GetSourceInfo(const xla::HloInstructionProto& instr,
OpSourceInfo GetSourceInfo(const xla::HloInstruction& instr) {
const auto stack_frame_id = instr.metadata().stack_frame_id();
const std::string stack_frame =
stack_frame_id != 0 ? GetOpLocationStack(stack_frame_id, instr) : "";
stack_frame_id != 0
? GetOpLocationStack(xla::StackFrameId{stack_frame_id}, instr)
: "";
return GetSourceInfo(instr.metadata().source_file(),
instr.metadata().source_line(), stack_frame);
}
Expand Down
16 changes: 15 additions & 1 deletion xprof/utils/hlo_module_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,19 @@ limitations under the License.
#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 {

Expand Down Expand Up @@ -85,11 +95,15 @@ inline std::string UncachedExpression(const xla::HloInstruction& instr,
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;
Expand Down
3 changes: 2 additions & 1 deletion xprof/utils/hlo_module_utils_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ limitations under the License.
#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"

Expand Down Expand Up @@ -129,7 +130,7 @@ TEST_F(HloModuleUtilsTest, TestGetLocationStack) {
GetModuleWithStackFrames());
const auto* root_instruction =
module_with_stack_frames->entry_computation()->root_instruction();
EXPECT_EQ(GetOpLocationStack(2, *root_instruction),
EXPECT_EQ(GetOpLocationStack(xla::StackFrameId{2}, *root_instruction),
"main.py:20:1\nmain.py:10:5\n");
}

Expand Down