RFC: Integrating oneDNN Graph Compiler into Inductor C++/OpenMP Backend for Enhanced Graph Fusion and Performance
Labels: triaged, oncall: pt2, module: inductor

### 🚀 The feature, motivation and pitch

Integrating oneDNN Graph Compiler into Inductor C++ Backend enables enhanced pattern fusion and performance for CPU. 
<img src="https://github.com/pytorch/pytorch/assets/19395079/47275ec4-f7ea-425b-9656-97a6a33b1844" alt="design" width="500"/>

### Motivation
Recent developments on the Inductor C++ backend have demonstrated promising performance on DL inference workloads with CPU, thanks to optimizations like Conv/GEMM + post-op fusions and vectorization (see [this](https://dev-discuss.pytorch.org/t/Inductor-update-4-cpu-backend-started-to-show-promising-performance-boost/874) and [this](https://dev-discuss.pytorch.org/t/torchinductor-update-5-cpu-backend-backend-performance-update-and-deep-dive-on-key-optimizations/1117)). 

[oneDNN Graph API](https://spec.oneapi.io/onednn-graph/latest/introduction.html) (codename LLGA) extents oneDNN with a high-level graph API. It goes beyond Conv/GEMM post-op fusions and supports [aggressive fusion patterns](http://oneapi-src.github.io/oneDNN/dev_guide_graph_fusion_patterns.html#aggressive-fusion-patterns) such as MultiheadAttention, MLP blocks, and more (with its [graph compiler backend](http://oneapi-src.github.io/oneDNN/dev_guide_graph_compiler.html)). Other features include  [low precision](http://oneapi-src.github.io/oneDNN/dev_guide_graph_low_precision.html). Since PyTorch 1.12, this API has been added in [TorchScript JIT fuser path](https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#use-onednn-graph-with-torchscript-for-inference) showing promising performance [#49444](https://github.com/pytorch/pytorch/issues/49444).  

Integrating the oneDNN Graph Compiler with Inductor C++ backend offers further performance enhancements. Additionally, adopting the oneDNN Graph API simplifies design and development.

### Plan
Our long-term goal is to use oneDNN Graph fusion by default, replacing post-op fusions. Starting as an experimental feature, we have implemented a `onednn_graph_fusion` pass into the Inductor post-graph passes, enabled by an inductor.cpp config option. We hope it will eventually become the default option after validation of no performance regression from Inductor C++ backend.

We start with CPU inference for float32 data type. In the future, we will add support for other PyTorch 2 features, like training, quantization, dynamic shape, BF16, etc.. 

### Implementation
#### Inductor Post-grad pass
We introduce the `onednn_graph_fusion` pass directly takes the FX graph from AOTAutograd as input. The FX graph is used to construct an LLGA Graph by lowering aten/prims IR to LLGA IR and by lowering `FakeTensor` to LLGA `LogicalTensor`. LLGA identifies fusion opportunities in the Graph and returns a list of LLGA `partition`s, each represents a set of fused operations. 

To enable the desired fusion, we rewrite the FX graph by fusing nodes within each LLGA partition into a call_function node. The target of the call_function node points to the corresponding LLGA kernel, which represents the compiled partitions. Other aten ops in the FX graph are executed by the Inductor C++ backend.

#### Other graph rewrites
Some graph rewrites are implemented to match current oneDNN fusion patterns and simplify the LLGA graph. For example, the LLGA op `MatMul` supports any number of dimensions while the `aten.bmm` op only supports 3D inputs, so we rewrite the FX graph to remove `ND -> 3D` and `3D -> ND` transformations before and after a call to `aten.bmm`. These graph rewrites will eventually be supported in the oneDNN Graph code, but currently remain part of our implementation. 

#### Additional context
To implement the backend in Python, we plan to add python binding for oneDNN graph API inside `jit/codegen/onednn` temporarily.

### User interface:
```
torch.compile(options={"cpp.onednn_graph": True})
```

### Preliminary performance
OneDNN Graph provides significant speedups on models that contain advanced fusion patterns such as transformer models, with potentially larger gains as the batch size increases. As the fuison-pattern coverage of oneDNN Graph increases, we expect speedups in more and more models. Currently, the oneDNN-graph pass has an accuracy rate of 56/68 models in the torchbench suite.

The following performance results are from a 32-core, single node test on Sapphire Rapids 

| Benchmark Set (Geometric Mean) | Inductor Speedup | Inductor + oneDNN Graph Speedup |
| -- | -- | -- |
| Torchbench | 1.29x | 1.13x |
| Torchbench  (BS=32) | 1.35x | 1.28x |
| Hugging Face in Torchbench (10 models) | 1.21x | 1.38x |

| Model | Batch Size | Inductor Speedup | Inductor + oneDNN Graph Speedup |
| -- | -- | -- | -- |
| Resnet50 | 32 | 1.72x | 1.81x |
| hf_GPT2_large | 1 | 1.29x | 1.40x |
| hf_GPT2_large | 4 | 1.14x | 1.29x |
| hf_GPT2_large | 32 | 1.22x | 1.47x |
| hf_T5_large | 1 | 1.22x | 2.03x |



<details><summary>Performance Details</summary>
<p>

Model | Batch Size | Accuracy (oneDNN) |   | oneDNN 3.2 | oneDNN+GC 3.2 | Inductor
-- | -- | -- | -- | -- | -- | --
BERT_pytorch | 2 | pass |   | 0.81 | 0.82 | 1.50
Background_Matting | 1 | pass |   | 0.90 | 0.90 | 1.13
LearningToPaint | 96 | pass |   | 1.21 | 1.22 | 1.27
Super_SloMo | 6 | pass |   | 1.11 | 1.11 | 1.17
alexnet | 128 | pass |   | 1.31 | 0.94 | 1.36
attention_is_all_you_need_pytorch | 32 | fail_accuracy |   | 0.82 | 0.16 | 0.91
basic_gnn_edgecnn | 1 | pass |   | 1.60 | 1.58 | 1.66
basic_gnn_gcn | 1 | pass |   | 0.37 | 0.37 | 0.49
basic_gnn_gin | 1 | pass |   | 0.55 | 0.56 | 0.53
basic_gnn_sage | 1 | pass |   | 0.36 | 0.08 | 0.33
cm3leon_generate | 0 | infra_error |   |   |   |  
dcgan | 256 | pass |   | 1.28 | 1.27 | 1.31
densenet121 | 64 | pass |   | 0.96 | 0.97 | 1.66
detectron2_fcos_r_50_fpn | 1 | pass |   | 0.96 | 0.95 | 1.07
dlrm | 2048 | pass |   | 0.90 | 0.96 | 1.17
doctr_det_predictor | 1 | pass |   | 1.34 | 1.35 | 1.40
doctr_reco_predictor | 1 | pass |   | 1.50 | 1.48 | 2.06
drq | 1 | pass |   | 0.61 | 0.09 | 0.99
fastNLP_Bert | 1 | pass |   | 1.24 | 0.08 | 1.27
functorch_dp_cifar10 | 64 | pass |   | 1.15 | 1.02 | 1.03
hf_Albert | 1 | pass |   | 1.19 | 0.06 | 1.29
hf_Bart | 1 | pass |   | 1.01 | 0.03 | 1.16
hf_Bert | 1 | pass |   | 1.27 | 0.02 | 1.25
hf_Bert_large | 1 | pass |   | 2.00 | 0.04 | 1.66
hf_BigBird | 1 | fail_accuracy |   | 0.94 | 0.94 | 1.31
hf_DistilBert | 1 | pass |   | 1.33 | 0.05 | 1.30
hf_GPT2 | 1 | pass |   | 1.12 | 1.10 | 1.01
hf_GPT2_large | 1 | pass_due_to_skip |   | 1.40 | 1.36 | 1.29
hf_Longformer | 0 | fail_to_run |   |   |   |  
hf_Reformer | 1 | fail_accuracy |   | 0.79 | 0.26 | 0.92
hf_T5 | 1 | pass |   | 1.30 | 1.30 | 1.01
hf_T5_base | 1 | pass |   | 1.49 | 1.49 | 1.01
hf_T5_generate | 1 | fail_to_run |   |   |   |  
hf_T5_large | 1 | pass_due_to_skip |   | 2.03 | 2.09 | 1.22
lennard_jones | 1000 | pass |   | 0.65 | 0.66 | 1.26
llama | 32 | pass |   | 0.61 | 0.61 | 0.49
mnasnet1_0 | 32 | pass |   | 1.73 | 1.83 | 2.27
mobilenet_v2 | 16 | pass |   | 1.77 | 1.76 | 2.40
mobilenet_v2_quantized_qat | 0 | fail_to_run |   |   |   |  
mobilenet_v3_large | 32 | pass |   | 2.63 | 2.67 | 3.10
nanogpt_generate | 0 | fail_to_run |   |   |   |  
nvidia_deeprecommender | 256 | fail_accuracy |   | 0.85 | 0.85 | 1.07
opacus_cifar10 | 64 | pass |   | 0.75 | 0.76 | 0.84
phlippe_densenet | 128 | pass |   | 0.71 | 0.72 | 1.80
phlippe_resnet | 128 | pass |   | 1.39 | 1.41 | 1.82
pytorch_CycleGAN_and_pix2pix | 1 | pass |   | 0.90 | 0.91 | 1.18
pytorch_stargan | 16 | pass |   | 0.96 | 0.97 | 0.97
pytorch_unet | 1 | pass |   | 1.13 | 1.14 | 1.06
resnet152 | 32 | pass |   | 1.44 | 1.45 | 1.52
resnet18 | 8 | pass |   | 1.60 | 1.53 | 1.75
resnet50 | 32 | pass |   | 1.81 | 1.82 | 1.72
resnet50_quantized_qat | 0 | fail_to_run |   |   |   |  
resnext50_32x4d | 8 | pass |   | 1.61 | 1.61 | 1.47
sam | 0 | infra_error |   |   |   |  
shufflenet_v2_x1_0 | 64 | pass |   | 1.63 | 1.72 | 2.14
soft_actor_critic | 256 | pass |   | 1.16 | 0.04 | 2.00
speech_transformer | 1 | pass |   | 1.01 | 0.06 | 1.00
squeezenet1_1 | 16 | pass |   | 1.39 | 1.39 | 2.61
timm_efficientnet | 64 | pass |   | 1.40 | 1.67 | 2.23
timm_nfnet | 128 | pass |   | 1.32 | 1.33 | 1.27
timm_regnet | 32 | pass |   | 1.35 | 1.36 | 1.55
timm_resnest | 32 | pass |   | 1.49 | 1.51 | 1.80
timm_vision_transformer | 32 | pass |   | 1.16 | 1.15 | 1.29
timm_vision_transformer_large | 32 | pass_due_to_skip |   | 1.15 | 1.19 | 1.16
timm_vovnet | 32 | pass |   | 1.12 | 1.13 | 1.43
vgg16 | 4 | pass |   | 1.50 | 0.70 | 1.42
vision_maskrcnn | 1 | fail_accuracy |   | 1.26 | 1.04 | 1.29
yolov3 | 8 | pass |   | 1.42 | 1.41 | 1.53
**Geometric Mean Speedup:** |   | 56/68 |   | 1.13 | 0.65 | 1.29



</p>
</details> 

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov