Add LFU evict strategy to dynamicemb (#52)

* Add LFU strategy to dynamicemb

* Add a test for LFU strategy

* Update test comment

* Delete some comments

* Precommit

* Update test and comments

---------

Co-authored-by: Runchu Zhao <runchuz@nvidia.com>

Author: z52527

corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py

@@ -666,6 +666,9 @@ def _create_score(self):
                 self._scores[table_name] = 1
             elif option.score_strategy == DynamicEmbScoreStrategy.CUSTOMIZED:
                 option.evict_strategy = DynamicEmbEvictStrategy.CUSTOMIZED
+            elif option.score_strategy == DynamicEmbScoreStrategy.LFU:
+                option.evict_strategy = DynamicEmbEvictStrategy.LFU
+                self._scores[table_name] = 1
 
     def _update_score(self):
         for table_name, option in zip(self._table_names, self._dynamicemb_options):
@@ -689,6 +692,8 @@ def _update_score(self):
                     self._scores[table_name] = 0
                 else:
                     self._scores[table_name] = new_score
+            elif option.score_strategy == DynamicEmbScoreStrategy.LFU:
+                self._scores[table_name] = 1
 
     def incremental_dump(
         self,

corelib/dynamicemb/dynamicemb/dynamicemb_config.py

@@ -208,6 +208,7 @@ class DynamicEmbScoreStrategy(enum.IntEnum):
     TIMESTAMP = 0
     STEP = 1
     CUSTOMIZED = 2
+    LFU = 3
 
 
 # Configs used as keys to group HKV variables(considering kernel behaviors, result type).

corelib/dynamicemb/src/dynamic_emb_op.cu

@@ -138,11 +138,11 @@ void insert_and_evict(
     bool unique_key = true,
     bool ignore_evict_strategy = false) {
 
-  if (not score and table->evict_strategy() == EvictStrategy::kCustomized) {
-    throw std::invalid_argument("Must specify the score when evict strategy is customized.");
+  if (not score and (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu)) {
+    throw std::invalid_argument("Must specify the score when evict strategy is customized or LFU.");
   }
   auto stream = at::cuda::getCurrentCUDAStream().stream();
-  if (table->evict_strategy() == EvictStrategy::kCustomized) {
+  if (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu) {
     auto&& option = at::TensorOptions().dtype(at::kUInt64).device(keys.device());
     // broadcast scores
     at::Tensor bc_scores = at::empty({static_cast<int64_t>(n)}, option);
@@ -188,8 +188,8 @@ void find_or_insert(std::shared_ptr<dyn_emb::DynamicVariableBase> table,
                   bool ignore_evict_strategy = false
                   )
 {
-  if (not score and table->evict_strategy() == EvictStrategy::kCustomized) {
-    throw std::invalid_argument("Must specify the score when evict strategy is customized.");
+  if (not score and (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu)) {
+    throw std::invalid_argument("Must specify the score when evict strategy is customized or LFU.");
   }
   if (n == 0) return;
   auto stream = at::cuda::getCurrentCUDAStream().stream();
@@ -203,7 +203,7 @@ void find_or_insert(std::shared_ptr<dyn_emb::DynamicVariableBase> table,
 
   auto found_tensor_data_ptr = found_tensor.data_ptr<bool>();
 
-  if (table->evict_strategy() == EvictStrategy::kCustomized) {
+  if (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu) {
     auto&& option = at::TensorOptions().dtype(at::kUInt64).device(keys.device());
     // broadcast scores
     at::Tensor bc_scores = at::empty({static_cast<int64_t>(n)}, option);
@@ -226,15 +226,15 @@ void find_or_insert_pointers(
   const std::optional<uint64_t> score = std::nullopt,
   bool unique_key = true,
   bool ignore_evict_strategy = false) {
-  if (not score and table->evict_strategy() == EvictStrategy::kCustomized) {
-    throw std::invalid_argument("Must specify the score when evict strategy is customized.");
+  if (not score and (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu)) {
+    throw std::invalid_argument("Must specify the score when evict strategy is customized or LFU.");
   }
   if (n == 0) return;
   auto stream = at::cuda::getCurrentCUDAStream().stream();
   auto values_data_ptr = reinterpret_cast<void**>(values.data_ptr<int64_t>());
   auto found_tensor_data_ptr = founds.data_ptr<bool>();
 
-  if (table->evict_strategy() == EvictStrategy::kCustomized) {
+  if (table->evict_strategy() == EvictStrategy::kCustomized || table->evict_strategy() == EvictStrategy::kLfu) {
     auto&& option = at::TensorOptions().dtype(at::kUInt64).device(keys.device());
     // broadcast scores
     at::Tensor bc_scores = at::empty({static_cast<int64_t>(n)}, option);

corelib/dynamicemb/src/dynamic_variable_base.cu

@@ -30,6 +30,7 @@ std::shared_ptr<DynamicVariableBase> VariableFactory::create(
     const OptimizerType optimizer_type = OptimizerType::Null) {
   // TODO:check input datetype , now only support key type is int64_t uint64_t ,
   // value type float
+  printf("create dynamic variable base\n");
   std::shared_ptr<DynamicVariableBase> table;
   DISPATCH_INTEGER_DATATYPE_FUNCTION(key_type, key_t, [&] {
     DISPATCH_FLOAT_DATATYPE_FUNCTION(value_type, value_t, [&] {

corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t___half_EvictStrategy__kLfu.cu

@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<int64_t, __half, EvictStrategy::kLfu>;
+} 

corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t___nv_bfloat16_EvictStrategy__kLfu.cu

@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<int64_t, __nv_bfloat16, EvictStrategy::kLfu>;
+} 

corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__int64_t_float_EvictStrategy__kLfu.cu

@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<int64_t, float, EvictStrategy::kLfu>;
+} 

corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t___half_EvictStrategy__kLfu.cu

@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<uint64_t, __half, EvictStrategy::kLfu>;
+} 

corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t___nv_bfloat16_EvictStrategy__kLfu.cu

@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<uint64_t, __nv_bfloat16, EvictStrategy::kLfu>;
+} 

corelib/dynamicemb/src/hkv_variable_instantiations/HKVVariable__uint64_t_float_EvictStrategy__kLfu.cu

@@ -0,0 +1,4 @@
+#include"hkv_variable.cuh" 
+namespace dyn_emb{
+    template class HKVVariable<uint64_t, float, EvictStrategy::kLfu>;
+}