Fixed bug where cache would never be allocated at the specified device_id

0e313a92 · tianxing wang · 164a18c7 · 0e313a92 · 0e313a92
Commit 0e313a92 authored May 13, 2024 by tianxing wang
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 2 deletions

src/cache.cuh
+4 -0

test/cache_test.cu
+2 -2

No files found.
--- a/src/cache.cuh
+++ b/src/cache.cuh
@@ -14,6 +14,10 @@ namespace gpucache {
 //            printf("Cache: keySize: %lu, valueSize: %u, dim: %u, capacity: %lu, "
 //                   "maxQueryNum: %u, deviceId: %u\n",
 //                   sizeof(KeyType), value_size, dim, capacity, max_query_num, device_id);
+            int device_count;
+            CUDA_CHECK(cudaGetDeviceCount(&device_count));
+            assert(device_id < device_count && "device_id should less than total device num");
+            CUDA_CHECK(cudaSetDevice(device_id));
            CUDA_CHECK(cudaMalloc((void **) &keys, capacity * key_size));
            CUDA_CHECK(cudaMalloc((void **) &values, capacity * value_size));
            CUDA_CHECK(cudaMalloc((void **) &timestamps, capacity * sizeof(uint8_t)));

--- a/test/cache_test.cu
+++ b/test/cache_test.cu
@@ -231,7 +231,7 @@ namespace gpucache{
        cfg.capacity = 4096 * 2;
        cfg.key_size = 4;
        cfg.max_query_num = 2048;
-        cfg.device_id = 0;
+        cfg.device_id = 1;
        cfg.dim = 172;
        FIFOCache<int32_t,uint32_t> *cache = NewFIFOCache<int32_t,uint32_t>(cfg);
@@ -247,7 +247,7 @@ namespace gpucache{
        cfg.capacity = 4096 * 2;
        cfg.key_size = 4;
        cfg.max_query_num = 2048;
-        cfg.device_id = 0;
+        cfg.device_id = 2;
        cfg.dim = 172;
        LFUCache<int32_t,uint32_t> *cache = NewLFUCache<int32_t,uint32_t>(cfg);