[GraphBolt] Cache policy replace fix.

dmlc · Jul 12, 2024 · 98658d7 · 98658d7
1 parent 945247e
commit 98658d7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/graphbolt/src/cache_policy.h b/graphbolt/src/cache_policy.h
@@ -86,7 +86,7 @@ struct CircularQueue {
 
 struct CacheKey {
   CacheKey(int64_t key, int64_t position)
-      : freq_(0), key_(key), position_in_cache_(position), reference_count_(0) {
+      : freq_(0), key_(key), position_in_cache_(position), reference_count_(1) {
     static_assert(sizeof(CacheKey) == 2 * sizeof(int64_t));
   }
 

diff --git a/python/dgl/graphbolt/impl/feature_cache.py b/python/dgl/graphbolt/impl/feature_cache.py
@@ -73,6 +73,7 @@ def replace(self, keys, values):
         """
         positions = self._policy.replace(keys)
         self._cache.replace(positions, values)
+        self._policy.reading_completed(keys)
 
     @property
     def miss_rate(self):