Repository: pytorch/torchdynamo

Issue:
TorchBench training error: MemoryError: std::bad_alloc
Repro:
```
benchmarks/torchbench.py -d cuda --inductor --training --float32 --no-skip -k timm_efficientdet
```
Error:
```
Traceback (most recent call last):
  File "/fsx/users/binbao/torchdynamo-tip/benchmarks/common.py", line 1236, in run_one_model
    new_result = optimized_model_iter_fn(model, example_inputs)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "benchmarks/torchbench.py", line 338, in forward_and_backward_pass
    def forward_and_backward_pass(self, mod, inputs, collect_outputs=True):
  File "benchmarks/torchbench.py", line 339, in <graph break in forward_and_backward_pass>
    cloned_inputs = clone_inputs(inputs)
  File "benchmarks/torchbench.py", line 344, in <graph break in forward_and_backward_pass>
    self.grad_scaler.scale(loss).backward()
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/monkey_patching.py", line 77, in _backward
    return _old_backward(*args, **kwargs)
  File "/fsx/users/binbao/pytorch-release/torch/_tensor.py", line 484, in backward
    torch.autograd.backward(
  File "/fsx/users/binbao/pytorch-release/torch/autograd/__init__.py", line 191, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/fsx/users/binbao/pytorch-release/torch/autograd/function.py", line 267, in apply
    return user_fn(self, *args)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 367, in backward
    CompiledFunction.compiled_bw = aot_config.bw_compiler(
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/optimizations/backends.py", line 542, in _wrapped_bw_compiler
    return torchdynamo.disable(bw_compiler(*args, **kwargs))
  File "/fsx/users/binbao/pytorch-release/functorch/functorch/_src/aot_autograd.py", line 237, in f
    out_f = compiler(fx_g, inps)
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/utils.py", line 75, in time_wrapper
    r = func(*args, **kwargs)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 184, in bw_compiler
    return compile_fx_inner(
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/debug_utils.py", line 329, in debug_wrapper
    compiled_fn = compiler(gm, example_inputs, **kwargs)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/debug.py", line 182, in inner
    return fn(*args, **kwargs)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 63, in compile_fx_inner
    compiled_fn = cudagraphify(
  File "/fsx/users/binbao/torchdynamo-tip/torchdynamo/utils.py", line 75, in time_wrapper
    r = func(*args, **kwargs)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/compile_fx.py", line 106, in cudagraphify
    model(*inputs)
  File "/tmp/torchinductor_binbao/7j/c7jxvonilvfwdeuue7g2acsacmhqsukshbbh7ux4wesdc654x5w3.py", line 14739, in call
    kernel172[grid(18)](buf651, sum_58, getitem_28, primals_34, buf106, buf103, buf91, buf697, sum_54, upsample_nearest2d_14, primals_32, buf102, buf653, buf655, buf657, buf659, buf661, buf663, buf699, buf701, buf703, buf705, 18, 7823)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 999, in __call__
    return self.kernel(*wargs, **kwargs, grid=self.grid)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/triton_ops/autotune.py", line 49, in __call__
    timings = {
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/triton_ops/autotune.py", line 50, in <dictcomp>
    config: self._bench(*args, config=config, **kwargs)
  File "/fsx/users/binbao/torchdynamo-tip/torchinductor/triton_ops/autotune.py", line 41, in _bench
    return super()._bench(*args, config=config, **kwargs)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1054, in _bench
    return triton.testing.do_bench(kernel_call)
  File "/fsx/users/binbao/triton/python/triton/testing.py", line 137, in do_bench
    fn()
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1053, in kernel_call
    self.kernel(*args, num_warps=config.num_warps, num_stages=config.num_stages, **current)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 988, in __call__
    return _triton.runtime.launch(wargs, self.fn.do_not_specialize, cache_key, self.fn.arg_names,
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 956, in add_to_cache
    return self.fn._warmup(key, arg_types=arg_types, device=device_idx, attributes=attributes, constants=constants, num_warps=num_warps, num_stages=num_stages,
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1285, in _warmup
    binary = self._compile(**compile)
  File "/fsx/users/binbao/triton/python/triton/code_gen.py", line 1320, in _compile
    name, asm, shared_mem = _triton.code_gen.compile_ttir(backend, generator.module, device, num_warps, num_stages, extern_libs)
MemoryError: std::bad_alloc
```


