Repository: pytorch/torchdynamo

Issue:
[inductor] Accuracy issue - volo_d1_224

Repro 
~~~

import torch
import torchdynamo
from torch import tensor, device
import torch.fx as fx
from torchdynamo.testing import rand_strided
from math import inf
from torchdynamo.debug_utils import run_fwd_maybe_bwd
from torchdynamo.debug_utils import same_two_models

args = [((2, 197, 384), (75648, 384, 1), torch.float32, 'cuda', True), ((2, 1, 384), (384, 384, 1), torch.float32, 'cuda', True), ((2, 1, 384), (384, 384, 1), torch.float32, 'cuda', True)]
args = [rand_strided(sh, st, dt, dev).requires_grad_(rg) for (sh, st, dt, dev, rg) in args]


from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.ln = LayerNorm((384,), eps=1e-05, elementwise_affine=True)
        self.linear1 = Linear(in_features=384, out_features=1000, bias=True)
        self.linear2 = Linear(in_features=384, out_features=1000, bias=True)



    def forward(self, cat_1, add_39, x):
        add_40 = add_39 + x
        getitem_51 = cat_1[(slice(None, None, None), slice(1, None, None))]
        cat_2 = torch.cat([add_40, getitem_51], dim = 1)
        ln = self.ln(cat_2)
        getitem_52 = ln[(slice(None, None, None), 0)]
        linear1 = self.linear1(getitem_52)
        getitem_53 = ln[(slice(None, None, None), slice(1, None, None))]
        linear2 = self.linear2(getitem_53)
        max_1 = linear2.max(1)
        getitem_54 = max_1[0]
        mul_20 = 0.5 * getitem_54
        add_41 = linear1 + mul_20
        return (add_41,)



mod = Repro().cuda()
opt_mod = torchdynamo.optimize("inductor")(mod)


mod.eval()
opt_mod.eval()
with torch.cuda.amp.autocast(enabled=False):
    assert same_two_models(mod, mod, args), "Eager itself failed"
    assert same_two_models(mod, opt_mod, args), "Dynamo failed"

~~~


