GPUCachedFeature

class dgl.graphbolt.GPUCachedFeature(fallback_feature: Feature, cache: GPUFeatureCache, offset: int = 0)

Bases: Feature

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])

count()

Get the count of the feature.

Returns:

The count of the feature.

Return type:

int

read(ids: Tensor | None = 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:

The read feature.

Return type:

torch.Tensor

read_async(ids: Tensor)

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:

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.

Return type:

A generator object.

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.

read_async_num_stages(ids_device: 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:

The number of stages of the read_async operation.

Return type:

int

size()

Get the size of the feature.

Returns:

The size of the feature.

Return type:

torch.Size

update(value: Tensor, ids: Tensor | None = 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.

property cache_size_in_bytes

Return the size taken by the cache in bytes.

property miss_rate

Returns the cache miss rate since creation.