Repository: pytorch
Issue: RFC: Integrating oneDNN Graph Compiler into Inductor C++/OpenMP Backend for Enhanced Graph Fusion and Performance
Labels: feature,triaged,oncall: pt2,module: inductor,oncall: cpu 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 would add 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.

In PyTorch 2.3, oneDNN Graph Inductor pass would be experimental, and would support FP32 & BF16 (with AMP).
PR - #120301
In PyTorch 2.4, we would add int8 support, as Quantization 2.0 would have matured by then.
In PyTorch 2.5, we may add support for generic dynamic-shape-supporting kernels, if oneDNN would have added their support by then.

### Implementation
#### Inductor joint graph passes
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 use the Inductor pattern-matcher to replacing a pattern with a target fusion, which calls an ATen op, such as `torch.ops.mkldnn._graph_sdpa_pattern_1`:

<img width="1069" alt="image" src="https://github.com/pytorch/pytorch/assets/90872293/7991c50c-29df-4a50-846d-1053196b9217">




### User interface:
```
torch.compile(options={"cpp.onednn_graph": True})
```