[Fix] benchmark 'unsupport API'

NVIDIA-Merlin · May 18, 2023 · dad5c7a · dad5c7a
1 parent 63582b5
commit dad5c7a
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/benchmark/merlin_hashtable_benchmark.cc.cu b/benchmark/merlin_hashtable_benchmark.cc.cu
@@ -318,6 +318,22 @@ float test_one_api(const API_Select api, const size_t dim,
         CUDA_CHECK(cudaFree(d_vectors_ptr));
         break;
       }
+      case API_Select::find_or_insert_ptr: {
+        V** d_vectors_ptr = nullptr;
+        bool* d_found;
+        CUDA_CHECK(cudaMalloc(&d_found, key_num_per_op_warmup * sizeof(bool)));
+        CUDA_CHECK(
+            cudaMalloc(&d_vectors_ptr, key_num_per_op_warmup * sizeof(V*)));
+        benchmark::array2ptr(d_vectors_ptr, d_vectors, options.dim,
+                             key_num_per_op_warmup, stream);
+        CUDA_CHECK(cudaStreamSynchronize(stream));
+        table->find_or_insert(key_num_per_op_warmup, d_keys, d_vectors_ptr,
+                              d_found, d_metas, stream);
+        CUDA_CHECK(cudaStreamSynchronize(stream));
+        CUDA_CHECK(cudaFree(d_vectors_ptr));
+        CUDA_CHECK(cudaFree(d_found));
+        break;
+      }
       default: {
         std::cout << "[Unsupport API]\n";
       }