wehub-resource-sync
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
"""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
|