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
(self, keys)
| 34 | self.total_queries = 0 |
| 35 | |
| 36 | def query(self, keys): |
| 37 | """Queries the GPU cache. |
| 38 | |
| 39 | Parameters |
| 40 | ---------- |
| 41 | keys : Tensor |
| 42 | The keys to query the GPU cache with. |
| 43 | |
| 44 | Returns |
| 45 | ------- |
| 46 | tuple(Tensor, Tensor, Tensor) |
| 47 | A tuple containing (values, missing_indices, missing_keys) where |
| 48 | values[missing_indices] corresponds to cache misses that should be |
| 49 | filled by quering another source with missing_keys. |
| 50 | """ |
| 51 | self.total_queries += keys.shape[0] |
| 52 | keys = F.astype(keys, self.idtype) |
| 53 | values, missing_index, missing_keys = _CAPI_DGLGpuCacheQuery( |
| 54 | self._cache, F.to_dgl_nd(keys) |
| 55 | ) |
| 56 | self.total_miss += missing_keys.shape[0] |
| 57 | return ( |
| 58 | F.from_dgl_nd(values), |
| 59 | F.from_dgl_nd(missing_index), |
| 60 | F.from_dgl_nd(missing_keys), |
| 61 | ) |
| 62 | |
| 63 | def replace(self, keys, values): |
| 64 | """Inserts key-value pairs into the GPU cache using the Least-Recently |