[Inductor] Fix the lowering of squeeze when input is not contiguous

leslie-fang-intel · leslie-fang-intel · commit 32a83bb4fc21 · 2025-02-07T21:09:28.000-08:00
ghstack-source-id: 1680d45 Pull Request resolved: #146746
diff --git a/test/inductor/test_unbacked_symints.py b/test/inductor/test_unbacked_symints.py
@@ -385,9 +385,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__()
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -1006,15 +1006,30 @@ def squeeze(x, dim=None):
     dims = OrderedSet((dim,) if not isinstance(dim, tuple) else dim)
 
     new_shape = []
+    new_stride = []
+    is_storage_and_layout = ir.is_storage_and_layout(x)
+    original_stride = x.get_stride() if is_storage_and_layout else None
+    new_offset = x.get_layout().offset if is_storage_and_layout else None
     for d, s in enumerate(x.get_size()):
         if not (
             d in dims
             and V.graph.sizevars.evaluate_expr(sympy.Eq(s, 1), size_oblivious=True)
         ):
             new_shape.append(s)
+            if is_storage_and_layout:
+                assert isinstance(original_stride, list)
+                new_stride.append(original_stride[d])
 
     # squeeze does nothing if the size isn't 1
-    return view(x, new_shape) if new_shape != x.get_size() else x
+    return (
+        (
+            as_strided(x, new_shape, new_stride, new_offset)
+            if is_storage_and_layout
+            else view(x, new_shape)
+        )
+        if new_shape != x.get_size()
+        else x
+    )
 
 
 @register_lowering(aten.squeeze_copy, type_promotion_kind=None)