Discussion:
Repro - backward only. Note, normally backward `fx_graph.py` would appear to be passing because indices would be generated as all zeros, when I change them to be random indices (in the repro below), the optimized and eager results start being different.  
```

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+gitcb4421b
# torch cuda version: 11.6
# torch git version: cb4421b19cf7aa3a1a6cfcc7e0677f2b2ba0a9b6


# 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 : 2


from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()



    def forward(self, primals_1, cat, select, _unsafe_view, unsqueeze_1, permute_4, permute_6, tangents_1):
        var_mean = torch.ops.aten.var_mean.correction(cat, [2], correction = 0, keepdim = True)
        getitem = var_mean[0]
        getitem_1 = var_mean[1];  var_mean = None
        add_1 = torch.ops.aten.add.Tensor(getitem, 1e-05);  getitem = None
        sqrt = torch.ops.aten.sqrt.default(add_1);  add_1 = None
        reciprocal = torch.ops.aten.reciprocal.default(sqrt);  sqrt = None
        sub = torch.ops.aten.sub.Tensor(cat, getitem_1);  cat = getitem_1 = None
        mul = torch.ops.aten.mul.Tensor(sub, reciprocal);  sub = None
        mul_3 = torch.ops.aten.mul.Tensor(tangents_1, 0.5)
        unsqueeze = torch.ops.aten.unsqueeze.default(mul_3, 1);  mul_3 = None
        zeros = torch.ops.aten.zeros.default([2, 196, 1000], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0))
        scatter = torch.ops.aten.scatter.src(zeros, 1, unsqueeze_1, unsqueeze);  zeros = unsqueeze_1 = unsqueeze = None
        sum_1 = torch.ops.aten.sum.dim_IntList(scatter, [0, 1], True)
        view = torch.ops.aten.view.default(sum_1, [1000]);  sum_1 = None
        view_1 = torch.ops.aten.view.default(scatter, [392, 1000]);  scatter = None
        permute_2 = torch.ops.aten.permute.default(view_1, [1, 0])
        mm_1 = torch.ops.aten.mm.default(permute_2, _unsafe_view);  permute_2 = _unsafe_view = None
        permute_3 = torch.ops.aten.permute.default(mm_1, [1, 0]);  mm_1 = None
        mm_2 = torch.ops.aten.mm.default(view_1, permute_4);  view_1 = permute_4 = None
        view_2 = torch.ops.aten.view.default(mm_2, [2, 196, 384]);  mm_2 = None
        permute_5 = torch.ops.aten.permute.default(permute_3, [1, 0]);  permute_3 = None
        new_zeros = torch.ops.aten.new_zeros.default(view_2, [2, 197, 384], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False)
        slice_scatter = torch.ops.aten.slice_scatter.default(new_zeros, view_2, 1, 1, 9223372036854775807);  new_zeros = view_2 = None
        new_zeros_1 = torch.ops.aten.new_zeros.default(slice_scatter, [2, 197, 384], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False)
        slice_scatter_1 = torch.ops.aten.slice_scatter.default(new_zeros_1, slice_scatter, 0, 0, 9223372036854775807);  new_zeros_1 = slice_scatter = None
        mm_3 = torch.ops.aten.mm.default(tangents_1, permute_6);  permute_6 = None
        permute_7 = torch.ops.aten.permute.default(tangents_1, [1, 0])
        mm_4 = torch.ops.aten.mm.default(permute_7, select);  permute_7 = select = None
        permute_8 = torch.ops.aten.permute.default(mm_4, [1, 0]);  mm_4 = None
        sum_2 = torch.ops.aten.sum.dim_IntList(tangents_1, [0], True);  tangents_1 = None
        view_3 = torch.ops.aten.view.default(sum_2, [1000]);  sum_2 = None
        permute_9 = torch.ops.aten.permute.default(permute_8, [1, 0]);  permute_8 = None
        new_zeros_2 = torch.ops.aten.new_zeros.default(mm_3, [2, 197, 384], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False)
        select_scatter = torch.ops.aten.select_scatter.default(new_zeros_2, mm_3, 1, 0);  new_zeros_2 = mm_3 = None
        new_zeros_3 = torch.ops.aten.new_zeros.default(select_scatter, [2, 197, 384], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False)
        slice_scatter_2 = torch.ops.aten.slice_scatter.default(new_zeros_3, select_scatter, 0, 0, 9223372036854775807);  new_zeros_3 = select_scatter = None
        add_5 = torch.ops.aten.add.Tensor(slice_scatter_1, slice_scatter_2);  slice_scatter_1 = slice_scatter_2 = None
        mul_5 = torch.ops.aten.mul.Tensor(add_5, primals_1);  primals_1 = None
        mul_6 = torch.ops.aten.mul.Tensor(mul_5, 384)
        sum_3 = torch.ops.aten.sum.dim_IntList(mul_5, [2], True)
        mul_7 = torch.ops.aten.mul.Tensor(mul_5, mul);  mul_5 = None
        sum_4 = torch.ops.aten.sum.dim_IntList(mul_7, [2], True);  mul_7 = None
        mul_8 = torch.ops.aten.mul.Tensor(mul, sum_4);  sum_4 = None
        sub_2 = torch.ops.aten.sub.Tensor(mul_6, sum_3);  mul_6 = sum_3 = None
        sub_3 = torch.ops.aten.sub.Tensor(sub_2, mul_8);  sub_2 = mul_8 = None
        div = torch.ops.aten.div.Tensor(reciprocal, 384);  reciprocal = None
        mul_9 = torch.ops.aten.mul.Tensor(div, sub_3);  div = sub_3 = None
        mul_10 = torch.ops.aten.mul.Tensor(add_5, mul);  mul = None
        sum_5 = torch.ops.aten.sum.dim_IntList(mul_10, [0, 1]);  mul_10 = None
        sum_6 = torch.ops.aten.sum.dim_IntList(add_5, [0, 1]);  add_5 = None
        slice_6 = torch.ops.aten.slice.Tensor(mul_9, 1, 0, 1)
        slice_7 = torch.ops.aten.slice.Tensor(mul_9, 1, 1, 197);  mul_9 = None
        new_zeros_4 = torch.ops.aten.new_zeros.default(slice_7, [2, 197, 384], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False)
        slice_scatter_3 = torch.ops.aten.slice_scatter.default(new_zeros_4, slice_7, 1, 1, 9223372036854775807);  new_zeros_4 = slice_7 = None
        new_zeros_5 = torch.ops.aten.new_zeros.default(slice_scatter_3, [2, 197, 384], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False)
        slice_scatter_4 = torch.ops.aten.slice_scatter.default(new_zeros_5, slice_scatter_3, 0, 0, 9223372036854775807);  new_zeros_5 = slice_scatter_3 = None
        return [sum_5, sum_6, permute_9, view_3, permute_5, view, slice_scatter_4, slice_6, slice_6]

args = [((384,), (1,), torch.float32, 'cuda'), ((2, 197, 384), (75648, 384, 1), torch.float32, 'cuda'), ((2, 384), (75648, 1), torch.float32, 'cuda'), ((392, 384), (384, 1), torch.float32, 'cuda'), ((2, 1, 1000), (1000, 1000, 1), torch.int64, 'cuda'), ((1000, 384), (384, 1), torch.float32, 'cuda'), ((1000, 384), (384, 1), torch.float32, 'cuda'), ((2, 1000), (1000, 1), torch.float32, 'cuda')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
for i, arg in enumerate(args):
    if arg.dtype == torch.int64:
        print(arg.size())
        arg_new = torch.randint(196, arg.size(), device=arg.device, dtype = arg.dtype)
        args[i] = arg_new
        print(args[i].max())
mod = make_fx(Repro())(*args)

from torchinductor.compile_fx import compile_fx_inner

compiled = compile_fx_inner(mod, args)
out_res = compiled(*args)
out_ref = mod(*args)
for o_res, o_ref in zip(out_res, out_ref):
    print("diff", (o_res-o_ref).abs().max(), o_res.max())
```
Ah, ok. For the `int64`, we should probably use `randint` in `rand_strided`. Good find!
> Ah, ok. For the `int64`, we should probably use `randint` in `rand_strided`. Good find!