Discussion:
This looks like an out of memory error
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.14.0a0+git25725fd
# torch cuda version: 11.6
# torch git version: 25725fd62448165b91647304c26d676db22b6955


# 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, add_108, relu_29, add_115, neg_21):
        select_69 = torch.ops.aten.select.int(relu_29, 0, 1);  relu_29 = None
        mul_264 = torch.ops.aten.mul.Tensor(add_108, select_69);  add_108 = select_69 = None
        div_69 = torch.ops.aten.div.Tensor(mul_264, add_115);  mul_264 = None
        div_93 = torch.ops.aten.div.Tensor(div_69, add_115);  div_69 = add_115 = None
        mul_704 = torch.ops.aten.mul.Tensor(neg_21, div_93);  neg_21 = div_93 = None
        sum_111 = torch.ops.aten.sum.default(mul_704);  mul_704 = None
        return (sum_111,)
        
args = [((2, 88, 20, 20), (35200, 400, 20, 1), torch.float32, 'cuda'), ((3,), (1,), torch.float32, 'cuda'), ((), (), torch.float32, 'cuda'), ((2, 88, 20, 20), (35200, 400, 20, 1), torch.float32, 'cuda')]
args = [rand_strided(sh, st, dt, dev) for (sh, st, dt, dev) in args]
mod = make_fx(Repro().to(device="cuda"))(*args)

from torchinductor.compile_fx import compile_fx_inner
from torchdynamo.debug_utils import same_two_models

compiled = compile_fx_inner(mod, args)
compiled(*args)
~~~
`tmp5` load in the reduction loop should have dense indexing:
```
        tmp3 = tl.load(in_ptr0 + ((r1 + (7823*x0)) % 70400 + tl.zeros([XBLOCK, RBLOCK], tl.int32)), xmask & rmask & tmp2, eviction_policy='evict_last', other=0)
        tmp4 = tl.load(in_ptr1 + ((r1 + (7823*x0)) % 70400 + tl.zeros([XBLOCK, RBLOCK], tl.int32)), xmask & rmask & tmp2, eviction_policy='evict_last', other=0)
        tmp5 = tl.load(in_ptr2 + (1), tmp2, other=0)
```
should be 
```
tmp5 = tl.load(in_ptr2 + (1 + tl.zeros([XBLOCK, RBLOCK], tl.int32)), tmp2, other=0)
```
instead
       

Resolving patch:
diff --git a/torchinductor/codegen/triton.py b/torchinductor/codegen/triton.py
--- a/torchinductor/codegen/triton.py
+++ b/torchinductor/codegen/triton.py
@@ -681,9 +681,11 @@ def indexing(
                 mask.append(f"{tree.prefix}mask")
             dense_mask.append(f"{tree.prefix}mask")
 
-        if (need_dense and not have_dense) or index == 0:
+        if (need_dense and not have_dense) or isinstance(
+            index, sympy.core.numbers.Integer
+        ):
             index_str = f"{index_str} + tl.zeros({self.dense_size_str()}, tl.int32)"
-            if index == 0:
+            if isinstance(index, sympy.core.numbers.Integer):
                 return index_str, "None"
             else:
                 mask = dense_mask


Test patch:
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
--- a/test/inductor/test_torchinductor.py
+++ b/test/inductor/test_torchinductor.py
@@ -3162,6 +3162,15 @@ def fn(a, dim, index, b):
             ],
         )
 
+    # issue #1150
+    def test_dense_mask_index(self):
+        def fn(x, y):
+            y = torch.ops.aten.select.int(y, 0, 2)
+            z = x * y
+            return z.sum()
+
+        self.common(fn, [torch.randn(102400), torch.randn(3)])
+
     def test_new_empty_strided(self):
         def fn(a):
             return aten.new_empty_strided(a, [1, 128, 128], [16384, 128, 1]).fill_(123)