Skip to content

Commit

Permalink
[Inductor] Fix the lowering of squeeze when input is not contiguous (p…
Browse files Browse the repository at this point in the history
…ytorch#146746)

**Summary**
Fix issue pytorch#143498. The issue happens when we lowering `select = torch.ops.aten.select.int(cat, 1, 0)`.

For example, when `cat` is contiguous with size[2, 2] stride[2,1]

- for eager, it returns a view of size[2,] stride[2,]
- for Inductor lowering, it returns wrong stride 1 instead of 2
```
TensorBox(
  ReinterpretView(
    StorageBox(
      ConcatKernel(name='buf10', layout=FixedLayout('cpu', torch.int64, size=[u0, 2], stride=[2, 1]), inputs=[ComputedBuffer(name='buf8', layout=NonOwningLayout('cpu', torch.int64, size=[u0, 1], stride=[2, 1]), data=Pointwise(device=device(type='cpu'), dtype=torch.int64, inner_fn=<function ReinterpretView.make_loader.<locals>.loader at 0x7f6b856449d0>, ranges=[u0, 1])), ComputedBuffer(name='buf9', layout=NonOwningLayout('cpu', torch.int64, size=[u0, 1], stride=[2, 1]), data=Pointwise(device=device(type='cpu'), dtype=torch.int64, inner_fn=<function ReinterpretView.make_loader.<locals>.loader at 0x7f6b85644790>, ranges=[u0, 1]))])
    ),
    FixedLayout('cpu', torch.int64, size=[u0], stride=[**1**]),
    origins=OrderedSet([select])
  )
)
```

To fix this issue, we give the right stride when lowering of `squeeze`.

**Test Plan**
```
python -u -m pytest -s -v test/inductor/test_unbacked_symints.py -k test_issue_143498
```

Pull Request resolved: pytorch#146746
Approved by: https://github.com/jgong5, https://github.com/sanchitintel, https://github.com/eellison
  • Loading branch information
leslie-fang-intel authored and Tytskiy committed Feb 18, 2025
1 parent e3f205d commit 404b797
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 4 deletions.
3 changes: 0 additions & 3 deletions test/inductor/test_unbacked_symints.py
Original file line number Diff line number Diff line change
Expand Up @@ -389,9 +389,6 @@ def fn(t, start, length):
@skipGPUIf(not HAS_GPU, "requires gpu and triton")
@dynamo_config.patch(capture_dynamic_output_shape_ops=True)
def test_issue_143498(self, device):
if device == "cpu":
raise unittest.SkipTest("CPU Failure")

class Model(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
Expand Down
4 changes: 3 additions & 1 deletion torch/_inductor/ir.py
Original file line number Diff line number Diff line change
Expand Up @@ -2777,7 +2777,9 @@ def fake_reindex(index): # type: ignore[no-untyped-def]
# TODO: unbacked should not diverge from backed in determining striding
# Need to require contiguous here instead of realize, see:
# https://github.com/pytorch/pytorch/issues/145561
x = ExternKernel.require_contiguous(x)
x = ExternKernel.require_exact_strides(
x, FlexibleLayout.contiguous_strides(x.get_size())
)

storage, old_layout = as_storage_and_layout(x, want_contiguous=True)
new_layout = FixedLayout(
Expand Down

0 comments on commit 404b797

Please sign in to comment.