Repository: pytorch/torchdynamo

Issue:
[inductor] AMP accuracy error - inception_v3
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+git071f875
# torch cuda version: 11.6
# torch git version: 071f875046202b87213865dfc180abdf8368f116


# 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, cat_11, _to_copy_373):
        avg_pool2d_8 = torch.ops.aten.avg_pool2d.default(cat_11, [3, 3], [1, 1], [1, 1]);  cat_11 = None
        convolution_93 = torch.ops.aten.convolution.default(avg_pool2d_8, _to_copy_373, None, [1, 1], [0, 0], [1, 1], False, [0, 0], 1);  avg_pool2d_8 = _to_copy_373 = None
        return (convolution_93,)
        
args = [((2, 2048, 8, 8), (131072, 64, 8, 1), torch.float16, 'cuda'), ((192, 2048, 1, 1), (2048, 1, 1, 1), torch.float16, 'cuda')]
args = [rand_strided(sh, st, dt, dev) for (sh, st, dt, dev) in args]
mod = make_fx(Repro())(*args)

from torchinductor.compile_fx import compile_fx_inner
from torchdynamo.debug_utils import same_two_models

compiled = compile_fx_inner(mod, args)
assert same_two_models(mod, compiled, args, only_fwd=True), "Accuracy failed"
~~~


Discussion:
This seems like a triton bug (and is a bug for fp32 and fp16, but fp32 just squeezes through based on some implicit 0 initialization that doesn't work for fp16). Currently, the following code is generated (I'm pasting simplified `where` generation that can be done with the latest triton, but with current `where` codegen results are the same, just the code is a bit messier):
```
    tmp34 = tl.load(in_ptr0 + (-1) + x4 + tl.zeros([XBLOCK], tl.int32), xmask & tmp33).to(tl.float32)
    tmp35 = tl.where(tmp33, tmp34, 0.0)
```
`xmask` is `True` everywhere. If instead we do 
```
tmp34 = tl.load(in_ptr0 + (-1) + x4 + tl.zeros([XBLOCK], tl.int32), xmask & tmp33, other=0.).to(tl.float32)
```
correct results are produced, even though it should be the same as snippet above. cc @jansel

Should we submit an upstream bug to Triton?

The `where` version is a bit more flexible, since it allows inline computation inside the `ops.masked()`. 
Yeah I agree `where` version is more flexible and it would be nice to keep it, I'll try to dig a bit more into what's happening. 
Filed a triton bug https://github.com/openai/triton/issues/737, for now I think the only workaround is to add `other=0.` kwarg to masked loads. 
Perhaps we should just always set `other=0.` until it is fixed?

Resolving patch:
diff --git a/torchinductor/codegen/triton.py b/torchinductor/codegen/triton.py
--- a/torchinductor/codegen/triton.py
+++ b/torchinductor/codegen/triton.py
@@ -766,7 +766,14 @@ def load(self, name: str, index: sympy.Expr):
             ep = ", eviction_policy='evict_last'"
         else:
             ep = ""
-        line = f"tl.load({var} + {index}, {mask}{ep})"
+        # "other" below is a workaround for https://github.com/openai/triton/issues/737
+        # for bool, even though it's likely subject to the same bug, setting `other` leads
+        # to LLVM errors so we are skipping it for now
+        if "tmp" in mask and V.graph.get_dtype(name) != torch.bool:
+            other = ", other=0"
+        else:
+            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)"
 


Test patch: