Repository: pytorch/torchdynamo

Issue:
[AMP] Bad decomp for softmax
Repro

~~~
from math import inf

import torch
import torch.fx as fx
from torch import device
from torch import tensor

import torchdynamo
from torchdynamo.debug_utils import run_fwd_maybe_bwd
from torchdynamo.debug_utils import same_two_models
from torchdynamo.testing import rand_strided

args = [
    ((1, 12, 128, 128), (196608, 16384, 128, 1), torch.float16, "cuda", True),
    ((1, 12, 128, 128), (128, 0, 0, 1), torch.bool, "cuda", False),
    ((), (), torch.float32, "cpu", False),
]
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__()

    def forward(self, matmul, expand_as, tensor):
        masked_fill = matmul.masked_fill(expand_as, tensor)
        softmax = torch.nn.functional.softmax(masked_fill, dim=-1)
        return (softmax,)


mod = Repro().cuda()
opt_mod = torchdynamo.optimize("aot_inductor_debug")(mod)


with torch.cuda.amp.autocast(enabled=True):
    ref = run_fwd_maybe_bwd(mod, args)
    res = run_fwd_maybe_bwd(opt_mod, args)


~~~


Discussion:

Resolving patch:
diff --git a/torchinductor/decomposition.py b/torchinductor/decomposition.py
--- a/torchinductor/decomposition.py
+++ b/torchinductor/decomposition.py
@@ -193,9 +193,12 @@ def rsub(a, b):
 def masked_fill(value, mask, other):
     if isinstance(other, numbers.Number):
         other = torch.tensor(other, dtype=value.dtype, device=value.device)
+    assert other.numel() == 1 and other.ndim == 0
     if other.device != value.device and other.numel() == 1:
         other = other.to(value.device)
-    value, mask, other = torch.broadcast_tensors(value, mask, other)
+    if other.dtype != value.dtype:
+        # TODO: error out on improper complex conversions
+        other = other.to(value.dtype)
     return torch.where(mask, other, value)
 
 
diff --git a/torchinductor/lowering.py b/torchinductor/lowering.py
--- a/torchinductor/lowering.py
+++ b/torchinductor/lowering.py
@@ -430,7 +430,7 @@ def register_pointwise(
     return fn
 
 
-@register_lowering(aten.where, broadcast=True, type_promotion_kind=None)
+@register_lowering(aten.where, broadcast=False, type_promotion_kind=None)
 def where(cond, a, b):
     def fn(*args):
         return ops.where(*args)
@@ -440,9 +440,18 @@ def fn(*args):
     if isinstance(b, (float, int)):
         b = constant_like(b)(a)
 
-    dtype = torch.promote_types(a.get_dtype(), b.get_dtype())
+    args = [cond, a, b]
+    dtype = get_promoted_dtype(
+        args[1], args[2], type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
+    )
+    indices = [i for i, x in enumerate(args) if isinstance(x, TensorBox)]
+    for i, x in zip(indices, broadcast_tensors(*[args[i] for i in indices])):
+        args[i] = x
+    for i in range(len(args)):
+        if isinstance(args[i], ir.Constant):
+            args[i] = ExpandView.create(args[i], list(args[indices[0]].get_size()))
     return make_pointwise(fn, override_return_dtype=dtype)(
-        cond, to_dtype(a, dtype), to_dtype(b, dtype)
+        args[0], to_dtype(args[1], dtype), to_dtype(args[2], dtype)
     )
 
 


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
@@ -1940,6 +1940,26 @@ def fn(mask, value):
             ),
         )
 
+    def test_masked_fill_promotion(self):
+        def fn(mask, value):
+            return aten.masked_fill(value, mask, torch.tensor(3.5))
+
+        opt_fn = torchdynamo.optimize("inductor")(fn)
+        for inp in (
+            torch.randn(
+                [16, 16],
+                dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                device=self.device,
+            ),
+            torch.randint(16, (16, 16), device=self.device),
+        ):
+
+            inputs = (
+                torch.randint(0, 1, [1, 16], dtype=torch.bool, device=self.device),
+                inp,
+            )
+            self.assertEqual(fn(*inputs), opt_fn(*inputs))
+
     def test_fill1(self):
         def fn(x):
             tmp = torch.ones_like(x)