Bugfix: rehashing in GPU hashtable is not enough when meeting large insert

Author: Lifann

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/cuckoo_hashtable_op_gpu.cu.cc

@@ -56,10 +56,10 @@ class CuckooHashTableOfTensorsGpu final : public LookupInterface {
       Status status = ReadInt64FromEnvVar("TF_HASHTABLE_INIT_SIZE",
                                           1024 * 8,  // 8192 KV pairs by default
                                           &env_var);
-      min_size_ = (size_t)env_var;
+      min_size_ = (size_t)env_var / 2;
       max_size_ = (size_t)env_var;
     } else {
-      min_size_ = init_size;
+      min_size_ = init_size / 2;
       max_size_ = init_size;
     }
 
@@ -176,47 +176,65 @@ class CuckooHashTableOfTensorsGpu final : public LookupInterface {
   }
 
   void RehashIfNeeded(cudaStream_t stream) {
-    K* d_keys;
-    gpu::ValueArrayBase<V>* d_values;
+    RehashIfNeeded(stream, min_size_);
+  }
+
+  void RehashIfNeeded(cudaStream_t stream, size_t expecting) {
+    K* d_keys = nullptr;
+    gpu::ValueArrayBase<V>* d_values = nullptr;
     size_t* d_dump_counter;
-    size_t new_max_size = max_size_;
+    size_t capacity = table_->get_capacity();
+
+    size_t cur_size = table_->get_size(stream);
+    expecting += cur_size;
+    expecting = max(min_size_, expecting);
 
-    size_t total_size = table_->get_size(stream);
     CUDA_CHECK(cudaStreamSynchronize(stream));
-    if (total_size >= 0.75 * max_size_) {
-      new_max_size = max_size_ * 2;
+    size_t new_max_size = capacity;
+    bool need_rehash = false;
+    while (expecting > 0.75 * new_max_size) {
+      new_max_size *= 2;
+      need_rehash = true;
     }
-    if (total_size < 0.25 * max_size_ && max_size_ > min_size_) {
-      new_max_size = max_size_ / 2;
+    if (expecting < 0.25 * capacity) {
+      new_max_size = capacity / 2;
+      need_rehash = true;
     }
-    if (new_max_size != max_size_) {  // rehash manually.
-      size_t capacity = table_->get_capacity();
+
+    if (need_rehash) {
+      LOG(INFO) << "Need to rehash GPU HashTable, where capacity=" << capacity
+                << ", current_size=" << cur_size << " and expecting "
+                << expecting;
       size_t h_dump_counter = 0;
-      CUDA_CHECK(cudaMallocManaged((void**)&d_dump_counter, sizeof(size_t)));
-      CUDA_CHECK(cudaMallocManaged((void**)&d_keys, sizeof(K) * capacity));
-      CUDA_CHECK(cudaMallocManaged((void**)&d_values,
-                                   sizeof(V) * runtime_dim_ * capacity));
-      table_->dump(d_keys, (gpu::ValueArrayBase<V>*)d_values, 0, capacity,
+
+      if (cur_size > 0) {
+        CUDA_CHECK(cudaMallocManaged((void**)&d_dump_counter, sizeof(size_t)));
+        CUDA_CHECK(cudaMallocManaged((void**)&d_keys, sizeof(K) * cur_size));
+        CUDA_CHECK(cudaMallocManaged((void**)&d_values, sizeof(V) * runtime_dim_ * cur_size));
+        table_->dump(d_keys, (gpu::ValueArrayBase<V>*)d_values, 0, capacity,
                    d_dump_counter, stream);
-      CUDA_CHECK(cudaStreamSynchronize(stream));
+        cudaMemcpyAsync(&h_dump_counter, d_dump_counter, sizeof(size_t), cudaMemcpyDeviceToHost, stream);
+        CUDA_CHECK(cudaStreamSynchronize(stream));
+      }
 
       delete table_;
       table_ = NULL;
       CreateTable(new_max_size, &table_);
-      CUDA_CHECK(cudaStreamSynchronize(stream));
-      CUDA_CHECK(cudaMemcpy((size_t*)&h_dump_counter, (size_t*)d_dump_counter,
-                            sizeof(size_t), cudaMemcpyDefault));
-      table_->upsert((const K*)d_keys, (const gpu::ValueArrayBase<V>*)d_values,
+
+      if (cur_size > 0) {
+        table_->upsert((const K*)d_keys, (const gpu::ValueArrayBase<V>*)d_values,
                      h_dump_counter, stream);
-      CUDA_CHECK(cudaStreamSynchronize(stream));
-      CUDA_CHECK(cudaFree(d_keys));
-      CUDA_CHECK(cudaFree(d_values));
-      CUDA_CHECK(cudaFree(d_dump_counter));
+        cudaStreamSynchronize(stream);
+        cudaFree(d_keys);
+        cudaFree(d_values);
+        cudaFree(d_dump_counter);
+      }
       max_size_ = new_max_size;
+
       LOG(INFO) << "HashTable on GPU changes to new status: [size="
-                << total_size << ", max_size=" << max_size_
+                << h_dump_counter << ", max_size=" << max_size_
                 << ", load factor=" << std::setprecision(2)
-                << (float)total_size / (float)max_size_ << "].";
+                << (float)h_dump_counter / (float)max_size_ << "].";
     }
   }
 
@@ -227,7 +245,7 @@ class CuckooHashTableOfTensorsGpu final : public LookupInterface {
     CUDA_CHECK(cudaStreamCreate(&_stream));
     {
       mutex_lock l(mu_);
-      RehashIfNeeded(_stream);
+      RehashIfNeeded(_stream, len);
       table_->upsert((const K*)keys.tensor_data().data(),
                      (const gpu::ValueArrayBase<V>*)values.tensor_data().data(),
                      len, _stream);
@@ -245,7 +263,7 @@ class CuckooHashTableOfTensorsGpu final : public LookupInterface {
     CUDA_CHECK(cudaStreamCreate(&_stream));
     {
       mutex_lock l(mu_);
-      RehashIfNeeded(_stream);
+      RehashIfNeeded(_stream, len);
       table_->accum(
           (const K*)keys.tensor_data().data(),
           (const gpu::ValueArrayBase<V>*)values_or_deltas.tensor_data().data(),
@@ -300,6 +318,7 @@ class CuckooHashTableOfTensorsGpu final : public LookupInterface {
     if (len > 0) {
       cudaStream_t _stream;
       CUDA_CHECK(cudaStreamCreate(&_stream));
+      RehashIfNeeded(_stream, len);
       CUDA_CHECK(cudaMallocManaged((void**)&d_keys, sizeof(K) * len));
       CUDA_CHECK(
           cudaMallocManaged((void**)&d_values, sizeof(V) * runtime_dim_ * len));