wehub-resource-sync
123 lines
3.8 KiB
C++
123 lines
3.8 KiB
C++
/*
|
|
* Copyright (c) 2023, NVIDIA CORPORATION.
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
#pragma once
|
|
|
|
#include <nv_util.h>
|
|
|
|
#include <cstdio>
|
|
#include <hash_functions.cuh>
|
|
#include <limits>
|
|
|
|
#include "gpu_cache_api.hpp"
|
|
#ifdef LIBCUDACXX_VERSION
|
|
#include <cuda/atomic>
|
|
#include <cuda/semaphore>
|
|
#endif
|
|
|
|
#define SET_ASSOCIATIVITY 2
|
|
#define SLAB_SIZE 32
|
|
#define TASK_PER_WARP_TILE_MACRO 1
|
|
|
|
namespace gpu_cache {
|
|
|
|
// slab for static slab list
|
|
template <typename key_type, int warp_size>
|
|
struct static_slab {
|
|
key_type slab_[warp_size];
|
|
};
|
|
|
|
// Static slablist(slabset) for GPU Cache
|
|
template <int set_associativity, typename key_type, int warp_size>
|
|
struct slab_set {
|
|
static_slab<key_type, warp_size> set_[set_associativity];
|
|
};
|
|
|
|
///////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
// GPU Cache
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher = MurmurHash3_32<key_type>,
|
|
typename slab_hasher = Mod_Hash<key_type, size_t>>
|
|
class gpu_cache : public gpu_cache_api<key_type> {
|
|
public:
|
|
// Ctor
|
|
gpu_cache(const size_t capacity_in_set, const size_t embedding_vec_size);
|
|
|
|
// Dtor
|
|
~gpu_cache();
|
|
|
|
// Query API, i.e. A single read from the cache
|
|
void Query(const key_type* d_keys, const size_t len, float* d_values, uint64_t* d_missing_index,
|
|
key_type* d_missing_keys, size_t* d_missing_len, cudaStream_t stream,
|
|
const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) override;
|
|
|
|
// Replace API, i.e. Follow the Query API to update the content of the cache to Most Recent
|
|
void Replace(const key_type* d_keys, const size_t len, const float* d_values, cudaStream_t stream,
|
|
const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) override;
|
|
|
|
// Update API, i.e. update the embeddings which exist in the cache
|
|
void Update(const key_type* d_keys, const size_t len, const float* d_values, cudaStream_t stream,
|
|
const size_t task_per_warp_tile = TASK_PER_WARP_TILE_MACRO) override;
|
|
|
|
// Dump API, i.e. dump some slabsets' keys from the cache
|
|
void Dump(key_type* d_keys, size_t* d_dump_counter, const size_t start_set_index,
|
|
const size_t end_set_index, cudaStream_t stream) override;
|
|
|
|
void Record(cudaStream_t stream) override {}
|
|
|
|
public:
|
|
using slabset = slab_set<set_associativity, key_type, warp_size>;
|
|
#ifdef LIBCUDACXX_VERSION
|
|
using atomic_ref_counter_type = cuda::atomic<ref_counter_type, cuda::thread_scope_device>;
|
|
using mutex = cuda::binary_semaphore<cuda::thread_scope_device>;
|
|
#endif
|
|
|
|
private:
|
|
static const size_t BLOCK_SIZE_ = 64;
|
|
|
|
// Cache data
|
|
slabset* keys_;
|
|
float* vals_;
|
|
ref_counter_type* slot_counter_;
|
|
|
|
// Global counter
|
|
#ifdef LIBCUDACXX_VERSION
|
|
atomic_ref_counter_type* global_counter_;
|
|
#else
|
|
ref_counter_type* global_counter_;
|
|
#endif
|
|
// CUDA device
|
|
int dev_;
|
|
|
|
// Cache capacity
|
|
size_t capacity_in_set_;
|
|
size_t num_slot_;
|
|
|
|
// Embedding vector size
|
|
size_t embedding_vec_size_;
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
// Array of mutex to protect (sub-)warp-level data structure, each mutex protect 1 slab set
|
|
mutex* set_mutex_;
|
|
#else
|
|
// Array of flag to protect (sub-)warp-level data structure, each flag act as a mutex and protect
|
|
// 1 slab set 1 for unlock, 0 for lock
|
|
int* set_mutex_;
|
|
#endif
|
|
};
|
|
|
|
} // namespace gpu_cache
|