Yes, but normally you wouldn't know what's the largest valid value, here I happened to know it's 195 because I looked at the model. 
Minified a bit further
```

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+gitcb4421b
# torch cuda version: 11.6
# torch git version: cb4421b19cf7aa3a1a6cfcc7e0677f2b2ba0a9b6


# 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 : 2


from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()



    def forward(self, primals_1, cat, select, _unsafe_view, unsqueeze_1, permute_4, permute_6, tangents_1):
        mul = cat
        mul_3 = tangents_1
        unsqueeze = torch.ops.aten.unsqueeze.default(mul_3, 1);  mul_3 = None
        zeros = torch.ops.aten.zeros.default([2, 196, 1000], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0))
        scatter = torch.ops.aten.scatter.src(zeros, 1, unsqueeze_1, unsqueeze);  zeros = unsqueeze_1 = unsqueeze = None
        view_1 = torch.ops.aten.view.default(scatter, [392, 1000]);  scatter = None
        mm_2 = torch.ops.aten.mm.default(view_1, permute_4);  view_1 = permute_4 = None
        view_2 = torch.ops.aten.view.default(mm_2, [2, 196, 384]);  mm_2 = None
        new_zeros = torch.ops.aten.new_zeros.default(view_2, [2, 197, 384], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False)
        slice_scatter = torch.ops.aten.slice_scatter.default(new_zeros, view_2, 1, 1, 9223372036854775807);  new_zeros = view_2 = None
        new_zeros_1 = torch.ops.aten.new_zeros.default(slice_scatter, [2, 197, 384], dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False)
        slice_scatter_1 = torch.ops.aten.slice_scatter.default(new_zeros_1, slice_scatter, 0, 0, 9223372036854775807);  new_zeros_1 = slice_scatter = None
        add_5 = slice_scatter_1
        mul_10 = torch.ops.aten.mul.Tensor(add_5, mul);  mul = None
        sum_5 = torch.ops.aten.sum.dim_IntList(mul_10, [0, 1]);  mul_10 = None
        return (sum_5,)

args = [((384,), (1,), torch.float32, 'cuda'), ((2, 197, 384), (75648, 384, 1), torch.float32, 'cuda'), ((2, 384), (75648, 1), torch.float32, 'cuda'), ((392, 384), (384, 1), torch.float32, 'cuda'), ((2, 1, 1000), (1000, 1000, 1), torch.int64, 'cuda'), ((1000, 384), (384, 1), torch.float32, 'cuda'), ((1000, 384), (384, 1), torch.float32, 'cuda'), ((2, 1000), (1000, 1), torch.float32, 'cuda')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
for i, arg in enumerate(args):
    if arg.dtype == torch.int64:
        arg_new = torch.randint(196, arg.size(), device=arg.device, dtype = arg.dtype)
        args[i] = arg_new
mod = make_fx(Repro())(*args)

from torchinductor.compile_fx import compile_fx_inner

compiled = compile_fx_inner(mod, args)
out_res = compiled(*args)
out_ref = mod(*args)
for o_res, o_ref in zip(out_res, out_ref):
    print("diff", (o_res-o_ref).abs().max(), o_res.max())
```
This one can be worked around by disabling `codegen_static_numels`, but today it has pretty bad perf consequences (more than 5% degradation), filed a triton bug with repro https://github.com/openai/triton/issues/741. We can see if marking xnumel/rnumel as divisible by 16 in the new runtime (with the corresponding guards) will resolve some or most of perf degradation, any volunteers to look into this? We would need to investigate this for dynamic shapes  anyway. 
cc @jansel

