[Bug] Partitioned scatter memory-budget pass misbehaves on XPU (skip/shared-budget logic)
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 9h
- Merged PRs (30d)
- 112
Description
## Bug Description
`TestPartitionedScatterOpt` memory-budget-aware scatter partitioning tests
fail on XPU. The partitioned-scatter Inductor pass is supposed to gate/limit
how many `index_put`/scatter rewrites it performs based on a computed memory
budget (headroom vs. `total_gpu`), but on XPU it (1) applies a partition even
when the test sets the floor so that headroom should be zero ("no headroom,
pass must skip" but a rewrite still happens), and (2) does not share the
per-invocation budget across multiple scatters in the same graph, rewriting
all 3 scatters when only 1 should fit the configured budget.
## Affected Tests
Cases:
op_ut,test/inductor/test_scatter_optimization.py,TestPartitionedScatterOpt.test_memory_aware_partition_count
op_ut,test/inductor/test_scatter_optimization.py,TestPartitionedScatterOpt.test_memory_budget_shared_across_scatters
## Error Message
```
AssertionError: Scalars are not equal!
Expected 0 but got 1.
Absolute difference: 1
Relative difference: inf
floor=total_gpu: no headroom, pass must skip
AssertionError: Scalars are not equal!
Expected 1 but got 3.
Absolute difference: 2
Relative difference: 2.0
budget fits one expanded buffer but 3 scatters were rewritten: the per-invocation budget is not shared
```
## Test Code Snippet
```python
# test/inductor/test_scatter_optimization.py:545-586 (test_memory_aware_partition_count)
def test_memory_aware_partition_count(self):
"""
Verify that live tensor memory constrains num_partitions.
...
Sub-test 1: headroom fits P-1 extra output buffers -> P=4 applied.
Sub-test 2: floor = total_gpu -> no headroom -> pass skips.
"""
torch.manual_seed(12)
N = 11_000_000
output_size = 10_000_000
persist_n = 100_000_000
def f(out, idx, vals, persistent):
scattered = out.index_put([idx], vals, accumulate=True)
return scattered + persistent.sum()
out = torch.zeros(output_size, dtype=torch.float32, device=GPU_TYPE)
idx = torch.randint(0, output_size, (N,), dtype=torch.int64, device=GPU_TYPE)
vals = torch.randn(N, dtype=torch.float32, device=GPU_TYPE)
persistent = torch.randn(persist_n, dtype=torch.float32, device=GPU_TYPE)
with torch.no_grad():
expected = f(out, idx, vals, persistent)
_, total_gpu = torch.xpu.mem_get_info() if GPU_TYPE == "xpu" else torch.cuda.mem_get_info()
...
# test/inductor/test_scatter_optimization.py:636-676 (test_memory_budget_shared_across_scatters)
def test_memory_budget_shared_across_scatters(self):
"""
Several scatters in one graph must share the budget, not each claim it.
...
Three 20 MB scatters here share room for one extra buffer, so exactly
one may be rewritten.
"""
torch.manual_seed(13)
S = 5_000_000
N = 20_000_000
def f(out0, out1, out2, idx, vals):
a = out0.index_put([idx], vals, accumulate=True)
b = out1.index_put([idx], vals, accumulate=True)
c = out2.index_put([idx], vals, accumulate=True)
return a, b, c
outs = [torch.zeros(S, dtype=torch.float32, device=GPU_TYPE) for _ in range(3)]
idx = torch.randint(0, S, (N,), dtype=torch.int64, device=GPU_TYPE)
vals = torch.randn(N, dtype=torch.float32, device=GPU_TYPE)
args = (*outs, idx, vals)
...
```
## Traceback
```
pytest_command:
pytest test/inductor/test_scatter_optimization.py -k TestPartitionedScatterOpt -v
Traceback:
_________ TestPartitionedScatterOpt.test_memory_aware_partition_count __________
File "test/inductor/test_scatter_optimization.py", line 620, in test_memory_aware_partition_count
self.assertEqual(
File ".../torch/_dynamo/test_case.py", line 127, in assertEqual
return super().assertEqual(x, y, *args, **kwargs)
File ".../torch/testing/_internal/common_utils.py", line 4844, in assertEqual
raise error_metas.pop()[0].to_error( # type: ignore[index]
AssertionError: Scalars are not equal!
Expected 0 but got 1.
Absolute difference: 1
Relative difference: inf
floor=total_gpu: no headroom, pass must skip
----------------------------- Captured stderr call -----------------------------
W0807 07:30:41.124000 3203240 site-packages/torch/_inductor/utils.py:3002] [0/0] get_device_tflops: no Triton fallback available for non-CUDA devices. Returning 0.0; roofline estimates will use memory bandwidth only.
I0807 07:30:41.125000 3203240 site-packages/torch/_inductor/compile_fx.py:1784] [0/0] [__inductor_metrics] Graph Metrics:
I0807 07:30:41.125000 3203240 site-packages/torch/_inductor/compile_fx.py:1784] [0/0] [__inductor_metrics] {'num_bytes_accessed': 1020006264, 'nodes_num_elem': [(SchedulerNode(name='op0'), 40000000), (SchedulerNode(name='op1'), 55000000), (SchedulerNode(name='op2'), 100000782), (SchedulerNode(name='op3'), 783), (SchedulerNode(name='op4'), 60000001)], 'node_runtimes': [(SchedulerNode(name='op0'), 0), (SchedulerNode(name='op1'), 0), (SchedulerNode(name='op2'), 0), (SchedulerNode(name='op3'), 0), (SchedulerNode(name='op4'), 0)]}
_____ TestPartitionedScatterOpt.test_memory_budget_shared_across_scatters ______
File "test/inductor/test_scatter_optimization.py", line 683, in test_memory_budget_shared_across_scatters
self.assertEqual(
File ".../torch/_dynamo/test_case.py", line 127, in assertEqual
return super().assertEqual(x, y, *args, **kwargs)
File ".../torch/testing/_internal/common_utils.py", line 4844, in assertEqual
raise error_metas.pop()[0].to_error( # type: ignore[index]
AssertionError: Scalars are not equal!
Expected 1 but got 3.
Absolute difference: 2
Relative difference: 2.0
budget fits one expanded buffer but 3 scatters were rewritten: the per-invocation budget is not shared
----------------------------- Captured stderr call -----------------------------
I0807 07:30:46.008000 3203240 site-packages/torch/_inductor/compile_fx.py:1784] [0/0] [__inductor_metrics] Graph Metrics:
I0807 07:30:46.008000 3203240 site-packages/torch/_inductor/compile_fx.py:1784] [0/0] [__inductor_metrics] {'num_bytes_accessed': 4200000000, 'nodes_num_elem': [(SchedulerNode(name='op0'), 80000000), (SchedulerNode(name='op4'), 80000000), (SchedulerNode(name='op8'), 80000000), (FusedSchedulerNode(nodes=op1_op5_op9), 540000000), (FusedSchedulerNode(nodes=op2_op3), 90000000), (FusedSchedulerNode(nodes=op6_op7), 90000000), (FusedSchedulerNode(nodes=op10_op11), 90000000)], 'node_runtimes': [(SchedulerNode(name='op0'), 0), (SchedulerNode(name='op1'), 0), (SchedulerNode(name='op2'), 0), (SchedulerNode(name='op3'), 0), (SchedulerNode(name='op4'), 0)]}
```
## Root Cause Analysis
Both tests use `torch.xpu.mem_get_info()` / `torch.cuda.mem_get_info()`
device-agnostically to compute `total_gpu`, then set
`config.partitioned_scatter_non_model_floor_bytes` to derive a specific
headroom for the partitioned-scatter Inductor pass to reason about. The
logged `node_runtimes` in the `[__inductor_metrics] Graph Metrics` output are
all `0` for every scheduler node on XPU (`'node_runtimes': [(SchedulerNode(...), 0), ...]`),
which is consistent with the `get_device_tflops: no Triton fallback
available for non-CUDA devices. Returning 0.0` warning seen right before it
(same underlying gap as the `get_device_tflops` XPU issue previously filed,
intel/torch-xpu-ops#4853). If the partitioned-scatter pass's memory/headroom
accounting or its cost model uses these zeroed-out per-node runtime
estimates (directly, or as part of a combined memory-bandwidth/compute
roofline heuristic) to decide how many partitions/rewrites fit the budget,
then on XPU the pass's decision logic can diverge from the CUDA-validated
expected behavior:
- Test 1: with a floor set so there should be no headroom at all, the pass
still performs 1 rewrite instead of skipping (expected 0, got 1).
- Test 2: with a budget sized for exactly 1 of 3 candidate scatters, the
pass rewrites all 3, i.e. it does not appear to be charging/decrementing a
shared budget across the multiple candidates evaluated in one graph
(expected 1, got 3).
These could either be (a) a downstream consequence of zeroed runtime
estimates changing which candidates the cost model considers "worth"
rewriting or how per-candidate cost is compared to the shared budget, or (b)
an independent scatter-partitioning budget-tracking bug that happens to
reproduce here on XPU. Further investigation inside
`torch/_inductor/fx_passes` (partitioned scatter/index_put pass) and its use
of `total_gpu` / `partitioned_scatter_non_model_floor_bytes` /
node runtime-based cost estimates is needed to confirm which.
## Related PyTorch Issues
- None found matching partitioned-scatter memory budget behavior on XPU.
## Related Intel/torch-xpu-ops Issues
- intel/torch-xpu-ops#4853 - Inductor flop counter reports 0 FLOPs on XPU
(get_device_tflops has no non-CUDA fallback) - same zeroed node-runtime
signature observed in this issue's captured stderr, possibly related root
cause.
## Versions
Branch: daisyden/inductor (repo: daisyden/pytorch)
Commit: 8f036e818c26e60f0a28fb51e89a5010315a27c3
Platform: linux, Python 3.10.20
Contributor guide
Assessment
This issue has not been assessed yet.