[patch] epoch API redesign

NVIDIA-Merlin · Sep 1, 2023 · d384e0b · d384e0b
1 parent a0b0db8
commit d384e0b
Show file tree

Hide file tree

Showing 7 changed files with 94 additions and 99 deletions.
diff --git a/README.md b/README.md
@@ -77,7 +77,7 @@ The `score_type` must be `uint64_t`. For more detail, please refer to [`class Ev
 | Name           | Definition of `Score`                                                                                                                                                                                           |
 |:---------------|:----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | __Lru__        | Device clock in a nanosecond, which could differ slightly from host clock.                                                                                                                                      |
-| __Lfu__        | Frequency provided by caller via the input parameter of `scores` of `insert-like` APIs as the increment of frequency.                                                                                           |
+| __Lfu__        | Frequency increment provided by caller via the input parameter of `scores` of `insert-like` APIs as the increment of frequency.                                                                                 |
 | __EpochLru__   | The high 32bits is the global epoch provided via the input parameter of `global_epoch`, <br>the low 32bits is equal to `(device_clock >> 20) & 0xffffffff` with granularity close to 1 ms.                      |
 | __EpochLfu__   | The high 32bits is the global epoch provided via the input parameter of `global_epoch`, <br>the low 32bits is the frequency, <br>the frequency will keep constant after reaching the max value of `0xffffffff`. |
 | __Customized__ | Fully provided by the caller via the input parameter of `scores` of `insert-like` APIs.                                                                                                                         |

diff --git a/include/merlin/core_kernels/kernel_utils.cuh b/include/merlin/core_kernels/kernel_utils.cuh
@@ -381,7 +381,7 @@ struct ScoreFunctor<K, V, S, EvictStrategyInternal::kEpochLru> {
   __forceinline__ __device__ static S desired_when_missed(
       const S* __restrict const input_scores, const int key_idx,
       const S& epoch) {
-    if (epoch == static_cast<S>(DEFAULT_GLOBAL_EPOCH) &&
+    if (epoch == static_cast<S>(IGNORED_GLOBAL_EPOCH) &&
         input_scores != nullptr) {
       return input_scores[key_idx];
     }
@@ -442,7 +442,7 @@ struct ScoreFunctor<K, V, S, EvictStrategyInternal::kEpochLfu> {
   __forceinline__ __device__ static S desired_when_missed(
       const S* __restrict const input_scores, const int key_idx,
       const S epoch) {
-    if (epoch == static_cast<S>(DEFAULT_GLOBAL_EPOCH)) {
+    if (epoch == static_cast<S>(IGNORED_GLOBAL_EPOCH)) {
       return input_scores[key_idx];
     }
     return make_epoch<S>(epoch) | (input_scores[key_idx] & SCORE_BITS_MASK);

diff --git a/include/merlin/types.cuh b/include/merlin/types.cuh
@@ -50,7 +50,7 @@ constexpr uint64_t LOCKED_KEY = UINT64_C(0xFFFFFFFFFFFFFFFD);
 constexpr uint64_t RESERVED_KEY_MASK = UINT64_C(0xFFFFFFFFFFFFFFFC);
 constexpr uint64_t MAX_SCORE = UINT64_C(0xFFFFFFFFFFFFFFFF);
 constexpr uint64_t EMPTY_SCORE = UINT64_C(0);
-constexpr uint64_t DEFAULT_GLOBAL_EPOCH = UINT64_C(0xFFFFFFFFFFFFFFFF);
+constexpr uint64_t IGNORED_GLOBAL_EPOCH = UINT64_C(0xFFFFFFFFFFFFFFFF);
 
 #define IS_RESERVED_KEY(key) ((RESERVED_KEY_MASK & (key)) == RESERVED_KEY_MASK)
 #define IS_VACANT_KEY(key) ((VACANT_KEY_MASK & (key)) == VACANT_KEY_MASK)
@@ -278,6 +278,11 @@ struct EvictStrategyInternal {
   constexpr static int kEpochLfu = 3;    ///< Epoch + LFU mode.
   constexpr static int kCustomized = 4;  ///< Customized mode.
 };
+union EvictStrategyParamUnion {  ///< Parameter for specialized strategy.
+  uint64_t global_epoch;
+};
+thread_local static EvictStrategyParamUnion EvictStrategyParam = {
+    static_cast<uint64_t>(IGNORED_GLOBAL_EPOCH)};
 
 }  // namespace merlin
 }  // namespace nv