Skip to content

Python Frame in ir dump is missing #4206

@JackCaoG

Description

@JackCaoG

🐛 Bug

after enabling XLA_IR_DEBUG and XLA_HLO_DEBUG, I would expect ir graph has the correct python frame and file location. However I am sseing

  %159 = f32[512]{0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0
  %160 = f32[512]{0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0
  %161 = f32[512]{0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0
  %162 = f32[512]{0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0
  %163 = f32[512,512,3,3]{3,2,1,0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0
  %164 = f32[512]{0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0
  %165 = f32[512]{0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0
  %166 = f32[512]{0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0
  %167 = f32[512]{0} xla::device_data(), location=@No Python Frames:-1, device=CPU:0

I took a look at upstream code(since this util has been upstreamed), it should be populate at
https://github.com/pytorch/pytorch/blob/master/torch/csrc/lazy/core/ir.cpp#L65

Node::Node(OpKind op, size_t num_outputs)
    : op_(op), num_outputs_(num_outputs), metadata_(GetMetaDataIfDebugging()) {}

I took a look at https://github.com/pytorch/pytorch/blob/6faa6c68e8b76fb68f3a2b2783685102d0e87c00/torch/csrc/lazy/core/ir_metadata.cpp#L96

MetaData GetMetaDataIfDebugging() {
  if (!FLAGS_torch_lazy_ir_debug) {
    return MetaData();
  }
  MetaData meta;
  meta.scope = GetCurrentScope();
  meta.frame_info = torch::lazy::GetPythonFramesFunction()();
  return meta;
}

and made sure that FLAGS_torch_lazy_ir_debug is currently mapped. We had a ampper in

static bool wants_frames = xla::sys_util::GetEnvBool("XLA_IR_DEBUG", false) |
xla::sys_util::GetEnvBool("XLA_HLO_DEBUG", false);
FLAGS_torch_lazy_ir_debug = wants_frames;
static bool no_scalars =
xla::sys_util::GetEnvBool("XLA_NO_SPECIAL_SCALARS", false);
FLAGS_torch_lazy_handle_special_scalars = !no_scalars;

However GetPythonFramesFunction seems to always return an empty python frame.
https://github.com/pytorch/pytorch/blob/3b8245ab12d54723b6e7bcceb176235f13f0348b/torch/csrc/lazy/core/debug_util.cpp#L58

std::function<std::vector<SourceLocation>()>& GetPythonFramesFunction() {
  static std::function<std::vector<SourceLocation>()> func_ = NoPythonFrames;
  return func_;
}

need to look into why and if I can fix this to unblock my debugging work.

FYI @alanwaketan @Krovatkin

Metadata

Metadata

Assignees

Labels

tracingLazy Tensor tracing

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions