[functionalization] functional tensors fail on torch.lstm

[davidberard98]

🐛 Describe the bug

Minimal repro & logs

torch.lstm fails here due to a mismatch between device type between the tensor.options() and torch.storage().

This issue occurs when running torch.lstm (specifically, demucs model from torchbench) with aot_eager backend in dynamo.

Versions

master branch

cc @bdhirsh @ezyang @soumith

[bdhirsh]

Ok, I think I tracked this down.

For some ops, functionalization tries to run the meta impl before running the rest of the functionalization logic - this is mostly to ensure that we raise any errors properly, if the underlying op would have raised errors.

for .set_(), the functionalization kernel looks like this:

    at::Tensor & set__source_Storage(c10::DispatchKeySet dispatchKeySet, at::Tensor & self, at::Storage source) {
      if (true) {
        // Before converting the mutable op to its functional variant, run meta tensors through the original op.
        // This will help us catch shape errors that apply to inplace ops that wouldn't apply to their functional variants
        // (We can only do this for inplace ops today though, because they technicaly all support meta tensors).
        at::AutoDispatchSkipFunctionalize func_guard;
        c10::impl::ExcludeDispatchKeyGuard guard(exclude_keys_for_meta_dispatch);
        auto self_meta = to_meta(self);
        at::_ops::set__source_Storage::call(self_meta, source);
      }
...

It tries to run the meta function by:
(1) Convert all input tensors into meta tensors
(2) pass converted inputs to the meta kernel

The problem is that we have a "storage" input, with a device that is now different from our meta tensors. We end up trying to run meta._tensor.set_(cuda_storage), which (correctly) raises an error.

The right fix is probably either to:
(1) Get this all to work with storages. Take the existing "convert to meta" logic, and apply it to the storage too.
(2) Skip the meta function in that situation. This is probably ok, since 90% of the time we're running functionalization with fake tensors underneath, which will end up running the meta logic anyway.

[ezyang]

I still do not understand why functionalization needs to run the meta logic lol

[bdhirsh]

So actually - for this specific case, we're running the meta tensor kernel for an inplace op, before redispatching to the out-of-place op.

You can have a situation where an inplace call should error, but converting to the functional call becomes ok - so we want to make sure we error (by running the inplace meta kernel first).

E.g.:

>>> a = torch.ones(2)
>>> b = torch.ones(2, 2)

# inplace fails
>>> a.add_(b)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: output with shape [2] doesn't match the broadcast shape [2, 2]

# out of place works due to broadcasting
>>> a.add(b)
tensor([[2., 2.],
        [2., 2.]])
>>>

[ezyang]

OK, but set_ isn't like conventional impact ops, right?

[bdhirsh]

yep, I agree this check isn't very useful for set_. I just figured it would be cleaner to keep the check around and try to make it work in more cases. But we could also, say, have that check not run for any ops that take in Storage as an argument. I can do that if you'd prefer - it's probably marginally beneficial to have fewer ops running a (pointless) meta kernel, although that will probably have a negligible impact on compile times.

[ezyang]

I don't see why it's cleaner. Set_ and add_ are completely different beasts and trying to shoehorn then into the same behavior makes the system more complicated, not less. You don't WANT them running the same code!

[bdhirsh]

Ok, I buy that argument :)