Repository: pytorch/torchdynamo

Issue:
[inductor] Accuracy for crossvit
https://github.com/pytorch/torchdynamo/issues/1190

~~~

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

# Set fake tensor off as there is a known bug with it.
import functorch.compile
functorch.compile.config.use_fake_tensor = False

from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.self_blocks_2__fusion_0__norm1 = LayerNorm((256,), eps=1e-06, elementwise_affine=True)
        self.self_blocks_2__fusion_0__attn_wq = Linear(in_features=256, out_features=256, bias=True)



    def forward(self, transpose_34, cat_10):
        self_blocks_2__fusion_0__norm1 = self.self_blocks_2__fusion_0__norm1(cat_10);  cat_10 = None
        getitem_64 = self_blocks_2__fusion_0__norm1[(slice(None, None, None), slice(0, 1, None), ...)];  self_blocks_2__fusion_0__norm1 = None
        self_blocks_2__fusion_0__attn_wq = self.self_blocks_2__fusion_0__attn_wq(getitem_64);  getitem_64 = None
        reshape_40 = self_blocks_2__fusion_0__attn_wq.reshape(2, 1, 4, 64);  self_blocks_2__fusion_0__attn_wq = None
        permute_24 = reshape_40.permute(0, 2, 1, 3);  reshape_40 = None
        matmul_32 = permute_24 @ transpose_34;  permute_24 = transpose_34 = None
        return (matmul_32,)


mod = Repro().cuda()
opt_mod = torchdynamo.optimize("inductor")(mod)

mod.eval()
opt_mod.eval()

args = [((2, 4, 64, 197), (50432, 64, 1, 256), torch.float32, 'cuda', True), ((2, 197, 256), (50432, 256, 1), torch.float32, 'cuda', True)]
args = [rand_strided(sh, st, dt, dev).requires_grad_(rg) for (sh, st, dt, dev, rg) in args]


with torch.cuda.amp.autocast(enabled=False):
    # ref = run_fwd_maybe_bwd(mod, args)
    # res = run_fwd_maybe_bwd(opt_mod, args)
    assert same_two_models(mod, mod, args), "Eager itself failed"
    assert same_two_models(mod, opt_mod, args), "Dynamo failed"
~~~


Discussion:
A smaller repro (fw functions only)
```
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_3, tangents_1):
        new_zeros = torch.ops.aten.new_zeros.default(tangents_1, [2, 197, 256], 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, tangents_1, 1, 0, 1);  new_zeros = tangents_1 = None
        mul_8 = torch.ops.aten.mul.Tensor(slice_scatter_1, primals_3);
        print(mul_8.size(), mul_8.stride(), slice_scatter_1.size(), slice_scatter_1.stride())
        sum_3 = torch.ops.aten.sum.dim_IntList(mul_8, [0, 1]);  mul_8 = None
        sum_4 = torch.ops.aten.sum.dim_IntList(slice_scatter_1, [0, 1]);  slice_scatter_1 = None
        return [sum_3, sum_4]

torch.manual_seed(12345)
args = [((2, 197, 256), (50432, 256, 1), torch.float32, 'cuda'), ((2, 1, 256), (256, 256, 1), torch.float32, 'cuda')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
#args[-1] = torch.ones_like(args[-1])
mod = make_fx(Repro())(*args)

from torchinductor.compile_fx import compile_fx_inner


compiled = compile_fx_inner(mod, args)
out = compiled(*args)
out_ref = mod(*args)
print((out[0] - out_ref[0])[:10], (out[1] - out_ref[1])[:10])
print(out[0][:10], out_ref[0][:10])
```

Resolving patch:
diff --git a/torchinductor/ir.py b/torchinductor/ir.py
--- a/torchinductor/ir.py
+++ b/torchinductor/ir.py
@@ -717,7 +717,11 @@ def fn(index):
                 reduction_type,
                 reduction_numel,
             )
-            reduction_hint = hint if hint != ReductionHint.DEFAULT else reduction_hint
+            # intermediate reduction in split can contain complex indexing,
+            # and num_splits will fail to correctly set the hint
+            # reuse the passed hint if available
+            if reduction_hint == ReductionHint.DEFAULT:
+                reduction_hint = hint
             if split > 1:
                 # triton doesn't support reduce to single element well, so break it up
                 return cls.create_multilayer(
diff --git a/torchinductor/sizevars.py b/torchinductor/sizevars.py
--- a/torchinductor/sizevars.py
+++ b/torchinductor/sizevars.py
@@ -130,11 +130,14 @@ def remove_zero_terms(base, divisor):
             for v in base.free_symbols:
                 if v in var_ranges:
                     # var smaller than divisor can be removed
+                    # if the rest is guaranteed to be multiple of divisor
                     rest = sympy.Wild("_rest", exclude=[v])
                     m = base.match(v + rest)
                     if m and v not in m[rest].free_symbols:
-                        if self.maybe_guard_leq(var_ranges[v], divisor):
-                            base = m[rest]
+                        gcd = sympy.gcd(m[rest], divisor)
+                        if gcd == divisor:
+                            if self.maybe_guard_leq(var_ranges[v], divisor):
+                                base = m[rest]
             return base
 
         def visit_indexing_div(base, divisor):


Test patch:
diff --git a/test/test_torchinductor.py b/test/test_torchinductor.py
--- a/test/test_torchinductor.py
+++ b/test/test_torchinductor.py
@@ -309,16 +309,10 @@ def test_indexing_simplification(self):
             sizevars.simplify_with_ranges(expr, var_ranges), i1 + 128 * i2 + 64 * r3
         )
 
-        # always-zero terms should be removed from IndexingDiv too
+        # small terms should be kept if the rest is not guaranteed to be divisible
         self.assertEqual(
             sizevars.simplify_with_ranges(IndexingDiv(r3 + i2 + i1, 32), var_ranges),
-            IndexingDiv(i1, 32),
-        )
-        self.assertEqual(
-            sizevars.simplify_with_ranges(
-                IndexingDiv(r3 + 2 * i2 + i1, 32), var_ranges
-            ),
-            IndexingDiv(i1 + 2 * i2, 32),
+            IndexingDiv(r3 + i2 + i1, 32),
         )
 
         expr = ModularIndexing(2 * i2 + r3, 1, 64)
@@ -592,6 +586,13 @@ def fn(a):
 
         self.common(fn, ((torch.rand((10, 3, 352, 352), dtype=torch.float16),)))
 
+    def test_expanded_reduction(self):
+        def fn(x, y):
+            z = x * y
+            return z.sum((0, 1))
+
+        self.common(fn, (torch.randn(2, 197, 256), torch.randn(2, 1, 256)))
+
     def test_min_max_reduction(self):
         def fn(a, b):
             return ((a + b).max(), (a + b).min(), torch.amax(a + 1, keepdim=True))