Repository: pytorch
Issue: [RFC] XPU device for PyTorch
Labels: feature,triaged,module: intel

## 🚀 Feature

This RFC proposes to add a new user visible 'XPU' device type and the corresponding Python device runtime API to PyTorch.

XPU is a device abstraction for Intel heterogeneous computation architectures, which can be mapped to CPU, GPU, FPGA, and accelerator.

For PyTorch users, XPU works as a normal torch device, just like CPU or CUDA device, but could run with Intel CPU optimization, GPU implementation, or a mixture of heterogeneous devices according to the device mapping. 

**The enabling of XPU contains two parts:**
1. adding a new 'XPU' device type and Python device runtime API to PyTorch, which is the focus of this RFC;
2. implementation of PyTorch ops and runtime of XPU, which will be supported via a PyTorch extension, i.e. Intel Extension for PyTorch (IPEX) that registers XPU implementation with C10 and C10D registration API when being imported by PyTorch users to enable the full functionality of XPU, similar to what other PyTorch extensions do. IPEX also provides additional Python API and implementation to configure the device implementation, for example, device mapping.

## Motivation

Intel is promoting oneAPI as the industry standard and delivers Intel oneAPI products (an implementation of oneAPI standard) for high-performance heterogeneous computing.

oneAPI brings a unified programming model for developers to have common developer experiences on XPU which stands for heterogeneous hardware architectures including CPUs, GPUs, FPGAs, and other accelerators.

Adding XPU as a PyTorch device is a fundamental step to allow PyTorch users to get the best performance of DL workloads on oneAPI-powered heterogeneous hardware architectures. 

## Pitch

For this RFC in particular, we propose the following changes: 

- Add a new device type 'XPU' ('xpu' for lower case) to PyTorch. Changes are needed for code related to device model and kernel dispatch, e.g. DeviceType, Backend and DispatchKey etc. 

- A modular design for different device types to register their runtime (e.g: CUDA or XPU) with a common set of APIs. Therefore, PyTorch frontend code can be shared among different devices if they rely on the common runtime API. 

- Decouple the PyTorch frontend codes hardcoded for CUDA runtime in long term. 

## Additional context

### Work with XPU

PyTorch users interface the XPU device as a normal device but need to import IPEX for full functionality.
```pytorch
import torch 
import ipex # Intel Extension for PyTorch 

xpu = torch.device('xpu') 

input = torch.randn([100]).to(xpu) 

model.to(xpu) 

output = model(input)
```

### Python API for Device Runtime

The IPEX will register its runtime modular into PyTorch during the first import. Therefore, a new API, `torch.register_runtime(name, module)`, will be added into PyTorch frontend, which will be introduced in another separate RFC soon.


<table>
<tr>
<td> Model code </td> <td> IPEX inside </td>
</tr>
<tr>
<td>

```pytorch
import torch
# XPU runtime will be registered at the first import
import ipex  

# Refer to the XPU runtime API 
torch.xpu.current_device() 
torch.xpu.current_stream()
...
```

</td>
<td>

```pytorch
# Register XPU runtime via PyTorch new API 
import sys
import torch

current_module = sys.modules[__name__] 
# New PyTorch API for runtime module registration 
torch.register_runtime('xpu', current_module)
...
```

</td>
</tr>
</table>

### XPU Device Mapping

The 'xpu' may map to different physical devices (xpu mapping). Additional APIs in IPEX, `ipex.set_xpu_mapping(mapping)`, are recommended so that the developer can configure the mapping.

```pytorch
import ipex 

# XPU == Explicit mapping 
# The explicit mapping exposes all physical devices as xpu 
# each with a different ordinal 
ipex.set_xpu_mapping(mapping="explicit") 
```

```pytorch
import ipex 

# XPU == Extended CPU optimization 
# The extended CPU mapping, which provides more specific optimizations 
# on CPU devices only. 

ipex.set_xpu_mapping(mapping="CPU")
```

```pytorch
import ipex 

# XPU == GPU implementation 
# The GPU device mapping, which maps XPU to GPU devices only. 

ipex.set_xpu_mapping(mapping="GPU") 
```

```pytorch
import ipex 

# XPU == AUTO implementation automatically schedules the computation on 
# heterogeneous devices, e.g. CPU/GPU combined. 

ipex.set_xpu_mapping(mapping="AUTO")
```

Together with above, if the xpu mapping is set to CPU, the xpu device index is always limited to 0, that is, `torch.device("xpu", 1)` will be invalid.