Repository: pytorch/torchdynamo

Issue:
Inductor fails on optimizer.step
Repro:
```
import torch
import torch.optim as optim
import torchdynamo

torchdynamo.config.fake_tensor_propagation = False

input = torch.ones([10, 10])
model = torch.nn.Sequential(*[torch.nn.Linear(10, 10) for _ in range(2)])
opt_model = torchdynamo.optimize("inductor")(model)
opt_model(input).sum().backward()

optimizer = optim.SGD(opt_model.parameters(), lr=5e-5)
opt_fn = torchdynamo.optimize("inductor")(optimizer.step)
opt_fn()
```

Error:
```
Traceback (most recent call last):
  File "/data/home/ybliang/Meta/torchdynamo/torchinductor/graph.py", line 234, in call_function
    return lowerings[target](*args, **kwargs)
  File "/data/home/ybliang/Meta/torchdynamo/torchinductor/lowering.py", line 185, in wrapped
    return decomp_fn(*args, **kwargs)
TypeError: fn() got an unexpected keyword argument 'alpha'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data/home/ybliang/Meta/torchdynamo/debug/debug3.py", line 14, in <module>
    opt_fn()
  File "/data/home/ybliang/Meta/torchdynamo/torchdynamo/eval_frame.py", line 154, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/torch/optim/optimizer.py", line 23, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/torch/optim/sgd.py", line 151, in step
    sgd(params_with_grad,
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/torch/optim/sgd.py", line 171, in sgd
    def sgd(params: List[Tensor],
  File "/data/home/ybliang/Meta/torchdynamo/torchdynamo/eval_frame.py", line 154, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 886, in forward
    return compiled_f(
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 872, in new_func
    compiled_fn = create_aot_dispatcher_function(
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 450, in create_aot_dispatcher_function
    return aot_dispatch_base(flat_fn, fake_flat_tensor_args, aot_config)
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 276, in aot_dispatch_base
    compiled_fw = aot_config.fw_compiler(fw_module, flat_args)
  File "/data/home/ybliang/Meta/torchdynamo/torchinductor/compile_fx.py", line 170, in fw_compiler
    return compile_fx_inner(
  File "/data/home/ybliang/Meta/torchdynamo/torchdynamo/debug_utils.py", line 253, in debug_wrapper
    raise e
  File "/data/home/ybliang/Meta/torchdynamo/torchdynamo/debug_utils.py", line 241, in debug_wrapper
    compiled_fn = compiler(gm, example_inputs, **kwargs)
  File "/data/home/ybliang/Meta/torchdynamo/torchinductor/debug.py", line 182, in inner
    return fn(*args, **kwargs)
  File "/data/home/ybliang/Meta/torchdynamo/torchinductor/compile_fx.py", line 56, in compile_fx_inner
    wrap(graph.run)(*example_inputs)
  File "/data/home/ybliang/Meta/torchdynamo/torchinductor/graph.py", line 129, in run
    return super().run(*args)
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/torch/fx/interpreter.py", line 129, in run
    self.env[node] = self.run_node(node)
  File "/data/home/ybliang/Meta/torchdynamo/torchinductor/graph.py", line 285, in run_node
    result = super().run_node(n)
  File "/fsx/users/ybliang/conda/envs/dynamo-gpu/lib/python3.9/site-packages/torch/fx/interpreter.py", line 162, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/data/home/ybliang/Meta/torchdynamo/torchinductor/graph.py", line 236, in call_function
    raise LoweringException(e, target, args, kwargs) from e
torchinductor.exc.LoweringException: TypeError: fn() got an unexpected keyword argument 'alpha'
  target: aten.add_.Tensor
  args[0]: TensorBox(StorageBox(
    InputBuffer(name='arg0_1', layout=FixedLayout('cpu', torch.float32, size=[s0, s0], stride=[s0, 1]))
  ))
  args[1]: TensorBox(StorageBox(
    InputBuffer(name='arg4_1', layout=FixedLayout('cpu', torch.float32, size=[s0, s0], stride=[s0, 1]))
  ))
  kwargs: {'alpha': -5e-05}

```
This failure occurs on both CPU and cuda.


