[TorchInductor] `cat_splitwithsizes_replace` misses equivalent negative dimension spelling
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.6k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
`cat_splitwithsizes_replace` removes a `cat` immediately followed by a `split_with_sizes` when the split recovers the original cat inputs. However, the pattern rejects an equivalent negative split dimension.
## Reproducer
```python
from __future__ import annotations
import torch
class Model(torch.nn.Module):
def forward(self, a: torch.Tensor, b: torch.Tensor):
joined = torch.cat([a, b], dim=1)
return torch.split(joined, [3, 5], dim=-1)
def main() -> None:
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required for this test")
torch.manual_seed(1875)
device = torch.device("cuda")
a = torch.randn(2, 3, device=device)
b = torch.randn(2, 5, device=device)
model = Model().to(device).eval()
with torch.no_grad():
expected = model(a, b)
actual = torch.compile(model, backend="inductor", fullgraph=True)(a, b)
torch.cuda.synchronize()
torch.testing.assert_close(actual, expected)
if __name__ == "__main__":
main()
```
For these rank-2 inputs, `dim=-1` and `dim=1` select the same logical axis, so the split exactly recovers the original `cat` inputs.
## Observed behavior
Before post-grad:
```text
cat = aten.cat.default([a, b], 1)
split_with_sizes = aten.split_with_sizes.default(cat, [3, 5], -1)
getitem = split_with_sizes[0]
getitem_1 = split_with_sizes[1]
return (getitem, getitem_1)
```
After post-grad, the graph is unchanged:
```text
cat = aten.cat.default([a, b], 1)
split_with_sizes = aten.split_with_sizes.default(cat, [3, 5], -1)
getitem = split_with_sizes[0]
getitem_1 = split_with_sizes[1]
return (getitem, getitem_1)
```
Since `dim=-1` refers to the same axis as the `cat` dimension here, this `cat` followed by `split_with_sizes` can be replaced with the original inputs, but the optimization is missed.
## Root cause
The check compares the raw integer dimension arguments:
https://github.com/pytorch/pytorch/blob/f07882e293662058329c946feb5615b1aaeedb82/torch/_inductor/fx_passes/post_grad.py#L1777-L1780
For this case, it compares `-1` with `1` and rejects the match even though they refer to the same logical axis. The dimensions need to be canonicalized before comparison.
The reverse `splitwithsizes_cat_replace` pattern appears to have the same raw-dimension comparison.
### Versions
PyTorch main commit `f07882e`
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo
Contributor guide
Research direction
Start in torch/_inductor/fx_passes/post_grad.py around lines 1777-1780 and inspect both the cat_splitwithsizes_replace and splitwithsizes_cat_replace dimension checks. Run the CUDA reproducer from the issue, then verify that equivalent positive and negative dimension spellings are recognized and the resulting compiled graph applies the replacement.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100