wehub-resource-sync
1646 lines
67 KiB
Plaintext
1646 lines
67 KiB
Plaintext
/*
|
|
* 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.
|
|
*/
|
|
|
|
#include <cooperative_groups.h>
|
|
|
|
#include <nv_gpu_cache.hpp>
|
|
|
|
namespace cg = cooperative_groups;
|
|
|
|
// Overload CUDA atomic for other 64bit unsigned/signed integer type
|
|
__forceinline__ __device__ long atomicAdd(long* address, long val) {
|
|
return (long)atomicAdd((unsigned long long*)address, (unsigned long long)val);
|
|
}
|
|
|
|
__forceinline__ __device__ long long atomicAdd(long long* address, long long val) {
|
|
return (long long)atomicAdd((unsigned long long*)address, (unsigned long long)val);
|
|
}
|
|
|
|
__forceinline__ __device__ unsigned long atomicAdd(unsigned long* address, unsigned long val) {
|
|
return (unsigned long)atomicAdd((unsigned long long*)address, (unsigned long long)val);
|
|
}
|
|
|
|
namespace gpu_cache {
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <int warp_size>
|
|
__forceinline__ __device__ void warp_tile_copy(const size_t lane_idx,
|
|
const size_t emb_vec_size_in_float, float* d_dst,
|
|
const float* d_src) {
|
|
#pragma unroll
|
|
for (size_t i = lane_idx; i < emb_vec_size_in_float; i += warp_size) {
|
|
d_dst[i] = d_src[i];
|
|
}
|
|
}
|
|
#else
|
|
template <int warp_size>
|
|
__forceinline__ __device__ void warp_tile_copy(const size_t lane_idx,
|
|
const size_t emb_vec_size_in_float,
|
|
volatile float* d_dst, volatile float* d_src) {
|
|
|
|
#pragma unroll
|
|
for (size_t i = lane_idx; i < emb_vec_size_in_float; i += warp_size) {
|
|
d_dst[i] = d_src[i];
|
|
}
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
// Will be called by multiple thread_block_tile((sub-)warp) on the same mutex
|
|
// Expect only one thread_block_tile return to execute critical section at any time
|
|
template <typename mutex, int warp_size>
|
|
__forceinline__ __device__ void warp_lock_mutex(const cg::thread_block_tile<warp_size>& warp_tile,
|
|
mutex& set_mutex) {
|
|
// The first thread of this (sub-)warp to acquire the lock
|
|
if (warp_tile.thread_rank() == 0) {
|
|
set_mutex.acquire();
|
|
}
|
|
warp_tile.sync(); // Synchronize the threads in the (sub-)warp. Execution barrier + memory fence
|
|
}
|
|
|
|
// The (sub-)warp holding the mutex will unlock the mutex after finishing the critical section on a
|
|
// set Expect any following (sub-)warp that acquire the mutex can see its modification done in the
|
|
// critical section
|
|
template <typename mutex, int warp_size>
|
|
__forceinline__ __device__ void warp_unlock_mutex(const cg::thread_block_tile<warp_size>& warp_tile,
|
|
mutex& set_mutex) {
|
|
warp_tile.sync(); // Synchronize the threads in the (sub-)warp. Execution barrier + memory fence
|
|
// The first thread of this (sub-)warp to release the lock
|
|
if (warp_tile.thread_rank() == 0) {
|
|
set_mutex.release();
|
|
}
|
|
}
|
|
#else
|
|
// Will be called by multiple thread_block_tile((sub-)warp) on the same mutex
|
|
// Expect only one thread_block_tile return to execute critical section at any time
|
|
template <int warp_size>
|
|
__forceinline__ __device__ void warp_lock_mutex(const cg::thread_block_tile<warp_size>& warp_tile,
|
|
volatile int& set_mutex) {
|
|
// The first thread of this (sub-)warp to acquire the lock
|
|
if (warp_tile.thread_rank() == 0) {
|
|
while (0 == atomicCAS((int*)&set_mutex, 1, 0))
|
|
;
|
|
}
|
|
__threadfence();
|
|
warp_tile.sync(); // Synchronize the threads in the (sub-)warp. Execution barrier + memory fence
|
|
}
|
|
|
|
// The (sub-)warp holding the mutex will unlock the mutex after finishing the critical section on a
|
|
// set Expect any following (sub-)warp that acquire the mutex can see its modification done in the
|
|
// critical section
|
|
template <int warp_size>
|
|
__forceinline__ __device__ void warp_unlock_mutex(const cg::thread_block_tile<warp_size>& warp_tile,
|
|
volatile int& set_mutex) {
|
|
__threadfence();
|
|
warp_tile.sync(); // Synchronize the threads in the (sub-)warp. Execution barrier + memory fence
|
|
// The first thread of this (sub-)warp to release the lock
|
|
if (warp_tile.thread_rank() == 0) {
|
|
atomicExch((int*)&set_mutex, 1);
|
|
}
|
|
}
|
|
#endif
|
|
|
|
// The (sub-)warp doing all reduction to find the slot with min slot_counter
|
|
// The slot with min slot_counter is the LR slot.
|
|
template <typename ref_counter_type, int warp_size>
|
|
__forceinline__ __device__ void warp_min_reduction(
|
|
const cg::thread_block_tile<warp_size>& warp_tile, ref_counter_type& min_slot_counter_val,
|
|
size_t& slab_distance, size_t& slot_distance) {
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
slot_distance = lane_idx;
|
|
|
|
for (size_t i = (warp_tile.size() >> 1); i > 0; i = i >> 1) {
|
|
ref_counter_type input_slot_counter_val = warp_tile.shfl_xor(min_slot_counter_val, (int)i);
|
|
size_t input_slab_distance = warp_tile.shfl_xor(slab_distance, (int)i);
|
|
size_t input_slot_distance = warp_tile.shfl_xor(slot_distance, (int)i);
|
|
|
|
if (input_slot_counter_val == min_slot_counter_val) {
|
|
if (input_slab_distance == slab_distance) {
|
|
if (input_slot_distance < slot_distance) {
|
|
slot_distance = input_slot_distance;
|
|
}
|
|
} else if (input_slab_distance < slab_distance) {
|
|
slab_distance = input_slab_distance;
|
|
slot_distance = input_slot_distance;
|
|
}
|
|
} else if (input_slot_counter_val < min_slot_counter_val) {
|
|
min_slot_counter_val = input_slot_counter_val;
|
|
slab_distance = input_slab_distance;
|
|
slot_distance = input_slot_distance;
|
|
}
|
|
}
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
// Kernel to initialize the GPU cache
|
|
// Init every entry of the cache with <unused_key, value> pair
|
|
template <typename slabset, typename ref_counter_type, typename atomic_ref_counter_type,
|
|
typename key_type, typename mutex>
|
|
__global__ void init_cache(slabset* keys, ref_counter_type* slot_counter,
|
|
atomic_ref_counter_type* global_counter, const size_t num_slot,
|
|
const key_type empty_key, mutex* set_mutex,
|
|
const size_t capacity_in_set) {
|
|
const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (idx < num_slot) {
|
|
// Set the key of this slot to unused key
|
|
// Flatten the cache
|
|
key_type* key_slot = (key_type*)keys;
|
|
key_slot[idx] = empty_key;
|
|
|
|
// Clear the counter for this slot
|
|
slot_counter[idx] = 0;
|
|
}
|
|
// First CUDA thread clear the global counter
|
|
if (idx == 0) {
|
|
new (global_counter) atomic_ref_counter_type(0);
|
|
}
|
|
|
|
// First capacity_in_set CUDA thread initialize mutex
|
|
if (idx < capacity_in_set) {
|
|
new (set_mutex + idx) mutex(1);
|
|
}
|
|
}
|
|
|
|
template <typename atomic_ref_counter_type, typename mutex>
|
|
__global__ void destruct_kernel(atomic_ref_counter_type* global_counter, mutex* set_mutex,
|
|
const size_t capacity_in_set) {
|
|
const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
// First CUDA thread destruct the global_counter
|
|
if (idx == 0) {
|
|
global_counter->~atomic_ref_counter_type();
|
|
}
|
|
// First capacity_in_set CUDA thread destruct the set mutex
|
|
if (idx < capacity_in_set) {
|
|
(set_mutex + idx)->~mutex();
|
|
}
|
|
}
|
|
#else
|
|
// Kernel to initialize the GPU cache
|
|
// Init every entry of the cache with <unused_key, value> pair
|
|
template <typename slabset, typename ref_counter_type, typename key_type>
|
|
__global__ void init_cache(slabset* keys, ref_counter_type* slot_counter,
|
|
ref_counter_type* global_counter, const size_t num_slot,
|
|
const key_type empty_key, int* set_mutex, const size_t capacity_in_set) {
|
|
const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (idx < num_slot) {
|
|
// Set the key of this slot to unused key
|
|
// Flatten the cache
|
|
key_type* key_slot = (key_type*)keys;
|
|
key_slot[idx] = empty_key;
|
|
|
|
// Clear the counter for this slot
|
|
slot_counter[idx] = 0;
|
|
}
|
|
// First CUDA thread clear the global counter
|
|
if (idx == 0) {
|
|
global_counter[idx] = 0;
|
|
}
|
|
|
|
// First capacity_in_set CUDA thread initialize mutex
|
|
if (idx < capacity_in_set) {
|
|
set_mutex[idx] = 1;
|
|
}
|
|
}
|
|
#endif
|
|
|
|
// Kernel to update global counter
|
|
// Resolve distance overflow issue as well
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <typename atomic_ref_counter_type>
|
|
__global__ void update_kernel_overflow_ignore(atomic_ref_counter_type* global_counter,
|
|
size_t* d_missing_len) {
|
|
// Update global counter
|
|
global_counter->fetch_add(1, cuda::std::memory_order_relaxed);
|
|
*d_missing_len = 0;
|
|
}
|
|
#else
|
|
template <typename ref_counter_type>
|
|
__global__ void update_kernel_overflow_ignore(ref_counter_type* global_counter,
|
|
size_t* d_missing_len) {
|
|
// Update global counter
|
|
atomicAdd(global_counter, 1);
|
|
*d_missing_len = 0;
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
// Kernel to read from cache
|
|
// Also update locality information for touched slot
|
|
template <typename key_type, typename ref_counter_type, typename atomic_ref_counter_type,
|
|
typename slabset, typename set_hasher, typename slab_hasher, typename mutex,
|
|
key_type empty_key, int set_associativity, int warp_size>
|
|
__global__ void get_kernel(const key_type* d_keys, const size_t len, float* d_values,
|
|
const size_t embedding_vec_size, uint64_t* d_missing_index,
|
|
key_type* d_missing_keys, size_t* d_missing_len,
|
|
const atomic_ref_counter_type* global_counter,
|
|
ref_counter_type* slot_counter, const size_t capacity_in_set,
|
|
const slabset* keys, const float* vals, mutex* set_mutex,
|
|
const size_t task_per_warp_tile) {
|
|
// Lane(thread) ID within a warp_tile
|
|
cg::thread_block_tile<warp_size> warp_tile =
|
|
cg::tiled_partition<warp_size>(cg::this_thread_block());
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
// Warp tile global ID
|
|
const size_t warp_tile_global_idx =
|
|
(blockIdx.x * (blockDim.x / warp_size)) + warp_tile.meta_group_rank();
|
|
// The index of key for this thread
|
|
const size_t key_idx = (warp_tile_global_idx * task_per_warp_tile) + lane_idx;
|
|
// The assigned key for this lane(thread)
|
|
key_type key;
|
|
// The dst slabset and the dst slab inside this set
|
|
size_t src_set;
|
|
size_t src_slab;
|
|
// The variable that contains the missing key
|
|
key_type missing_key;
|
|
// The variable that contains the index for the missing key
|
|
uint64_t missing_index;
|
|
// The counter for counting the missing key in this warp
|
|
uint8_t warp_missing_counter = 0;
|
|
// Active flag: whether current lane(thread) has unfinished task
|
|
bool active = false;
|
|
if (lane_idx < task_per_warp_tile) {
|
|
if (key_idx < len) {
|
|
active = true;
|
|
key = d_keys[key_idx];
|
|
src_set = set_hasher::hash(key) % capacity_in_set;
|
|
src_slab = slab_hasher::hash(key) % set_associativity;
|
|
}
|
|
}
|
|
|
|
// Lane participate in warp_tile ballot to produce warp-level work queue
|
|
unsigned active_mask = warp_tile.ballot(active);
|
|
|
|
// The warp-level outer loop: finish all the tasks within the work queue
|
|
while (active_mask != 0) {
|
|
// Next task in the work quere, start from lower index lane(thread)
|
|
int next_lane = __ffs(active_mask) - 1;
|
|
// Broadcast the task and the global index to all lane in the warp_tile
|
|
key_type next_key = warp_tile.shfl(key, next_lane);
|
|
size_t next_idx = warp_tile.shfl(key_idx, next_lane);
|
|
size_t next_set = warp_tile.shfl(src_set, next_lane);
|
|
size_t next_slab = warp_tile.shfl(src_slab, next_lane);
|
|
|
|
// Counter to record how many slab have been searched
|
|
size_t counter = 0;
|
|
|
|
// Working queue before task started
|
|
const unsigned old_active_mask = active_mask;
|
|
|
|
// Lock the slabset before operating the slabset
|
|
warp_lock_mutex<mutex, warp_size>(warp_tile, set_mutex[next_set]);
|
|
|
|
// The warp-level inner loop: finish a single task in the work queue
|
|
while (active_mask == old_active_mask) {
|
|
// When all the slabs inside a slabset have been searched, mark missing task, task is
|
|
// completed
|
|
if (counter >= set_associativity) {
|
|
if (lane_idx == warp_missing_counter) {
|
|
missing_key = next_key;
|
|
missing_index = next_idx;
|
|
}
|
|
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
warp_missing_counter++;
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// The warp_tile read out the slab
|
|
key_type read_key = keys[next_set].set_[next_slab].slab_[lane_idx];
|
|
|
|
// Compare the slab data with the target key
|
|
int found_lane = __ffs(warp_tile.ballot(read_key == next_key)) - 1;
|
|
|
|
// If found, mark hit task, copy the founded data, the task is completed
|
|
if (found_lane >= 0) {
|
|
size_t found_offset = (next_set * set_associativity + next_slab) * warp_size + found_lane;
|
|
if (lane_idx == (size_t)next_lane) {
|
|
slot_counter[found_offset] = global_counter->load(cuda::std::memory_order_relaxed);
|
|
active = false;
|
|
}
|
|
|
|
warp_tile_copy<warp_size>(lane_idx, embedding_vec_size,
|
|
d_values + next_idx * embedding_vec_size,
|
|
vals + found_offset * embedding_vec_size);
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Compare the slab data with empty key, if found empty key, mark missing task, task is
|
|
// completed
|
|
if (warp_tile.ballot(read_key == empty_key) != 0) {
|
|
if (lane_idx == warp_missing_counter) {
|
|
missing_key = next_key;
|
|
missing_index = next_idx;
|
|
}
|
|
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
warp_missing_counter++;
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Not found in this slab, the task is not completed, goto searching next slab
|
|
counter++;
|
|
next_slab = (next_slab + 1) % set_associativity;
|
|
}
|
|
|
|
// Unlock the slabset after operating the slabset
|
|
warp_unlock_mutex<mutex, warp_size>(warp_tile, set_mutex[next_set]);
|
|
}
|
|
|
|
// After warp_tile complete the working queue, save the result for output
|
|
// First thread of the warp_tile accumulate the missing length to global variable
|
|
size_t warp_position;
|
|
if (lane_idx == 0) {
|
|
warp_position = atomicAdd(d_missing_len, (size_t)warp_missing_counter);
|
|
}
|
|
warp_position = warp_tile.shfl(warp_position, 0);
|
|
|
|
if (lane_idx < warp_missing_counter) {
|
|
d_missing_keys[warp_position + lane_idx] = missing_key;
|
|
d_missing_index[warp_position + lane_idx] = missing_index;
|
|
}
|
|
}
|
|
#else
|
|
// Kernel to read from cache
|
|
// Also update locality information for touched slot
|
|
template <typename key_type, typename ref_counter_type, typename slabset, typename set_hasher,
|
|
typename slab_hasher, key_type empty_key, int set_associativity, int warp_size>
|
|
__global__ void get_kernel(const key_type* d_keys, const size_t len, float* d_values,
|
|
const size_t embedding_vec_size, uint64_t* d_missing_index,
|
|
key_type* d_missing_keys, size_t* d_missing_len,
|
|
ref_counter_type* global_counter,
|
|
volatile ref_counter_type* slot_counter, const size_t capacity_in_set,
|
|
volatile slabset* keys, volatile float* vals, volatile int* set_mutex,
|
|
const size_t task_per_warp_tile) {
|
|
// Lane(thread) ID within a warp_tile
|
|
cg::thread_block_tile<warp_size> warp_tile =
|
|
cg::tiled_partition<warp_size>(cg::this_thread_block());
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
// Warp tile global ID
|
|
const size_t warp_tile_global_idx =
|
|
(blockIdx.x * (blockDim.x / warp_size)) + warp_tile.meta_group_rank();
|
|
// The index of key for this thread
|
|
const size_t key_idx = (warp_tile_global_idx * task_per_warp_tile) + lane_idx;
|
|
// The assigned key for this lane(thread)
|
|
key_type key;
|
|
// The dst slabset and the dst slab inside this set
|
|
size_t src_set;
|
|
size_t src_slab;
|
|
// The variable that contains the missing key
|
|
key_type missing_key;
|
|
// The variable that contains the index for the missing key
|
|
uint64_t missing_index;
|
|
// The counter for counting the missing key in this warp
|
|
uint8_t warp_missing_counter = 0;
|
|
// Active flag: whether current lane(thread) has unfinished task
|
|
bool active = false;
|
|
if (lane_idx < task_per_warp_tile) {
|
|
if (key_idx < len) {
|
|
active = true;
|
|
key = d_keys[key_idx];
|
|
src_set = set_hasher::hash(key) % capacity_in_set;
|
|
src_slab = slab_hasher::hash(key) % set_associativity;
|
|
}
|
|
}
|
|
|
|
// Lane participate in warp_tile ballot to produce warp-level work queue
|
|
unsigned active_mask = warp_tile.ballot(active);
|
|
|
|
// The warp-level outer loop: finish all the tasks within the work queue
|
|
while (active_mask != 0) {
|
|
// Next task in the work quere, start from lower index lane(thread)
|
|
int next_lane = __ffs(active_mask) - 1;
|
|
// Broadcast the task and the global index to all lane in the warp_tile
|
|
key_type next_key = warp_tile.shfl(key, next_lane);
|
|
size_t next_idx = warp_tile.shfl(key_idx, next_lane);
|
|
size_t next_set = warp_tile.shfl(src_set, next_lane);
|
|
size_t next_slab = warp_tile.shfl(src_slab, next_lane);
|
|
|
|
// Counter to record how many slab have been searched
|
|
size_t counter = 0;
|
|
|
|
// Working queue before task started
|
|
const unsigned old_active_mask = active_mask;
|
|
|
|
// Lock the slabset before operating the slabset
|
|
warp_lock_mutex<warp_size>(warp_tile, set_mutex[next_set]);
|
|
|
|
// The warp-level inner loop: finish a single task in the work queue
|
|
while (active_mask == old_active_mask) {
|
|
// When all the slabs inside a slabset have been searched, mark missing task, task is
|
|
// completed
|
|
if (counter >= set_associativity) {
|
|
if (lane_idx == warp_missing_counter) {
|
|
missing_key = next_key;
|
|
missing_index = next_idx;
|
|
}
|
|
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
warp_missing_counter++;
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// The warp_tile read out the slab
|
|
key_type read_key = ((volatile key_type*)(keys[next_set].set_[next_slab].slab_))[lane_idx];
|
|
|
|
// Compare the slab data with the target key
|
|
int found_lane = __ffs(warp_tile.ballot(read_key == next_key)) - 1;
|
|
|
|
// If found, mark hit task, copy the founded data, the task is completed
|
|
if (found_lane >= 0) {
|
|
size_t found_offset = (next_set * set_associativity + next_slab) * warp_size + found_lane;
|
|
if (lane_idx == (size_t)next_lane) {
|
|
slot_counter[found_offset] = atomicAdd(global_counter, 0);
|
|
active = false;
|
|
}
|
|
|
|
warp_tile_copy<warp_size>(lane_idx, embedding_vec_size,
|
|
(volatile float*)(d_values + next_idx * embedding_vec_size),
|
|
(volatile float*)(vals + found_offset * embedding_vec_size));
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Compare the slab data with empty key, if found empty key, mark missing task, task is
|
|
// completed
|
|
if (warp_tile.ballot(read_key == empty_key) != 0) {
|
|
if (lane_idx == warp_missing_counter) {
|
|
missing_key = next_key;
|
|
missing_index = next_idx;
|
|
}
|
|
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
warp_missing_counter++;
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Not found in this slab, the task is not completed, goto searching next slab
|
|
counter++;
|
|
next_slab = (next_slab + 1) % set_associativity;
|
|
}
|
|
|
|
// Unlock the slabset after operating the slabset
|
|
warp_unlock_mutex<warp_size>(warp_tile, set_mutex[next_set]);
|
|
}
|
|
|
|
// After warp_tile complete the working queue, save the result for output
|
|
// First thread of the warp_tile accumulate the missing length to global variable
|
|
size_t warp_position;
|
|
if (lane_idx == 0) {
|
|
warp_position = atomicAdd(d_missing_len, (size_t)warp_missing_counter);
|
|
}
|
|
warp_position = warp_tile.shfl(warp_position, 0);
|
|
|
|
if (lane_idx < warp_missing_counter) {
|
|
d_missing_keys[warp_position + lane_idx] = missing_key;
|
|
d_missing_index[warp_position + lane_idx] = missing_index;
|
|
}
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
// Kernel to insert or replace the <k,v> pairs into the cache
|
|
template <typename key_type, typename slabset, typename ref_counter_type, typename mutex,
|
|
typename atomic_ref_counter_type, typename set_hasher, typename slab_hasher,
|
|
key_type empty_key, int set_associativity, int warp_size,
|
|
ref_counter_type max_ref_counter_type = std::numeric_limits<ref_counter_type>::max(),
|
|
size_t max_slab_distance = std::numeric_limits<size_t>::max()>
|
|
__global__ void insert_replace_kernel(const key_type* d_keys, const float* d_values,
|
|
const size_t embedding_vec_size, const size_t len,
|
|
slabset* keys, float* vals, ref_counter_type* slot_counter,
|
|
mutex* set_mutex,
|
|
const atomic_ref_counter_type* global_counter,
|
|
const size_t capacity_in_set,
|
|
const size_t task_per_warp_tile) {
|
|
// Lane(thread) ID within a warp_tile
|
|
cg::thread_block_tile<warp_size> warp_tile =
|
|
cg::tiled_partition<warp_size>(cg::this_thread_block());
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
// Warp tile global ID
|
|
const size_t warp_tile_global_idx =
|
|
(blockIdx.x * (blockDim.x / warp_size)) + warp_tile.meta_group_rank();
|
|
// The index of key for this thread
|
|
const size_t key_idx = (warp_tile_global_idx * task_per_warp_tile) + lane_idx;
|
|
// The assigned key for this lane(thread)
|
|
key_type key;
|
|
// The dst slabset and the dst slab inside this set
|
|
size_t src_set;
|
|
size_t src_slab;
|
|
// Active flag: whether current lane(thread) has unfinished task
|
|
bool active = false;
|
|
if (lane_idx < task_per_warp_tile) {
|
|
if (key_idx < len) {
|
|
active = true;
|
|
key = d_keys[key_idx];
|
|
src_set = set_hasher::hash(key) % capacity_in_set;
|
|
src_slab = slab_hasher::hash(key) % set_associativity;
|
|
}
|
|
}
|
|
|
|
// Lane participate in warp_tile ballot to produce warp-level work queue
|
|
unsigned active_mask = warp_tile.ballot(active);
|
|
|
|
// The warp-level outer loop: finish all the tasks within the work queue
|
|
while (active_mask != 0) {
|
|
// Next task in the work quere, start from lower index lane(thread)
|
|
int next_lane = __ffs(active_mask) - 1;
|
|
// Broadcast the task, the global index and the src slabset and slab to all lane in a warp_tile
|
|
key_type next_key = warp_tile.shfl(key, next_lane);
|
|
size_t next_idx = warp_tile.shfl(key_idx, next_lane);
|
|
size_t next_set = warp_tile.shfl(src_set, next_lane);
|
|
size_t next_slab = warp_tile.shfl(src_slab, next_lane);
|
|
size_t first_slab = next_slab;
|
|
|
|
// Counter to record how many slab have been searched
|
|
size_t counter = 0;
|
|
|
|
// Variable to keep the min slot counter during the probing
|
|
ref_counter_type min_slot_counter_val = max_ref_counter_type;
|
|
// Variable to keep the slab distance for slot with min counter
|
|
size_t slab_distance = max_slab_distance;
|
|
// Variable to keep the slot distance for slot with min counter within the slab
|
|
size_t slot_distance;
|
|
// Working queue before task started
|
|
const unsigned old_active_mask = active_mask;
|
|
|
|
// Lock the slabset before operating the slabset
|
|
warp_lock_mutex<mutex, warp_size>(warp_tile, set_mutex[next_set]);
|
|
|
|
// The warp-level inner loop: finish a single task in the work queue
|
|
while (active_mask == old_active_mask) {
|
|
// When all the slabs inside a slabset have been searched
|
|
// and no empty slots or target slots are found. Replace with LRU
|
|
if (counter >= set_associativity) {
|
|
// (sub)Warp all-reduction, the reduction result store in all threads
|
|
warp_min_reduction<ref_counter_type, warp_size>(warp_tile, min_slot_counter_val,
|
|
slab_distance, slot_distance);
|
|
|
|
// Calculate the position of LR slot
|
|
size_t target_slab = (first_slab + slab_distance) % set_associativity;
|
|
size_t slot_index =
|
|
(next_set * set_associativity + target_slab) * warp_size + slot_distance;
|
|
|
|
// Replace the LR slot
|
|
if (lane_idx == (size_t)next_lane) {
|
|
keys[next_set].set_[target_slab].slab_[slot_distance] = key;
|
|
slot_counter[slot_index] = global_counter->load(cuda::std::memory_order_relaxed);
|
|
}
|
|
|
|
warp_tile_copy<warp_size>(lane_idx, embedding_vec_size,
|
|
vals + slot_index * embedding_vec_size,
|
|
d_values + next_idx * embedding_vec_size);
|
|
|
|
// Replace complete, mark this task completed
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// The warp_tile read out the slab
|
|
key_type read_key = keys[next_set].set_[next_slab].slab_[lane_idx];
|
|
|
|
// Compare the slab data with the target key
|
|
int found_lane = __ffs(warp_tile.ballot(read_key == next_key)) - 1;
|
|
|
|
// If found target key, the insertion/replace is no longer needed.
|
|
// Refresh the slot, the task is completed
|
|
if (found_lane >= 0) {
|
|
size_t found_offset = (next_set * set_associativity + next_slab) * warp_size + found_lane;
|
|
if (lane_idx == (size_t)next_lane) {
|
|
slot_counter[found_offset] = global_counter->load(cuda::std::memory_order_relaxed);
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Compare the slab data with empty key.
|
|
// If found empty key, do insertion,the task is complete
|
|
found_lane = __ffs(warp_tile.ballot(read_key == empty_key)) - 1;
|
|
if (found_lane >= 0) {
|
|
size_t found_offset = (next_set * set_associativity + next_slab) * warp_size + found_lane;
|
|
|
|
if (lane_idx == (size_t)next_lane) {
|
|
keys[next_set].set_[next_slab].slab_[found_lane] = key;
|
|
slot_counter[found_offset] = global_counter->load(cuda::std::memory_order_relaxed);
|
|
}
|
|
|
|
warp_tile_copy<warp_size>(lane_idx, embedding_vec_size,
|
|
vals + found_offset * embedding_vec_size,
|
|
d_values + next_idx * embedding_vec_size);
|
|
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// If no target or unused slot found in this slab,
|
|
// Refresh LR info, continue probing
|
|
ref_counter_type read_slot_counter =
|
|
slot_counter[(next_set * set_associativity + next_slab) * warp_size + lane_idx];
|
|
if (read_slot_counter < min_slot_counter_val) {
|
|
min_slot_counter_val = read_slot_counter;
|
|
slab_distance = counter;
|
|
}
|
|
|
|
counter++;
|
|
next_slab = (next_slab + 1) % set_associativity;
|
|
}
|
|
|
|
// Unlock the slabset after operating the slabset
|
|
warp_unlock_mutex<mutex, warp_size>(warp_tile, set_mutex[next_set]);
|
|
}
|
|
}
|
|
#else
|
|
// Kernel to insert or replace the <k,v> pairs into the cache
|
|
template <typename key_type, typename slabset, typename ref_counter_type, typename set_hasher,
|
|
typename slab_hasher, key_type empty_key, int set_associativity, int warp_size,
|
|
ref_counter_type max_ref_counter_type = std::numeric_limits<ref_counter_type>::max(),
|
|
size_t max_slab_distance = std::numeric_limits<size_t>::max()>
|
|
__global__ void insert_replace_kernel(const key_type* d_keys, const float* d_values,
|
|
const size_t embedding_vec_size, const size_t len,
|
|
volatile slabset* keys, volatile float* vals,
|
|
volatile ref_counter_type* slot_counter,
|
|
volatile int* set_mutex, ref_counter_type* global_counter,
|
|
const size_t capacity_in_set,
|
|
const size_t task_per_warp_tile) {
|
|
// Lane(thread) ID within a warp_tile
|
|
cg::thread_block_tile<warp_size> warp_tile =
|
|
cg::tiled_partition<warp_size>(cg::this_thread_block());
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
// Warp tile global ID
|
|
const size_t warp_tile_global_idx =
|
|
(blockIdx.x * (blockDim.x / warp_size)) + warp_tile.meta_group_rank();
|
|
// The index of key for this thread
|
|
const size_t key_idx = (warp_tile_global_idx * task_per_warp_tile) + lane_idx;
|
|
// The assigned key for this lane(thread)
|
|
key_type key;
|
|
// The dst slabset and the dst slab inside this set
|
|
size_t src_set;
|
|
size_t src_slab;
|
|
// Active flag: whether current lane(thread) has unfinished task
|
|
bool active = false;
|
|
if (lane_idx < task_per_warp_tile) {
|
|
if (key_idx < len) {
|
|
active = true;
|
|
key = d_keys[key_idx];
|
|
src_set = set_hasher::hash(key) % capacity_in_set;
|
|
src_slab = slab_hasher::hash(key) % set_associativity;
|
|
}
|
|
}
|
|
|
|
// Lane participate in warp_tile ballot to produce warp-level work queue
|
|
unsigned active_mask = warp_tile.ballot(active);
|
|
|
|
// The warp-level outer loop: finish all the tasks within the work queue
|
|
while (active_mask != 0) {
|
|
// Next task in the work quere, start from lower index lane(thread)
|
|
int next_lane = __ffs(active_mask) - 1;
|
|
// Broadcast the task, the global index and the src slabset and slab to all lane in a warp_tile
|
|
key_type next_key = warp_tile.shfl(key, next_lane);
|
|
size_t next_idx = warp_tile.shfl(key_idx, next_lane);
|
|
size_t next_set = warp_tile.shfl(src_set, next_lane);
|
|
size_t next_slab = warp_tile.shfl(src_slab, next_lane);
|
|
size_t first_slab = next_slab;
|
|
|
|
// Counter to record how many slab have been searched
|
|
size_t counter = 0;
|
|
|
|
// Variable to keep the min slot counter during the probing
|
|
ref_counter_type min_slot_counter_val = max_ref_counter_type;
|
|
// Variable to keep the slab distance for slot with min counter
|
|
size_t slab_distance = max_slab_distance;
|
|
// Variable to keep the slot distance for slot with min counter within the slab
|
|
size_t slot_distance;
|
|
// Working queue before task started
|
|
const unsigned old_active_mask = active_mask;
|
|
|
|
// Lock the slabset before operating the slabset
|
|
warp_lock_mutex<warp_size>(warp_tile, set_mutex[next_set]);
|
|
|
|
// The warp-level inner loop: finish a single task in the work queue
|
|
while (active_mask == old_active_mask) {
|
|
// When all the slabs inside a slabset have been searched
|
|
// and no empty slots or target slots are found. Replace with LRU
|
|
if (counter >= set_associativity) {
|
|
// (sub)Warp all-reduction, the reduction result store in all threads
|
|
warp_min_reduction<ref_counter_type, warp_size>(warp_tile, min_slot_counter_val,
|
|
slab_distance, slot_distance);
|
|
|
|
// Calculate the position of LR slot
|
|
size_t target_slab = (first_slab + slab_distance) % set_associativity;
|
|
size_t slot_index =
|
|
(next_set * set_associativity + target_slab) * warp_size + slot_distance;
|
|
|
|
// Replace the LR slot
|
|
if (lane_idx == (size_t)next_lane) {
|
|
((volatile key_type*)(keys[next_set].set_[target_slab].slab_))[slot_distance] = key;
|
|
slot_counter[slot_index] = atomicAdd(global_counter, 0);
|
|
}
|
|
|
|
warp_tile_copy<warp_size>(lane_idx, embedding_vec_size,
|
|
(volatile float*)(vals + slot_index * embedding_vec_size),
|
|
(volatile float*)(d_values + next_idx * embedding_vec_size));
|
|
|
|
// Replace complete, mark this task completed
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// The warp_tile read out the slab
|
|
key_type read_key = ((volatile key_type*)(keys[next_set].set_[next_slab].slab_))[lane_idx];
|
|
|
|
// Compare the slab data with the target key
|
|
int found_lane = __ffs(warp_tile.ballot(read_key == next_key)) - 1;
|
|
|
|
// If found target key, the insertion/replace is no longer needed.
|
|
// Refresh the slot, the task is completed
|
|
if (found_lane >= 0) {
|
|
size_t found_offset = (next_set * set_associativity + next_slab) * warp_size + found_lane;
|
|
if (lane_idx == (size_t)next_lane) {
|
|
slot_counter[found_offset] = atomicAdd(global_counter, 0);
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Compare the slab data with empty key.
|
|
// If found empty key, do insertion,the task is complete
|
|
found_lane = __ffs(warp_tile.ballot(read_key == empty_key)) - 1;
|
|
if (found_lane >= 0) {
|
|
size_t found_offset = (next_set * set_associativity + next_slab) * warp_size + found_lane;
|
|
|
|
if (lane_idx == (size_t)next_lane) {
|
|
((volatile key_type*)(keys[next_set].set_[next_slab].slab_))[found_lane] = key;
|
|
slot_counter[found_offset] = atomicAdd(global_counter, 0);
|
|
}
|
|
|
|
warp_tile_copy<warp_size>(lane_idx, embedding_vec_size,
|
|
(volatile float*)(vals + found_offset * embedding_vec_size),
|
|
(volatile float*)(d_values + next_idx * embedding_vec_size));
|
|
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// If no target or unused slot found in this slab,
|
|
// Refresh LR info, continue probing
|
|
ref_counter_type read_slot_counter =
|
|
slot_counter[(next_set * set_associativity + next_slab) * warp_size + lane_idx];
|
|
if (read_slot_counter < min_slot_counter_val) {
|
|
min_slot_counter_val = read_slot_counter;
|
|
slab_distance = counter;
|
|
}
|
|
|
|
counter++;
|
|
next_slab = (next_slab + 1) % set_associativity;
|
|
}
|
|
|
|
// Unlock the slabset after operating the slabset
|
|
warp_unlock_mutex<warp_size>(warp_tile, set_mutex[next_set]);
|
|
}
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
// Kernel to update the existing keys in the cache
|
|
// Will not change the locality information
|
|
template <typename key_type, typename slabset, typename set_hasher, typename slab_hasher,
|
|
typename mutex, key_type empty_key, int set_associativity, int warp_size>
|
|
__global__ void update_kernel(const key_type* d_keys, const size_t len, const float* d_values,
|
|
const size_t embedding_vec_size, const size_t capacity_in_set,
|
|
const slabset* keys, float* vals, mutex* set_mutex,
|
|
const size_t task_per_warp_tile) {
|
|
// Lane(thread) ID within a warp_tile
|
|
cg::thread_block_tile<warp_size> warp_tile =
|
|
cg::tiled_partition<warp_size>(cg::this_thread_block());
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
// Warp tile global ID
|
|
const size_t warp_tile_global_idx =
|
|
(blockIdx.x * (blockDim.x / warp_size)) + warp_tile.meta_group_rank();
|
|
// The index of key for this thread
|
|
const size_t key_idx = (warp_tile_global_idx * task_per_warp_tile) + lane_idx;
|
|
// The assigned key for this lane(thread)
|
|
key_type key;
|
|
// The dst slabset and the dst slab inside this set
|
|
size_t src_set;
|
|
size_t src_slab;
|
|
// Active flag: whether current lane(thread) has unfinished task
|
|
bool active = false;
|
|
if (lane_idx < task_per_warp_tile) {
|
|
if (key_idx < len) {
|
|
active = true;
|
|
key = d_keys[key_idx];
|
|
src_set = set_hasher::hash(key) % capacity_in_set;
|
|
src_slab = slab_hasher::hash(key) % set_associativity;
|
|
}
|
|
}
|
|
|
|
// Lane participate in warp_tile ballot to produce warp-level work queue
|
|
unsigned active_mask = warp_tile.ballot(active);
|
|
|
|
// The warp-level outer loop: finish all the tasks within the work queue
|
|
while (active_mask != 0) {
|
|
// Next task in the work quere, start from lower index lane(thread)
|
|
int next_lane = __ffs(active_mask) - 1;
|
|
// Broadcast the task and the global index to all lane in the warp_tile
|
|
key_type next_key = warp_tile.shfl(key, next_lane);
|
|
size_t next_idx = warp_tile.shfl(key_idx, next_lane);
|
|
size_t next_set = warp_tile.shfl(src_set, next_lane);
|
|
size_t next_slab = warp_tile.shfl(src_slab, next_lane);
|
|
|
|
// Counter to record how many slab have been searched
|
|
size_t counter = 0;
|
|
|
|
// Working queue before task started
|
|
const unsigned old_active_mask = active_mask;
|
|
|
|
// Lock the slabset before operating the slabset
|
|
warp_lock_mutex<mutex, warp_size>(warp_tile, set_mutex[next_set]);
|
|
|
|
// The warp-level inner loop: finish a single task in the work queue
|
|
while (active_mask == old_active_mask) {
|
|
// When all the slabs inside a slabset have been searched, mark missing task, do nothing, task
|
|
// complete
|
|
if (counter >= set_associativity) {
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// The warp_tile read out the slab
|
|
key_type read_key = keys[next_set].set_[next_slab].slab_[lane_idx];
|
|
|
|
// Compare the slab data with the target key
|
|
int found_lane = __ffs(warp_tile.ballot(read_key == next_key)) - 1;
|
|
|
|
// If found, mark hit task, update the value, the task is completed
|
|
if (found_lane >= 0) {
|
|
size_t found_offset = (next_set * set_associativity + next_slab) * warp_size + found_lane;
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
warp_tile_copy<warp_size>(lane_idx, embedding_vec_size,
|
|
vals + found_offset * embedding_vec_size,
|
|
d_values + next_idx * embedding_vec_size);
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Compare the slab data with empty key, if found empty key, mark missing task, do nothing,
|
|
// task is completed
|
|
if (warp_tile.ballot(read_key == empty_key) != 0) {
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Not found in this slab, the task is not completed, goto searching next slab
|
|
counter++;
|
|
next_slab = (next_slab + 1) % set_associativity;
|
|
}
|
|
|
|
// Unlock the slabset after operating the slabset
|
|
warp_unlock_mutex<mutex, warp_size>(warp_tile, set_mutex[next_set]);
|
|
}
|
|
}
|
|
#else
|
|
// Kernel to update the existing keys in the cache
|
|
// Will not change the locality information
|
|
template <typename key_type, typename slabset, typename set_hasher, typename slab_hasher,
|
|
key_type empty_key, int set_associativity, int warp_size>
|
|
__global__ void update_kernel(const key_type* d_keys, const size_t len, const float* d_values,
|
|
const size_t embedding_vec_size, const size_t capacity_in_set,
|
|
volatile slabset* keys, volatile float* vals, volatile int* set_mutex,
|
|
const size_t task_per_warp_tile) {
|
|
// Lane(thread) ID within a warp_tile
|
|
cg::thread_block_tile<warp_size> warp_tile =
|
|
cg::tiled_partition<warp_size>(cg::this_thread_block());
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
// Warp tile global ID
|
|
const size_t warp_tile_global_idx =
|
|
(blockIdx.x * (blockDim.x / warp_size)) + warp_tile.meta_group_rank();
|
|
// The index of key for this thread
|
|
const size_t key_idx = (warp_tile_global_idx * task_per_warp_tile) + lane_idx;
|
|
// The assigned key for this lane(thread)
|
|
key_type key;
|
|
// The dst slabset and the dst slab inside this set
|
|
size_t src_set;
|
|
size_t src_slab;
|
|
// Active flag: whether current lane(thread) has unfinished task
|
|
bool active = false;
|
|
if (lane_idx < task_per_warp_tile) {
|
|
if (key_idx < len) {
|
|
active = true;
|
|
key = d_keys[key_idx];
|
|
src_set = set_hasher::hash(key) % capacity_in_set;
|
|
src_slab = slab_hasher::hash(key) % set_associativity;
|
|
}
|
|
}
|
|
|
|
// Lane participate in warp_tile ballot to produce warp-level work queue
|
|
unsigned active_mask = warp_tile.ballot(active);
|
|
|
|
// The warp-level outer loop: finish all the tasks within the work queue
|
|
while (active_mask != 0) {
|
|
// Next task in the work quere, start from lower index lane(thread)
|
|
int next_lane = __ffs(active_mask) - 1;
|
|
// Broadcast the task and the global index to all lane in the warp_tile
|
|
key_type next_key = warp_tile.shfl(key, next_lane);
|
|
size_t next_idx = warp_tile.shfl(key_idx, next_lane);
|
|
size_t next_set = warp_tile.shfl(src_set, next_lane);
|
|
size_t next_slab = warp_tile.shfl(src_slab, next_lane);
|
|
|
|
// Counter to record how many slab have been searched
|
|
size_t counter = 0;
|
|
|
|
// Working queue before task started
|
|
const unsigned old_active_mask = active_mask;
|
|
|
|
// Lock the slabset before operating the slabset
|
|
warp_lock_mutex<warp_size>(warp_tile, set_mutex[next_set]);
|
|
|
|
// The warp-level inner loop: finish a single task in the work queue
|
|
while (active_mask == old_active_mask) {
|
|
// When all the slabs inside a slabset have been searched, mark missing task, do nothing, task
|
|
// complete
|
|
if (counter >= set_associativity) {
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// The warp_tile read out the slab
|
|
key_type read_key = ((volatile key_type*)(keys[next_set].set_[next_slab].slab_))[lane_idx];
|
|
|
|
// Compare the slab data with the target key
|
|
int found_lane = __ffs(warp_tile.ballot(read_key == next_key)) - 1;
|
|
|
|
// If found, mark hit task, update the value, the task is completed
|
|
if (found_lane >= 0) {
|
|
size_t found_offset = (next_set * set_associativity + next_slab) * warp_size + found_lane;
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
warp_tile_copy<warp_size>(lane_idx, embedding_vec_size,
|
|
(volatile float*)(vals + found_offset * embedding_vec_size),
|
|
(volatile float*)(d_values + next_idx * embedding_vec_size));
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Compare the slab data with empty key, if found empty key, mark missing task, do nothing,
|
|
// task is completed
|
|
if (warp_tile.ballot(read_key == empty_key) != 0) {
|
|
if (lane_idx == (size_t)next_lane) {
|
|
active = false;
|
|
}
|
|
|
|
active_mask = warp_tile.ballot(active);
|
|
break;
|
|
}
|
|
|
|
// Not found in this slab, the task is not completed, goto searching next slab
|
|
counter++;
|
|
next_slab = (next_slab + 1) % set_associativity;
|
|
}
|
|
|
|
// Unlock the slabset after operating the slabset
|
|
warp_unlock_mutex<warp_size>(warp_tile, set_mutex[next_set]);
|
|
}
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <typename key_type, typename slabset, typename mutex, key_type empty_key,
|
|
int set_associativity, int warp_size>
|
|
__global__ void dump_kernel(key_type* d_keys, size_t* d_dump_counter, const slabset* keys,
|
|
mutex* set_mutex, const size_t start_set_index,
|
|
const size_t end_set_index) {
|
|
// Block-level counter used by all warp tiles within a block
|
|
__shared__ uint32_t block_acc;
|
|
// Initialize block-level counter
|
|
if (threadIdx.x == 0) {
|
|
block_acc = 0;
|
|
}
|
|
__syncthreads();
|
|
// Lane(thread) ID within a warp tile
|
|
cg::thread_block_tile<warp_size> warp_tile =
|
|
cg::tiled_partition<warp_size>(cg::this_thread_block());
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
// Warp tile target slabset id
|
|
const size_t set_idx =
|
|
((blockIdx.x * (blockDim.x / warp_size)) + warp_tile.meta_group_rank()) + start_set_index;
|
|
// Keys dump from cache
|
|
key_type read_key[set_associativity];
|
|
// Lane(thread) offset for storing each key
|
|
uint32_t thread_key_offset[set_associativity];
|
|
// Warp offset for storing each key
|
|
uint32_t warp_key_offset;
|
|
// Block offset for storing each key
|
|
__shared__ size_t block_key_offset;
|
|
|
|
// Warp tile dump target slabset
|
|
if (set_idx < end_set_index) {
|
|
// Lock the slabset before operating the slabset
|
|
warp_lock_mutex<mutex, warp_size>(warp_tile, set_mutex[set_idx]);
|
|
|
|
// The warp tile read out the slabset
|
|
for (unsigned slab_id = 0; slab_id < set_associativity; slab_id++) {
|
|
// The warp tile read out a slab
|
|
read_key[slab_id] = keys[set_idx].set_[slab_id].slab_[lane_idx];
|
|
}
|
|
|
|
// Finish dumping the slabset, unlock the slabset
|
|
warp_unlock_mutex<mutex, warp_size>(warp_tile, set_mutex[set_idx]);
|
|
|
|
// Each lane(thread) within the warp tile calculate the offset to store its keys
|
|
uint32_t warp_tile_total_keys = 0;
|
|
for (unsigned slab_id = 0; slab_id < set_associativity; slab_id++) {
|
|
unsigned valid_mask = warp_tile.ballot(read_key[slab_id] != empty_key);
|
|
thread_key_offset[slab_id] =
|
|
__popc(valid_mask & ((1U << lane_idx) - 1U)) + warp_tile_total_keys;
|
|
warp_tile_total_keys = warp_tile_total_keys + __popc(valid_mask);
|
|
}
|
|
|
|
// Each warp tile request a unique place from the block-level counter
|
|
if (lane_idx == 0) {
|
|
warp_key_offset = atomicAdd(&block_acc, warp_tile_total_keys);
|
|
}
|
|
warp_key_offset = warp_tile.shfl(warp_key_offset, 0);
|
|
}
|
|
|
|
// Each block request a unique place in global memory output buffer
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
block_key_offset = atomicAdd(d_dump_counter, (size_t)block_acc);
|
|
}
|
|
__syncthreads();
|
|
|
|
// Warp tile store the (non-empty)keys back to output buffer
|
|
if (set_idx < end_set_index) {
|
|
for (unsigned slab_id = 0; slab_id < set_associativity; slab_id++) {
|
|
if (read_key[slab_id] != empty_key) {
|
|
d_keys[block_key_offset + warp_key_offset + thread_key_offset[slab_id]] = read_key[slab_id];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
#else
|
|
template <typename key_type, typename slabset, key_type empty_key, int set_associativity,
|
|
int warp_size>
|
|
__global__ void dump_kernel(key_type* d_keys, size_t* d_dump_counter, volatile slabset* keys,
|
|
volatile int* set_mutex, const size_t start_set_index,
|
|
const size_t end_set_index) {
|
|
// Block-level counter used by all warp tiles within a block
|
|
__shared__ uint32_t block_acc;
|
|
// Initialize block-level counter
|
|
if (threadIdx.x == 0) {
|
|
block_acc = 0;
|
|
}
|
|
__syncthreads();
|
|
// Lane(thread) ID within a warp tile
|
|
cg::thread_block_tile<warp_size> warp_tile =
|
|
cg::tiled_partition<warp_size>(cg::this_thread_block());
|
|
const size_t lane_idx = warp_tile.thread_rank();
|
|
// Warp tile target slabset id
|
|
const size_t set_idx =
|
|
((blockIdx.x * (blockDim.x / warp_size)) + warp_tile.meta_group_rank()) + start_set_index;
|
|
// Keys dump from cache
|
|
key_type read_key[set_associativity];
|
|
// Lane(thread) offset for storing each key
|
|
uint32_t thread_key_offset[set_associativity];
|
|
// Warp offset for storing each key
|
|
uint32_t warp_key_offset;
|
|
// Block offset for storing each key
|
|
__shared__ size_t block_key_offset;
|
|
|
|
// Warp tile dump target slabset
|
|
if (set_idx < end_set_index) {
|
|
// Lock the slabset before operating the slabset
|
|
warp_lock_mutex<warp_size>(warp_tile, set_mutex[set_idx]);
|
|
|
|
// The warp tile read out the slabset
|
|
for (unsigned slab_id = 0; slab_id < set_associativity; slab_id++) {
|
|
// The warp tile read out a slab
|
|
read_key[slab_id] = ((volatile key_type*)(keys[set_idx].set_[slab_id].slab_))[lane_idx];
|
|
}
|
|
|
|
// Finish dumping the slabset, unlock the slabset
|
|
warp_unlock_mutex<warp_size>(warp_tile, set_mutex[set_idx]);
|
|
|
|
// Each lane(thread) within the warp tile calculate the offset to store its keys
|
|
uint32_t warp_tile_total_keys = 0;
|
|
for (unsigned slab_id = 0; slab_id < set_associativity; slab_id++) {
|
|
unsigned valid_mask = warp_tile.ballot(read_key[slab_id] != empty_key);
|
|
thread_key_offset[slab_id] =
|
|
__popc(valid_mask & ((1U << lane_idx) - 1U)) + warp_tile_total_keys;
|
|
warp_tile_total_keys = warp_tile_total_keys + __popc(valid_mask);
|
|
}
|
|
|
|
// Each warp tile request a unique place from the block-level counter
|
|
if (lane_idx == 0) {
|
|
warp_key_offset = atomicAdd(&block_acc, warp_tile_total_keys);
|
|
}
|
|
warp_key_offset = warp_tile.shfl(warp_key_offset, 0);
|
|
}
|
|
|
|
// Each block request a unique place in global memory output buffer
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
block_key_offset = atomicAdd(d_dump_counter, (size_t)block_acc);
|
|
}
|
|
__syncthreads();
|
|
|
|
// Warp tile store the (non-empty)keys back to output buffer
|
|
if (set_idx < end_set_index) {
|
|
for (unsigned slab_id = 0; slab_id < set_associativity; slab_id++) {
|
|
if (read_key[slab_id] != empty_key) {
|
|
d_keys[block_key_offset + warp_key_offset + thread_key_offset[slab_id]] = read_key[slab_id];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
#endif
|
|
///////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::gpu_cache(const size_t capacity_in_set, const size_t embedding_vec_size)
|
|
: capacity_in_set_(capacity_in_set), embedding_vec_size_(embedding_vec_size) {
|
|
// Check parameter
|
|
if (capacity_in_set_ == 0) {
|
|
printf("Error: Invalid value for capacity_in_set.\n");
|
|
return;
|
|
}
|
|
if (embedding_vec_size_ == 0) {
|
|
printf("Error: Invalid value for embedding_vec_size.\n");
|
|
return;
|
|
}
|
|
if (set_associativity <= 0) {
|
|
printf("Error: Invalid value for set_associativity.\n");
|
|
return;
|
|
}
|
|
if (warp_size != 1 && warp_size != 2 && warp_size != 4 && warp_size != 8 && warp_size != 16 &&
|
|
warp_size != 32) {
|
|
printf("Error: Invalid value for warp_size.\n");
|
|
return;
|
|
}
|
|
|
|
// Get the current CUDA dev
|
|
CUDA_CHECK(cudaGetDevice(&dev_));
|
|
|
|
// Calculate # of slot
|
|
num_slot_ = capacity_in_set_ * set_associativity * warp_size;
|
|
|
|
// Allocate GPU memory for cache
|
|
CUDA_CHECK(cudaMalloc((void**)&keys_, sizeof(slabset) * capacity_in_set_));
|
|
CUDA_CHECK(cudaMalloc((void**)&vals_, sizeof(float) * embedding_vec_size_ * num_slot_));
|
|
CUDA_CHECK(cudaMalloc((void**)&slot_counter_, sizeof(ref_counter_type) * num_slot_));
|
|
CUDA_CHECK(cudaMalloc((void**)&global_counter_, sizeof(atomic_ref_counter_type)));
|
|
|
|
// Allocate GPU memory for set mutex
|
|
CUDA_CHECK(cudaMalloc((void**)&set_mutex_, sizeof(mutex) * capacity_in_set_));
|
|
|
|
// Initialize the cache, set all entry to unused <K,V>
|
|
init_cache<<<((num_slot_ - 1) / BLOCK_SIZE_) + 1, BLOCK_SIZE_>>>(
|
|
keys_, slot_counter_, global_counter_, num_slot_, empty_key, set_mutex_, capacity_in_set_);
|
|
|
|
// Wait for initialization to finish
|
|
CUDA_CHECK(cudaStreamSynchronize(0));
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#else
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::gpu_cache(const size_t capacity_in_set, const size_t embedding_vec_size)
|
|
: capacity_in_set_(capacity_in_set), embedding_vec_size_(embedding_vec_size) {
|
|
// Check parameter
|
|
if (capacity_in_set_ == 0) {
|
|
printf("Error: Invalid value for capacity_in_set.\n");
|
|
return;
|
|
}
|
|
if (embedding_vec_size_ == 0) {
|
|
printf("Error: Invalid value for embedding_vec_size.\n");
|
|
return;
|
|
}
|
|
if (set_associativity <= 0) {
|
|
printf("Error: Invalid value for set_associativity.\n");
|
|
return;
|
|
}
|
|
if (warp_size != 1 && warp_size != 2 && warp_size != 4 && warp_size != 8 && warp_size != 16 &&
|
|
warp_size != 32) {
|
|
printf("Error: Invalid value for warp_size.\n");
|
|
return;
|
|
}
|
|
|
|
// Get the current CUDA dev
|
|
CUDA_CHECK(cudaGetDevice(&dev_));
|
|
|
|
// Calculate # of slot
|
|
num_slot_ = capacity_in_set_ * set_associativity * warp_size;
|
|
|
|
// Allocate GPU memory for cache
|
|
CUDA_CHECK(cudaMalloc((void**)&keys_, sizeof(slabset) * capacity_in_set_));
|
|
CUDA_CHECK(cudaMalloc((void**)&vals_, sizeof(float) * embedding_vec_size_ * num_slot_));
|
|
CUDA_CHECK(cudaMalloc((void**)&slot_counter_, sizeof(ref_counter_type) * num_slot_));
|
|
CUDA_CHECK(cudaMalloc((void**)&global_counter_, sizeof(ref_counter_type)));
|
|
|
|
// Allocate GPU memory for set mutex
|
|
CUDA_CHECK(cudaMalloc((void**)&set_mutex_, sizeof(int) * capacity_in_set_));
|
|
|
|
// Initialize the cache, set all entry to unused <K,V>
|
|
init_cache<<<((num_slot_ - 1) / BLOCK_SIZE_) + 1, BLOCK_SIZE_>>>(
|
|
keys_, slot_counter_, global_counter_, num_slot_, empty_key, set_mutex_, capacity_in_set_);
|
|
|
|
// Wait for initialization to finish
|
|
CUDA_CHECK(cudaStreamSynchronize(0));
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::~gpu_cache() {
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Destruct CUDA std object
|
|
destruct_kernel<<<((capacity_in_set_ - 1) / BLOCK_SIZE_) + 1, BLOCK_SIZE_>>>(
|
|
global_counter_, set_mutex_, capacity_in_set_);
|
|
// Wait for destruction to finish
|
|
CUDA_CHECK(cudaStreamSynchronize(0));
|
|
|
|
// Free GPU memory for cache
|
|
CUDA_CHECK(cudaFree(keys_));
|
|
CUDA_CHECK(cudaFree(vals_));
|
|
CUDA_CHECK(cudaFree(slot_counter_));
|
|
CUDA_CHECK(cudaFree(global_counter_));
|
|
// Free GPU memory for set mutex
|
|
CUDA_CHECK(cudaFree(set_mutex_));
|
|
}
|
|
#else
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::~gpu_cache() noexcept(false) {
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Free GPU memory for cache
|
|
CUDA_CHECK(cudaFree(keys_));
|
|
CUDA_CHECK(cudaFree(vals_));
|
|
CUDA_CHECK(cudaFree(slot_counter_));
|
|
CUDA_CHECK(cudaFree(global_counter_));
|
|
// Free GPU memory for set mutex
|
|
CUDA_CHECK(cudaFree(set_mutex_));
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
void gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::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) {
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Check if it is a valid query
|
|
if (len == 0) {
|
|
// Set the d_missing_len to 0 before return
|
|
CUDA_CHECK(cudaMemsetAsync(d_missing_len, 0, sizeof(size_t), stream));
|
|
return;
|
|
}
|
|
|
|
// Update the global counter as user perform a new(most recent) read operation to the cache
|
|
// Resolve distance overflow issue as well.
|
|
update_kernel_overflow_ignore<atomic_ref_counter_type>
|
|
<<<1, 1, 0, stream>>>(global_counter_, d_missing_len);
|
|
|
|
// Read from the cache
|
|
// Touch and refresh the hitting slot
|
|
const size_t keys_per_block = (BLOCK_SIZE_ / warp_size) * task_per_warp_tile;
|
|
const size_t grid_size = ((len - 1) / keys_per_block) + 1;
|
|
get_kernel<key_type, ref_counter_type, atomic_ref_counter_type, slabset, set_hasher, slab_hasher,
|
|
mutex, empty_key, set_associativity, warp_size><<<grid_size, BLOCK_SIZE_, 0, stream>>>(
|
|
d_keys, len, d_values, embedding_vec_size_, d_missing_index, d_missing_keys, d_missing_len,
|
|
global_counter_, slot_counter_, capacity_in_set_, keys_, vals_, set_mutex_,
|
|
task_per_warp_tile);
|
|
|
|
// Check for GPU error before return
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#else
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
void gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::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) {
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Check if it is a valid query
|
|
if (len == 0) {
|
|
// Set the d_missing_len to 0 before return
|
|
CUDA_CHECK(cudaMemsetAsync(d_missing_len, 0, sizeof(size_t), stream));
|
|
return;
|
|
}
|
|
|
|
// Update the global counter as user perform a new(most recent) read operation to the cache
|
|
// Resolve distance overflow issue as well.
|
|
update_kernel_overflow_ignore<ref_counter_type>
|
|
<<<1, 1, 0, stream>>>(global_counter_, d_missing_len);
|
|
|
|
// Read from the cache
|
|
// Touch and refresh the hitting slot
|
|
const size_t keys_per_block = (BLOCK_SIZE_ / warp_size) * task_per_warp_tile;
|
|
const size_t grid_size = ((len - 1) / keys_per_block) + 1;
|
|
get_kernel<key_type, ref_counter_type, slabset, set_hasher, slab_hasher, empty_key,
|
|
set_associativity, warp_size><<<grid_size, BLOCK_SIZE_, 0, stream>>>(
|
|
d_keys, len, d_values, embedding_vec_size_, d_missing_index, d_missing_keys, d_missing_len,
|
|
global_counter_, slot_counter_, capacity_in_set_, keys_, vals_, set_mutex_,
|
|
task_per_warp_tile);
|
|
|
|
// Check for GPU error before return
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
void gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::Replace(const key_type* d_keys, const size_t len,
|
|
const float* d_values, cudaStream_t stream,
|
|
const size_t task_per_warp_tile) {
|
|
// Check if it is a valid replacement
|
|
if (len == 0) {
|
|
return;
|
|
}
|
|
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Try to insert the <k,v> paris into the cache as long as there are unused slot
|
|
// Then replace the <k,v> pairs into the cache
|
|
const size_t keys_per_block = (BLOCK_SIZE_ / warp_size) * task_per_warp_tile;
|
|
const size_t grid_size = ((len - 1) / keys_per_block) + 1;
|
|
insert_replace_kernel<key_type, slabset, ref_counter_type, mutex, atomic_ref_counter_type,
|
|
set_hasher, slab_hasher, empty_key, set_associativity, warp_size>
|
|
<<<grid_size, BLOCK_SIZE_, 0, stream>>>(d_keys, d_values, embedding_vec_size_, len, keys_,
|
|
vals_, slot_counter_, set_mutex_, global_counter_,
|
|
capacity_in_set_, task_per_warp_tile);
|
|
|
|
// Check for GPU error before return
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#else
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
void gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::Replace(const key_type* d_keys, const size_t len,
|
|
const float* d_values, cudaStream_t stream,
|
|
const size_t task_per_warp_tile) {
|
|
// Check if it is a valid replacement
|
|
if (len == 0) {
|
|
return;
|
|
}
|
|
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Try to insert the <k,v> paris into the cache as long as there are unused slot
|
|
// Then replace the <k,v> pairs into the cache
|
|
const size_t keys_per_block = (BLOCK_SIZE_ / warp_size) * task_per_warp_tile;
|
|
const size_t grid_size = ((len - 1) / keys_per_block) + 1;
|
|
insert_replace_kernel<key_type, slabset, ref_counter_type, set_hasher, slab_hasher, empty_key,
|
|
set_associativity, warp_size><<<grid_size, BLOCK_SIZE_, 0, stream>>>(
|
|
d_keys, d_values, embedding_vec_size_, len, keys_, vals_, slot_counter_, set_mutex_,
|
|
global_counter_, capacity_in_set_, task_per_warp_tile);
|
|
|
|
// Check for GPU error before return
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
void gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::Update(const key_type* d_keys, const size_t len, const float* d_values,
|
|
cudaStream_t stream, const size_t task_per_warp_tile) {
|
|
// Check if it is a valid update request
|
|
if (len == 0) {
|
|
return;
|
|
}
|
|
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Update the value of input keys that are existed in the cache
|
|
const size_t keys_per_block = (BLOCK_SIZE_ / warp_size) * task_per_warp_tile;
|
|
const size_t grid_size = ((len - 1) / keys_per_block) + 1;
|
|
update_kernel<key_type, slabset, set_hasher, slab_hasher, mutex, empty_key, set_associativity,
|
|
warp_size><<<grid_size, BLOCK_SIZE_, 0, stream>>>(
|
|
d_keys, len, d_values, embedding_vec_size_, capacity_in_set_, keys_, vals_, set_mutex_,
|
|
task_per_warp_tile);
|
|
|
|
// Check for GPU error before return
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#else
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
void gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::Update(const key_type* d_keys, const size_t len, const float* d_values,
|
|
cudaStream_t stream, const size_t task_per_warp_tile) {
|
|
// Check if it is a valid update request
|
|
if (len == 0) {
|
|
return;
|
|
}
|
|
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Update the value of input keys that are existed in the cache
|
|
const size_t keys_per_block = (BLOCK_SIZE_ / warp_size) * task_per_warp_tile;
|
|
const size_t grid_size = ((len - 1) / keys_per_block) + 1;
|
|
update_kernel<key_type, slabset, set_hasher, slab_hasher, empty_key, set_associativity, warp_size>
|
|
<<<grid_size, BLOCK_SIZE_, 0, stream>>>(d_keys, len, d_values, embedding_vec_size_,
|
|
capacity_in_set_, keys_, vals_, set_mutex_,
|
|
task_per_warp_tile);
|
|
|
|
// Check for GPU error before return
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#endif
|
|
|
|
#ifdef LIBCUDACXX_VERSION
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
void gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::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) {
|
|
// Check if it is a valid dump request
|
|
if (start_set_index >= capacity_in_set_) {
|
|
printf("Error: Invalid value for start_set_index. Nothing dumped.\n");
|
|
return;
|
|
}
|
|
if (end_set_index <= start_set_index || end_set_index > capacity_in_set_) {
|
|
printf("Error: Invalid value for end_set_index. Nothing dumped.\n");
|
|
return;
|
|
}
|
|
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Set the global counter to 0 first
|
|
CUDA_CHECK(cudaMemsetAsync(d_dump_counter, 0, sizeof(size_t), stream));
|
|
|
|
// Dump keys from the cache
|
|
const size_t grid_size =
|
|
(((end_set_index - start_set_index) - 1) / (BLOCK_SIZE_ / warp_size)) + 1;
|
|
dump_kernel<key_type, slabset, mutex, empty_key, set_associativity, warp_size>
|
|
<<<grid_size, BLOCK_SIZE_, 0, stream>>>(d_keys, d_dump_counter, keys_, set_mutex_,
|
|
start_set_index, end_set_index);
|
|
|
|
// Check for GPU error before return
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#else
|
|
template <typename key_type, typename ref_counter_type, key_type empty_key, int set_associativity,
|
|
int warp_size, typename set_hasher, typename slab_hasher>
|
|
void gpu_cache<key_type, ref_counter_type, empty_key, set_associativity, warp_size, set_hasher,
|
|
slab_hasher>::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) {
|
|
// Check if it is a valid dump request
|
|
if (start_set_index >= capacity_in_set_) {
|
|
printf("Error: Invalid value for start_set_index. Nothing dumped.\n");
|
|
return;
|
|
}
|
|
if (end_set_index <= start_set_index || end_set_index > capacity_in_set_) {
|
|
printf("Error: Invalid value for end_set_index. Nothing dumped.\n");
|
|
return;
|
|
}
|
|
|
|
// Device Restorer
|
|
nv::CudaDeviceRestorer dev_restorer;
|
|
// Check device
|
|
dev_restorer.check_device(dev_);
|
|
|
|
// Set the global counter to 0 first
|
|
CUDA_CHECK(cudaMemsetAsync(d_dump_counter, 0, sizeof(size_t), stream));
|
|
|
|
// Dump keys from the cache
|
|
const size_t grid_size =
|
|
(((end_set_index - start_set_index) - 1) / (BLOCK_SIZE_ / warp_size)) + 1;
|
|
dump_kernel<key_type, slabset, empty_key, set_associativity, warp_size>
|
|
<<<grid_size, BLOCK_SIZE_, 0, stream>>>(d_keys, d_dump_counter, keys_, set_mutex_,
|
|
start_set_index, end_set_index);
|
|
|
|
// Check for GPU error before return
|
|
CUDA_CHECK(cudaGetLastError());
|
|
}
|
|
#endif
|
|
|
|
template class gpu_cache<unsigned int, uint64_t, std::numeric_limits<unsigned int>::max(),
|
|
SET_ASSOCIATIVITY, SLAB_SIZE>;
|
|
template class gpu_cache<long long, uint64_t, std::numeric_limits<long long>::max(),
|
|
SET_ASSOCIATIVITY, SLAB_SIZE>;
|
|
} // namespace gpu_cache
|