Resolving patch:
diff --git a/torchinductor/codegen/triton.py b/torchinductor/codegen/triton.py
--- a/torchinductor/codegen/triton.py
+++ b/torchinductor/codegen/triton.py
@@ -773,7 +773,7 @@ def load(self, name: str, index: sympy.Expr):
             other = ", other=0"
         else:
             other = ""
-        line = f"tl.load({var} + {index}, {mask}{ep}{other})"
+        line = f"tl.load({var} + ({index}), {mask}{ep}{other})"
         if V.graph.get_dtype(name) in (torch.float16, torch.bfloat16):
             line += ".to(tl.float32)"
 
@@ -797,9 +797,9 @@ def store(self, name, index, value, mode=None):
         var = self.args.output(name)
         index, mask = self.indexing(index, value, dense_indexing=True)
         if mode is None:
-            line = f"tl.store({var} + {index}, {value}, {mask})"
+            line = f"tl.store({var} + ({index}), {value}, {mask})"
         elif mode == "atomic_add":
-            line = f"tl.atomic_add({var} + {index}, {value}, {mask})"
+            line = f"tl.atomic_add({var} + ({index}), {value}, {mask})"
         else:
             raise NotImplementedError(f"store mode={mode}")
         self.stores.writeline(name, line)


Test patch:
diff --git a/test/test_torchinductor.py b/test/test_torchinductor.py
--- a/test/test_torchinductor.py
+++ b/test/test_torchinductor.py
@@ -2918,6 +2918,15 @@ def fn(a, dim, index, b):
             rtol=1e-3,
         )
 
+    def test_scatter4(self):
+        def fn(x, ind, src):
+            return torch.scatter(x, 0, ind, src)
+
+        self.common(
+            fn,
+            (torch.randn(196, 992), torch.randint(196, (1, 992)), torch.randn(1, 992)),
+        )
+
     @unittest.skip("Flaky test, needs debugging")
     def test_scatter_add1(self):
         def fn(a, dim, index, b):