chore: import upstream snapshot with attribution

python/dgl/graphbolt/impl/__init__.py

@@ -0,0 +1,18 @@
+"""Implementation of GraphBolt."""
+from .basic_feature_store import *
+from .fused_csc_sampling_graph import *
+from .gpu_feature_cache import *
+from .gpu_cached_feature import *
+from .in_subgraph_sampler import *
+from .legacy_dataset import *
+from .neighbor_sampler import *
+from .temporal_neighbor_sampler import *
+from .ondisk_dataset import *
+from .ondisk_metadata import *
+from .sampled_subgraph_impl import *
+from .torch_based_feature_store import *
+from .uniform_negative_sampler import *
+from .gpu_graph_cache import *
+from .cpu_feature_cache import *
+from .cpu_cached_feature import *
+from .cooperative_conv import *

python/dgl/graphbolt/impl/basic_feature_store.py

@@ -0,0 +1,60 @@
+"""Basic feature store for GraphBolt."""
+
+from typing import Dict, Tuple
+
+from ..feature_store import Feature, FeatureKey, FeatureStore
+
+__all__ = ["BasicFeatureStore"]
+
+
+class BasicFeatureStore(FeatureStore):
+    r"""A basic feature store to manage multiple features for access."""
+
+    def __init__(self, features: Dict[Tuple[str, str, str], Feature]):
+        r"""Initiate a basic feature store.
+
+
+        Parameters
+        ----------
+        features : Dict[Tuple[str, str, str], Feature]
+            The dict of features served by the feature store, in which the key
+            is tuple of (domain, type_name, feature_name).
+
+        Returns
+        -------
+        The feature stores.
+        """
+        super().__init__()
+        self._features = features
+
+    def __getitem__(self, feature_key: FeatureKey) -> Feature:
+        """Access the underlying `Feature` with its (domain, type, name) as
+        the feature_key.
+        """
+        return self._features[feature_key]
+
+    def __setitem__(self, feature_key: FeatureKey, feature: Feature):
+        """Set the underlying `Feature` with its (domain, type, name) as
+        the feature_key and feature as the value.
+        """
+        self._features[feature_key] = feature
+
+    def __contains__(self, feature_key: FeatureKey) -> bool:
+        """Checks whether the provided (domain, type, name) as the feature_key
+        is container in the BasicFeatureStore."""
+        return feature_key in self._features
+
+    def __len__(self):
+        """Return the number of features."""
+        return len(self._features)
+
+    def keys(self):
+        """Get the keys of the features.
+
+        Returns
+        -------
+        List[tuple]
+            The keys of the features. The tuples are in `(domain, type_name,
+            feat_name)` format.
+        """
+        return list(self._features.keys())

python/dgl/graphbolt/impl/cooperative_conv.py

