[Inductor] test_var_mean_uses_two_step_for_non_split_reductions fails on XPU: two-step variance heuristic is gated on device.type == "cuda"
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 9h
- Merged PRs (30d)
- 112
Description
### 🐛 Describe the bug
`test/inductor/test_cooperative_reductions.py::TestVarianceReductionHeuristic::test_var_mean_uses_two_step_for_non_split_reductions`
fails on XPU because Inductor's two-step variance heuristic is gated on
`device.type == "cuda"`, so XPU always falls through to Welford.
#### Root cause
`torch/_inductor/lowering.py::use_two_step_variance` selects the Triton
two-step (sum-of-squares) path only for CUDA:
```python
if device and device.type == "cpu":
threshold = config.cpp.use_two_step_variance_threshold
elif device and device.type == "cuda" and is_triton(x) and is_cuda_two_step_dtype:
min_numel = config.triton.use_two_step_variance_min_numel
threshold = config.triton.use_two_step_variance_threshold
check_for_split = True
else:
threshold = config.unroll_reductions_threshold
```
On XPU the `else` branch is taken, so `threshold` becomes
`config.unroll_reductions_threshold` (8). A 4096-element reduction is far above
that, `use_two_step_variance` returns `False`, and Inductor emits
`triton_helpers.welford_reduce` instead of `tl.sum`. The heuristic was added for
CUDA only in pytorch/pytorch#184383 ("[inductor] Use two-step variance for small
CUDA reductions"); XPU uses the same Triton codegen path and the same
`config.triton.*` knobs, so the gate is too narrow rather than intentionally
CUDA-specific.
Note the three sibling tests in the same class (`..._keeps_welford_for_float32_reductions`,
`..._keeps_welford_for_small_reductions`, `..._keeps_welford_for_split_reductions`)
currently pass on XPU only vacuously: they assert Welford *is* used, which is
always true on XPU since the two-step path is unreachable.
#### Reproducer
```bash
python test/inductor/test_cooperative_reductions.py \
TestVarianceReductionHeuristic.test_var_mean_uses_two_step_for_non_split_reductions
```
(The test currently skips unless `GPU_TYPE` is `cuda`; the XPU run above is with
the skip relaxed to `GPU_TYPE not in ["cuda", "xpu"]`.)
Standalone:
```python
import torch
from torch._inductor.utils import run_and_get_code
def fn(x):
return torch.var_mean(x, dim=-1, correction=0)
x = torch.randn([4, 4096], device="xpu", dtype=torch.float16)
_, (code,) = run_and_get_code(torch.compile(fn, fullgraph=True), x)
print("tl.sum:", "tl.sum" in code, "| welford_reduce:", "welford_reduce" in code)
# XPU: tl.sum: False | welford_reduce: True
# CUDA: tl.sum: True | welford_reduce: False
```
#### Error log
```
======================================================================
FAIL: test_var_mean_uses_two_step_for_non_split_reductions (__main__.TestVarianceReductionHeuristic)
----------------------------------------------------------------------
Traceback (most recent call last):
File "torch/testing/_internal/common_utils.py", line 3886, in wrapper
method(*args, **kwargs)
File "test/inductor/test_cooperative_reductions.py", line 51, in test_var_mean_uses_two_step_for_non_split_reductions
self.assertIn("tl.sum", source_code)
AssertionError: 'tl.sum' not found in '...triton_red_fused_var_mean_0...'
```
The generated kernel uses `triton_helpers.welford_reduce` / `triton_helpers.welford`:
```python
tmp3_mean = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
tmp3_m2 = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
tmp3_weight = tl.zeros([XBLOCK, R0_BLOCK], tl.float32)
for r0_offset in tl.range(0, r0_numel, R0_BLOCK):
...
tmp3_mean_next, tmp3_m2_next, tmp3_weight_next = triton_helpers.welford_reduce(
tmp2, tmp3_mean, tmp3_m2, tmp3_weight, roffset == 0
)
tmp4, tmp5, tmp6 = triton_helpers.welford(tmp3_mean, tmp3_m2, tmp3_weight, 1)
```
Numerics are correct (`assertEqual(result, expected)` passes); only the codegen
strategy differs, so this is a missed-optimization / device-coverage gap rather
than an accuracy bug. Welford does more work per element, so fp16/bf16
`var_mean` over non-split reductions is slower on XPU than it needs to be.
#### Proposed fix
Extend the gate to XPU in `torch/_inductor/lowering.py`:
```python
elif (
device
and device.type in ("cuda", "xpu")
and is_triton(x)
and is_cuda_two_step_dtype
):
```
Verified locally: with this change the kernel emits `tl.sum` and no
`welford_reduce`, and `assertEqual(result, expected)` still passes for fp16 and
bf16. The fix belongs in `pytorch/pytorch`; this issue tracks it on the XPU side.
### Versions
```
PyTorch: 2.15.0.dev20260903+xpu
Triton: 3.6.0
Device: Intel(R) Data Center GPU Max 1550 (PVC)
Driver: 1.6.33578+57
OS: Linux
```
branch: https://github.com/daisyden/pytorch/tree/daisyden/inductor2
---
> This report (root cause analysis, reproducer, and proposed fix) was drafted
> with an AI assistant. The failure, the root cause, and the proposed fix were
> reproduced and reviewed by a human on Intel Data Center GPU Max 1550 before
> submission.
Contributor guide
Research direction
Start with torch/_inductor/lowering.py::use_two_step_variance and test/inductor/test_cooperative_reductions.py::TestVarianceReductionHeuristic. Run the named test with its XPU skip condition relaxed, then verify that the non-split fp16/bf16 case emits tl.sum rather than welford_reduce while the sibling Welford tests and numerical assertions still pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100