intel / intel/torch-xpu-ops

[xpu-aglinment] Fix NLLLoss2d backward 32-bit offset overflow

Open
#4,723 0 comments 0 reactions 1 assignee Claimed by @laifenxiawucha View on GitHub
ai_generated
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.