Placeholder tracing segmentation faults

[rpsilva-aws]

🐛 Bug

There is a segmentation fault when attempting to stage the HLO dump that involve placeholder tensors. The issue seems to originate from the PjRt computation client, as it attempts to execute the computation. However, for placeholder tensors, it requires the buffers to be present from a given PjRt data, which is not the case for placeholders. We should not be doing any materialization when tracing and returning the HLO, so this might be an undesired behavior.

To Reproduce

import torch
import torch_xla
from torch_xla.core.xla_builder import create_placeholder_tensor

shape = (32,32)
dtype = torch.float32
p = create_placeholder_tensor(shape, dtype)
result = p + 1.0
print(torch_xla._XLAC._get_xla_tensors_hlo(result))

Environment

• Reproducible on XLA backend [CPU/TPU/CUDA]: CPU
• torch_xla version: 2.6+

Additional context

[ysiraichi]

I'm not really familiar with the innards of PjRt, but is this undesired behavior, though? My guess is that _get_xla_tensors_hlo() was implemented under the assumption that the parameters were not placeholders.

Also, I don't think we are materializing anything when tracing. That said, I agree that, in theory, we wouldn't need to materialize the tensors just for generating the HLO. In this sense, I think this is more a feature request than a bug.

[rpsilva-aws]

I meant that it is part of the computation execution, which was somewhat surprising since the intention was to only trace all participating XLA tensors and their respective IDs. There is probably a good reason to also exercise this path. Do you have a possible idea?

We fail here:

xla/torch_xla/csrc/runtime/pjrt_computation_client.cc

Line 741 in c4b45a9

XLA_CHECK(pjrt_device == pjrt_data->buffer->device())

Since the placeholder tensors do not have a PjRt buffer (even empty might not work, as it derives it from the shape), then there is some undesired semantics here imo. This would require a bit more investigation.

I agree and understand the point, but at the same time, it is a new feature that is hitting a seg fault when used with existing APIs, which is more concerning. Perhaps the question is more on why do we execute the computation in the API. I think we can shift it to a feature request, but adding an xcheck that captures the unsupported case would be good. I could raise a PR with a simple: XLA_CHECK(pjrt_data->buffer != nullptr)

[rpsilva-aws]

@ysiraichi #9062 to at least prevent the seg fault, since the use case here seems valid.