@@ -0,0 +1,117 @@
+"""Graphbolt cooperative convolution."""
+from typing import Dict, Union
+
+import torch
+
+from ..sampled_subgraph import SampledSubgraph
+from ..subgraph_sampler import all_to_all, convert_to_hetero, revert_to_homo
+
+__all__ = ["CooperativeConvFunction", "CooperativeConv"]
+
+
+class CooperativeConvFunction(torch.autograd.Function):
+    """Cooperative convolution operation from Cooperative Minibatching.
+
+    Implements the `all-to-all` message passing algorithm
+    in Cooperative Minibatching, which was initially proposed in
+    `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
+    was later first fully described in
+    `Cooperative Minibatching in Graph Neural Networks
+    <https://arxiv.org/abs/2310.12403>`__.
+    Cooperation between the GPUs eliminates duplicate work performed across the
+    GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
+    performing GNN minibatching. This reduces the redundant computations across
+    GPUs at the expense of communication.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        subgraph: SampledSubgraph,
+        tensor: Union[torch.Tensor, Dict[str, torch.Tensor]],
+    ):
+        """Implements the forward pass."""
+        counts_sent = convert_to_hetero(subgraph._counts_sent)
+        counts_received = convert_to_hetero(subgraph._counts_received)
+        seed_inverse_ids = convert_to_hetero(subgraph._seed_inverse_ids)
+        seed_sizes = convert_to_hetero(subgraph._seed_sizes)
+        ctx.communication_variables = (
+            counts_sent,
+            counts_received,
+            seed_inverse_ids,
+            seed_sizes,
+        )
+        outs = {}
+        for ntype, typed_tensor in convert_to_hetero(tensor).items():
+            out = typed_tensor.new_empty(
+                (sum(counts_sent[ntype]),) + typed_tensor.shape[1:]
+            )
+            all_to_all(
+                torch.split(out, counts_sent[ntype]),
+                torch.split(
+                    typed_tensor[seed_inverse_ids[ntype]],
+                    counts_received[ntype],
+                ),
+            )
+            outs[ntype] = out
+        return revert_to_homo(out)
+
+    @staticmethod
+    def backward(
+        ctx, grad_output: Union[torch.Tensor, Dict[str, torch.Tensor]]
+    ):
+        """Implements the backward pass."""
+        (
+            counts_sent,
+            counts_received,
+            seed_inverse_ids,
+            seed_sizes,
+        ) = ctx.communication_variables
+        delattr(ctx, "communication_variables")
+        outs = {}
+        for ntype, typed_grad_output in convert_to_hetero(grad_output).items():
+            out = typed_grad_output.new_empty(
+                (sum(counts_received[ntype]),) + typed_grad_output.shape[1:]
+            )
+            all_to_all(
+                torch.split(out, counts_received[ntype]),
+                torch.split(typed_grad_output, counts_sent[ntype]),
+            )
+            i = out.new_empty(2, out.shape[0], dtype=torch.int64)
+            i[0] = seed_inverse_ids[ntype]  # src
+            i[1] = torch.arange(
+                out.shape[0], device=typed_grad_output.device
+            )  # dst
+            coo = torch.sparse_coo_tensor(
+                i,
+                torch.ones(
+                    i.shape[1], dtype=grad_output.dtype, device=i.device
+                ),
+                size=(seed_sizes[ntype], i.shape[1]),
+            )
+            outs[ntype] = torch.sparse.mm(coo, out)
+        return None, revert_to_homo(outs)
+
+
+class CooperativeConv(torch.nn.Module):
+    """Cooperative convolution operation from Cooperative Minibatching.
+
+    Implements the `all-to-all` message passing algorithm
+    in Cooperative Minibatching, which was initially proposed in
+    `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
+    was later first fully described in
+    `Cooperative Minibatching in Graph Neural Networks
+    <https://arxiv.org/abs/2310.12403>`__.
+    Cooperation between the GPUs eliminates duplicate work performed across the
+    GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
+    performing GNN minibatching. This reduces the redundant computations across
+    GPUs at the expense of communication.
+    """
+
+    def forward(
+        self,
+        subgraph: SampledSubgraph,
+        x: Union[torch.Tensor, Dict[str, torch.Tensor]],
+    ):
+        """Implements the forward pass."""
+        return CooperativeConvFunction.apply(subgraph, x)

python/dgl/graphbolt/impl/cpu_cached_feature.py

@@ -0,0 +1,499 @@
+"""CPU cached feature for GraphBolt."""
+from typing import Dict, Optional, Union
+
+import torch
+
+from ..base import get_device_to_host_uva_stream, get_host_to_device_uva_stream
+from ..feature_store import (
+    bytes_to_number_of_items,
+    Feature,
+    FeatureKey,
+    wrap_with_cached_feature,
+)
+
+from .cpu_feature_cache import CPUFeatureCache
+
+__all__ = ["CPUCachedFeature", "cpu_cached_feature"]
+
+
+class CPUCachedFeature(Feature):
+    r"""CPU cached feature wrapping a fallback feature. Use `cpu_cached_feature`
+    to construct an instance of this class.
+
+    Parameters
+    ----------
+    fallback_feature : Feature
+        The fallback feature.
+    cache : CPUFeatureCache
+        A CPUFeatureCache instance to serve as the cache backend.
+    offset : int, optional
+        The offset value to add to the given ids before using the cache. This
+        parameter is useful if multiple `CPUCachedFeature`s are sharing a single
+        CPUFeatureCache object.
+    """
+
+    _cache_type = CPUFeatureCache
+
+    def __init__(
+        self,
+        fallback_feature: Feature,
+        cache: CPUFeatureCache,
+        offset: int = 0,
+    ):
+        super(CPUCachedFeature, self).__init__()
+        assert isinstance(fallback_feature, Feature), (
+            f"The fallback_feature must be an instance of Feature, but got "
+            f"{type(fallback_feature)}."
+        )
+        self._fallback_feature = fallback_feature
+        self._feature = cache
+        self._offset = offset
+
+    def read(self, ids: torch.Tensor = None):
+        """Read the feature by index.
+
+        Parameters
+        ----------
+        ids : torch.Tensor, optional
+            The index of the feature. If specified, only the specified indices
+            of the feature are read. If None, the entire feature is returned.
+
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        if ids is None:
+            return self._fallback_feature.read()
+        return self._feature.query_and_replace(
+            ids.cpu(), self._fallback_feature.read, self._offset
+        ).to(ids.device)
+
+    def read_async(self, ids: torch.Tensor):
+        r"""Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            ``read_async_num_stages(ids.device)``\ th invocation. The return result
+            can be accessed by calling ``.wait()``. on the returned future object.
+            It is undefined behavior to call ``.wait()`` more than once.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        policy = self._feature._policy
+        cache = self._feature._cache
+        if ids.is_cuda and self.is_pinned():
+            ids_device = ids.device
+            current_stream = torch.cuda.current_stream()
+            device_to_host_stream = get_device_to_host_uva_stream()
+            device_to_host_stream.wait_stream(current_stream)
+            with torch.cuda.stream(device_to_host_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                ids = ids.to("cpu", non_blocking=True)
+                ids_copy_event = torch.cuda.Event()
+                ids_copy_event.record()
+
+            yield  # first stage is done.
+
+            ids_copy_event.synchronize()
+            policy_future = policy.query_and_replace_async(ids, self._offset)
+
+            yield
+
+            (
+                positions,
+                index,
+                pointers,
+                missing_keys,
+                found_offsets,
+                missing_offsets,
+            ) = policy_future.wait()
+            self._feature.total_queries += ids.shape[0]
+            self._feature.total_miss += missing_keys.shape[0]
+            found_cnt = ids.size(0) - missing_keys.size(0)
+            found_positions = positions[:found_cnt]
+            missing_positions = positions[found_cnt:]
+            found_pointers = pointers[:found_cnt]
+            missing_pointers = pointers[found_cnt:]
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                found_positions = found_positions.to(
+                    ids_device, non_blocking=True
+                )
+                values_from_cpu = cache.index_select(found_positions)
+                values_from_cpu.record_stream(current_stream)
+                values_from_cpu_copy_event = torch.cuda.Event()
+                values_from_cpu_copy_event.record()
+
+            fallback_reader = self._fallback_feature.read_async(missing_keys)
+            for _ in range(
+                self._fallback_feature.read_async_num_stages(
+                    missing_keys.device
+                )
+            ):
+                missing_values_future = next(fallback_reader, None)
+                yield  # fallback feature stages.
+
+            values_from_cpu_copy_event.synchronize()
+            reading_completed = policy.reading_completed_async(
+                found_pointers, found_offsets
+            )
+
+            missing_values = missing_values_future.wait()
+            replace_future = cache.replace_async(
+                missing_positions, missing_values
+            )
+
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                index = index.to(ids_device, non_blocking=True)
+                missing_values = missing_values.to(
+                    ids_device, non_blocking=True
+                )
+                index.record_stream(current_stream)
+                missing_values.record_stream(current_stream)
+                missing_values_copy_event = torch.cuda.Event()
+                missing_values_copy_event.record()
+
+            yield
+
+            reading_completed.wait()
+            replace_future.wait()
+            writing_completed = policy.writing_completed_async(
+                missing_pointers, missing_offsets
+            )
+
+            class _Waiter:
+                def __init__(self, events, existing, missing, index):
+                    self.events = events
+                    self.existing = existing
+                    self.missing = missing
+                    self.index = index
+
+                def wait(self):
+                    """Returns the stored value when invoked."""
+                    for event in self.events:
+                        event.wait()
+                    values = torch.empty(
+                        (self.index.shape[0],) + self.missing.shape[1:],
+                        dtype=self.missing.dtype,
+                        device=ids_device,
+                    )
+                    num_found = self.existing.size(0)
+                    found_index = self.index[:num_found]
+                    missing_index = self.index[num_found:]
+                    values[found_index] = self.existing
+                    values[missing_index] = self.missing
+                    # Ensure there is no memory leak.
+                    self.events = self.existing = None
+                    self.missing = self.index = None
+                    return values
+
+            yield _Waiter(
+                [
+                    writing_completed,
+                    values_from_cpu_copy_event,
+                    missing_values_copy_event,
+                ],
+                values_from_cpu,
+                missing_values,
+                index,
+            )
+        elif ids.is_cuda:
+            ids_device = ids.device
+            current_stream = torch.cuda.current_stream()
+            device_to_host_stream = get_device_to_host_uva_stream()
+            device_to_host_stream.wait_stream(current_stream)
+            with torch.cuda.stream(device_to_host_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                ids = ids.to("cpu", non_blocking=True)
+                ids_copy_event = torch.cuda.Event()
+                ids_copy_event.record()
+
+            yield  # first stage is done.
+
+            ids_copy_event.synchronize()
+            policy_future = policy.query_and_replace_async(ids, self._offset)
+
+            yield
+
+            (
+                positions,
+                index,
+                pointers,
+                missing_keys,
+                found_offsets,
+                missing_offsets,
+            ) = policy_future.wait()
+            self._feature.total_queries += ids.shape[0]
+            self._feature.total_miss += missing_keys.shape[0]
+            found_cnt = ids.size(0) - missing_keys.size(0)
+            found_positions = positions[:found_cnt]
+            missing_positions = positions[found_cnt:]
+            found_pointers = pointers[:found_cnt]
+            missing_pointers = pointers[found_cnt:]
+            values_future = cache.query_async(
+                found_positions, index, ids.shape[0]
+            )
+
+            fallback_reader = self._fallback_feature.read_async(missing_keys)
+            for _ in range(
+                self._fallback_feature.read_async_num_stages(
+                    missing_keys.device
+                )
+            ):
+                missing_values_future = next(fallback_reader, None)
+                yield  # fallback feature stages.
+
+            values = values_future.wait()
+            reading_completed = policy.reading_completed_async(
+                found_pointers, found_offsets
+            )
+
+            missing_index = index[found_cnt:]
+
+            missing_values = missing_values_future.wait()
+            replace_future = cache.replace_async(
+                missing_positions, missing_values
+            )
+            values = torch.ops.graphbolt.scatter_async(
+                values, missing_index, missing_values
+            )
+
+            yield
+
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                values = values.wait().to(ids_device, non_blocking=True)
+                values.record_stream(current_stream)
+                values_copy_event = torch.cuda.Event()
+                values_copy_event.record()
+
+            reading_completed.wait()
+            replace_future.wait()
+            writing_completed = policy.writing_completed_async(
+                missing_pointers, missing_offsets
+            )
+
+            class _Waiter:
+                def __init__(self, events, values):
+                    self.events = events
+                    self.values = values
+
+                def wait(self):
+                    """Returns the stored value when invoked."""
+                    for event in self.events:
+                        event.wait()
+                    values = self.values
+                    # Ensure there is no memory leak.
+                    self.events = self.values = None
+                    return values
+
+            yield _Waiter([values_copy_event, writing_completed], values)
+        else:
+            policy_future = policy.query_and_replace_async(ids, self._offset)
+
+            yield
+
+            (
+                positions,
+                index,
+                pointers,
+                missing_keys,
+                found_offsets,
+                missing_offsets,
+            ) = policy_future.wait()
+            self._feature.total_queries += ids.shape[0]
+            self._feature.total_miss += missing_keys.shape[0]
+            found_cnt = ids.size(0) - missing_keys.size(0)
+            found_positions = positions[:found_cnt]
+            missing_positions = positions[found_cnt:]
+            found_pointers = pointers[:found_cnt]
+            missing_pointers = pointers[found_cnt:]
+            values_future = cache.query_async(
+                found_positions, index, ids.shape[0]
+            )
+
+            fallback_reader = self._fallback_feature.read_async(missing_keys)
+            for _ in range(
+                self._fallback_feature.read_async_num_stages(
+                    missing_keys.device
+                )
+            ):
+                missing_values_future = next(fallback_reader, None)
+                yield  # fallback feature stages.
+
+            values = values_future.wait()
+            reading_completed = policy.reading_completed_async(
+                found_pointers, found_offsets
+            )
+
+            missing_index = index[found_cnt:]
+
+            missing_values = missing_values_future.wait()
+            replace_future = cache.replace_async(
+                missing_positions, missing_values
+            )
+            values = torch.ops.graphbolt.scatter_async(
+                values, missing_index, missing_values
+            )
+
+            yield
+
+            reading_completed.wait()
+            replace_future.wait()
+            writing_completed = policy.writing_completed_async(
+                missing_pointers, missing_offsets
+            )
+
+            class _Waiter:
+                def __init__(self, event, values):
+                    self.event = event
+                    self.values = values
+
+                def wait(self):
+                    """Returns the stored value when invoked."""
+                    self.event.wait()
+                    values = self.values.wait()
+                    # Ensure there is no memory leak.
+                    self.event = self.values = None
+                    return values
+
+            yield _Waiter(writing_completed, values)
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        if ids_device.type == "cuda":
+            return 4 + self._fallback_feature.read_async_num_stages(
+                torch.device("cpu")
+            )
+        else:
+            return 3 + self._fallback_feature.read_async_num_stages(ids_device)
+
+    def size(self):
+        """Get the size of the feature.
+
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        return self._fallback_feature.size()
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        return self._fallback_feature.count()
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Update the feature.
+
+        Parameters
+        ----------
+        value : torch.Tensor
+            The updated value of the feature.
+        ids : torch.Tensor, optional
+            The indices of the feature to update. If specified, only the
+            specified indices of the feature will be updated. For the feature,
+            the `ids[i]` row is updated to `value[i]`. So the indices and value
+            must have the same length. If None, the entire feature will be
+            updated.
+        """
+        if ids is None:
+            feat0 = value[:1]
+            self._fallback_feature.update(value)
+            cache_size = min(
+                bytes_to_number_of_items(self.cache_size_in_bytes, feat0),
+                value.shape[0],
+            )
+            self._feature = None  # Destroy the existing cache first.
+            self._feature = self._cache_type(
+                (cache_size,) + feat0.shape[1:], feat0.dtype
+            )
+        else:
+            self._fallback_feature.update(value, ids)
+            self._feature.replace(ids, value, None, self._offset)
+
+    def is_pinned(self):
+        """Returns True if the cache storage is pinned."""
+        return self._feature.is_pinned()
+
+    @property
+    def cache_size_in_bytes(self):
+        """Return the size taken by the cache in bytes."""
+        return self._feature.max_size_in_bytes
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self._feature.miss_rate
+
+
+def cpu_cached_feature(
+    fallback_features: Union[Feature, Dict[FeatureKey, Feature]],
+    max_cache_size_in_bytes: int,
+    policy: Optional[str] = None,
+    pin_memory: bool = False,
+) -> Union[CPUCachedFeature, Dict[FeatureKey, CPUCachedFeature]]:
+    r"""CPU cached feature wrapping a fallback feature.
+
+    Parameters
+    ----------
+    fallback_features : Union[Feature, Dict[FeatureKey, Feature]]
+        The fallback feature(s).
+    max_cache_size_in_bytes : int
+        The capacity of the cache in bytes. The size should be a few factors
+        larger than the size of each read request. Otherwise, the caching policy
+        will hang due to all cache entries being read and/or write locked,
+        resulting in a deadlock.
+    policy : str, optional
+        The cache eviction policy algorithm name. The available policies are
+        ["s3-fifo", "sieve", "lru", "clock"]. Default is "sieve".
+    pin_memory : bool, optional
+        Whether the cache storage should be allocated on system pinned memory.
+        Default is False.
+    Returns
+    -------
+    Union[CPUCachedFeature, Dict[FeatureKey, CPUCachedFeature]]
+        New feature(s) wrapped with CPUCachedFeature.
+    """
+    return wrap_with_cached_feature(
+        CPUCachedFeature,
+        fallback_features,
+        max_cache_size_in_bytes,
+        policy=policy,
+        pin_memory=pin_memory,
+    )

python/dgl/graphbolt/impl/cpu_feature_cache.py

@@ -0,0 +1,176 @@
+"""CPU Feature Cache implementation wrapper for graphbolt."""
+import torch
+
+__all__ = ["CPUFeatureCache"]
+
+caching_policies = {
+    "s3-fifo": torch.ops.graphbolt.s3_fifo_cache_policy,
+    "sieve": torch.ops.graphbolt.sieve_cache_policy,
+    "lru": torch.ops.graphbolt.lru_cache_policy,
+    "clock": torch.ops.graphbolt.clock_cache_policy,
+}
+
+
+class CPUFeatureCache(object):
+    r"""High level wrapper for the CPU feature cache.
+
+    Parameters
+    ----------
+    cache_shape : List[int]
+        The shape of the cache. cache_shape[0] gives us the capacity.
+    dtype : torch.dtype
+        The data type of the elements stored in the cache.
+    policy: str, optional
+        The cache policy. Default is "sieve". "s3-fifo", "lru" and "clock" are
+        also available.
+    num_parts: int, optional
+        The number of cache partitions for parallelism. Default is
+        `torch.get_num_threads()`.
+    pin_memory: bool, optional
+        Whether the cache storage should be pinned.
+    """
+
+    def __init__(
+        self,
+        cache_shape,
+        dtype,
+        policy=None,
+        num_parts=None,
+        pin_memory=False,
+    ):
+        if policy is None:
+            policy = "sieve"
+        assert (
+            policy in caching_policies
+        ), f"{list(caching_policies.keys())} are the available caching policies."
+        if num_parts is None:
+            num_parts = torch.get_num_threads()
+        min_num_cache_items = num_parts * (10 if policy == "s3-fifo" else 1)
+        # Since we partition the cache, each partition needs to have a positive
+        # number of slots. In addition, each "s3-fifo" partition needs at least
+        # 10 slots since the small queue is 10% and the small queue needs a
+        # positive size.
+        if cache_shape[0] < min_num_cache_items:
+            cache_shape = (min_num_cache_items,) + cache_shape[1:]
+        self._policy = caching_policies[policy](cache_shape[0], num_parts)
+        self._cache = torch.ops.graphbolt.feature_cache(
+            cache_shape, dtype, pin_memory
+        )
+        self.total_miss = 0
+        self.total_queries = 0
+
+    def is_pinned(self):
+        """Returns True if the cache storage is pinned."""
+        return self._cache.is_pinned()
+
+    @property
+    def max_size_in_bytes(self):
+        """Return the size taken by the cache in bytes."""
+        return self._cache.nbytes
+
+    def query(self, keys, offset=0):
+        """Queries the cache.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the cache with.
+        offset : int
+            The offset to be added to the keys. Default is 0.
+
+        Returns
+        -------
+        tuple(Tensor, Tensor, Tensor, Tensor)
+            A tuple containing
+            (values, missing_indices, missing_keys, missing_offsets) where
+            values[missing_indices] corresponds to cache misses that should be
+            filled by quering another source with missing_keys. If keys is
+            pinned, then the returned values tensor is pinned as well. The
+            missing_offsets tensor has the partition offsets of missing_keys.
+        """
+        self.total_queries += keys.shape[0]
+        (
+            positions,
+            index,
+            missing_keys,
+            found_pointers,
+            found_offsets,
+            missing_offsets,
+        ) = self._policy.query(keys, offset)
+        values = self._cache.query(positions, index, keys.shape[0])
+        self._policy.reading_completed(found_pointers, found_offsets)
+        self.total_miss += missing_keys.shape[0]
+        missing_index = index[positions.size(0) :]
+        return values, missing_index, missing_keys, missing_offsets
+
+    def query_and_replace(self, keys, reader_fn, offset=0):
+        """Queries the cache. Then inserts the keys that are not found by
+        reading them by calling `reader_fn(missing_keys)`, which are then
+        inserted into the cache using the selected caching policy algorithm
+        to remove the old entries if it is full.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the cache with.
+        reader_fn : reader_fn(keys: torch.Tensor) -> torch.Tensor
+            A function that will take a missing keys tensor and will return
+            their values.
+        offset : int
+            The offset to be added to the keys. Default is 0.
+
+        Returns
+        -------
+        Tensor
+            A tensor containing values corresponding to the keys. Should equal
+            `reader_fn(keys)`, computed in a faster way.
+        """
+        self.total_queries += keys.shape[0]
+        (
+            positions,
+            index,
+            pointers,
+            missing_keys,
+            found_offsets,
+            missing_offsets,
+        ) = self._policy.query_and_replace(keys, offset)
+        found_cnt = keys.size(0) - missing_keys.size(0)
+        found_positions = positions[:found_cnt]
+        values = self._cache.query(found_positions, index, keys.shape[0])
+        found_pointers = pointers[:found_cnt]
+        self._policy.reading_completed(found_pointers, found_offsets)
+        self.total_miss += missing_keys.shape[0]
+        missing_index = index[found_cnt:]
+        missing_values = reader_fn(missing_keys)
+        values[missing_index] = missing_values
+        missing_positions = positions[found_cnt:]
+        self._cache.replace(missing_positions, missing_values)
+        missing_pointers = pointers[found_cnt:]
+        self._policy.writing_completed(missing_pointers, missing_offsets)
+        return values
+
+    def replace(self, keys, values, offsets=None, offset=0):
+        """Inserts key-value pairs into the cache using the selected caching
+        policy algorithm to remove old key-value pairs if it is full.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to insert to the cache.
+        values : Tensor
+            The values to insert to the cache.
+        offsets : Tensor, optional
+            The partition offsets of the keys.
+        offset : int
+            The offset to be added to the keys. Default is 0.
+        """
+        positions, pointers, offsets = self._policy.replace(
+            keys, offsets, offset
+        )
+        self._cache.replace(positions, values)
+        self._policy.writing_completed(pointers, offsets)
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self.total_miss / self.total_queries

python/dgl/graphbolt/impl/gpu_cached_feature.py

@@ -0,0 +1,283 @@
+"""GPU cached feature for GraphBolt."""
+from typing import Dict, Union
+
+import torch
+
+from ..feature_store import (
+    bytes_to_number_of_items,
+    Feature,
+    FeatureKey,
+    wrap_with_cached_feature,
+)
+
+from .gpu_feature_cache import GPUFeatureCache
+
+__all__ = ["GPUCachedFeature", "gpu_cached_feature"]
+
+
+class GPUCachedFeature(Feature):
+    r"""GPU cached feature wrapping a fallback feature. It uses the least
+    recently used (LRU) algorithm as the cache eviction policy. Use
+    `gpu_cached_feature` to construct an instance of this class.
+
+    Places the GPU cache to torch.cuda.current_device().
+
+    Parameters
+    ----------
+    fallback_feature : Feature
+        The fallback feature.
+    cache : GPUFeatureCache
+        A GPUFeatureCache instance to serve as the cache backend.
+    offset : int, optional
+        The offset value to add to the given ids before using the cache. This
+        parameter is useful if multiple `GPUCachedFeature`s are sharing a single
+        GPUFeatureCache object.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+    >>> torch_feat = torch.arange(10).reshape(2, -1).to("cuda")
+    >>> cache_size = 5
+    >>> fallback_feature = gb.TorchBasedFeature(torch_feat)
+    >>> feature = gb.gpu_cached_feature(fallback_feature, cache_size)
+    >>> feature.read()
+    tensor([[0, 1, 2, 3, 4],
+            [5, 6, 7, 8, 9]], device='cuda:0')
+    >>> feature.read(torch.tensor([0]).to("cuda"))
+    tensor([[0, 1, 2, 3, 4]], device='cuda:0')
+    >>> feature.update(torch.tensor([[1 for _ in range(5)]]).to("cuda"),
+    ...                torch.tensor([1]).to("cuda"))
+    >>> feature.read(torch.tensor([0, 1]).to("cuda"))
+    tensor([[0, 1, 2, 3, 4],
+            [1, 1, 1, 1, 1]], device='cuda:0')
+    >>> feature.size()
+    torch.Size([5])
+    """
+
+    _cache_type = GPUFeatureCache
+
+    def __init__(
+        self,
+        fallback_feature: Feature,
+        cache: GPUFeatureCache,
+        offset: int = 0,
+    ):
+        super(GPUCachedFeature, self).__init__()
+        assert isinstance(fallback_feature, Feature), (
+            f"The fallback_feature must be an instance of Feature, but got "
+            f"{type(fallback_feature)}."
+        )
+        self._fallback_feature = fallback_feature
+        self._feature = cache
+        self._offset = offset
+
+    def read(self, ids: torch.Tensor = None):
+        """Read the feature by index.
+
+        The returned tensor is always in GPU memory, no matter whether the
+        fallback feature is in memory or on disk.
+
+        Parameters
+        ----------
+        ids : torch.Tensor, optional
+            The index of the feature. If specified, only the specified indices
+            of the feature are read. If None, the entire feature is returned.
+
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        if ids is None:
+            return self._fallback_feature.read()
+        values, missing_index, missing_keys = self._feature.query(
+            ids if self._offset == 0 else ids + self._offset
+        )
+        missing_values = self._fallback_feature.read(
+            missing_keys if self._offset == 0 else missing_keys - self._offset
+        )
+        values[missing_index] = missing_values
+        self._feature.replace(missing_keys, missing_values)
+        return values
+
+    def read_async(self, ids: torch.Tensor):
+        r"""Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            ``read_async_num_stages(ids.device)``\ th invocation. The return result
+            can be accessed by calling ``.wait()``. on the returned future object.
+            It is undefined behavior to call ``.wait()`` more than once.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        future = self._feature.query(
+            ids if self._offset == 0 else ids + self._offset, async_op=True
+        )
+
+        yield
+
+        values, missing_index, missing_keys = future.wait()
+
+        fallback_reader = self._fallback_feature.read_async(
+            missing_keys if self._offset == 0 else missing_keys - self._offset
+        )
+        fallback_num_stages = self._fallback_feature.read_async_num_stages(
+            missing_keys.device
+        )
+        for i in range(fallback_num_stages):
+            missing_values_future = next(fallback_reader, None)
+            if i < fallback_num_stages - 1:
+                yield  # fallback feature stages.
+
+        class _Waiter:
+            def __init__(
+                self,
+                feature,
+                values,
+                missing_index,
+                missing_keys,
+                missing_values_future,
+            ):
+                self.feature = feature
+                self.values = values
+                self.missing_index = missing_index
+                self.missing_keys = missing_keys
+                self.missing_values_future = missing_values_future
+
+            def wait(self):
+                """Returns the stored value when invoked."""
+                missing_values = self.missing_values_future.wait()
+                self.feature.replace(self.missing_keys, missing_values)
+                self.values[self.missing_index] = missing_values
+                values = self.values
+                # Ensure there is no memory leak.
+                self.feature = self.values = self.missing_index = None
+                self.missing_keys = self.missing_values_future = None
+                return values
+
+        yield _Waiter(
+            self._feature,
+            values,
+            missing_index,
+            missing_keys,
+            missing_values_future,
+        )
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        assert ids_device.type == "cuda"
+        return 1 + self._fallback_feature.read_async_num_stages(ids_device)
+
+    def size(self):
+        """Get the size of the feature.
+
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        return self._fallback_feature.size()
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        return self._fallback_feature.count()
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Update the feature.
+
+        Parameters
+        ----------
+        value : torch.Tensor
+            The updated value of the feature.
+        ids : torch.Tensor, optional
+            The indices of the feature to update. If specified, only the
+            specified indices of the feature will be updated. For the feature,
+            the `ids[i]` row is updated to `value[i]`. So the indices and value
+            must have the same length. If None, the entire feature will be
+            updated.
+        """
+        if ids is None:
+            feat0 = value[:1]
+            self._fallback_feature.update(value)
+            cache_size = min(
+                bytes_to_number_of_items(self.cache_size_in_bytes, feat0),
+                value.shape[0],
+            )
+            self._feature = None  # Destroy the existing cache first.
+            self._feature = self._cache_type(
+                (cache_size,) + feat0.shape[1:], feat0.dtype
+            )
+        else:
+            self._fallback_feature.update(value, ids)
+            self._feature.replace(ids, value)
+
+    @property
+    def cache_size_in_bytes(self):
+        """Return the size taken by the cache in bytes."""
+        return self._feature.max_size_in_bytes
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self._feature.miss_rate
+
+
+def gpu_cached_feature(
+    fallback_features: Union[Feature, Dict[FeatureKey, Feature]],
+    max_cache_size_in_bytes: int,
+) -> Union[GPUCachedFeature, Dict[FeatureKey, GPUCachedFeature]]:
+    r"""GPU cached feature wrapping a fallback feature. It uses the least
+    recently used (LRU) algorithm as the cache eviction policy.
+
+    Places the GPU cache to torch.cuda.current_device().
+
+    Parameters
+    ----------
+    fallback_features : Union[Feature, Dict[FeatureKey, Feature]]
+        The fallback feature(s).
+    max_cache_size_in_bytes : int
+        The capacity of the GPU cache in bytes.
+    Returns
+    -------
+    Union[GPUCachedFeature, Dict[FeatureKey, GPUCachedFeature]]
+        The feature(s) wrapped with GPUCachedFeature.
+    """
+    return wrap_with_cached_feature(
+        GPUCachedFeature, fallback_features, max_cache_size_in_bytes
+    )

python/dgl/graphbolt/impl/gpu_feature_cache.py

@@ -0,0 +1,80 @@
+"""HugeCTR gpu_cache wrapper for graphbolt."""
+from functools import reduce
+from operator import mul
+
+import torch
+
+
+class GPUFeatureCache(object):
+    """High-level wrapper for GPU embedding cache"""
+
+    def __init__(self, cache_shape, dtype):
+        major, _ = torch.cuda.get_device_capability()
+        assert (
+            major >= 7
+        ), "GPUFeatureCache is supported only on CUDA compute capability >= 70 (Volta)."
+        self._cache = torch.ops.graphbolt.gpu_cache(cache_shape, dtype)
+        element_size = torch.tensor([], dtype=dtype).element_size()
+        self.max_size_in_bytes = reduce(mul, cache_shape) * element_size
+        self.total_miss = 0
+        self.total_queries = 0
+
+    def query(self, keys, async_op=False):
+        """Queries the GPU cache.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the GPU cache with.
+        async_op: bool
+            Boolean indicating whether the call is asynchronous. If so, the
+            result can be obtained by calling wait on the returned future.
+
+        Returns
+        -------
+        tuple(Tensor, Tensor, Tensor)
+            A tuple containing (values, missing_indices, missing_keys) where
+            values[missing_indices] corresponds to cache misses that should be
+            filled by quering another source with missing_keys.
+        """
+
+        class _Waiter:
+            def __init__(self, gpu_cache, future):
+                self.gpu_cache = gpu_cache
+                self.future = future
+
+            def wait(self):
+                """Returns the stored value when invoked."""
+                gpu_cache = self.gpu_cache
+                values, missing_index, missing_keys = (
+                    self.future.wait() if async_op else self.future
+                )
+                # Ensure there is no leak.
+                self.gpu_cache = self.future = None
+
+                gpu_cache.total_queries += values.shape[0]
+                gpu_cache.total_miss += missing_keys.shape[0]
+                return values, missing_index, missing_keys
+
+        if async_op:
+            return _Waiter(self, self._cache.query_async(keys))
+        else:
+            return _Waiter(self, self._cache.query(keys)).wait()
+
+    def replace(self, keys, values):
+        """Inserts key-value pairs into the GPU cache using the Least-Recently
+        Used (LRU) algorithm to remove old key-value pairs if it is full.
+
+        Parameters
+        ----------
+        keys: Tensor
+            The keys to insert to the GPU cache.
+        values: Tensor
+            The values to insert to the GPU cache.
+        """
+        self._cache.replace(keys, values)
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self.total_miss / self.total_queries

python/dgl/graphbolt/impl/gpu_graph_cache.py

@@ -0,0 +1,118 @@
+"""HugeCTR gpu_cache wrapper for graphbolt."""
+import torch
+
+
+class GPUGraphCache(object):
+    r"""High-level wrapper for GPU graph cache.
+
+    Places the GPU graph cache to torch.cuda.current_device().
+
+    Parameters
+    ----------
+    num_edges : int
+        Upperbound on number of edges to cache.
+    threshold : int
+        The number of accesses before the neighborhood of a vertex is cached.
+    indptr_dtype : torch.dtype
+        The dtype of the indptr tensor of the graph.
+    dtypes : list[torch.dtype]
+        The dtypes of the edge tensors that are going to be cached.
+    has_original_edge_ids : bool
+        Whether the graph to be cached has original edge ids.
+    """
+
+    def __init__(
+        self, num_edges, threshold, indptr_dtype, dtypes, has_original_edge_ids
+    ):
+        major, _ = torch.cuda.get_device_capability()
+        assert (
+            major >= 7
+        ), "GPUGraphCache is supported only on CUDA compute capability >= 70 (Volta)."
+        self._cache = torch.ops.graphbolt.gpu_graph_cache(
+            num_edges, threshold, indptr_dtype, dtypes, has_original_edge_ids
+        )
+        self.total_miss = 0
+        self.total_queries = 0
+
+    def query(self, keys):
+        """Queries the GPU cache.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the GPU graph cache with.
+
+        Returns
+        -------
+        tuple(Tensor, func)
+            A tuple containing (missing_keys, replace_fn) where replace_fn is a
+            function that should be called with the graph structure
+            corresponding to the missing keys. Its arguments are
+            (Tensor, list(Tensor)), where the first tensor is the missing indptr
+            and the second list is the missing edge tensors.
+        """
+        self.total_queries += keys.shape[0]
+        (
+            index,
+            position,
+            num_hit,
+            num_threshold,
+        ) = self._cache.query(keys)
+        self.total_miss += keys.shape[0] - num_hit
+
+        def replace_functional(missing_indptr, missing_edge_tensors):
+            return self._cache.replace(
+                keys,
+                index,
+                position,
+                num_hit,
+                num_threshold,
+                missing_indptr,
+                missing_edge_tensors,
+            )
+
+        return keys[index[num_hit:]], replace_functional
+
+    def query_async(self, keys):
+        """Queries the GPU cache asynchronously.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the GPU graph cache with.
+
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns the missing keys on the second
+            invocation and expects the fetched indptr and edge tensors on the
+            next invocation. The third and last invocation returns a future
+            object and the return result can be accessed by calling `.wait()`
+            on the returned future object. It is undefined behavior to call
+            `.wait()` more than once.
+        """
+        future = self._cache.query_async(keys)
+
+        yield
+
+        index, position, num_hit, num_threshold = future.wait()
+
+        self.total_queries += keys.shape[0]
+        self.total_miss += keys.shape[0] - num_hit
+
+        missing_indptr, missing_edge_tensors = yield keys[index[num_hit:]]
+
+        yield self._cache.replace_async(
+            keys,
+            index,
+            position,
+            num_hit,
+            num_threshold,
+            missing_indptr,
+            missing_edge_tensors,
+        )
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self.total_miss / self.total_queries

python/dgl/graphbolt/impl/in_subgraph_sampler.py

@@ -0,0 +1,86 @@
+"""In-subgraph sampler for GraphBolt."""
+
+from torch.utils.data import functional_datapipe
+
+from ..internal import unique_and_compact_csc_formats
+
+from ..subgraph_sampler import SubgraphSampler
+from .sampled_subgraph_impl import SampledSubgraphImpl
+
+
+__all__ = ["InSubgraphSampler"]
+
+
+@functional_datapipe("sample_in_subgraph")
+class InSubgraphSampler(SubgraphSampler):
+    """Sample the subgraph induced on the inbound edges of the given nodes.
+
+    Functional name: :obj:`sample_in_subgraph`.
+
+    In-subgraph sampler is responsible for sampling a subgraph from given data,
+    returning an induced subgraph along with compacted information.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    graph : FusedCSCSamplingGraph
+        The graph on which to perform in_subgraph sampling.
+
+    Examples
+    -------
+    >>> import dgl.graphbolt as gb
+    >>> import torch
+    >>> indptr = torch.LongTensor([0, 3, 5, 7, 9, 12, 14])
+    >>> indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4])
+    >>> graph = gb.fused_csc_sampling_graph(indptr, indices)
+    >>> item_set = gb.ItemSet(len(indptr) - 1, names="seeds")
+    >>> item_sampler = gb.ItemSampler(item_set, batch_size=2)
+    >>> insubgraph_sampler = gb.InSubgraphSampler(item_sampler, graph)
+    >>> for _, data in enumerate(insubgraph_sampler):
+    ...     print(data.sampled_subgraphs[0].sampled_csc)
+    ...     print(data.sampled_subgraphs[0].original_row_node_ids)
+    ...     print(data.sampled_subgraphs[0].original_column_node_ids)
+    CSCFormatBase(indptr=tensor([0, 3, 5]),
+                indices=tensor([0, 1, 2, 3, 4]),
+    )
+    tensor([0, 1, 4, 2, 3])
+    tensor([0, 1])
+    CSCFormatBase(indptr=tensor([0, 2, 4]),
+                indices=tensor([2, 3, 4, 0]),
+    )
+    tensor([2, 3, 0, 5, 1])
+    tensor([2, 3])
+    CSCFormatBase(indptr=tensor([0, 3, 5]),
+                indices=tensor([2, 3, 1, 4, 0]),
+    )
+    tensor([4, 5, 0, 3, 1])
+    tensor([4, 5])
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+    ):
+        super().__init__(datapipe)
+        self.graph = graph
+        self.sampler = graph.in_subgraph
+
+    def sample_subgraphs(
+        self, seeds, seeds_timestamp, seeds_pre_time_window=None
+    ):
+        subgraph = self.sampler(seeds)
+        (
+            original_row_node_ids,
+            compacted_csc_formats,
+            _,
+        ) = unique_and_compact_csc_formats(subgraph.sampled_csc, seeds)
+        subgraph = SampledSubgraphImpl(
+            sampled_csc=compacted_csc_formats,
+            original_column_node_ids=seeds,
+            original_row_node_ids=original_row_node_ids,
+            original_edge_ids=subgraph.original_edge_ids,
+        )
+        seeds = original_row_node_ids
+        return (seeds, [subgraph])

