chore: import upstream snapshot with attribution

python/dgl/cuda/gpu_cache.py

@@ -0,0 +1,86 @@
+"""API wrapping HugeCTR gpu_cache."""
+#    Copyright (c) 2022, NVIDIA Corporation
+#    All rights reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+#  @file gpu_cache.py
+#  @brief API for managing a GPU Cache
+
+from .. import backend as F
+from .._ffi.function import _init_api
+
+
+class GPUCache(object):
+    """High-level wrapper for GPU embedding cache"""
+
+    def __init__(self, num_items, num_feats, idtype=F.int64):
+        assert idtype in [F.int32, F.int64]
+        self._cache = _CAPI_DGLGpuCacheCreate(
+            num_items, num_feats, 32 if idtype == F.int32 else 64
+        )
+        self.idtype = idtype
+        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 cache with.
+
+        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.
+        """
+        self.total_queries += keys.shape[0]
+        keys = F.astype(keys, self.idtype)
+        values, missing_index, missing_keys = _CAPI_DGLGpuCacheQuery(
+            self._cache, F.to_dgl_nd(keys)
+        )
+        self.total_miss += missing_keys.shape[0]
+        return (
+            F.from_dgl_nd(values),
+            F.from_dgl_nd(missing_index),
+            F.from_dgl_nd(missing_keys),
+        )
+
+    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.
+        """
+        keys = F.astype(keys, self.idtype)
+        values = F.astype(values, F.float32)
+        _CAPI_DGLGpuCacheReplace(
+            self._cache, F.to_dgl_nd(keys), F.to_dgl_nd(values)
+        )
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self.total_miss / self.total_queries
+
+
+_init_api("dgl.cuda", __name__)