xpu: support torch.utils.data.DataLoader(pin_memory_device='xpu')

[dvrogozh]

torch.utils.data.DataLoader(pin_memory_device='xpu') is currently not supported with upstream PyTorch XPU backend. I know that with IPEX this feature was supported. Please, support the feature if it's relevant for XPU or update PyTorch documentation with optimization notice that feature is not needed for XPU.

Relevant environment:

• torch - a0429c0
• torchvision - pytorch/vision@947ae1d

This can be reproduced with:

$ cat data_loader_mempin.py
import torch
import torchvision

print('torch.__version__:' + torch.__version__)
print('torch.xpu.is_available(): ' + str(torch.xpu.is_available()))

dataset = torchvision.datasets.FakeData(size=10, transform=torchvision.transforms.ToTensor())
loader = torch.utils.data.DataLoader(dataset, pin_memory=True, pin_memory_device='xpu')

with torch.inference_mode():
    for batch_idx, (sample, target) in enumerate(loader):
        print(batch_idx)
        print(sample.is_pinned())

Gives this output:

$ python3 data_loader_mempin.py
torch.__version__:2.4.0a0+gita0429c0
torch.xpu.is_available(): True
Traceback (most recent call last):
  File "/home/dvrogozh/tests/data_loader_mempin.py", line 11, in <module>
    for batch_idx, (sample, target) in enumerate(loader):
  File "/home/dvrogozh/git/pytorch/pytorch/torch/utils/data/dataloader.py", line 629, in __next__
    data = self._next_data()
  File "/home/dvrogozh/git/pytorch/pytorch/torch/utils/data/dataloader.py", line 674, in _next_data
    data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)
  File "/home/dvrogozh/git/pytorch/pytorch/torch/utils/data/_utils/pin_memory.py", line 88, in pin_memory
    clone[i] = pin_memory(item, device)
  File "/home/dvrogozh/git/pytorch/pytorch/torch/utils/data/_utils/pin_memory.py", line 58, in pin_memory
    return data.pin_memory(device)
NotImplementedError: The operator 'aten::_pin_memory' is not currently implemented for the XPU device. Please open a feature on https://github.com/intel/torch-xpu-ops/issues. You can set the environment variable `PYTORCH_ENABLE_XPU_FALLBACK=1` to use the CPU implementation as a fallback for XPU unimplemented operators. WARNING: this will bring unexpected performance compared with running natively on XPU.

CC: @jgong5 @mingfeima @XiaobingSuper @ashokei @jingxu10 @gujinghui @EikanWang @fengyuan14 @guangyey

cc @andrewkho @gokulavasan @ssnl @VitalyFedyunin @dzhulgakov @gujinghui @EikanWang @fengyuan14 @guangyey

[dvrogozh]

Issue at torch-xpu-ops (filed as noted in the log): intel/torch-xpu-ops#261

[mikaylagawarecki]

cc @albanD just in case this is related to unified API for accelerator backends

[gujinghui]

The unified API for pin memory is already in PyTorch master, right?
The only missed part for this issue is pin_memory op impl in torch_xpu_ops, which is already in progress.
One issue in torch_xpu_ops is good enough? I don't think it's necessary to create issue in PyTorch here.

[dvrogozh]

The unified API for pin memory is already in PyTorch master, right?
The only missed part for this issue is pin_memory op impl in torch_xpu_ops, which is already in progress.

According a printed call stack - yes. Once pin_memory op will be implemented in torch_xpu_ops, then someone will need to recheck whether pin_memory API works for XPU (I will do that once fix or PR will be available).

One issue in torch_xpu_ops is good enough? I don't think it's necessary to create issue in PyTorch here.

Pytorch has better visibility than torch_xpu_ops. To me it does make sense to have issue here in pytorch rather than in torch_xpu_ops. Besides this issue affect xpu backend at pytorch API level - another reason to have it as pytorch issue. This way it can be tracked towards xpu backend readiness for specific pytorch target.

[gujinghui]

If need more attention, pls. add @EikanWang , @guangyey , @fengyuan14 and me as assignee or reviewer for all xpu related PR and issues. It's more helpful to solve the issue in short-term if urgent. Thanks.

[dvrogozh]

If need more attention, pls. add @EikanWang , @guangyey , @fengyuan14 and me as assignee or reviewer for all xpu related PR and issues.

Good to know. However, I don't have the rights to add add assignees or reviewers. But I will be adding to CC in comments. As of now I do not request special attention to this issue from my side and can wait for the fix. Mostly it's to give heads up to other people who might be trying XPU backend that this feature is missing.

[gujinghui]

If need more attention, pls. add @EikanWang , @guangyey , @fengyuan14 and me as assignee or reviewer for all xpu related PR and issues.

Good to know. However, I don't have the rights to add add assignees or reviewers. But I will be adding to CC in comments. As of now I do not request special attention to this issue from my side and can wait for the fix. Mostly it's to give heads up to other people who might be trying XPU backend that this feature is missing.

It's already in our plan, actually. We will take attention on it. Thanks.