Repository: pytorch/torchdynamo

Issue:
diffusion inductor errors
## Repro

The [lucidrains](https://github.com/lucidrains) models are probably a good idea to add to test suite.

`pip install denoising_diffusion_pytorch`

```python
import torch
import torchdynamo
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
)

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,   # number of steps
    loss_type = 'l1'    # L1 or L2
).cuda()

@torchdynamo.optimize("inductor")
def train():
    training_images = torch.randn(8, 3, 128, 128).cuda() # images are normalized from 0 to 1
    loss = diffusion(training_images)
    loss.backward()

train()
# after a lot of training

sampled_images = diffusion.sample(batch_size = 4)
sampled_images.shape # (4, 3, 128, 128)```
```

## Logs

```
(dynamo) ubuntu@ip-172-31-31-152:~/tests$ python denoise.py 
torchdynamo.symbolic_convert: [WARNING] Graph break: call_function BuiltinVariable(callable) [NestedUserFunctionVariable()] {} from user code at   File "denoise.py", line 20, in train
    loss = diffusion(training_images)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py", line 689, in forward
    return self.p_losses(img, t, *args, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py", line 650, in p_losses
    noise = default(noise, lambda: torch.randn_like(x_start))
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py", line 38, in default
    return d() if callable(d) else d

torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.graph: [WARNING] Creating implicit fallback for:
  target: aten.randint.low
  args[0]: 0
  args[1]: 1000
  args[2]: [8]
  kwargs: {'device': device(type='cuda', index=0), 'pin_memory': False}
torchinductor.ir: [WARNING] Using FallbackKernel: torch.ops.aten.randint.low
torchinductor.graph: [WARNING] Creating implicit fallback for:
  target: aten.randn_like.default
  args[0]: TensorBox(StorageBox(
    InputBuffer(name='arg0_1', layout=FixedLayout('cuda', torch.float32, size=[s0, s1, s2, s2], stride=[s1*s2**2, s2**2, s2, 1]))
  ))
  kwargs: {'dtype': torch.float32, 'layout': torch.strided, 'device': device(type='cuda', index=0), 'pin_memory': False}
torchinductor.ir: [WARNING] Using FallbackKernel: torch.ops.aten.randn_like.default
torchdynamo.symbolic_convert: [WARNING] Graph break: UNPACK_SEQUENCE NNModuleVariable() from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py", line 668, in <graph break in p_losses>
    model_out = self.model(x, t, x_self_cond)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py", line 362, in forward
    for block1, block2, attn, downsample in self.downs:

Traceback (most recent call last):
  File "/home/ubuntu/torchdynamo/torchinductor/graph.py", line 255, in call_function
    return lowerings[target](*args, **kwargs)
  File "/home/ubuntu/torchdynamo/torchinductor/lowering.py", line 187, in wrapped
    return decomp_fn(*args, **kwargs)
  File "/home/ubuntu/torchdynamo/torchinductor/lowering.py", line 1456, in gather
    assert 0 <= dim < len(x.get_size())
AssertionError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "denoise.py", line 23, in <module>
    train()
  File "/home/ubuntu/torchdynamo/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "denoise.py", line 17, in train
    @torchdynamo.optimize("inductor")
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py", line 683, in forward
    def forward(self, img, *args, **kwargs):
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py", line 648, in p_losses
    def p_losses(self, x_start, t, noise = None):
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py", line 650, in <graph break in p_losses>
    noise = default(noise, lambda: torch.randn_like(x_start))
  File "/home/ubuntu/torchdynamo/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 904, in forward
    return compiled_f(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 890, in new_func
    compiled_fn = create_aot_dispatcher_function(
  File "/home/ubuntu/torchdynamo/torchdynamo/utils.py", line 76, in time_wrapper
    r = func(*args, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 468, in create_aot_dispatcher_function
    return aot_dispatch_base(flat_fn, fake_flat_tensor_args, aot_config)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 278, in aot_dispatch_base
    compiled_fw = aot_config.fw_compiler(fw_module, flat_args)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 237, in f
    out_f = compiler(fx_g, inps)
  File "/home/ubuntu/torchdynamo/torchdynamo/utils.py", line 76, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/ubuntu/torchdynamo/torchinductor/compile_fx.py", line 246, in fw_compiler
    return compile_fx_inner(
  File "/home/ubuntu/torchdynamo/torchdynamo/debug_utils.py", line 360, in debug_wrapper
    compiled_fn = compiler(gm, example_inputs, **kwargs)
  File "/home/ubuntu/torchdynamo/torchinductor/debug.py", line 182, in inner
    return fn(*args, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/home/ubuntu/torchdynamo/torchinductor/compile_fx.py", line 83, in compile_fx_inner
    wrap(graph.run)(*example_inputs)
  File "/home/ubuntu/torchdynamo/torchdynamo/utils.py", line 76, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/ubuntu/torchdynamo/torchinductor/graph.py", line 145, in run
    return super().run(*args)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/fx/interpreter.py", line 129, in run
    self.env[node] = self.run_node(node)
  File "/home/ubuntu/torchdynamo/torchinductor/graph.py", line 312, in run_node
    result = super().run_node(n)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/fx/interpreter.py", line 162, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/home/ubuntu/torchdynamo/torchinductor/graph.py", line 257, in call_function
    raise LoweringException(e, target, args, kwargs) from e
torchinductor.exc.LoweringException: AssertionError: 
  target: aten.gather.default
  args[0]: TensorBox(StorageBox(
    InputBuffer(name='arg0_1', layout=FixedLayout('cuda', torch.float32, size=[s0], stride=[1]))
  ))
  args[1]: -1
  args[2]: TensorBox(StorageBox(
    InputBuffer(name='arg3_1', layout=FixedLayout('cuda', torch.int64, size=[s1], stride=[1]))
  ))
```


Discussion:
From the error log I can say it's trying to pass negative ```dim``` to ```gather``` function, which is supported by eager mode:
```
>>> t = torch.tensor([[1, 2], [3, 4]])
>>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
tensor([[1, 1],
        [4, 3]])
>>> torch.gather(t, -1, torch.tensor([[0, 0], [1, 0]]))
tensor([[1, 1],
        [4, 3]])
```
I think inductor should support this, will send a fix soon.

Resolving patch:
diff --git a/torchinductor/lowering.py b/torchinductor/lowering.py
--- a/torchinductor/lowering.py
+++ b/torchinductor/lowering.py
@@ -1453,16 +1453,17 @@ def full(size, fill_value, **kwargs):
 @register_lowering(aten.gather, type_promote=False)
 def gather(x, dim, index):
     assert isinstance(x, TensorBox)
-    assert isinstance(dim, int)
-    assert "int" in str(index.get_dtype())
-    assert 0 <= dim < len(x.get_size())
+    assert index.get_dtype() == torch.int64
+    offset = len(x.get_size()) == 0
+    dim = _validate_dim(x, dim, offset)
 
     x_loader = x.make_loader()
     index_loader = index.make_loader()
 
     def fn(idx):
         idx = list(idx)
-        idx[dim] = ops.indirect_indexing(index_loader(idx))
+        if len(idx) != 0:
+            idx[dim] = ops.indirect_indexing(index_loader(idx))
         return x_loader(idx)
 
     return Pointwise.create(


Test patch:
diff --git a/test/test_torchinductor.py b/test/test_torchinductor.py
--- a/test/test_torchinductor.py
+++ b/test/test_torchinductor.py
@@ -1108,9 +1108,12 @@ def fn(a, b):
             check_lowp=False,
         )
 
-    def test_gather(self):
+    def test_gather1(self):
         def fn(a, b):
-            return (torch.gather(a.expand([4, 5, 10, 6]), 3, b + 1),)
+            return (
+                torch.gather(a.expand([4, 5, 10, 6]), 3, b + 1),
+                torch.gather(a.expand([4, 5, 10, 6]), -1, b + 1),
+            )
 
         self.common(
             fn,
@@ -1120,6 +1123,15 @@ def fn(a, b):
             ),
         )
 
+    def test_gather2(self):
+        # 0d tensor
+        def fn(a, b):
+            return torch.gather(a, 0, b) + torch.gather(a, -1, b)
+
+        x = torch.tensor(123)
+        y = torch.tensor(0)
+        self.assertEqual(fn(x, y), x + x)
+
     def test_slice1(self):
         def fn(a):
             return (