pytorch / pytorch/pytorch

[TorchInductor] `cat_splitwithsizes_replace` misses equivalent negative dimension spelling

Open
#196,905 2 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: inductor module: performance oncall: pt2 triaged
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.