| 47 | } |
| 48 | |
| 49 | void TestKeyValueStore(KeyValueStore* store, size_t num_embeddings, size_t test_embeddings, |
| 50 | size_t embedding_vec_size) { |
| 51 | auto device = Singleton<ep::DeviceManagerRegistry>::Get()->GetDevice(DeviceType::kCUDA, 0); |
| 52 | ep::Stream* stream = device->CreateStream(); |
| 53 | |
| 54 | store->SaveSnapshot("init"); |
| 55 | |
| 56 | uint64_t* keys = nullptr; |
| 57 | float* values = nullptr; |
| 58 | float* values1 = nullptr; |
| 59 | uint64_t* keys_host = nullptr; |
| 60 | float* values_host = nullptr; |
| 61 | uint64_t* context = nullptr; |
| 62 | uint32_t* n_missing = nullptr; |
| 63 | uint32_t* host_n_missing = nullptr; |
| 64 | uint64_t* missing_keys = nullptr; |
| 65 | uint32_t* missing_indices = nullptr; |
| 66 | size_t keys_size = sizeof(uint64_t) * num_embeddings; |
| 67 | size_t values_size = sizeof(float) * embedding_vec_size * num_embeddings; |
| 68 | size_t context_size = sizeof(uint64_t) * num_embeddings; |
| 69 | const size_t batch_size = 128; |
| 70 | OF_CUDA_CHECK(cudaMalloc(&keys, keys_size)); |
| 71 | OF_CUDA_CHECK(cudaMalloc(&values, values_size)); |
| 72 | OF_CUDA_CHECK(cudaMalloc(&values1, values_size)); |
| 73 | OF_CUDA_CHECK(cudaMalloc(&context, context_size)); |
| 74 | OF_CUDA_CHECK(cudaMallocHost(&keys_host, keys_size)); |
| 75 | OF_CUDA_CHECK(cudaMallocHost(&values_host, values_size)); |
| 76 | OF_CUDA_CHECK(cudaMallocHost(&host_n_missing, sizeof(uint32_t))); |
| 77 | OF_CUDA_CHECK(cudaMalloc(&missing_keys, batch_size * sizeof(uint64_t))); |
| 78 | OF_CUDA_CHECK(cudaMalloc(&missing_indices, batch_size * sizeof(uint32_t))); |
| 79 | OF_CUDA_CHECK(cudaMalloc(&n_missing, sizeof(uint32_t))); |
| 80 | for (size_t i = 0; i < num_embeddings; ++i) { |
| 81 | uint64_t key = i + 1; |
| 82 | keys_host[i] = key; |
| 83 | for (size_t j = 0; j < embedding_vec_size; j++) { |
| 84 | values_host[i * embedding_vec_size + j] = key; |
| 85 | } |
| 86 | } |
| 87 | OF_CUDA_CHECK(cudaMemcpy(keys, keys_host, keys_size, cudaMemcpyDefault)); |
| 88 | OF_CUDA_CHECK(cudaMemcpy(values, values_host, values_size, cudaMemcpyDefault)); |
| 89 | |
| 90 | store->Put(stream, 0, keys, values); |
| 91 | OF_CUDA_CHECK(cudaDeviceSynchronize()); |
| 92 | OF_CUDA_CHECK(cudaGetLastError()); |
| 93 | |
| 94 | for (size_t offset = 0; offset < test_embeddings; offset += batch_size) { |
| 95 | const size_t num_keys = std::min(batch_size, test_embeddings - offset); |
| 96 | store->Get(stream, num_keys, keys + offset, values1 + offset * embedding_vec_size, n_missing, |
| 97 | missing_indices); |
| 98 | OF_CUDA_CHECK(cudaMemcpy(host_n_missing, n_missing, sizeof(uint32_t), cudaMemcpyDefault)); |
| 99 | OF_CUDA_CHECK(cudaDeviceSynchronize()); |
| 100 | ASSERT_EQ(*host_n_missing, num_keys); |
| 101 | store->Put(stream, num_keys, keys + offset, values + offset * embedding_vec_size); |
| 102 | } |
| 103 | |
| 104 | OF_CUDA_CHECK(cudaDeviceSynchronize()); |
| 105 | |
| 106 | store->SaveSnapshot("final"); |
no test coverage detected