Repository: pytorch
Issue: [RFC] Add JIT graph fuser for oneDNN Graph API
Labels: oncall: jit

## Feature

This RFC proposes to add the infrastructure of a new PyTorch JIT graph fuser based on [oneDNN Graph API](https://spec.oneapi.com/onednn-graph/latest/introduction.html) (codename: LLGA).

**oneDNN Graph API** _extends_ oneDNN with a high-level graph API, as a complement to the **oneDNN Primitive API**. It can benefit the PyTorch community by maximizing the optimization opportunity with target-specific optimization and code generation in a larger scope than individual DL ops across a variety of Intel and non-Intel HW.

oneDNN Graph API accepts a deep learning computation graph based on its [Opset](https://spec.oneapi.com/onednn-graph/latest/ops/index.html) as input and performs graph partitioning, where nodes that are candidates for fusion are grouped together. oneDNN Graph compiles and executes a group of deep learning operations in a graph partition as a fused operation. [oneDNN Graph programming model](https://spec.oneapi.com/onednn-graph/latest/programming_model.html) is designed for easy integration with deep learning frameworks.

The PyTorch integration will be mainly composed of two parts: 1) **Graph rewrite** for graph partitioning and fusion; 2) **Runtime** for partition compilation and execution.

1. Graph Rewrite

    oneDNN Graph rewrite pass applies fusion on TorchScript IR according to oneDNN Graph partitioning. The rewrite pass is composed of the following stages:

    - **Passing graph**: Translate TorchScript IR into oneDNN Graph IR represented by the `graph` object. All PyTorch `Node`s are mapped to oneDNN Graph `op`s and passed with `graph::add_op(op)`.

    - **Partitioning**: oneDNN Graph picks out regions of interest and returns `std::vector<partition>`. Each `partition` corresponds to an unordered list of ops to form a fusion op.

    - **Rewrite**: Cluster JIT nodes into fusion ops based on the partitions returned from the partitioning stage.

2. Runtime

    Each oneDNN Graph partition is [compiled](https://spec.oneapi.com/onednn-graph/latest/programming_model.html#partition) and [executed](https://spec.oneapi.com/onednn-graph/latest/programming_model.html#compiled-partition) in the implementation of the fusion op.

## Pitch

Add the infrastructure of a new JIT graph fuser on top of oneDNN Graph API.

## Additional context

We are structuring a new graph fuser, that follows many of the designs of TensorExpression and CudaFuser.

1. PyTorch JIT Integration and User Opt-in API

    Most of the source code lives in `torch/csrc/jit/codegen/onednn/`. Since oneDNN Graph supports fusion, code generation and execution with a unified graph API, it results in a very light-weight bridge code (< 2k LOC) for the new fuser.

    We plan to register a new custom pre-pass for graph rewrite and a new JIT op called `prim::LlgaFusionGroup` for runtime implementation. The optimization pass will be disabled by default at first, and users should enable it by calling

    ```python
    torch._C._jit_set_llga_enabled(True)
    ```

2. Alias and Mutation Reduction

    oneDNN Graph Ops are designed to be pure functional in order to simplify the graph analysis. PyTorch Ops with alias and mutation semantics are treated as [unsupported Ops](https://spec.oneapi.com/onednn-graph/latest/ops/misc/Wildcard_1.html) and will be excluded from the fusion. In order to maximize the fusion opportunities, there has to be a graph rewrite pass beforehand to convert PyTorch Ops with alias and mutation semantics to Ops with functional semantics with the best effort. We notice that there is a `MutationRemover` utilily class that can help. We are also aware that PyTorch is introducing the functional transformation (FX) tool to facilitate such a graph rewrite pass. We want to ask for architectural guidance on what is the recommended path to follow.

    oneDNN Graph API does not support control flow yet and the current implementation only supports analyzing the non-control-flow portion of the graph.

@jgong5 @Jianhui-li

cc @gmagogsfm