[upstream] Failing test case `TestExportOpInfoCPU.test_fake_export_nn_functional_max_unpool2d_grad_target_device_xpu:0_cpu_float32`
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 13h
- Merged PRs (30d)
- 107
Description
### 🐛 Describe the bug
On a CPU-only build, propagating fake tensors whose device is `xpu:0` through
`max_unpool2d` fails with:
```
RuntimeError: PyTorch is not linked with support for xpu devices
```
Repro (needs a build with neither XPU nor CUDA linked; no `torch.export`
required, plain fake propagation is enough):
```python
import torch
from torch._subclasses.fake_tensor import FakeTensorMode
x, idx = torch.nn.functional.max_pool2d(torch.randn(1, 2, 3, 6), 3, return_indices=True)
x = x.contiguous(memory_format=torch.channels_last)
idx = idx.contiguous(memory_format=torch.channels_last)
with FakeTensorMode(allow_non_fake_inputs=True):
torch.nn.functional.max_unpool2d(x.to("xpu:0"), idx.to("xpu:0"), 3, output_size=(1, 2, 3, 6))
```
In CI it shows up as:
```
PYTORCH_OPINFO_SAMPLE_INPUT_INDEX=8 python test/export/test_export_opinfo.py \
TestExportOpInfoCPU.test_fake_export_nn_functional_max_unpool2d_grad_target_device_xpu:0_cpu_float32
```
## Root cause
`_max_unpoolnd` restrides its output to match the input's memory format, for
`cpu` and `xpu` only, because those are the backends whose native kernels
preserve memory format (`torch/_decomp/decompositions.py`, `_restride`):
```python
def _restride(t: TensorLike) -> TensorLike:
if dim == 2 and self.device.type in ("cpu", "xpu"):
return t.contiguous(memory_format=utils.suggest_memory_format(self))
return t
```
`Tensor.contiguous` is not guard-free. Its generated Python binding wraps the
call in a device guard on the tensor's own device
(`torch/csrc/autograd/generated/python_variable_methods.cpp`):
```cpp
static Tensor dispatch_contiguous(const Tensor & self, at::MemoryFormat memory_format) {
pybind11::gil_scoped_release no_gil;
OptionalDeviceGuard device_guard(device_of(self));
return self.contiguous(memory_format);
}
```
A `FakeTensor` reports its *fake* device, so `device_of(self)` is `xpu:0` and
constructing the guard calls `getDeviceGuardImpl(DeviceType::XPU)`, which throws
on a build with no XPU support linked in.
Two things mask this:
- `THPVariable_contiguous` early-returns before `dispatch_contiguous` when the
tensor is already contiguous in the requested format, so only samples with a
non-contiguous (channels-last) input construct the guard. Index 8 is the first
such OpInfo sample.
- The `cuda` target never enters `_restride` at all, so the pre-existing
`cuda:0` parametrization of this test has always passed.
There is also an asymmetry worth deciding on here: `FakeTensorMode` calls
`torch._C._ensureCUDADeviceGuardSet()` when `avoid_device_init` is set, but
`c10::impl::ensureCUDADeviceGuardSet()` only installs a `FakeGuardImpl`
when the CUDA guard is already registered and reports 0 devices. It is a no-op
on a CPU-only build, and there is no XPU equivalent.
## Possible fixes
- Make the decomposition guard-free. `torch.ops.aten.contiguous.default(t,
memory_format=...)`, `t.clone(memory_format=...)` and `aten.clone` all work on
a fake tensor whose device type is not linked into the build; only the
`Tensor.contiguous` Python binding does not.
- Or, more generally, register a fake device guard for the device types a
`FakeTensorMode` is targeting, so any guard-taking binding is safe under fake
propagation rather than fixing decompositions one at a time.
### Versions
CPU-only build without XPU nor CUDA available.
Contributor guide
Assessment
This issue has not been assessed yet.