Extend TorchInductor to support more backends
Labels: triaged, oncall: pt2, module: inductor

### 🚀 The feature, motivation and pitch

This RFC proposes supporting TorchInductor to add a new backend at runtime to generate code for a particular device.

### Motivation
Currently, there are two integration levels to add a new backend for PT compiler: at the AtenIR/PrimsIR level as a new backend of Dynamo, and at the Inductor loop IR level as a new codegen backend of the Inductor.

- Integrating a new backend at the AtenIR/PrimsIR level is straightforward and provides more freedom to the new backend to optimize the captured graph. But it might be suboptimal performance if the backend lacks the DL compiler capability because the decomposed operations for a single operation might need more memory and worse data locality compared with the single operation if the decomposed operations can't be fused.

- Integrating at the Inductor loop IR level can significantly simplify the complexity of design and implementation by leveraging the Inductor's fusion capability and other optimizations directly. The new backend just needs to focus on how to generate optimal code for a particular device.

This RFC focuses on the latter one.

### Proposal
There are currently two backends for the Inductor - C++/OpenMP and Triton. And we propose providing a dynamic registration mechanism on the Inductor side for a new backend. The new backend can be out-of-tree and register its codegen for a particular device at runtime.

It can provide more flexibility as a new backend just needs to customize three fundamental structures: `Scheduling`, `Kernel` and `WrapperCodegen`. The sample code could be as follows:

```python

# in-tree/out-of-tree
class ExtensionKernel(Kernel):
    pass

# in-tree/out-of-tree
class ExtensionScheduling(Scheduling):
    ...
    def codegen_nodes(self, nodes):
        with ExtensionKernel(...):
            pass
    ...

# in-tree/out-of-tree
class ExtensionWrapperCodeGen(WrapperCodeGen):
    ...

```
The `ExtensionKernel` class and `ExtensionScheduling` class are used to define the operations and scheduling, respectively, for the new backend. The `ExtensionWrapperCodeGen` class generates the Python wrapper code for the backend to glue the kernel. All three classes can be out-of-tree as they can be registered to the Inductor at runtime.

Besides the key structures, we will provide two utility functions to support the runtime registration - `register_scheduling_for_device` and `get_scheduling_for_device`. The two utility functions are used to register and get the backend for a particular device, respectively.
```python
# in-tree
def register_scheduling_for_device(device:str, ext_scheduling: Scheduling):
    pass

# in-tree
def get_scheduling_for_device(device:str) -> Scheduling:
    pass
```

After that, we just need to add a default path to route to the newly added backend according to the device type.
```python
# in-tree
def create_backend(self, device: torch.device):
    ...
    if device.type == "cpu":
        from .codegen.cpp import CppScheduling
        return CppScheduling(self)
    elif device.type == "cuda":
        from .codegen.triton import TritonScheduling
        return TritonScheduling(self)
    else:
        return get_scheduling_for_device(device.type)(self)
```

### Alternatives

Another option is to add a new backend-related code to Inductor directly. Take the backend creation as an example, the code sample could be as follows:

```python
# in-tree
class NewBackendKernel(Kernel):
    pass

# in-tree
class NewBackendScheduling(Scheduling):
    ...
    def codegen_nodes(self, nodes):
        with NewBackendKernel(...):
            pass
    ...

# in-tree
class NewBackendWrapperCodeGen(WrapperCodeGen):
    ...

# in-tree
def create_backend(self, device: torch.device):
    ...
    if device.type == "cpu":
        from .codegen.cpp import CppScheduling
        return CppScheduling(self)
    elif device.type == "xpu":
        from .codegen.new_device import NewBackendScheduling
        return NewBackendScheduling(self)
    else:
        from .codegen.triton import TritonScheduling
        return TritonScheduling(self)
```

Compared to the former option, the latter is more straightforward and simple. But it requires all the new backend-related code to be in-tree. In other words, the new backend must be landed the PyTorch. It is a high bar to support other devices.

Currently, the common practice to support new devices is through PyTorch Extension(out of the pytorch code tree), where hardware vendors implement all the necessary operations in their own PyTorch extensions. And it might take longer to finalize the design, maturize the feature, and some other things to meet the bar of upstreaming. Before it is a part of PyTorch, it will conflict with this alternative's "in-tree" requirement.

Besides that, it's worth noting that the CI/CD pipeline for the latter option may be more challenging since new devices are typically not available on AWS.

### Additional context

_No response_

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @voznesenskym @penguinwu @anijain2305 @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @desertfire