[xpu-aglinment] Fix NLLLoss2d backward 32-bit offset overflow
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 9h
- Merged PRs (30d)
- 112
Description
### 🐛 Describe the bug
`F.nll_loss(...).backward()` produces a silently wrong gradient on XPU when the
flattened offset reaches 2^31. The XPU SYCL backward kernel still uses 32-bit
`sample`, `toffset`, and `ioffset` values, so the offset wraps around instead of
addressing the full tensor.
This was found during XPU alignment for upstream PyTorch issue
https://github.com/pytorch/pytorch/issues/190139. The upstream CUDA path was
already fixed by selecting a 64-bit index type when 32-bit indexing is unsafe;
the analogous XPU kernel was never updated, so XPU still exhibits the original
overflow.
### Reproducer
```python
import torch
import torch.nn.functional as F
print("torch", torch.__version__)
print("xpu available", torch.xpu.is_available())
x = torch.zeros((2**16 + 1, 2**15, 1, 1), device="xpu", dtype=torch.float16, requires_grad=True)
target = torch.zeros((2**16 + 1, 1, 1), device="xpu", dtype=torch.long)
F.nll_loss(x, target, reduction="sum").backward()
assert x.grad[0, 0, 0, 0].item() == -1.0
assert x.grad[-1, 0, 0, 0].item() == -1.0
```
### Actual behavior
The first gradient is correct while the final sample gradient is zero, matching
32-bit offset wraparound at `sample * map_nelem * n_classes == 2**31`:
```text
x.grad[0,0,0,0] = -1.0
x.grad[-1,0,0,0] = 0.0
RESULT: confirmed
```
### Expected behavior
Every valid target position receives its gradient regardless of tensor size,
matching CPU/CUDA eager semantics:
```text
x.grad[0,0,0,0] = -1.0
x.grad[-1,0,0,0] = -1.0
```
### Alignment metadata
- Upstream source: https://github.com/pytorch/pytorch/issues/190139
- Source type: upstream issue (CUDA fixed; XPU unaligned)
- Local XPU result: confirmed
- Routed area: `intel/torch-xpu-ops`, XPU NLLLoss2d backward kernel offset typing
- Target repository: `intel-sandbox/torch-xpu-ops-exp`
### Notes
The upstream CUDA issue was closed after its kernel selected a 64-bit index type
when 32-bit indexing is unsafe. The XPU fix should independently align the SYCL
implementation with that guard (widen `sample`/`toffset`/`ioffset` to 64-bit when
the flattened extent exceeds `INT_MAX`) rather than modifying shared PyTorch code.
A regression test should be added under `test/repro/` and run with `pytest`.
### Versions
```text
PyTorch version: 2.14.0.dev20260713+xpu
OS: Ubuntu 22.04.5 LTS (x86_64)
Is XPU available: True
Intel GPU: Intel(R) Data Center GPU Max 1100
Relevant path: src/ATen/native/xpu/sycl/*Loss*.cpp (NLLLoss2d backward)
```
Contributor guide
Assessment
This issue has not been assessed yet.