Repository: pytorch
Issue: [RFC] Add common device runtime abstraction
Labels: feature,triaged

## 🚀 Feature
This RFC proposes to add common frontend device runtime API and a device runtime plugin design to PyTorch. They align with CUDA runtime (torch.cuda package) and aim to simplify the work of adding new devices to PyTorch. Common code and algorithms can be shared based on this new device runtime abstraction. To enable new HWs, minimal changes are required in PyTorch packages which keep the user code unchanged. 

Below are the proposed changes at the high level. Details are described in the “additional context”.

- Common Python frontend device runtime API, abstracted from torch.cuda, under the `torch.<device>` where `<device>` corresponds to c10::DeviceType, e.g., “torch.xpu” for the XPU device. 
- C++ device runtime plugin interfaces, abstracted from CUDA Hooks to runtime, dispatched with c10::DeviceType. 
- Runtime plugin registration interface for python frondend runtime API and C++ runtime interface with c10::DeviceType as the key.
- A few modifications to some PyTorch python modules (e.g., DDP) to generalize the device runtime calls with new frontend runtime API.

## Motivation

The CUDA runtime is currently the only accelerator device runtime supported in the stock PyTorch. Whenever the Pytorch modules need to interact with the device runtime, the Pytorch calls the CUDA interface directly with the assumption that all the devices, other than the CPU, are CUDA devices. This makes it hard for adding a new device to PyTorch. One has to replicate a lot of common code (like DDP in XLA) for the new device to work around the CUDA specialization. 
As the Pytorch supports more and more new devices other than CUDA, like XPU proposed in #48246, the abstraction of runtime interface is required to decouple the Pytorch modules from CUDA. 

## Pitch

Add general frontend device runtime API and device runtime plugin design for PyTorch.

## Additional context

### Runtime plug-in registration interface

- Python frontend runtime API: register_runtime
- C++ runtime interface: register_hooks

### Python Frontend Device Runtime API

<table>
<tr>
<td> Runtime API </td> <td> CUDA API </td>
</tr>
<tr>
<td>
is_avaliable
</td>
<td>
torch.cuda.is_avaliable
</td>
</tr>

<tr>
<td>
current_stream
</td>
<td>
torch.cuda.current_stream
</td>
</tr>

<tr>
<td>
default_stream
</td>
<td>
torch.cuda.default_stream
</td>
</tr>

<tr>
<td>
device_count
</td>
<td>
torch.cuda.device_count
</td>
</tr>

<tr>
<td>
current_device
</td>
<td>
torch.cuda.current_device
</td>
</tr>

<tr>
<td>
set_device
</td>
<td>
torch.cuda.set_device
</td>
</tr>

<tr>
<td>
synchronize
</td>
<td>
torch.cuda.synchronize
</td>
</tr>

<tr>
<td>
get_device_name
</td>
<td>
torch.cuda.get_device_name
</td>
</tr>

<tr>
<td>
class Stream
</td>
<td>
torch.cuda.Stream
</td>
</tr>

<tr>
<td>
class Event
</td>
<td>
torch.cuda.Event
</td>
</tr>

<tr>
<td>
Random Generator related API
</td>
<td>
torch.cuda.random
</td>
</tr>

<tr>
<td>
Memory Management related API
</td>
<td>
torch.cuda.memory
</td>
</tr>
</table>

### C++ Runtime Interface

<table>
<tr>
<td> Runtime C++ Interface</td> <td> CUDA Hooks Interface</td>
</tr>
<tr>
<td>
init
</td>
<td>
initCUDA
</td>
</tr>

<tr>
<td>
isAvaliable
</td>
<td>
hasCUDA
</td>
</tr>

<tr>
<td>
getDefaultGenerator
</td>
<td>
getDefaultCUDAGenerator
</td>
</tr>

<tr>
<td>
getDeviceFromPtr
</td>
<td>
getDeviceFromPtr
</td>
</tr>

<tr>
<td>
getDeviceAllocator
</td>
<td>
getCUDADeviceAllocator
</td>
</tr>

<tr>
<td>
isPinnedPtr
</td>
<td>
isPinnedPtr
</td>
</tr>

<tr>
<td>
getPinnedMemoryAllocator
</td>
<td>
getPinnedMemoryAllocator
</td>
</tr>

<tr>
<td>
showConfig
</td>
<td>
showConfig
</td>
</tr>

<tr>
<td>
getNumDevices
</td>
<td>
getNumGPUs
</td>
</tr>

<tr>
<td>
current_device
</td>
<td>
current_device
</td>
</tr>
</table>

### Changes in python example

<table>
<tr>
<td> Coupled with CUDA </td> <td> General code with runtime API </td>
</tr>
<tr>
<td>

```python

# get device type from model, tensor or global default one.
dev_type = ...

# Refer to the cuda API 
if dev_type == "cuda"  and torch.cuda.is_available(): 
  ...
```

</td>
<td>

```python

# get device type from model, tensor or global default one.
dev_type = ... 

# General code to refer the runtime API 
if torch.runtime[dev_type].is_available():
  ...
```

</td>
</tr>
</table>

### Changes in C++ example

<table>
<tr>
<td> Coupled with CUDA Hooks </td> <td> General code with runtime Interface </td>
</tr>
<tr>
<td>

```c++
Device getDeviceFromPtr(void* data, DeviceType device_type) {
  initDeviceIfNeeded(device_type);
  if (device_type == at::kCPU) {
    return DeviceType::CPU;
  } else if (device_type == at::kCUDA) {
    return at::detail::getCUDAHooks().getDeviceFromPtr(data);
  } else {
    AT_ERROR(DeviceTypeName(device_type), " device type not enabled.");
  }
}
```

</td>
<td>

```c++
Device getDeviceFromPtr(void* data, DeviceType device_type) {
  initDeviceIfNeeded(device_type);
  if (device_type == at::kCPU) {
    return DeviceType::CPU;
  } else {
    auto* runtime = c10::runtime::getRuntime(device_type);
    if (runtime) {
      return runtime->getDeviceFromPtr(data);
    } else {
      AT_ERROR(DeviceTypeName(device_type), " device type not enabled.");
    }
  }
}
```

</td>
</tr>
</table>

### User code changes

<table>
<tr>
<td> Pytorch CUDA </td> <td> Pytorch with XPU Plug-in</td>
</tr>
<tr>
<td>

```python
import torch

# Refer the CUDA API
torch.cuda.seed_all()
stream = torch.cuda.current_stream()
...
```

</td>
<td>

```python
import torch
import ipex

# Refer the XPU API
torch.xpu.seed_all()
stream = torch.xpu.current_stream()
...
```

</td>
</tr>
</table>