python/dgl/graphbolt/impl/legacy_dataset.py

@@ -0,0 +1,158 @@
+"""Graphbolt dataset for legacy DGLDataset."""
+
+from typing import List, Union
+
+from ..base import etype_tuple_to_str
+from ..dataset import Dataset, Task
+from ..itemset import HeteroItemSet, ItemSet
+from ..sampling_graph import SamplingGraph
+from .basic_feature_store import BasicFeatureStore
+from .fused_csc_sampling_graph import from_dglgraph
+from .ondisk_dataset import OnDiskTask
+from .torch_based_feature_store import TorchBasedFeature
+
+
+class LegacyDataset(Dataset):
+    """A Graphbolt dataset for legacy DGLDataset."""
+
+    def __init__(self, legacy):
+        # Only supports single graph cases.
+        assert len(legacy) == 1
+        graph = legacy[0]
+        # Handle OGB Dataset.
+        if isinstance(graph, tuple):
+            graph, _ = graph
+        if graph.is_homogeneous:
+            self._init_as_homogeneous_node_pred(legacy)
+        else:
+            self._init_as_heterogeneous_node_pred(legacy)
+
+    def _init_as_heterogeneous_node_pred(self, legacy):
+        def _init_item_set_dict(idx, labels):
+            item_set_dict = {}
+            for key in idx.keys():
+                item_set = ItemSet(
+                    (idx[key], labels[key][idx[key]]),
+                    names=("seeds", "labels"),
+                )
+                item_set_dict[key] = item_set
+            return HeteroItemSet(item_set_dict)
+
+        # OGB Dataset has the idx split.
+        if hasattr(legacy, "get_idx_split"):
+            graph, labels = legacy[0]
+            split_idx = legacy.get_idx_split()
+
+            # Initialize tasks.
+            tasks = []
+            metadata = {
+                "num_classes": legacy.num_classes,
+                "name": "node_classification",
+            }
+            train_set = _init_item_set_dict(split_idx["train"], labels)
+            validation_set = _init_item_set_dict(split_idx["valid"], labels)
+            test_set = _init_item_set_dict(split_idx["test"], labels)
+            task = OnDiskTask(metadata, train_set, validation_set, test_set)
+            tasks.append(task)
+            self._tasks = tasks
+
+            item_set_dict = {}
+            for ntype in graph.ntypes:
+                item_set = ItemSet(graph.num_nodes(ntype), names="seeds")
+                item_set_dict[ntype] = item_set
+            self._all_nodes_set = HeteroItemSet(item_set_dict)
+
+            features = {}
+            for ntype in graph.ntypes:
+                for name in graph.nodes[ntype].data.keys():
+                    tensor = graph.nodes[ntype].data[name]
+                    if tensor.dim() == 1:
+                        tensor = tensor.view(-1, 1)
+                    features[("node", ntype, name)] = TorchBasedFeature(tensor)
+            for etype in graph.canonical_etypes:
+                for name in graph.edges[etype].data.keys():
+                    tensor = graph.edges[etype].data[name]
+                    if tensor.dim() == 1:
+                        tensor = tensor.view(-1, 1)
+                    gb_etype = etype_tuple_to_str(etype)
+                    features[("edge", gb_etype, name)] = TorchBasedFeature(
+                        tensor
+                    )
+            self._feature = BasicFeatureStore(features)
+            self._graph = from_dglgraph(graph, is_homogeneous=False)
+            self._dataset_name = legacy.name
+        else:
+            raise NotImplementedError(
+                "Only support heterogeneous ogn node pred dataset"
+            )
+
+    def _init_as_homogeneous_node_pred(self, legacy):
+        from dgl.data import AsNodePredDataset
+
+        legacy = AsNodePredDataset(legacy)
+
+        # Initialize tasks.
+        tasks = []
+        metadata = {
+            "num_classes": legacy.num_classes,
+            "name": "node_classification",
+        }
+        train_labels = legacy[0].ndata["label"][legacy.train_idx]
+        validation_labels = legacy[0].ndata["label"][legacy.val_idx]
+        test_labels = legacy[0].ndata["label"][legacy.test_idx]
+        train_set = ItemSet(
+            (legacy.train_idx, train_labels),
+            names=("seeds", "labels"),
+        )
+        validation_set = ItemSet(
+            (legacy.val_idx, validation_labels),
+            names=("seeds", "labels"),
+        )
+        test_set = ItemSet(
+            (legacy.test_idx, test_labels), names=("seeds", "labels")
+        )
+        task = OnDiskTask(metadata, train_set, validation_set, test_set)
+        tasks.append(task)
+        self._tasks = tasks
+
+        num_nodes = legacy[0].num_nodes()
+        self._all_nodes_set = ItemSet(num_nodes, names="seeds")
+        features = {}
+        for name in legacy[0].ndata.keys():
+            tensor = legacy[0].ndata[name]
+            if tensor.dim() == 1:
+                tensor = tensor.view(-1, 1)
+            features[("node", None, name)] = TorchBasedFeature(tensor)
+        for name in legacy[0].edata.keys():
+            tensor = legacy[0].edata[name]
+            if tensor.dim() == 1:
+                tensor = tensor.view(-1, 1)
+            features[("edge", None, name)] = TorchBasedFeature(tensor)
+        self._feature = BasicFeatureStore(features)
+        self._graph = from_dglgraph(legacy[0], is_homogeneous=True)
+        self._dataset_name = legacy.name
+
+    @property
+    def tasks(self) -> List[Task]:
+        """Return the tasks."""
+        return self._tasks
+
+    @property
+    def graph(self) -> SamplingGraph:
+        """Return the graph."""
+        return self._graph
+
+    @property
+    def feature(self) -> BasicFeatureStore:
+        """Return the feature."""
+        return self._feature
+
+    @property
+    def dataset_name(self) -> str:
+        """Return the dataset name."""
+        return self._dataset_name
+
+    @property
+    def all_nodes_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the itemset containing all nodes."""
+        return self._all_nodes_set

python/dgl/graphbolt/impl/ondisk_metadata.py

@@ -0,0 +1,129 @@
+"""Ondisk metadata of GraphBolt."""
+
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+import pydantic
+
+from ..internal_utils import version
+
+
+__all__ = [
+    "OnDiskFeatureDataFormat",
+    "OnDiskTVTSetData",
+    "OnDiskTVTSet",
+    "OnDiskFeatureDataDomain",
+    "OnDiskFeatureData",
+    "OnDiskMetaData",
+    "OnDiskGraphTopologyType",
+    "OnDiskGraphTopology",
+    "OnDiskTaskData",
+]
+
+
+class ExtraMetaData(pydantic.BaseModel, extra="allow"):
+    """Group extra fields into metadata. Internal use only."""
+
+    extra_fields: Optional[Dict[str, Any]] = {}
+
+    # As pydantic 2.0 has changed the API of validators, we need to use
+    # different validators for different versions to be compatible with
+    # previous versions.
+    if version.parse(pydantic.__version__) >= version.parse("2.0"):
+
+        @pydantic.model_validator(mode="before")
+        @classmethod
+        def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+            """Build extra fields."""
+            for key in list(values.keys()):
+                if key not in cls.model_fields:
+                    values["extra_fields"] = values.get("extra_fields", {})
+                    values["extra_fields"][key] = values.pop(key)
+            return values
+
+    else:
+
+        @pydantic.root_validator(pre=True)
+        @classmethod
+        def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+            """Build extra fields."""
+            for key in list(values.keys()):
+                if key not in cls.__fields__:
+                    values["extra_fields"] = values.get("extra_fields", {})
+                    values["extra_fields"][key] = values.pop(key)
+            return values
+
+
+class OnDiskFeatureDataFormat(str, Enum):
+    """Enum of data format."""
+
+    TORCH = "torch"
+    NUMPY = "numpy"
+
+
+class OnDiskTVTSetData(pydantic.BaseModel):
+    """Train-Validation-Test set data."""
+
+    name: Optional[str] = None
+    format: OnDiskFeatureDataFormat
+    in_memory: Optional[bool] = True
+    path: str
+
+
+class OnDiskTVTSet(pydantic.BaseModel):
+    """Train-Validation-Test set."""
+
+    type: Optional[str] = None
+    data: List[OnDiskTVTSetData]
+
+
+class OnDiskFeatureDataDomain(str, Enum):
+    """Enum of feature data domain."""
+
+    NODE = "node"
+    EDGE = "edge"
+    GRAPH = "graph"
+
+
+class OnDiskFeatureData(ExtraMetaData):
+    r"""The description of an on-disk feature."""
+    domain: OnDiskFeatureDataDomain
+    type: Optional[str] = None
+    name: str
+    format: OnDiskFeatureDataFormat
+    path: str
+    in_memory: Optional[bool] = True
+
+
+class OnDiskGraphTopologyType(str, Enum):
+    """Enum of graph topology type."""
+
+    FUSED_CSC_SAMPLING = "FusedCSCSamplingGraph"
+
+
+class OnDiskGraphTopology(pydantic.BaseModel):
+    """The description of an on-disk graph topology."""
+
+    type: OnDiskGraphTopologyType
+    path: str
+
+
+class OnDiskTaskData(ExtraMetaData):
+    """Task specification in YAML."""
+
+    train_set: Optional[List[OnDiskTVTSet]] = []
+    validation_set: Optional[List[OnDiskTVTSet]] = []
+    test_set: Optional[List[OnDiskTVTSet]] = []
+
+
+class OnDiskMetaData(pydantic.BaseModel):
+    """Metadata specification in YAML.
+
+    As multiple node/edge types and multiple splits are supported, each TVT set
+    is a list of list of ``OnDiskTVTSet``.
+    """
+
+    dataset_name: Optional[str] = None
+    graph_topology: Optional[OnDiskGraphTopology] = None
+    feature_data: Optional[List[OnDiskFeatureData]] = []
+    tasks: Optional[List[OnDiskTaskData]] = []

python/dgl/graphbolt/impl/sampled_subgraph_impl.py

@@ -0,0 +1,118 @@
+"""Sampled subgraph for FusedCSCSamplingGraph."""
+# pylint: disable= invalid-name
+from dataclasses import dataclass
+from typing import Dict, Union
+
+import torch
+
+from ..base import CSCFormatBase, etype_str_to_tuple
+from ..internal_utils import get_attributes
+from ..sampled_subgraph import SampledSubgraph
+
+__all__ = ["SampledSubgraphImpl"]
+
+
+@dataclass
+class SampledSubgraphImpl(SampledSubgraph):
+    r"""Sampled subgraph of CSCSamplingGraph.
+
+    Examples
+    --------
+    >>> sampled_csc = {"A:relation:B": CSCFormatBase(indptr=torch.tensor([0, 1, 2, 3]),
+    ... indices=torch.tensor([0, 1, 2]))}
+    >>> original_column_node_ids = {'B': torch.tensor([10, 11, 12])}
+    >>> original_row_node_ids = {'A': torch.tensor([13, 14, 15])}
+    >>> original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
+    >>> subgraph = gb.SampledSubgraphImpl(
+    ... sampled_csc=sampled_csc,
+    ... original_column_node_ids=original_column_node_ids,
+    ... original_row_node_ids=original_row_node_ids,
+    ... original_edge_ids=original_edge_ids
+    ... )
+    >>> print(subgraph.sampled_csc)
+    {"A:relation:B": CSCForamtBase(indptr=torch.tensor([0, 1, 2, 3]),
+    ... indices=torch.tensor([0, 1, 2]))}
+    >>> print(subgraph.original_column_node_ids)
+    {'B': tensor([10, 11, 12])}
+    >>> print(subgraph.original_row_node_ids)
+    {'A': tensor([13, 14, 15])}
+    >>> print(subgraph.original_edge_ids)
+    {"A:relation:B": tensor([19, 20, 21])}
+    """
+    sampled_csc: Union[CSCFormatBase, Dict[str, CSCFormatBase]] = None
+    original_column_node_ids: Union[
+        Dict[str, torch.Tensor], torch.Tensor
+    ] = None
+    original_row_node_ids: Union[Dict[str, torch.Tensor], torch.Tensor] = None
+    original_edge_ids: Union[Dict[str, torch.Tensor], torch.Tensor] = None
+    # Used to fetch sampled_csc.indices if it is missing.
+    _edge_ids_in_fused_csc_sampling_graph: Union[
+        Dict[str, torch.Tensor], torch.Tensor
+    ] = None
+
+    def __post_init__(self):
+        if isinstance(self.sampled_csc, dict):
+            for etype, pair in self.sampled_csc.items():
+                assert (
+                    isinstance(etype, str)
+                    and len(etype_str_to_tuple(etype)) == 3
+                ), "Edge type should be a string in format of str:str:str."
+                assert pair.indptr is not None and isinstance(
+                    pair.indptr, torch.Tensor
+                ), "Node pair should be have indptr of type torch.Tensor."
+                # For CUDA, indices may be None because it will be fetched later.
+                if not pair.indptr.is_cuda or pair.indices is not None:
+                    assert isinstance(
+                        pair.indices, torch.Tensor
+                    ), "Node pair should be have indices of type torch.Tensor."
+                else:
+                    assert isinstance(
+                        self._edge_ids_in_fused_csc_sampling_graph.get(
+                            etype, None
+                        ),
+                        torch.Tensor,
+                    ), "When indices is missing, sampled edge ids needs to be provided."
+        else:
+            assert self.sampled_csc.indptr is not None and isinstance(
+                self.sampled_csc.indptr, torch.Tensor
+            ), "Node pair should be have torch.Tensor indptr."
+            # For CUDA, indices may be None because it will be fetched later.
+            if (
+                not self.sampled_csc.indptr.is_cuda
+                or self.sampled_csc.indices is not None
+            ):
+                assert isinstance(
+                    self.sampled_csc.indices, torch.Tensor
+                ), "Node pair should have a torch.Tensor indices."
+            else:
+                assert isinstance(
+                    self._edge_ids_in_fused_csc_sampling_graph, torch.Tensor
+                ), "When indices is missing, sampled edge ids needs to be provided."
+
+    def __repr__(self) -> str:
+        return _sampled_subgraph_str(self, "SampledSubgraphImpl")
+
+
+def _sampled_subgraph_str(sampled_subgraph: SampledSubgraph, classname) -> str:
+    final_str = classname + "("
+
+    attributes = get_attributes(sampled_subgraph)
+    attributes.reverse()
+
+    for name in attributes:
+        if name in "_edge_ids_in_fused_csc_sampling_graph":
+            continue
+        val = getattr(sampled_subgraph, name)
+
+        def _add_indent(_str, indent):
+            lines = _str.split("\n")
+            lines = [lines[0]] + [" " * indent + line for line in lines[1:]]
+            return "\n".join(lines)
+
+        val = str(val)
+        final_str = (
+            final_str
+            + f"{name}={_add_indent(val, len(name) + len(classname) + 1)},\n"
+            + " " * len(classname)
+        )
+    return final_str[: -len(classname)] + ")"

python/dgl/graphbolt/impl/temporal_neighbor_sampler.py

@@ -0,0 +1,270 @@
+"""Temporal neighbor subgraph samplers for GraphBolt."""
+import torch
+from torch.utils.data import functional_datapipe
+
+from ..internal import compact_csc_format
+
+from ..subgraph_sampler import SubgraphSampler
+from .sampled_subgraph_impl import SampledSubgraphImpl
+
+
+__all__ = ["TemporalNeighborSampler", "TemporalLayerNeighborSampler"]
+
+
+class TemporalNeighborSamplerImpl(SubgraphSampler):
+    """Base class for TemporalNeighborSamplers."""
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+        fanouts,
+        replace,
+        prob_name,
+        node_timestamp_attr_name,
+        edge_timestamp_attr_name,
+        sampler,
+    ):
+        super().__init__(datapipe)
+        self.graph = graph
+        # Convert fanouts to a list of tensors.
+        self.fanouts = []
+        for fanout in fanouts:
+            if not isinstance(fanout, torch.Tensor):
+                fanout = torch.LongTensor([int(fanout)])
+            self.fanouts.insert(0, fanout)
+        self.replace = replace
+        self.prob_name = prob_name
+        self.node_timestamp_attr_name = node_timestamp_attr_name
+        self.edge_timestamp_attr_name = edge_timestamp_attr_name
+        self.sampler = sampler
+
+    def sample_subgraphs(
+        self, seeds, seeds_timestamp, seeds_pre_time_window=None
+    ):
+        assert (
+            seeds_timestamp is not None
+        ), "seeds_timestamp must be provided for temporal neighbor sampling."
+        subgraphs = []
+        num_layers = len(self.fanouts)
+        # Enrich seeds with all node types. Ensure that the dtype and device
+        # remain consistent with those of the existing seeds.
+        if isinstance(seeds, dict):
+            first_val = next(iter(seeds.items()))[1]
+            ntypes = list(self.graph.node_type_to_id.keys())
+            seeds = {
+                ntype: seeds.get(
+                    ntype,
+                    torch.tensor(
+                        [], dtype=first_val.dtype, device=first_val.device
+                    ),
+                )
+                for ntype in ntypes
+            }
+            empty_tensor = torch.tensor(
+                [], dtype=torch.int64, device=first_val.device
+            )
+            seeds_timestamp = {
+                ntype: seeds_timestamp.get(ntype, empty_tensor)
+                for ntype in ntypes
+            }
+            if seeds_pre_time_window:
+                seeds_pre_time_window = {
+                    ntype: seeds_pre_time_window.get(ntype, empty_tensor)
+                    for ntype in ntypes
+                }
+        for hop in range(num_layers):
+            subgraph = self.sampler(
+                seeds,
+                seeds_timestamp,
+                self.fanouts[hop],
+                self.replace,
+                seeds_pre_time_window,
+                self.prob_name,
+                self.node_timestamp_attr_name,
+                self.edge_timestamp_attr_name,
+            )
+            (
+                original_row_node_ids,
+                compacted_csc_formats,
+                row_timestamps,
+            ) = compact_csc_format(subgraph.sampled_csc, seeds, seeds_timestamp)
+
+            subgraph = SampledSubgraphImpl(
+                sampled_csc=compacted_csc_formats,
+                original_column_node_ids=seeds,
+                original_row_node_ids=original_row_node_ids,
+                original_edge_ids=subgraph.original_edge_ids,
+            )
+
+            subgraphs.insert(0, subgraph)
+            seeds = original_row_node_ids
+            seeds_timestamp = row_timestamps
+        return seeds, subgraphs
+
+
+@functional_datapipe("temporal_sample_neighbor")
+class TemporalNeighborSampler(TemporalNeighborSamplerImpl):
+    """Temporally sample neighbor edges from a graph and return sampled
+    subgraphs.
+
+    Functional name: :obj:`temporal_sample_neighbor`.
+
+    Neighbor sampler is responsible for sampling a subgraph from given data. It
+    returns an induced subgraph along with compacted information. In the
+    context of a node classification task, the neighbor sampler directly
+    utilizes the nodes provided as seed nodes. However, in scenarios involving
+    link prediction, the process needs another pre-peocess operation. That is,
+    gathering unique nodes from the given node pairs, encompassing both
+    positive and negative node pairs, and employs these nodes as the seed nodes
+    for subsequent steps.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    graph : FusedCSCSamplingGraph
+        The graph on which to perform subgraph sampling.
+    fanouts: list[torch.Tensor] or list[int]
+        The number of edges to be sampled for each node with or without
+        considering edge types. The length of this parameter implicitly
+        signifies the layer of sampling being conducted.
+        Note: The fanout order is from the outermost layer to innermost layer.
+        For example, the fanout '[15, 10, 5]' means that 15 to the outermost
+        layer, 10 to the intermediate layer and 5 corresponds to the innermost
+        layer.
+    replace: bool
+        Boolean indicating whether the sample is preformed with or
+        without replacement. If True, a value can be selected multiple
+        times. Otherwise, each value can be selected only once.
+    prob_name: str, optional
+        The name of an edge attribute used as the weights of sampling for
+        each node. This attribute tensor should contain (unnormalized)
+        probabilities corresponding to each neighboring edge of a node.
+        It must be a 1D floating-point or boolean tensor, with the number
+        of elements equalling the total number of edges.
+    node_timestamp_attr_name: str, optional
+        The name of an node attribute used as the timestamps of nodes.
+        It must be a 1D integer tensor, with the number of elements
+        equalling the total number of nodes.
+    edge_timestamp_attr_name: str, optional
+        The name of an edge attribute used as the timestamps of edges.
+        It must be a 1D integer tensor, with the number of elements
+        equalling the total number of edges.
+
+    Examples
+    -------
+    TODO(zhenkun) : Add an example after the API to pass timestamps is finalized.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+        fanouts,
+        replace=False,
+        prob_name=None,
+        node_timestamp_attr_name=None,
+        edge_timestamp_attr_name=None,
+    ):
+        super().__init__(
+            datapipe,
+            graph,
+            fanouts,
+            replace,
+            prob_name,
+            node_timestamp_attr_name,
+            edge_timestamp_attr_name,
+            graph.temporal_sample_neighbors,
+        )
+
+
+@functional_datapipe("temporal_sample_layer_neighbor")
+class TemporalLayerNeighborSampler(TemporalNeighborSamplerImpl):
+    """Temporally sample neighbor edges from a graph and return sampled
+    subgraphs.
+
+    Functional name: :obj:`temporal_sample_layer_neighbor`.
+
+    Sampler that builds computational dependency of node representations via
+    labor sampling for multilayer GNN from the NeurIPS 2023 paper
+    `Layer-Neighbor Sampling -- Defusing Neighborhood Explosion in GNNs
+    <https://proceedings.neurips.cc/paper_files/paper/2023/file/51f9036d5e7ae822da8f6d4adda1fb39-Paper-Conference.pdf>`__
+
+    Layer-Neighbor sampler is responsible for sampling a subgraph from given
+    data. It returns an induced subgraph along with compacted information. In
+    the context of a node classification task, the neighbor sampler directly
+    utilizes the nodes provided as seed nodes. However, in scenarios involving
+    link prediction, the process needs another pre-process operation. That is,
+    gathering unique nodes from the given node pairs, encompassing both
+    positive and negative node pairs, and employs these nodes as the seed nodes
+    for subsequent steps. When the graph is hetero, sampled subgraphs in
+    minibatch will contain every edge type even though it is empty after
+    sampling.
+
+    Implements the approach described in Appendix A.3 of the paper. Similar to
+    dgl.dataloading.LaborSampler but this uses sequential poisson sampling
+    instead of poisson sampling to keep the count of sampled edges per vertex
+    deterministic like NeighborSampler. Thus, it is a drop-in replacement for
+    NeighborSampler. However, unlike NeighborSampler, it samples fewer vertices
+    and edges for multilayer GNN scenario without harming convergence speed with
+    respect to training iterations.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    graph : FusedCSCSamplingGraph
+        The graph on which to perform subgraph sampling.
+    fanouts: list[torch.Tensor] or list[int]
+        The number of edges to be sampled for each node with or without
+        considering edge types. The length of this parameter implicitly
+        signifies the layer of sampling being conducted.
+        Note: The fanout order is from the outermost layer to innermost layer.
+        For example, the fanout '[15, 10, 5]' means that 15 to the outermost
+        layer, 10 to the intermediate layer and 5 corresponds to the innermost
+        layer.
+    replace: bool
+        Boolean indicating whether the sample is preformed with or
+        without replacement. If True, a value can be selected multiple
+        times. Otherwise, each value can be selected only once.
+    prob_name: str, optional
+        The name of an edge attribute used as the weights of sampling for
+        each node. This attribute tensor should contain (unnormalized)
+        probabilities corresponding to each neighboring edge of a node.
+        It must be a 1D floating-point or boolean tensor, with the number
+        of elements equalling the total number of edges.
+    node_timestamp_attr_name: str, optional
+        The name of an node attribute used as the timestamps of nodes.
+        It must be a 1D integer tensor, with the number of elements
+        equalling the total number of nodes.
+    edge_timestamp_attr_name: str, optional
+        The name of an edge attribute used as the timestamps of edges.
+        It must be a 1D integer tensor, with the number of elements
+        equalling the total number of edges.
+
+    Examples
+    -------
+    TODO(zhenkun) : Add an example after the API to pass timestamps is finalized.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+        fanouts,
+        replace=False,
+        prob_name=None,
+        node_timestamp_attr_name=None,
+        edge_timestamp_attr_name=None,
+    ):
+        super().__init__(
+            datapipe,
+            graph,
+            fanouts,
+            replace,
+            prob_name,
+            node_timestamp_attr_name,
+            edge_timestamp_attr_name,
+            graph.temporal_sample_layer_neighbors,
+        )

python/dgl/graphbolt/impl/torch_based_feature_store.py

@@ -0,0 +1,665 @@
+"""Torch-based feature store for GraphBolt."""
+
+import copy
+import textwrap
+from typing import Dict, List
+
+import numpy as np
+import torch
+
+from ..base import (
+    get_device_to_host_uva_stream,
+    get_host_to_device_uva_stream,
+    index_select,
+)
+from ..feature_store import Feature
+from ..internal_utils import gb_warning, is_wsl
+from .basic_feature_store import BasicFeatureStore
+from .ondisk_metadata import OnDiskFeatureData
+
+__all__ = ["TorchBasedFeature", "DiskBasedFeature", "TorchBasedFeatureStore"]
+
+
+class _Waiter:
+    def __init__(self, event, values):
+        self.event = event
+        self.values = values
+
+    def wait(self):
+        """Returns the stored value when invoked."""
+        self.event.wait()
+        values = self.values
+        # Ensure there is no memory leak.
+        self.event = self.values = None
+        return values
+
+
+class TorchBasedFeature(Feature):
+    r"""A wrapper of pytorch based feature.
+
+    Initialize a torch based feature store by a torch feature.
+    Note that the feature can be either in memory or on disk.
+
+    Parameters
+    ----------
+    torch_feature : torch.Tensor
+        The torch feature.
+        Note that the dimension of the tensor should be greater than 1.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+
+    1. The feature is in memory.
+
+    >>> torch_feat = torch.arange(10).reshape(2, -1)
+    >>> feature = gb.TorchBasedFeature(torch_feat)
+    >>> feature.read()
+    tensor([[0, 1, 2, 3, 4],
+            [5, 6, 7, 8, 9]])
+    >>> feature.read(torch.tensor([0]))
+    tensor([[0, 1, 2, 3, 4]])
+    >>> feature.update(torch.tensor([[1 for _ in range(5)]]),
+    ...                      torch.tensor([1]))
+    >>> feature.read(torch.tensor([0, 1]))
+    tensor([[0, 1, 2, 3, 4],
+            [1, 1, 1, 1, 1]])
+    >>> feature.size()
+    torch.Size([5])
+
+    2. The feature is on disk. Note that you can use gb.numpy_save_aligned as a
+    replacement for np.save to potentially get increased performance.
+
+    >>> import numpy as np
+    >>> arr = np.array([[1, 2], [3, 4]])
+    >>> np.save("/tmp/arr.npy", arr)
+    >>> torch_feat = torch.from_numpy(np.load("/tmp/arr.npy", mmap_mode="r+"))
+    >>> feature = gb.TorchBasedFeature(torch_feat)
+    >>> feature.read()
+    tensor([[1, 2],
+            [3, 4]])
+    >>> feature.read(torch.tensor([0]))
+    tensor([[1, 2]])
+
+    3. Pinned CPU feature.
+
+    >>> torch_feat = torch.arange(10).reshape(2, -1).pin_memory()
+    >>> feature = gb.TorchBasedFeature(torch_feat)
+    >>> feature.read().device
+    device(type='cuda', index=0)
+    >>> feature.read(torch.tensor([0]).cuda()).device
+    device(type='cuda', index=0)
+    """
+
+    def __init__(self, torch_feature: torch.Tensor, metadata: Dict = None):
+        super().__init__()
+        self._is_inplace_pinned = set()
+        assert isinstance(torch_feature, torch.Tensor), (
+            f"torch_feature in TorchBasedFeature must be torch.Tensor, "
+            f"but got {type(torch_feature)}."
+        )
+        assert torch_feature.dim() > 1, (
+            f"dimension of torch_feature in TorchBasedFeature must be greater "
+            f"than 1, but got {torch_feature.dim()} dimension."
+        )
+        # Make sure the tensor is contiguous.
+        self._tensor = torch_feature.contiguous()
+        self._metadata = metadata
+
+    def __del__(self):
+        # torch.Tensor.pin_memory() is not an inplace operation. To make it
+        # truly in-place, we need to use cudaHostRegister. Then, we need to use
+        # cudaHostUnregister to unpin the tensor in the destructor.
+        # https://github.com/pytorch/pytorch/issues/32167#issuecomment-753551842
+        for tensor in self._is_inplace_pinned:
+            assert self._inplace_unpinner(tensor.data_ptr()) == 0
+
+    def read(self, ids: torch.Tensor = None):
+        """Read the feature by index.
+
+        If the feature is on pinned CPU memory and `ids` is on GPU or pinned CPU
+        memory, it will be read by GPU and the returned tensor will be on GPU.
+        Otherwise, the returned tensor will be on CPU.
+
+        Parameters
+        ----------
+        ids : torch.Tensor, optional
+            The index of the feature. If specified, only the specified indices
+            of the feature are read. If None, the entire feature is returned.
+
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        if ids is None:
+            if self._tensor.is_pinned():
+                return self._tensor.cuda()
+            return self._tensor
+        return index_select(self._tensor, ids)
+
+    def read_async(self, ids: torch.Tensor):
+        r"""Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            ``read_async_num_stages(ids.device)``\ th invocation. The return result
+            can be accessed by calling ``.wait()``. on the returned future object.
+            It is undefined behavior to call ``.wait()`` more than once.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        assert self._tensor.device.type == "cpu"
+        if ids.is_cuda and self.is_pinned():
+            current_stream = torch.cuda.current_stream()
+            host_to_device_stream = get_host_to_device_uva_stream()
+            host_to_device_stream.wait_stream(current_stream)
+            with torch.cuda.stream(host_to_device_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                values = index_select(self._tensor, ids)
+                values.record_stream(current_stream)
+                values_copy_event = torch.cuda.Event()
+                values_copy_event.record()
+
+            yield _Waiter(values_copy_event, values)
+        elif ids.is_cuda:
+            ids_device = ids.device
+            current_stream = torch.cuda.current_stream()
+            device_to_host_stream = get_device_to_host_uva_stream()
+            device_to_host_stream.wait_stream(current_stream)
+            with torch.cuda.stream(device_to_host_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                ids = ids.to(self._tensor.device, non_blocking=True)
+                ids_copy_event = torch.cuda.Event()
+                ids_copy_event.record()
+
+            yield  # first stage is done.
+
+            ids_copy_event.synchronize()
+            values = torch.ops.graphbolt.index_select_async(self._tensor, ids)
+            yield
+
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                values_cuda = values.wait().to(ids_device, non_blocking=True)
+                values_cuda.record_stream(current_stream)
+                values_copy_event = torch.cuda.Event()
+                values_copy_event.record()
+
+            yield _Waiter(values_copy_event, values_cuda)
+        else:
+            yield torch.ops.graphbolt.index_select_async(self._tensor, ids)
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        if ids_device.type == "cuda":
+            if self._tensor.is_cuda:
+                # If the ids and the tensor are on cuda, no need for async.
+                return 0
+            return 1 if self.is_pinned() else 3
+        else:
+            return 1
+
+    def size(self):
+        """Get the size of the feature.
+
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        return self._tensor.size()[1:]
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        return self._tensor.size()[0]
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Update the feature store.
+
+        Parameters
+        ----------
+        value : torch.Tensor
+            The updated value of the feature.
+        ids : torch.Tensor, optional
+            The indices of the feature to update. If specified, only the
+            specified indices of the feature will be updated. For the feature,
+            the `ids[i]` row is updated to `value[i]`. So the indices and value
+            must have the same length. If None, the entire feature will be
+            updated.
+        """
+        if ids is None:
+            self._tensor = value
+        else:
+            assert ids.shape[0] == value.shape[0], (
+                f"ids and value must have the same length, "
+                f"but got {ids.shape[0]} and {value.shape[0]}."
+            )
+            assert self.size() == value.size()[1:], (
+                f"The size of the feature is {self.size()}, "
+                f"while the size of the value is {value.size()[1:]}."
+            )
+            if self._tensor.is_pinned() and value.is_cuda and ids.is_cuda:
+                raise NotImplementedError(
+                    "Update the feature on pinned CPU memory by GPU is not "
+                    "supported yet."
+                )
+            self._tensor[ids] = value
+
+    def metadata(self):
+        """Get the metadata of the feature.
+
+        Returns
+        -------
+        Dict
+            The metadata of the feature.
+        """
+        return (
+            self._metadata if self._metadata is not None else super().metadata()
+        )
+
+    def pin_memory_(self):
+        """In-place operation to copy the feature to pinned memory. Returns the
+        same object modified in-place."""
+        # torch.Tensor.pin_memory() is not an inplace operation. To make it
+        # truly in-place, we need to use cudaHostRegister. Then, we need to use
+        # cudaHostUnregister to unpin the tensor in the destructor.
+        # https://github.com/pytorch/pytorch/issues/32167#issuecomment-753551842
+        x = self._tensor
+        if not x.is_pinned() and x.device.type == "cpu":
+            assert (
+                x.is_contiguous()
+            ), "Tensor pinning is only supported for contiguous tensors."
+            cudart = torch.cuda.cudart()
+            assert (
+                cudart.cudaHostRegister(
+                    x.data_ptr(), x.numel() * x.element_size(), 0
+                )
+                == 0
+            )
+
+            self._is_inplace_pinned.add(x)
+            self._inplace_unpinner = cudart.cudaHostUnregister
+
+        return self
+
+    def is_pinned(self):
+        """Returns True if the stored feature is pinned."""
+        return self._tensor.is_pinned()
+
+    def to(self, device):  # pylint: disable=invalid-name
+        """Copy `TorchBasedFeature` to the specified device."""
+        # copy.copy is a shallow copy so it does not copy tensor memory.
+        self2 = copy.copy(self)
+        if device == "pinned":
+            self2._tensor = self2._tensor.pin_memory()
+        else:
+            self2._tensor = self2._tensor.to(device)
+        return self2
+
+    def __repr__(self) -> str:
+        ret = (
+            "{Classname}(\n"
+            "    feature={feature},\n"
+            "    metadata={metadata},\n"
+            ")"
+        )
+
+        feature_str = textwrap.indent(
+            str(self._tensor), " " * len("    feature=")
+        ).strip()
+        metadata_str = textwrap.indent(
+            str(self.metadata()), " " * len("    metadata=")
+        ).strip()
+
+        return ret.format(
+            Classname=self.__class__.__name__,
+            feature=feature_str,
+            metadata=metadata_str,
+        )
+
+
+class DiskBasedFeature(Feature):
+    r"""A wrapper of disk based feature.
+
+    Initialize a disk based feature fetcher by a numpy file. Note that you can
+    use gb.numpy_save_aligned as a replacement for np.save to potentially get
+    increased performance.
+
+    Parameters
+    ----------
+    path : string
+        The path to the numpy feature file.
+        Note that the dimension of the numpy should be greater than 1.
+    metadata : Dict
+        The metadata of the feature.
+    num_threads : int
+        The number of threads driving io_uring queues.
+    Examples
+    --------
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+    >>> torch_feat = torch.arange(10).reshape(2, -1)
+    >>> pth = "path/to/feat.npy"
+    >>> np.save(pth, torch_feat)
+    >>> feature = gb.DiskBasedFeature(pth)
+    >>> feature.read(torch.tensor([0]))
+    tensor([[0, 1, 2, 3, 4]])
+    >>> feature.size()
+    torch.Size([5])
+    """
+
+    def __init__(self, path: str, metadata: Dict = None, num_threads=None):
+        super().__init__()
+        mmap_mode = "r+"
+        ondisk_data = np.load(path, mmap_mode=mmap_mode)
+        assert ondisk_data.flags[
+            "C_CONTIGUOUS"
+        ], "DiskBasedFeature only supports C_CONTIGUOUS array."
+        self._tensor = torch.from_numpy(ondisk_data)
+
+        self._metadata = metadata
+        if torch.ops.graphbolt.detect_io_uring():
+            self._ondisk_npy_array = torch.ops.graphbolt.ondisk_npy_array(
+                path, self._tensor.dtype, self._tensor.shape, num_threads
+            )
+
+    def read(self, ids: torch.Tensor = None):
+        """Read the feature by index.
+        The returned tensor will be on CPU.
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        if ids is None:
+            return self._tensor
+        elif torch.ops.graphbolt.detect_io_uring():
+            try:
+                return self._ondisk_npy_array.index_select(ids).wait()
+            except RuntimeError:
+                raise IndexError
+        else:
+            return index_select(self._tensor, ids)
+
+    def read_async(self, ids: torch.Tensor):
+        r"""Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            ``read_async_num_stages(ids.device)``\ th invocation. The return result
+            can be accessed by calling ``.wait()``. on the returned future object.
+            It is undefined behavior to call ``.wait()`` more than once.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        assert torch.ops.graphbolt.detect_io_uring()
+        if ids.is_cuda:
+            ids_device = ids.device
+            current_stream = torch.cuda.current_stream()
+            device_to_host_stream = get_device_to_host_uva_stream()
+            device_to_host_stream.wait_stream(current_stream)
+            with torch.cuda.stream(device_to_host_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                ids = ids.to(self._tensor.device, non_blocking=True)
+                ids_copy_event = torch.cuda.Event()
+                ids_copy_event.record()
+
+            yield  # first stage is done.
+
+            ids_copy_event.synchronize()
+            values = self._ondisk_npy_array.index_select(ids)
+            yield
+
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                values_cuda = values.wait().to(ids_device, non_blocking=True)
+                values_cuda.record_stream(current_stream)
+                values_copy_event = torch.cuda.Event()
+                values_copy_event.record()
+
+            yield _Waiter(values_copy_event, values_cuda)
+        else:
+            yield self._ondisk_npy_array.index_select(ids)
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        return 3 if ids_device.type == "cuda" else 1
+
+    def size(self):
+        """Get the size of the feature.
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        return self._tensor.size()[1:]
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        return self._tensor.size()[0]
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Disk based feature does not support update for now."""
+        raise NotImplementedError
+
+    def metadata(self):
+        """Get the metadata of the feature.
+        Returns
+        -------
+        Dict
+            The metadata of the feature.
+        """
+        return (
+            self._metadata if self._metadata is not None else super().metadata()
+        )
+
+    def read_into_memory(self) -> TorchBasedFeature:
+        """Change disk-based feature to torch-based feature."""
+        return TorchBasedFeature(self._tensor, self._metadata)
+
+    def to(self, _):  # pylint: disable=invalid-name
+        """Placeholder `DiskBasedFeature` to implementation. It is a no-op."""
+        gb_warning(
+            "`DiskBasedFeature.to(device)` is not supported. Leaving unmodified."
+        )
+        return self
+
+    def pin_memory_(self):  # pylint: disable=invalid-name
+        r"""Placeholder `DiskBasedFeature` pin_memory_ implementation. It is a no-op."""
+        gb_warning(
+            "`DiskBasedFeature.pin_memory_()` is not supported. Leaving unmodified."
+        )
+        return self
+
+    def __repr__(self) -> str:
+        ret = (
+            "{Classname}(\n"
+            "    feature={feature},\n"
+            "    metadata={metadata},\n"
+            ")"
+        )
+
+        feature_str = textwrap.indent(
+            str(self._tensor), " " * len("    feature=")
+        ).strip()
+        metadata_str = textwrap.indent(
+            str(self.metadata()), " " * len("    metadata=")
+        ).strip()
+
+        return ret.format(
+            Classname=self.__class__.__name__,
+            feature=feature_str,
+            metadata=metadata_str,
+        )
+
+
+class TorchBasedFeatureStore(BasicFeatureStore):
+    r"""A store to manage multiple pytorch based feature for access.
+
+    The feature stores are described by the `feat_data`. The `feat_data` is a
+    list of `OnDiskFeatureData`.
+
+    For a feature store, its format must be either "pt" or "npy" for Pytorch or
+    Numpy formats. If the format is "pt", the feature store must be loaded in
+    memory. If the format is "npy", the feature store can be loaded in memory or
+    on disk. Note that you can use gb.numpy_save_aligned as a replacement for
+    np.save to potentially get increased performance.
+
+    Parameters
+    ----------
+    feat_data : List[OnDiskFeatureData]
+        The description of the feature stores.
+
+    Examples
+    --------
+    >>> import torch
+    >>> import numpy as np
+    >>> from dgl import graphbolt as gb
+    >>> edge_label = torch.tensor([[1], [2], [3]])
+    >>> node_feat = torch.tensor([[1, 2, 3], [4, 5, 6]])
+    >>> torch.save(edge_label, "/tmp/edge_label.pt")
+    >>> gb.numpy_save_aligned("/tmp/node_feat.npy", node_feat.numpy())
+    >>> feat_data = [
+    ...     gb.OnDiskFeatureData(domain="edge", type="author:writes:paper",
+    ...         name="label", format="torch", path="/tmp/edge_label.pt",
+    ...         in_memory=True),
+    ...     gb.OnDiskFeatureData(domain="node", type="paper", name="feat",
+    ...         format="numpy", path="/tmp/node_feat.npy", in_memory=False),
+    ... ]
+    >>> feature_store = gb.TorchBasedFeatureStore(feat_data)
+    """
+
+    def __init__(self, feat_data: List[OnDiskFeatureData]):
+        features = {}
+        for spec in feat_data:
+            key = (spec.domain, spec.type, spec.name)
+            metadata = spec.extra_fields
+            if spec.format == "torch":
+                assert spec.in_memory, (
+                    f"Pytorch tensor can only be loaded in memory, "
+                    f"but the feature {key} is loaded on disk."
+                )
+                features[key] = TorchBasedFeature(
+                    torch.load(spec.path, weights_only=False), metadata=metadata
+                )
+            elif spec.format == "numpy":
+                if spec.in_memory:
+                    # TorchBasedFeature is always in memory by default.
+                    features[key] = TorchBasedFeature(
+                        torch.as_tensor(np.load(spec.path)), metadata=metadata
+                    )
+                else:
+                    # DiskBasedFeature is always out of memory by default.
+                    features[key] = DiskBasedFeature(
+                        spec.path, metadata=metadata
+                    )
+            else:
+                raise ValueError(f"Unknown feature format {spec.format}")
+        super().__init__(features)
+
+    def pin_memory_(self):
+        """In-place operation to copy the feature store to pinned memory.
+        Returns the same object modified in-place."""
+        if is_wsl():
+            gb_warning(
+                "In place pinning is not supported on WSL. "
+                "Returning the out of place pinned `TorchBasedFeatureStore`."
+            )
+            return self.to("pinned")
+        for feature in self._features.values():
+            feature.pin_memory_()
+
+        return self
+
+    def is_pinned(self):
+        """Returns True if all the stored features are pinned."""
+        return all(feature.is_pinned() for feature in self._features.values())
+
+    def to(self, device):  # pylint: disable=invalid-name
+        """Copy `TorchBasedFeatureStore` to the specified device."""
+        # copy.copy is a shallow copy so it does not copy tensor memory.
+        self2 = copy.copy(self)
+        self2._features = {k: v.to(device) for k, v in self2._features.items()}
+        return self2
+
+    def __repr__(self) -> str:
+        ret = "{Classname}(\n" + "    {features}\n" + ")"
+        features_str = textwrap.indent(str(self._features), "    ").strip()
+        return ret.format(
+            Classname=self.__class__.__name__, features=features_str
+        )

python/dgl/graphbolt/impl/uniform_negative_sampler.py

@@ -0,0 +1,86 @@
+"""Uniform negative sampler for GraphBolt."""
+
+import torch
+from torch.utils.data import functional_datapipe
+
+from ..negative_sampler import NegativeSampler
+
+__all__ = ["UniformNegativeSampler"]
+
+
+@functional_datapipe("sample_uniform_negative")
+class UniformNegativeSampler(NegativeSampler):
+    """Sample negative destination nodes for each source node based on a uniform
+    distribution.
+
+    Functional name: :obj:`sample_uniform_negative`.
+
+    It's important to note that the term 'negative' refers to false negatives,
+    indicating that the sampled pairs are not ensured to be absent in the graph.
+    For each edge ``(u, v)``, it is supposed to generate `negative_ratio` pairs
+    of negative edges ``(u, v')``, where ``v'`` is chosen uniformly from all
+    the nodes in the graph.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    graph : FusedCSCSamplingGraph
+        The graph on which to perform negative sampling.
+    negative_ratio : int
+        The proportion of negative samples to positive samples.
+
+    Examples
+    --------
+    >>> from dgl import graphbolt as gb
+    >>> indptr = torch.LongTensor([0, 1, 2, 3, 4])
+    >>> indices = torch.LongTensor([1, 2, 3, 0])
+    >>> graph = gb.fused_csc_sampling_graph(indptr, indices)
+    >>> seeds = torch.tensor([[0, 1], [1, 2], [2, 3], [3, 0]])
+    >>> item_set = gb.ItemSet(seeds, names="seeds")
+    >>> item_sampler = gb.ItemSampler(
+    ...     item_set, batch_size=4,)
+    >>> neg_sampler = gb.UniformNegativeSampler(
+    ...     item_sampler, graph, 2)
+    >>> for minibatch in neg_sampler:
+    ...       print(minibatch.seeds)
+    ...       print(minibatch.labels)
+    ...       print(minibatch.indexes)
+    tensor([[0, 1], [1, 2], [2, 3], [3, 0], [0, 1], [0, 3], [1, 1], [1, 2],
+        [2, 1], [2, 0], [3, 0], [3, 2]])
+    tensor([1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.])
+    tensor([0, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3, 3])
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+        negative_ratio,
+    ):
+        super().__init__(datapipe, negative_ratio)
+        self.graph = graph
+
+    def _sample_with_etype(self, seeds, etype=None):
+        assert seeds.ndim == 2 and seeds.shape[1] == 2, (
+            "Only tensor with shape N*2 is supported for negative"
+            + f" sampling, but got {seeds.shape}."
+        )
+        # Sample negative edges, and concatenate positive edges with them.
+        all_seeds = self.graph.sample_negative_edges_uniform(
+            etype,
+            seeds,
+            self.negative_ratio,
+        )
+        # Construct indexes for all node pairs.
+        pos_num = seeds.shape[0]
+        negative_ratio = self.negative_ratio
+        pos_indexes = torch.arange(0, pos_num, device=all_seeds.device)
+        neg_indexes = pos_indexes.repeat_interleave(negative_ratio)
+        indexes = torch.cat((pos_indexes, neg_indexes))
+        # Construct labels for all node pairs.
+        neg_num = all_seeds.shape[0] - pos_num
+        labels = torch.empty(pos_num + neg_num, device=all_seeds.device)
+        labels[:pos_num] = 1
+        labels[pos_num:] = 0
+        return all_seeds, labels, indexes