Repository: pytorch
Issue: [JIT] [Autocast] JIT Autocast Pass operations' list should be extendable and consistent with imperative path
Labels: oncall: jit

### Tasks:
- [ ] Make autocast ops extensible, e.g. via a function that modifies a global list of ops & their behavior, **per device**
- [ ] To better copy eager behavior, we could try using codegen to generate both eager and jit autocasting rules

### 🐛 Describe the bug

Currently, the operations' list is hardcoded in `JIT Autocast pass`   by different categories (https://github.com/pytorch/pytorch/blob/ef50186a7d05cd65fc605598480b52375cadd4b2/torch/csrc/jit/passes/autocast.cpp#L322-L440). 

Regardless of the limitation mentioned in the document (https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/JIT-AUTOCAST.md). We think the `JIT Autocast operation list` should aligne with `Imperative path Autocast`:
1. The imperative path Autocast Operation list is registered by `Autocast` dispatch key, which means it's extensible by third party to register same operation into different operation list (For the case, third party such as IPEX has optimization on this operation). So `JIT Autocast operation` list should aligned with `Imperative path Autocast` to be extensible.
2. The imperative path Autocast Operation list is distinguishable by different devices by dispatch key of `Autocast` and `AutocastCPU`. The JIT Autocast should also align with it.

If the `JIT Autocast operation list` is not aligned with `Imperative path Autocast`, we will probably see the issue when using Autocast with JIT Trace with third party such as IPEX. Take the below case for example:
```
import torch
import torch.nn as nn
import intel_extension_for_pytorch as ipex

class AtenSoftmaxRepalce(nn.Module):
    def __init__(self, dim=-1):
        super(AtenSoftmaxRepalce, self).__init__()
        self.conv = torch.nn.Conv2d(3, 64, (3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        self.softmax = torch.nn.Softmax(dim)

    def forward(self, x):
        return self.softmax(x)

model = AtenSoftmaxRepalce()
model.eval()

x = torch.rand(1, 3, 224, 224).to(torch.bfloat16)

with torch.no_grad():
    with torch.cpu.amp.autocast(cache_enabled=False):
        model = torch.jit.trace(model, x).eval()
        model = torch.jit.freeze(model)

        print("------start the run--------------")
        print(model.graph_for(x))
        print(model(x).dtype)
```

In IPEX, we have optimization on `Softmax` to support bf16 input and output. So we extend the imperative path Autocast operation list by change `softmax` from `black list` into `fall through list` inside IPEX. However, after going through the `Autocast JIT Pass`, it will insert the `aten::_autocast_to_full_precision` node to cast the `bf16 input` into `fp32` which changes the original graph generated by `jit trace`.

Here is the new graph after we enable the `Autocast JIT Pass`:
```
graph(%self.1 : __torch__.___torch_mangle_3.AtenSoftmaxRepalce,
      %x : Tensor):
  %2 : int = prim::Constant[value=-1](), scope: __module.softmax # /home/lesliefang/pytorch_1_7_1/ssd-rn34/frameworks.ai.pytorch.private-cpu/torch/nn/functional.py:1783:0
  %3 : NoneType = prim::Constant(), scope: __module.softmax
  %4 : bool = prim::Constant[value=0]()
  %5 : bool = prim::Constant[value=1]()
  %8 : Tensor = prim::profile[profiled_type=BFloat16(1, 3, 224, 224, strides=[150528, 50176, 224, 1], requires_grad=0, device=cpu), seen_none=0](%x)
  %6 : Tensor = aten::_autocast_to_full_precision(%8, %4, %5)
  %9 : Tensor = prim::profile[profiled_type=Float(1, 3, 224, 224, strides=[150528, 50176, 224, 1], requires_grad=0, device=cpu), seen_none=0](%6)
  %7 : Tensor = aten::softmax(%9, %2, %3), scope: __module.softmax # /home/lesliefang/pytorch_1_7_1/ssd-rn34/frameworks.ai.pytorch.private-cpu/torch/nn/functional.py:1783:0
  %10 : Tensor = prim::profile[profiled_type=Float(1, 3, 224, 224, strides=[150528, 50176, 224, 1], requires_grad=0, device=cpu), seen_none=0](%7)
   = prim::profile()
  return (%10)
``` 

Here is the original graph if we disable the `Autocast JIT Pass`:
```
graph(%self.1 : __torch__.___torch_mangle_3.AtenSoftmaxRepalce,
      %x : Tensor):
  %2 : int = prim::Constant[value=-1](), scope: __module.softmax # /home/lesliefang/pytorch_1_7_1/ssd-rn34/frameworks.ai.pytorch.private-cpu/torch/nn/functional.py:1783:0
  %3 : NoneType = prim::Constant(), scope: __module.softmax
  %5 : Tensor = prim::profile[profiled_type=BFloat16(1, 3, 224, 224, strides=[150528, 50176, 224, 1], requires_grad=0, device=cpu), seen_none=0](%x)
  %4 : Tensor = aten::softmax(%5, %2, %3), scope: __module.softmax # /home/lesliefang/pytorch_1_7_1/ssd-rn34/frameworks.ai.pytorch.private-cpu/torch/nn/functional.py:1783:0
  %6 : Tensor = prim::profile[profiled_type=BFloat16(1, 3, 224, 224, strides=[150528, 50176, 224, 1], requires_grad=0, device=cpu), seen_none=0](%4)
   = prim::profile()
  return (%6)

torch.bfloat16
```

### Versions

And we are using stock pytorch with commit: https://github.com/pytorch/pytorch/commit/123297a8c02e0ebbf8b0ae3d3cb16d0dc2e350b4.