chore: import upstream snapshot with attribution

python/dgl/storages/__init__.py

@@ -0,0 +1,10 @@
+"""Feature storage classes for DataLoading"""
+from .. import backend as F
+from .base import *
+from .numpy import *
+
+# Defines the name TensorStorage
+if F.get_preferred_backend() == "pytorch":
+    from .pytorch_tensor import PyTorchTensorStorage as TensorStorage
+else:
+    from .tensor import BaseTensorStorage as TensorStorage

python/dgl/storages/base.py

@@ -0,0 +1,93 @@
+"""Base classes and functionalities for feature storages."""
+
+import threading
+
+STORAGE_WRAPPERS = {}
+
+
+def register_storage_wrapper(type_):
+    """Decorator that associates a type to a ``FeatureStorage`` object."""
+
+    def deco(cls):
+        STORAGE_WRAPPERS[type_] = cls
+        return cls
+
+    return deco
+
+
+def wrap_storage(storage):
+    """Wrap an object into a FeatureStorage as specified by the ``register_storage_wrapper``
+    decorators.
+    """
+    for type_, storage_cls in STORAGE_WRAPPERS.items():
+        if isinstance(storage, type_):
+            return storage_cls(storage)
+
+    assert isinstance(
+        storage, FeatureStorage
+    ), "The frame column must be a tensor or a FeatureStorage object, got {}".format(
+        type(storage)
+    )
+    return storage
+
+
+class _FuncWrapper(object):
+    def __init__(self, func):
+        self.func = func
+
+    def __call__(self, buf, *args):
+        buf[0] = self.func(*args)
+
+
+class ThreadedFuture(object):
+    """Wraps a function into a future asynchronously executed by a Python
+    ``threading.Thread`.  The function is being executed upon instantiation of
+    this object.
+    """
+
+    def __init__(self, target, args):
+        self.buf = [None]
+
+        thread = threading.Thread(
+            target=_FuncWrapper(target),
+            args=[self.buf] + list(args),
+            daemon=True,
+        )
+        thread.start()
+        self.thread = thread
+
+    def wait(self):
+        """Blocks the current thread until the result becomes available and returns it."""
+        self.thread.join()
+        return self.buf[0]
+
+
+class FeatureStorage(object):
+    """Feature storage object which should support a fetch() operation.  It is the
+    counterpart of a tensor for homogeneous graphs, or a dict of tensor for heterogeneous
+    graphs where the keys are node/edge types.
+    """
+
+    def requires_ddp(self):
+        """Whether the FeatureStorage requires the DataLoader to set use_ddp."""
+        return False
+
+    def fetch(self, indices, device, pin_memory=False, **kwargs):
+        """Retrieve the features at the given indices.
+
+        If :attr:`indices` is a tensor, this is equivalent to
+
+        .. code::
+
+           storage[indices]
+
+        If :attr:`indices` is a dict of tensor, this is equivalent to
+
+        .. code::
+
+           {k: storage[k][indices[k]] for k in indices.keys()}
+
+        The subclasses can choose to utilize or ignore the flag :attr:`pin_memory`
+        depending on the underlying framework.
+        """
+        raise NotImplementedError

python/dgl/storages/numpy.py

@@ -0,0 +1,25 @@
+"""Feature storage for ``numpy.memmap`` object."""
+import numpy as np
+
+from .. import backend as F
+from .base import FeatureStorage, register_storage_wrapper, ThreadedFuture
+
+
+@register_storage_wrapper(np.memmap)
+class NumpyStorage(FeatureStorage):
+    """FeatureStorage that asynchronously reads features from a ``numpy.memmap`` object."""
+
+    def __init__(self, arr):
+        self.arr = arr
+
+    # pylint: disable=unused-argument
+    def _fetch(self, indices, device, pin_memory=False):
+        result = F.zerocopy_from_numpy(self.arr[indices])
+        result = F.copy_to(result, device)
+        return result
+
+    # pylint: disable=unused-argument
+    def fetch(self, indices, device, pin_memory=False, **kwargs):
+        return ThreadedFuture(
+            target=self._fetch, args=(indices, device, pin_memory)
+        )

python/dgl/storages/pytorch_tensor.py

@@ -0,0 +1,58 @@
+"""Feature storages for PyTorch tensors."""
+
+import torch
+
+from ..utils import gather_pinned_tensor_rows
+from .base import register_storage_wrapper
+from .tensor import BaseTensorStorage
+
+
+def _fetch_cpu(indices, tensor, feature_shape, device, pin_memory, **kwargs):
+    result = torch.empty(
+        indices.shape[0],
+        *feature_shape,
+        dtype=tensor.dtype,
+        pin_memory=pin_memory,
+    )
+    torch.index_select(tensor, 0, indices, out=result)
+    kwargs["non_blocking"] = pin_memory
+    result = result.to(device, **kwargs)
+    return result
+
+
+def _fetch_cuda(indices, tensor, device, **kwargs):
+    return torch.index_select(tensor, 0, indices).to(device, **kwargs)
+
+
+@register_storage_wrapper(torch.Tensor)
+class PyTorchTensorStorage(BaseTensorStorage):
+    """Feature storages for slicing a PyTorch tensor."""
+
+    def fetch(self, indices, device, pin_memory=False, **kwargs):
+        device = torch.device(device)
+        storage_device_type = self.storage.device.type
+        indices_device_type = indices.device.type
+        if storage_device_type != "cuda":
+            if indices_device_type == "cuda":
+                if self.storage.is_pinned():
+                    return gather_pinned_tensor_rows(self.storage, indices)
+                else:
+                    raise ValueError(
+                        f"Got indices on device {indices.device} whereas the feature tensor "
+                        f"is on {self.storage.device}. Please either (1) move the graph "
+                        f"to GPU with to() method, or (2) pin the graph with "
+                        f"pin_memory_() method."
+                    )
+            # CPU to CPU or CUDA - use pin_memory and async transfer if possible
+            else:
+                return _fetch_cpu(
+                    indices,
+                    self.storage,
+                    self.storage.shape[1:],
+                    device,
+                    pin_memory,
+                    **kwargs,
+                )
+        else:
+            # CUDA to CUDA or CPU
+            return _fetch_cuda(indices, self.storage, device, **kwargs)

python/dgl/storages/tensor.py

@@ -0,0 +1,17 @@
+"""Feature storages for tensors across different frameworks."""
+from .. import backend as F
+from .base import FeatureStorage
+
+
+class BaseTensorStorage(FeatureStorage):
+    """FeatureStorage that synchronously slices features from a tensor and transfers
+    it to the given device.
+    """
+
+    def __init__(self, tensor):
+        self.storage = tensor
+
+    def fetch(
+        self, indices, device, pin_memory=False, **kwargs
+    ):  # pylint: disable=unused-argument
+        return F.copy_to(F.gather_row(self.storage, indices), device, **kwargs)