Compare oneDNN and OpenBLAS backend of PyTorch on arm64 architecture
Labels: module: cpu, triaged, module: bfloat16, module: arm

### 🚀 The feature, motivation and pitch

I am optimizing PyTorch performance on ARM server because the customer has this requirement.
Currently, I am focusing on inference task, and I have found that the current version of PyTorch can use two backends: one is oneDNN, which is connected to the Arm Compute Library (ACL), and the other is using the BLAS library. On ARM server devices, it is typically the OpenBLAS library. 
I am using the ARM Neoverse-N2 architecture, which supports BF16 instruction extension. Both ACL and OpenBLAS have implemented optimization for BF16 GEMM, and if PyTorch users can use this optimization, it will significantly improve performance.

In fact, users can already use BF16 in the following way:
1. If build with BLAS, the way of use `.to(torch.bfloat16)`
```python
# Assuming we are using the ResNet-50 model.
model = model.to(torch.bfloat16)
x = x.to(torch.bfloat16)
pred = model(x)
```
This approach will cause all operations to be performed in bf16 type.
This will speed up gemm, but other ops such as activation layers and addition op will perform worse than fp32, and some ops do not support bf16, such as the roi-head commonly used in object detection networks.

2. the way of use `ONEDNN_DEFAULT_FPMATH_MODE`
oneDNN + ACL has a better solution, where they set an environment variable that only converts ops with improvements to bf16, while others continue to use fp32.
```bash
ONEDNN_DEFAULT_FPMATH_MODE=BF16 python resnet.py
```
Users do not need to change the code, they just specify that the value of the environment variable is BF16 before the command.

So, I tested resnet50 with pytorch profiler.

1. If build with BLAS, the way of use `.to(torch.bfloat16)`
```bash
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                             Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                     aten::conv2d         0.14%     247.000us        57.84%     102.213ms       1.929ms            53
                aten::convolution         0.21%     371.000us        57.70%     101.966ms       1.924ms            53
               aten::_convolution         0.30%     525.000us        57.49%     101.595ms       1.917ms            53
                aten::thnn_conv2d         0.11%     189.000us        57.19%     101.056ms       1.907ms            53
       aten::_slow_conv2d_forward        56.86%     100.473ms        57.08%     100.867ms       1.903ms            53
                 aten::batch_norm         0.25%     436.000us        19.34%      34.178ms     644.868us            53
     aten::_batch_norm_impl_index         0.28%     493.000us        19.23%      33.980ms     641.132us            53
          aten::native_batch_norm        18.72%      33.078ms        18.92%      33.425ms     630.660us            53
                       aten::add_         9.58%      16.927ms         9.58%      16.927ms       1.058ms            16
                      aten::relu_         0.30%     529.000us         9.38%      16.571ms     338.184us            49
                 aten::clamp_min_         9.25%      16.348ms         9.25%      16.348ms     333.633us            49
                 aten::max_pool2d         0.01%      23.000us         2.95%       5.219ms       5.219ms             1
    aten::max_pool2d_with_indices         2.94%       5.196ms         2.94%       5.196ms       5.196ms             1
                     aten::linear         0.02%      28.000us         0.36%     642.000us     642.000us             1
                      aten::addmm         0.33%     576.000us         0.34%     593.000us     593.000us             1
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 176.709ms
```

2. the way of use `ONEDNN_DEFAULT_FPMATH_MODE`
```bash
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                             Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                     aten::conv2d         0.08%     296.000us        95.89%     343.604ms       6.483ms            53
                aten::convolution         0.11%     392.000us        95.81%     343.308ms       6.478ms            53
               aten::_convolution         0.16%     586.000us        95.70%     342.916ms       6.470ms            53
         aten::mkldnn_convolution        68.22%     244.440ms        69.49%     249.009ms      12.450ms            20
                aten::thnn_conv2d         0.04%     129.000us        26.03%      93.264ms       2.826ms            33
       aten::_slow_conv2d_forward        25.92%      92.882ms        25.99%      93.135ms       2.822ms            33
                 aten::batch_norm         0.05%     187.000us         1.39%       4.996ms      94.264us            53
     aten::_batch_norm_impl_index         0.15%     549.000us         1.33%       4.773ms      90.057us            53
                 aten::max_pool2d         0.01%      22.000us         1.27%       4.553ms       4.553ms             1
    aten::max_pool2d_with_indices         1.26%       4.531ms         1.26%       4.531ms       4.531ms             1
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 358.335ms
```
So far, the OpenBLAS backend is significantly faster than the ACL.

So now there is a problem, better performance, but not easy to use. Why can't we provide convenient bf16 optimization support for the pytorch+BLAS solution? pytorch does not read environment variables, so we can provide a similar interface:
```python
# Add this line before user's code
torch.set_float32_fast_math_mode("BF16")
```
The user adds this line in front of his code, and then PyTorch can use the BF16 optimization when calling gemm, while the rest of the operations still use fp32 by default, thus getting the best performance.



### Alternatives

_No response_

### Additional context

_No response_

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @malfet