Repository: pytorch
Issue: Codegen issues with Tensor(a!)? arguments in native_functions.yaml schemas
Labels: triaged,module: codegen

## 🐛 Bug: Codegen issues with Tensor(a!)? arguments in native_functions.yaml schemas

Currently, there are no function schemas in [native_functions.yaml](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml#L746) that take any argument which is both mutable *and* optional. When trying to add/modify a function to take a `Tensor(a!)?`, a number of compilation issues occur.

In general the issues are:
1. Codegen step fails due to inconsistent handling of the optional, mutable type
2. Issues with efficiently handling/unwrapping non-const `c10::optional<Tensor>&`
  a. Specifically, this causes issues with `cached_cast` in the `autocast_mode.cpp` macros

## To Reproduce

See: https://github.com/pytorch/pytorch/pull/65754

1. Modify `batch_norm` and `instance_norm` in [native_functions.yaml](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml#L746) to take `Tensor(a!)?` for `running_mean` and `running_var`. **At this point compilation will fail due to inconsistent expected types**. See [here](https://github.com/pytorch/pytorch/blob/master/tools/codegen/api/cpp.py#L83).
2. Fix cpp.py ^ and modify the c++ implementations of batch_norm, instance_norm, etc. to take `c10::optional<Tensor>&` instead of `const c10::optional<Tensor>&`.
  a. Note that there are some related issues here: when extracting the values from `optional<Tensor>`, some functions use `at::borrow_from_optional_tensor` which will return `const Tensor&`; if we want the extracted type to be a non-const `Tensor`, I think a copy will be introduced (or the `MaybeOwned` class would need to be changed)
3. Fix the function signatures in autocast_mode.cpp (i.e. `c10::optional<Tensor>&` instead of `const c10::optional<Tensor>&`)
4. `python setup.py develop` and get some compilation issues. (There are other issues, but this is the relevant one):

```
...
../aten/src/ATen/autocast_mode.cpp: In instantiation of ‘static Ret at::autocast::WrapFunction_<(at::autocast::CastPolicy)1, device_type, Redispatch, F, Ret, c10::guts::typelist::typelist<Args ...> >::call(Args ...) [with c10::DeviceType device_type = (c10::DeviceType)0; Redispatch = at::Tensor(const at::Tensor&, const c10::optional<at::Tensor>&, const c10::optional<at::Tensor>&, c10::optional<at::Tensor>&, c10::optional<at::Tensor>&, bool, double, double, bool); Redispatch* F = at::batch_norm; Ret = at::Tensor; Args = {const at::Tensor&, const c10::optional<at::Tensor>&, const c10::optional<at::Tensor>&, c10::optional<at::Tensor>&, c10::optional<at::Tensor>&, bool, double, double, bool}]’:
../aten/src/ATen/autocast_mode.cpp:487:3:   required from here
../aten/src/ATen/autocast_mode.cpp:183:16: error: cannot bind non-const lvalue reference of type ‘c10::optional<at::Tensor>&’ to an rvalue of type ‘c10::optional<at::Tensor>’
     return (*F)(cached_cast(at::kFloat, args, device_type)...);
            ~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
../aten/src/ATen/autocast_mode.cpp: In instantiation of ‘static Ret at::autocast::WrapFunction_<(at::autocast::CastPolicy)1, device_type, Redispatch, F, Ret, c10::guts::typelist::typelist<Args ...> >::call(Args ...) [with c10::DeviceType device_type = (c10::DeviceType)0; Redispatch = at::Tensor(const at::Tensor&, const c10::optional<at::Tensor>&, const c10::optional<at::Tensor>&, c10::optional<at::Tensor>&, c10::optional<at::Tensor>&, bool, double, double, bool); Redispatch* F = at::instance_norm; Ret = at::Tensor; Args = {const at::Tensor&, const c10::optional<at::Tensor>&, const c10::optional<at::Tensor>&, c10::optional<at::Tensor>&, c10::optional<at::Tensor>&, bool, double, double, bool}]’:
../aten/src/ATen/autocast_mode.cpp:509:3:   required from here
../aten/src/ATen/autocast_mode.cpp:183:16: error: cannot bind non-const lvalue reference of type ‘c10::optional<at::Tensor>&’ to an rvalue of type ‘c10::optional<at::Tensor>’
In file included from ../aten/src/ATen/core/boxing/KernelFunction_impl.h:2,
                 from ../aten/src/ATen/core/boxing/KernelFunction.h:273,
                 from ../aten/src/ATen/core/op_registration/op_registration.h:11,
                 from ../torch/library.h:69,
                 from ../aten/src/ATen/autocast_mode.cpp:2:
../aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h: In instantiation of ‘std::
...
```

This issue is caused when expanding [the macro for registering the batch_norm kernel](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/autocast_mode.cpp#L487) by the [cached_cast](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/autocast_mode.cpp#L111) function. In the implementation for `c10::optional<Tensor>` arguments, a `c10::optional<Tensor>` rvalue is returned and passed into the batch_norm function (which takes a `c10::optional<Tensor>&`, so it can't take an rvalue).

To reproduce this last issue, use this change https://github.com/pytorch/pytorch/pull/65754

cc @eellison