pytorch

acc9f9af - enable bf16 mkldnn path for gemm (#61891)

enable bf16 mkldnn path for gemm (#61891) Summary: # Goal: Integrate mkldnn bf16 Gemm to pytorch ## BF16 Suport for mm, addmm, bmm, addbmm, baddbmm, mv, addmv, dot (with mkldnn matmul primitive): https://oneapi-src.github.io/oneDNN/group__dnnl__api__matmul.html For gemm related ops, we keep all inputs under plain format. So we will not introduce opaque tensor for these ops to save mem copy here.  The minimized integration is only dispatch to mkldnn in addmm, but for gemm with 3-D input (with additional dim for"batch") this will call mkldnn gemm for "batch" times. Since mkldnn matmul support input with multiple dims, we directly dispatch to mkldnn gemm in {bmm, addbmm, baddbmm} to reduce the time to create mkldnn memory desc, primitive, etc. For the different definition for "bias" between mkldnn(which must be shape of (1, N)) and pytorch (which can be same shape with gemm result (M, N)), we use a fused sum to handle it. ## User Case: User case is exactly same with before because no opaque tensor's is introduced. Since the pytorch has already support bf16 data type with CPU tensor before, we can leverage the existed bf16 gemm UT. ## Gemm performance gain on CPX 28Cores/Socket: Note: data is collected using PyTorch operator benchmarks: https://github.com/pytorch/pytorch/tree/master/benchmarks/operator_benchmark (with adding bfloat16 dtype) ### use 1 thread on 1 core ### torch.addmm (M, N) * (N, K) + (M, K) | impl |16x16x16|32x32x32| 64x64x64 | 128x128x128| 256x256x256| 512x512x512|1024x1024x1024| |:---:|:---:| :---: | :---: | :---: | :---: | :---: | :---: | | aten-fp32| 4.115us|4.583us|8.230us|26.972us|211.857us|1.458ms|11.258ms| | aten-bf16 | 15.812us| 105.087us|801.787us|3.767ms|20.274ms|122.440ms|836.453ms| | mkldnn-bf16 |20.561us |22.510us|24.551us|37.709us|143.571us|0.835ms|5.76ms| We can see mkldnn-bf16 are better than aten bf16, but for smaller shapes, mkldnn bf16 are not better than aten fp32. This is because onednn overhead, this overhead more like a "constant" overhead and while problems get larger, we can ignore it. Also we are continue optimize the kernel efficiency and decrease the overhead as well. More shapes | impl |1x2048x2048|2048x1x2048| 2048x2048x1 | |:---:|:---:| :---: | :---: | | aten-fp32| 0.640ms|3.794ms|0.641ms| | aten-bf16 | 2.924ms| 3.868ms|23.413ms| | mkldnn-bf16 |0.335ms |4.490ms|0.368ms| ### use 1 socket (28 thread, 28 core) | impl | 256x256x256| 512x512x512|1024x1024x1024| 2048x2048x2048|4096x4096x4096| |:---:| :---: | :---: | :---: | :---: | :---: | | aten-fp32| 35.943us |140.315us|643.510us|5.827ms|41.761ms| | mkldnn-bf16 |53.432us|114.716us|421.858us|2.863ms|23.029ms| More shapes | impl |128x2048x2048|2048x128x2048| 2048x2048x128 | |:---:|:---:| :---: | :---: | | aten-fp32| 0.561ms|0.458ms|0.406ms| | mkldnn-bf16 |0.369ms |0.331ms|0.239ms| We dose not show aten-bf16 for this case since aten-bf16 always compute as single thread and the performance is extreme poor. The trend for this case is similar for 1 thread on 1 core. Pull Request resolved: https://github.com/pytorch/pytorch/pull/61891 Reviewed By: iramazanli Differential Revision: D29998114 Pulled By: VitalyFedyunin fbshipit-source-id: 459dc5874c638d62f290c96684ca0a694ded4b5a