Repository: pytorch
Issue: [RFC] CPU float16 performance optimization on eager mode.
Labels: feature,module: cpu,triaged,module: half

### 🚀 The feature, motivation and pitch

The RFC is to improve float16 performance as well as Op coverage on PyTorch CPU backend on eager mode.

Float16 and BFloat16 are both commonly used reduced floating point types for performance improvement in neural network inference/training. On the CPU side, previous optimization efforts have been placed more on BFloat16, which leaves float16 at a relatively primitive status.

On the 4th generation Intel® Xeon® Scalable processor (Sapphire Rapids), a new fp16 instruction set architecture for Intel® AVX-512 has been added, e.g. **avx512-fp16**. The instruction set supports a wide range of general-purpose numeric operations for fp16. One the next generation of Xeon, Intel® Advanced Matrix Extensions (AMX) will have fp16 support, e.g. **amx-fp16**.

This proposal would help the scenario when the model is pre trained on GPU with mixed precision of float16/float32 and users intend to do deployment on the CPU side without modifying the model weights, for instance, many HuggingFace models belong to this scenario.

**This project will be targeting at:**
* Improve float16 Op coverage on PT from 52% to ~80% (for reference, BFloat16 Op coverage ratio is 83%). Torchbench models and HuggingFace models will be prioritized.
* Improve float16 performance： Since the current fp16 performance on CPU is very low, we can use fp16 v.s. fp32 speedup as the metrics here. Performance speedup will match bf16, on average a 4x-5x performance speedup will be achieved on hardware with amx-fp16 support.
* Add Automatic Mixed Precision (AMP) support for float16 on CPU.
* Improve float16 numerical stability: use fp32 as accumulate type in reduction Ops, e.g. mean.

**Technically, the optimization will be carried out as below:**

Compute intensive Ops (e.g. Convolution, Gemm, and RNN):

* Rely on oneDNN for optimal performance when the hardware has float16 acceleration.
* Functional coverage will be added for hardware without float16 acceleration, no performance gain.

Generic ATen kernels:

* Extend avx256 and avx512 and vectorization utils for dtype `Half`. Add native convert intrinsics:  `_mm256_cvtph_ps`/`_mm256_cvtps_ph` (Rounding mode: RNE).
* Unary and binary Op kernels: map to fp32 for computation.
* Non arithmetic Ops: do direct memory copy (no dtype conversion), e.g. cat, index_select.
* Reduction Ops: use fp32 as accumulate type.
* NN Ops: reuse kernels of BFloat16.

**Test Plan:**
* Extend `vec_test_all_types_AVX2` and `vec_test_all_types_AVX512` for float16.
* Add OpInfo at `torch/testing/_internal/common_methods_invocations.py`.
* Provide specific test cases for reduced floating point.

### Alternatives

_No response_

### Additional context

**Previous RFC on extending AMP fp16 on CPU:**
* https://github.com/pytorch/pytorch/issues/96093

**Float16 support in torch inductor** working in parallel (implemented in similar method as BFloat16 support), has dependency on explicit vectorization utils of `at::vec::Vectorized<Half>`.

**Pull Requests related to this feature requests:**
* https://github.com/pytorch/pytorch/pull/96076
* https://github.com/pytorch/pytorch/pull/96077
* https://github.com/pytorch/pytorch/pull/96078
* https://github.com/pytorch/pytorch/pull/96079
* https://github.com/pytorch/pytorch/pull/96080
* https://github.com/pytorch/pytorch/pull/96081
* https://github.com/pytorch/pytorch/pull/96082



cc @jgong5 @XiaobingSuper @sanchitintel @ashokei @jingxu10