Repository: pytorch
Issue: RFC: Enabling AVX512 dispatch for compute-intensive ATen ops
Labels: module: performance,module: cpu,triaged,module: intel

# 🚀 The feature, motivation and pitch

## Summary
On some more recent x86-64 architectures,  AVX512 performs better than AVX2 on compute-bound workloads, and AVX512 instructions do not cause as much throttling as they did on older CPU architectures. Eager mode performance can thus be improved on such machines.

Targeting for PyTorch 2.1, we propose enabling the dispatch of AVX512 ATen kernels if they'd be _expected_ to perform better than their AVX2 counterparts. The default ATen CPU capability would thus be AVX512.

## Approach

### Naïve solution
First, we would extend & enhance the coverage of individual ATen ops in [OpBench](https://github.com/pytorch/pytorch/tree/master/benchmarks/operator_benchmark), which would offer us more insights into ATen kernel performance characteristics by varying certain factors such as vectorization ISA, thread-pool size, memory layout, dtype, etc. OpBench already benchmarks for various input-sizes. #104655 is an example PR towards this endeavor.

[We enabled AVX512 dispatch for kernels that _always_ perform well with AVX512](https://github.com/pytorch/pytorch/pull/104165). This solution was easier to implement, as it entailed only enabling AVX512 dispatch for compute-bound ATen ops, and disabling AVX512 dispatch for memory-bound ATen ops. We do not compile AVX512 kernels of memory-bound ATen ops to reduce the binary-size. This approach is quite restrictive & disables AVX512 dispatch for most ATen ops. This set of kernels is quite small, as we are aiming for _high precision & low recall_. 

### Can we do better?
Rather than disable AVX512 dispatch for all kernels that perform poorly with AVX512 in certain cases, we can enable AVX512 dispatch for the cases in which AVX512 performance would be better. We would have to analyze oodles of data gleaned with OpBench (n factors such as CPU generation, dtype, input size, number of threads, etc) to drive this task.

The second solution requires more fine-grained analysis of ATen ops' characteristics, so we are starting with the first one, and plan to gradually move towards the second one.

### Tasks completed
- [x] Enable AVX512 dispatch of those AVX512 ATen kernels that would always perform better than their AVX2 counterparts. - #104165

### Ongoing tasks (have open PRs)
- [ ] Extend OpBench coverage - #104655 (Stalled because there doesn't seem to be a way to only run those ops' benchmarks whose kernels are being modified in a PR)


### Tasks under investigation
- [ ] Enable AVX512 dispatch for BF16/FP16 kernels even if other dtypes don't perform as well with AVX512
- [ ] A new design PoC for AVX-n kernel dispatch.
- [ ] Adoption of new design.


We welcome your comments & are open to changing our approach after discussion. Thanks!

cc @ngimel @jgong5 @mingfeima @XiaobingSuper @ashokei @jingxu10 @frank-wei @malfet 


### Alternatives

_No response_

### Additional context

Currently, the default ATen CPU capability is AVX2 but users can use the environment variable `ATEN_CPU_CAPABILITY=avx512` to use AVX512 ATen kernels