Discussion:

Resolving patch:
diff --git a/torchinductor/lowering.py b/torchinductor/lowering.py
--- a/torchinductor/lowering.py
+++ b/torchinductor/lowering.py
@@ -249,9 +249,17 @@ def promote_constants(inputs):
     ]
 
 
-def make_pointwise(fn, override_dtype=None, override_device=None, override_bool=None):
-    def inner(*inputs: List[TensorBox]):
+def make_pointwise(
+    fn, override_dtype=None, override_device=None, override_bool=None, allow_alpha=False
+):
+    def inner(*inputs: List[TensorBox], alpha=None):
         inputs = promote_constants(inputs)
+        if allow_alpha:
+            if alpha is not None and alpha != 1:
+                inputs = list(inputs)
+                inputs[-1] = mul(inputs[-1], alpha)
+        else:
+            assert alpha is None
         loaders = [x.make_loader() for x in inputs]
         ranges = inputs[0].get_size()
         dtype = override_dtype or inputs[0].get_dtype()
@@ -366,6 +374,7 @@ def register_pointwise(
     override_dtype=None,
     override_device=None,
     override_bool=None,
+    allow_alpha=False,
 ):
     """A pointwise function that maps ops.{name} to inputs"""
     name = name or aten_fn.__name__
@@ -378,6 +387,7 @@ def register_pointwise(
         override_dtype=override_dtype,
         override_device=override_device,
         override_bool=override_bool,
+        allow_alpha=allow_alpha,
     )
     fn = register_lowering(aten_fn, broadcast=broadcast, type_promote=type_promote)(fn)
 
@@ -2887,7 +2897,7 @@ def sum_(x, axis=None, keepdims=False, *, dtype=None):
     make_reduction("argmin", override_dtype=torch.int64)
 )
 
-add = register_pointwise(aten.add)
+add = register_pointwise(aten.add, allow_alpha=True)
 exp = register_pointwise(aten.exp)
 floor = register_pointwise(aten.floor)
 mul = register_pointwise(aten.mul)
@@ -2895,7 +2905,7 @@ def sum_(x, axis=None, keepdims=False, *, dtype=None):
 sigmoid = register_pointwise(aten.sigmoid)
 sqrt = register_pointwise(aten.sqrt)
 square = register_pointwise(aten.square)
-sub = register_pointwise(aten.sub)
+sub = register_pointwise(aten.sub, allow_alpha=True)
 trunc = register_pointwise(aten.trunc)
 
 register_pointwise(aten.cos)
@@ -2935,8 +2945,8 @@ def sum_(x, axis=None, keepdims=False, *, dtype=None):
 
 def register_inplace(aten_op, outplace_op):
     @register_lowering(aten_op, type_promote=False)
-    def fn(*args):
-        result = outplace_op(*args)
+    def fn(*args, **kwargs):
+        result = outplace_op(*args, **kwargs)
         return mutate_to(args[0], result)
 
     return fn


Test patch:
diff --git a/test/test_torchinductor.py b/test/test_torchinductor.py
--- a/test/test_torchinductor.py
+++ b/test/test_torchinductor.py
@@ -3109,6 +3109,21 @@ def fn(x):
             result = fn(v)
             assert same(result, mod(v))
 
+        def test_inplace_add_alpha(self):
+            def fn(x, y):
+                aten.add_.Tensor(x, y, alpha=0.55)
+                return (x,)
+
+            x1 = torch.zeros(10)
+            x2 = torch.zeros(10)
+            x3 = torch.zeros(10)
+            y = torch.randn(10)
+            fn_fx = make_fx(fn)(x1, y)
+            fn_compiled = compile_fx_inner(fn_fx, [x1, y])
+            fn(x2, y)
+            fn_compiled(x3, y)
+            assert same(x2, x3)
+
 
 if HAS_CUDA: