intel / intel/torch-xpu-ops

[Inductor] test_var_mean_uses_two_step_for_non_split_reductions fails on XPU: two-step variance heuristic is gated on device.type == "cuda"

Open Beginner friendly
#5,221 2 comments 0 reactions 0 assignees View on GitHub
agent:active module: inductor test: ut
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.