Repository: pytorch
Issue: [RFC] Support MemoryView for Tensors
Labels: module: internals,feature,module: cuda,triaged,module: numpy,module: viewing and reshaping

## 🚀 Feature
[RFC] Support MemoryView for Tensors

## Motivation
When we are checkpointing the state dict to remote storage, we currently move entire tensors from GPU to CPU before uploading the bytes which is pretty memory intensive. Instead of moving entire tensors to CPU before uploading it to remote storage, we are considering to do this via streaming — taking a tensor from GPU to CPU chunk by chunk and prevent copying the entire tensor at once. On the checkpoint loading path, we hope to do something similar and prevent caching the entire tensor in CPU before moving the bytes to GPU.

An approach to do this is to use the Tensor's underlying Storage buffer and pass an offset and length for that buffer to a library responsible for writing out bytes over the wire or to disk which avoids making a copy of the data. However, we were not able to find something similar that is possible at the Python level without avoiding the numpy dependency. This also has been brought up previously [here](https://github.com/pytorch/pytorch/issues/19143).

## Pitch

We are proposing a MemoryView class for tensors which is similar to the memoryview available in python. The couple of functionalities it supports are:

* From tensor to MemoryView:
    * For a tensor in CPU or GPU memory, given (tensor, offset, size), return the memory view so that the caller can get a view of the bytes. Users can only move MemoryView across hosts and devices (in the above example, users will construct a memory view for the tensors in GPU, move it to CPU and stream out bytes to storage).
* From MemoryView to tensor:
    * Assuming that the source memory view and the destination tensor are residing on the same device (GPU/DRAM). Given a memory view and a tensor, populate the tensor using the bytes in the memory view. Users can later on move the tensor across host/device or fill a part of another tensor using this tensor.

**Note that we are not planning to support non-contiguous tensors in the first version.**


## Interfaces

1) We will expose the torch.MemoryView class which is similar to Python’s memory_view. The init function takes in either a tensor or a python memory view. 
```
class MemoryView:
    def __init__(
        self, 
        input_view: Optional[memory_view], 
        input_tensor: Optional[Tensor],
        input_tensor_offset: Optional[int],
        memory_view_size: Optional[int],
    ) -> None:
```
The init function caches the python memory view if a memory view is passed in, otherwise it will convert the tensor into a memory view and cache the result. 

2) We expose a variable for users to understand if the memory view is on cuda or in DRAM.
```
memory_view.is_cuda: bool
```
3) We expose an API for users to move memory view to device from host.
```
def to_device(self, device) -> MemoryView
```
4) Returns a copy of this object in CPU memory. If this object is already in CPU memory and on the correct device, then no copy is performed and the original object is returned.
```
def cpu(self) -> MemoryView
```
5) Fills the content of the tensor using bytes in the memory view.
```
def fill_tensor(self, tensor) -> None
```
6) Returns the length of the memory view
```
def nbytes(self) -> int
```
7) Returns the reference to the Python memory view and enable users to interact with the memory view object itself. (This is designed to work for CPU memory view only as of now because Python’s memoryview has limited support for GPU)
```
def get_py_memoryview(self) -> memoryview
```
## Additional Context / Implementation Details

**Python Layer:** In the Python layer, we will mainly implement the MemoryView class which will interact with C++ APIs. If the user passed in a memory view to the init function, we will cache the view along with its metadata (device, nbytes etc.) If users passed in a tensor instead of a memory view, we will convert it to a python memory view and cache it in the MemoryView class along with its metadata (device, nbytes etc.). The utility functions can be implemented by simply returning the cached metadata in the class (e.g. nbytes, is_cuda, get_py_memoryview). The cpu() operation can be implemented by allocating some bytes in cpu and copying the bytes over, while the to_device() operation can be implemented by converting the memory view to a CPU based tensor, moving the tensor to GPU and finally constructing a MemoryView object using the tensor on device. The initialization function and fill_tensor requires us to implement a couple of C++ interfaces described in the following section.

**C++ Layer:** In the C++ layer, we will need two functions to convert between memory views and tensors:
```
m.def("to_py_memory_view", [](const at::Tensor& tensor, int64_t offset, int64_t size) -> py::memoryview {
        return py::memoryview::from_memory(
            (void*)tensor.data() + (offset + tensor.storage_offset()) * tensor.element_size(), 
            size, 
            true));
      });

m.def("from_py_memory_view", [](
    const at::Tensor& tensor, const py::memoryview& memoryview) -> void {
        populate(tensor, memoryview);
      });
```

cc @ezyang @bhosmer @smessmer @ljk53 @bdhirsh @ngimel @mruberry @rgommers