Repository: pytorch/torchdynamo

Issue:
[inductor] [tacotron2] torchinductor.exc.LoweringException: AssertionError: s1
Repro

~~~

import torch
from torch import tensor, device
import torch.fx as fx
from torchdynamo.testing import rand_strided
from math import inf
from torch.fx.experimental.proxy_tensor import make_fx

# torch version: 1.13.0a0+git6d6e04d
# torch cuda version: 11.6
# torch git version: 6d6e04d6cc9a3d7cf5d9a2eda5baafd5c3ee75c0


# CUDA Info:
# nvcc: NVIDIA (R) Cuda compiler driver
# Copyright (c) 2005-2022 NVIDIA Corporation
# Built on Thu_Feb_10_18:23:41_PST_2022
# Cuda compilation tools, release 11.6, V11.6.112
# Build cuda_11.6.r11.6/compiler.30978841_0

# GPU Hardware Info:
# NVIDIA A100-SXM4-40GB : 8


from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()



    def forward(self, arg0_1):
        squeeze = torch.ops.aten.squeeze.dim(arg0_1, 1);  arg0_1 = None
        return [squeeze]

args = [((64, 80), (80, 1), torch.float32, 'cuda')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
mod = make_fx(Repro())(*args)

from torchinductor.compile_fx import compile_fx_inner

compiled = compile_fx_inner(mod, args)
compiled(*args)
~~~


~~~
Traceback (most recent call last):
  File "/scratch/anijain/work/torchdynamo/torchinductor/graph.py", line 246, in call_function
    return lowerings[target](*args, **kwargs)
  File "/scratch/anijain/work/torchdynamo/torchinductor/lowering.py", line 185, in wrapped
    return decomp_fn(*args, **kwargs)
  File "/scratch/anijain/work/torchdynamo/torchinductor/lowering.py", line 451, in squeeze
    assert removed == 1, removed
AssertionError: s1

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/scratch/anijain/work/torchdynamo/repro.py", line 42, in <module>
    compiled = compile_fx_inner(mod, args)
  File "/scratch/anijain/work/torchdynamo/torchdynamo/debug_utils.py", line 330, in debug_wrapper
    compiled_fn = compiler(gm, example_inputs, **kwargs)
  File "/scratch/anijain/work/torchdynamo/torchinductor/debug.py", line 182, in inner
    return fn(*args, **kwargs)
  File "/scratch/anijain/work/torchdynamo/torchinductor/compile_fx.py", line 59, in compile_fx_inner
    wrap(graph.run)(*example_inputs)
  File "/scratch/anijain/work/torchdynamo/torchdynamo/utils.py", line 75, in time_wrapper
    r = func(*args, **kwargs)
  File "/scratch/anijain/work/torchdynamo/torchinductor/graph.py", line 141, in run
    return super().run(*args)
  File "/scratch/anijain/work/pytorch/torch/fx/interpreter.py", line 129, in run
    self.env[node] = self.run_node(node)
  File "/scratch/anijain/work/torchdynamo/torchinductor/graph.py", line 296, in run_node
    result = super().run_node(n)
  File "/scratch/anijain/work/pytorch/torch/fx/interpreter.py", line 162, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/scratch/anijain/work/torchdynamo/torchinductor/graph.py", line 248, in call_function
    raise LoweringException(e, target, args, kwargs) from e
torchinductor.exc.LoweringException: AssertionError: s1
  target: aten.squeeze.dim
  args[0]: TensorBox(StorageBox(
    InputBuffer(name='arg0_1', layout=FixedLayout('cuda', torch.float32, size=[s0, s1], stride=[s1, 1]))
  ))
  args[1]: 1
~~~


Discussion:

Resolving patch:
diff --git a/torchinductor/lowering.py b/torchinductor/lowering.py
--- a/torchinductor/lowering.py
+++ b/torchinductor/lowering.py
@@ -448,8 +448,11 @@ def squeeze(x, dim=None):
     dim = _validate_dim(x, dim, 0)
     new_shape = list(x.get_size())
     removed = new_shape.pop(dim)
-    assert removed == 1, removed
-    return view(x, new_shape)
+    if V.graph.sizevars.maybe_guard_equals(removed, 1):
+        return view(x, new_shape)
+
+    # squeeze does nothing if the size isn't 1
+    return x
 
 
 @register_lowering([aten.squeeze_])
diff --git a/torchinductor/sizevars.py b/torchinductor/sizevars.py
--- a/torchinductor/sizevars.py
+++ b/torchinductor/sizevars.py
@@ -181,6 +181,8 @@ def guard_equals(self, left: Expr, right: Expr) -> Expr:
 
     def maybe_guard_equals(self, left: Expr, right: Expr) -> bool:
         """if left==right, guard on that fact and return true"""
+        if left == right:
+            return True
         if self.size_hint(left - right) == 0:
             self.guard_equals(left, right)
             return True


Test patch:
diff --git a/test/test_torchinductor.py b/test/test_torchinductor.py
--- a/test/test_torchinductor.py
+++ b/test/test_torchinductor.py
@@ -3137,6 +3137,14 @@ def fn(x, y):
             fn_compiled(x3, y)
             assert same(x2, x3)
 
+        def test_no_op_squeeze(self):
+            @torchdynamo.optimize("inductor")
+            def forward(arg0_1):
+                return torch.ops.aten.squeeze.dim(arg0_1, 1)
+
+            x = torch.randn((10, 20))
+            assert same(x, forward(x))
+
 
 if HAS_CUDA: