Aluode/PerceptionLabPortable

1// File from https://github.com/mlpen/YOSO/blob/main/encoders/backbones/efficient_attentions/yoso/yoso_v1/cuda/fast_lsh_cumulation_cuda.cu2 3#include "fast_lsh_cumulation_cuda.h"4#include "common_cuda_device.h"5#include "common_cuda.h"6#include "common.h"7#include <stdio.h>8//////////////////////////////////////////////////////////////////////////////////////////////////9//////////////////////////////////////////////////////////////////////////////////////////////////10 11inline __device__ void fast_hadamard_transform(float *vector_buffer, int vector_dim, int dim_idx) {12  int stride = vector_dim / 2;13  while (stride > (WARP_SIZE / 2)) {14    __syncthreads();15    int sign = 1 - ((dim_idx / stride) % 2) * 2;16    float val1 = vector_buffer[dim_idx];17    float val2 = vector_buffer[dim_idx + sign * stride];18    __syncthreads();19    vector_buffer[dim_idx] = float(sign) * val1 + val2;20    stride = stride / 2;21  }22 23  float val = vector_buffer[dim_idx];24  #pragma unroll25  for (stride = (WARP_SIZE / 2); stride > 0; stride = stride / 2) {26    int sign = 1 - ((dim_idx / stride) % 2) * 2;27    val = float(sign) * val + __shfl_xor_sync(FULL_MASK, val, stride);28  }29  vector_buffer[dim_idx] = val;30}31 32__global__ void fast_hash_ver1_cuda_kernel(33  int *mask,        // [batch_size, num_vector]34  float *vector,    // [batch_size, num_vector, vector_dim]35  int *Dmat,        // [batch_size, 3, num_part, vector_dim]36  int *hash_code,   // [batch_size, num_vector, num_hash_f]37  int batch_size,38  int num_vector,39  int vector_dim,40  int num_part,41  int num_hash_f,42  int hash_code_len43) {44 45  int batch_idx = blockIdx.z;46  int vector_idx = blockIdx.y;47  int part_idx = blockIdx.x;48 49  int dim_idx = threadIdx.x;50 51  int batch_idx__vector_idx = batch_idx * num_vector + vector_idx;52  if (mask[batch_idx__vector_idx] == 0) {53    return;54  }55 56  extern __shared__ float buffer[];57  float *vector_buffer = buffer;58 59  vector_buffer[dim_idx] = vector[batch_idx__vector_idx * vector_dim + dim_idx];60 61  vector_buffer[dim_idx] = vector_buffer[dim_idx] * (float)Dmat[((batch_idx * 3 + 0) * num_part + part_idx) * vector_dim + dim_idx];62  fast_hadamard_transform(vector_buffer, vector_dim, dim_idx);63  vector_buffer[dim_idx] = vector_buffer[dim_idx] * (float)Dmat[((batch_idx * 3 + 1) * num_part + part_idx) * vector_dim + dim_idx];64  fast_hadamard_transform(vector_buffer, vector_dim, dim_idx);65  vector_buffer[dim_idx] = vector_buffer[dim_idx] * (float)Dmat[((batch_idx * 3 + 2) * num_part + part_idx) * vector_dim + dim_idx];66  fast_hadamard_transform(vector_buffer, vector_dim, dim_idx);67 68  int num_hash_per_part = vector_dim / hash_code_len;69  if (hash_code_len == 8 || hash_code_len == 16) {70    int code = select(vector_buffer[dim_idx] > 0, 1 << (dim_idx % hash_code_len), 0);71    for (int offset = 1; offset < hash_code_len; offset = offset * 2) {72      code += __shfl_xor_sync(FULL_MASK, code, offset);73    }74    if (dim_idx % hash_code_len == 0) {75      int hash_f_idx = part_idx * num_hash_per_part + dim_idx / hash_code_len;76      if (hash_f_idx < num_hash_f) {77        hash_code[batch_idx__vector_idx * num_hash_f + hash_f_idx] = code;78      }79    }80  } else {81    vector_buffer[dim_idx] = select(vector_buffer[dim_idx] > 0, 1 << (dim_idx % hash_code_len), 0);82    __syncthreads();83    if (dim_idx < num_hash_per_part) {84      int code = 0;85      for (int i = 0; i < hash_code_len; i++) {86        code += vector_buffer[dim_idx * hash_code_len + i];87      }88      int hash_f_idx = part_idx * num_hash_per_part + dim_idx;89      if (hash_f_idx < num_hash_f) {90        hash_code[batch_idx__vector_idx * num_hash_f + hash_f_idx] = code;91      }92    }93  }94}95 96__global__ void lsh_cumulation_ver1_step1_cuda_kernel(97  int *key_mask,           // [batch_size, num_key]98  int *key_hash_code,      // [batch_size, num_key, num_hash_f]99  float *value,            // [batch_size, num_key, value_dim]100  float *hashtable_value,  // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE]101  int batch_size,102  int num_hash_f,103  int hashtable_capacity,104  int num_key,105  int value_dim,106  int offset_warp107) {108 109  int warp_thread_idx = threadIdx.x;110 111  int batch_idx = blockIdx.y;112  int key_idx = blockIdx.x * blockDim.y + threadIdx.y;113 114  int batch_idx__key_idx = batch_idx * num_key + key_idx;115  if (key_mask[batch_idx__key_idx] == 0) {116    return;117  }118 119  if (num_hash_f > WARP_SIZE) {120    float warp_value = value[batch_idx__key_idx * value_dim + offset_warp + warp_thread_idx];121    for (int hash_f_start = 0; hash_f_start < num_hash_f; hash_f_start = hash_f_start + WARP_SIZE) {122      int warp_hashcode = key_hash_code[batch_idx__key_idx * num_hash_f + hash_f_start + warp_thread_idx];123      #pragma unroll124      for (int hash_f_offset = 0; hash_f_offset < WARP_SIZE; hash_f_offset++) {125        int current_hashcode = warp_hashcode;126        current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_offset);127        int hashtable_idx = (batch_idx * num_hash_f + (hash_f_start + hash_f_offset)) * hashtable_capacity + current_hashcode;128        atomicAdd(&hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx], warp_value);129      }130    }131  } else {132    float warp_value = value[batch_idx__key_idx * value_dim + offset_warp + warp_thread_idx];133    int warp_hashcode = 0;134    if (warp_thread_idx < num_hash_f) {135      warp_hashcode = key_hash_code[batch_idx__key_idx * num_hash_f + warp_thread_idx];136    }137    for (int hash_f_idx = 0; hash_f_idx < num_hash_f; hash_f_idx++) {138      int current_hashcode = warp_hashcode;139      current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_idx);140      int hashtable_idx = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + current_hashcode;141      atomicAdd(&hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx], warp_value);142    }143  }144 145}146 147__global__ void lsh_cumulation_ver1_step2_cuda_kernel(148  int *query_mask,         // [batch_size, num_query]149  int *query_hash_code,    // [batch_size, num_query, num_hash_f]150  float *hashtable_value,  // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE]151  float *cumulation_value, // [batch_size, num_query, value_dim]152  int batch_size,153  int num_hash_f,154  int hashtable_capacity,155  int num_query,156  int value_dim,157  int offset_warp158) {159 160  int warp_thread_idx = threadIdx.x;161 162  int batch_idx = blockIdx.y;163  int query_idx = blockIdx.x * blockDim.y + threadIdx.y;164 165  int batch_idx__query_idx = batch_idx * num_query + query_idx;166  if (query_mask[batch_idx__query_idx] == 0) {167    return;168  }169 170  if (num_hash_f > WARP_SIZE) {171    float warp_value = 0;172    for (int hash_f_start = 0; hash_f_start < num_hash_f; hash_f_start = hash_f_start + WARP_SIZE) {173      int warp_hashcode = query_hash_code[batch_idx__query_idx * num_hash_f + hash_f_start + warp_thread_idx];174      #pragma unroll175      for (int hash_f_offset = 0; hash_f_offset < WARP_SIZE; hash_f_offset++) {176        int current_hashcode = warp_hashcode;177        current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_offset);178        int hashtable_idx = (batch_idx * num_hash_f + (hash_f_start + hash_f_offset)) * hashtable_capacity + current_hashcode;179        warp_value = warp_value + hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx];180      }181    }182    cumulation_value[batch_idx__query_idx * value_dim + offset_warp + warp_thread_idx] = warp_value / float(num_hash_f);183  } else {184    float warp_value = 0;185    int warp_hashcode = 0;186    if (warp_thread_idx < num_hash_f) {187      warp_hashcode = query_hash_code[batch_idx__query_idx * num_hash_f + warp_thread_idx];188    }189    for (int hash_f_idx = 0; hash_f_idx < num_hash_f; hash_f_idx++) {190      int current_hashcode = warp_hashcode;191      current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_idx);192      int hashtable_idx = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + current_hashcode;193      warp_value = warp_value + hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx];194    }195    cumulation_value[batch_idx__query_idx * value_dim + offset_warp + warp_thread_idx] = warp_value / float(num_hash_f);196  }197 198}199 200__global__ void lsh_weighted_cumulation_ver1_step1_cuda_kernel(201  int *key_mask,            // [batch_size, num_key]202  int *key_hash_code,       // [batch_size, num_key, num_hash_f]203  float *key_weight,        // [batch_size, num_key, weight_dim]204  float *value,             // [batch_size, num_key, value_dim]205  float *hashtable_value,   // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE]206  int batch_size,207  int num_hash_f,208  int hashtable_capacity,209  int num_key,210  int value_dim,211  int weight_dim,212  int offset_warp,213  int weight_idx214) {215 216  int warp_thread_idx = threadIdx.x;217 218  int batch_idx = blockIdx.y;219  int key_idx = blockIdx.x * blockDim.y + threadIdx.y;220 221  int batch_idx__key_idx = batch_idx * num_key + key_idx;222  if (key_mask[batch_idx__key_idx] == 0) {223    return;224  }225 226  if (num_hash_f > WARP_SIZE) {227    float warp_value = key_weight[batch_idx__key_idx * weight_dim + weight_idx] * value[batch_idx__key_idx * value_dim + offset_warp + warp_thread_idx];228    for (int hash_f_start = 0; hash_f_start < num_hash_f; hash_f_start = hash_f_start + WARP_SIZE) {229      int warp_hashcode = key_hash_code[batch_idx__key_idx * num_hash_f + hash_f_start + warp_thread_idx];230      #pragma unroll231      for (int hash_f_offset = 0; hash_f_offset < WARP_SIZE; hash_f_offset++) {232        int current_hashcode = warp_hashcode;233        current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_offset);234        int hashtable_idx = (batch_idx * num_hash_f + (hash_f_start + hash_f_offset)) * hashtable_capacity + current_hashcode;235        atomicAdd(&hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx], warp_value);236      }237    }238  } else {239    float warp_value = key_weight[batch_idx__key_idx * weight_dim + weight_idx] * value[batch_idx__key_idx * value_dim + offset_warp + warp_thread_idx];240    int warp_hashcode = 0;241    if (warp_thread_idx < num_hash_f) {242      warp_hashcode = key_hash_code[batch_idx__key_idx * num_hash_f + warp_thread_idx];243    }244    for (int hash_f_idx = 0; hash_f_idx < num_hash_f; hash_f_idx++) {245      int current_hashcode = warp_hashcode;246      current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_idx);247      int hashtable_idx = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + current_hashcode;248      atomicAdd(&hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx], warp_value);249    }250  }251 252}253 254__global__ void lsh_weighted_cumulation_ver1_step2_cuda_kernel(255  int *query_mask,          // [batch_size, num_query]256  int *query_hash_code,     // [batch_size, num_query, num_hash_f]257  float *query_weight,      // [batch_size, num_query, weight_dim]258  float *hashtable_value,   // [batch_size, num_hash_f, hashtable_capacity, WARP_SIZE]259  float *cumulation_value,  // [batch_size, num_query, value_dim]260  int batch_size,261  int num_hash_f,262  int hashtable_capacity,263  int num_query,264  int value_dim,265  int weight_dim,266  int offset_warp,267  int weight_idx268) {269 270  int warp_thread_idx = threadIdx.x;271 272  int batch_idx = blockIdx.y;273  int query_idx = blockIdx.x * blockDim.y + threadIdx.y;274 275  int batch_idx__query_idx = batch_idx * num_query + query_idx;276  if (query_mask[batch_idx__query_idx] == 0) {277    return;278  }279 280  if (num_hash_f > WARP_SIZE) {281    float warp_value = 0;282    for (int hash_f_start = 0; hash_f_start < num_hash_f; hash_f_start = hash_f_start + WARP_SIZE) {283      int warp_hashcode = query_hash_code[batch_idx__query_idx * num_hash_f + hash_f_start + warp_thread_idx];284      #pragma unroll285      for (int hash_f_offset = 0; hash_f_offset < WARP_SIZE; hash_f_offset++) {286        int current_hashcode = warp_hashcode;287        current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_offset);288        int hashtable_idx = (batch_idx * num_hash_f + (hash_f_start + hash_f_offset)) * hashtable_capacity + current_hashcode;289        warp_value = warp_value + hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx];290      }291    }292    float warp_weight = query_weight[batch_idx__query_idx * weight_dim + weight_idx];293    cumulation_value[batch_idx__query_idx * value_dim + offset_warp + warp_thread_idx] += warp_weight * warp_value / float(num_hash_f);294  } else {295    float warp_value = 0;296    int warp_hashcode = 0;297    if (warp_thread_idx < num_hash_f) {298      warp_hashcode = query_hash_code[batch_idx__query_idx * num_hash_f + warp_thread_idx];299    }300    for (int hash_f_idx = 0; hash_f_idx < num_hash_f; hash_f_idx++) {301      int current_hashcode = warp_hashcode;302      current_hashcode = __shfl_sync(FULL_MASK, current_hashcode, hash_f_idx);303      int hashtable_idx = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + current_hashcode;304      warp_value = warp_value + hashtable_value[hashtable_idx * WARP_SIZE + warp_thread_idx];305    }306    float warp_weight = query_weight[batch_idx__query_idx * weight_dim + weight_idx];307    cumulation_value[batch_idx__query_idx * value_dim + offset_warp + warp_thread_idx] += warp_weight * warp_value / float(num_hash_f);308  }309 310}311 312__global__ void count_sort_step1_cuda_kernel(313  int *key_mask,         // [batch_size, num_key]314  int *key_hash_code,    // [batch_size, num_key, num_hash_f]315  int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity]316  int batch_size,317  int num_hash_f,318  int hashtable_capacity,319  int num_key320) {321 322  int batch_idx = blockIdx.y;323  int key_idx = blockIdx.x * blockDim.y + threadIdx.y;324  int hash_f_idx = threadIdx.x;325 326  int batch_idx__key_idx = batch_idx * num_key + key_idx;327  if (key_mask[batch_idx__key_idx] == 0) {328    return;329  }330 331  int hash_code = key_hash_code[batch_idx__key_idx * num_hash_f + hash_f_idx];332  atomicAdd(&count_sort_table[(batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + hash_code], 1);333 334}335 336__global__ void count_sort_step2_cuda_kernel(337  int *count_sort_table,  // [batch_size, num_hash_f, hashtable_capacity]338  int batch_size,339  int num_hash_f,340  int hashtable_capacity341) {342 343  int batch_idx = blockIdx.y;344  int hash_f_idx = blockIdx.x;345 346  int num_threads = blockDim.x;347  int thread_id = threadIdx.x;348 349  int batch_idx__hash_f_idx = batch_idx * num_hash_f + hash_f_idx;350 351  extern __shared__ float buffer[];352  int *table_buffer = (int*)buffer;353 354  if (thread_id == 0) {355    table_buffer[0] = 0;356  }357  copy_data<int>(&count_sort_table[batch_idx__hash_f_idx * hashtable_capacity], &table_buffer[1], hashtable_capacity - 1, num_threads, thread_id);358 359  for (int table_idx_start = 0; table_idx_start < hashtable_capacity; table_idx_start = table_idx_start + num_threads) {360    int thread_value = table_buffer[table_idx_start + thread_id];361    int next_thread_value = 0;362    for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) {363      next_thread_value = __shfl_up_sync(FULL_MASK, thread_value, offset);364      if (thread_id % WARP_SIZE >= offset) {365        thread_value = thread_value + next_thread_value;366      }367    }368    table_buffer[table_idx_start + thread_id] = thread_value;369  }370  __syncthreads();371 372  if (hashtable_capacity > WARP_SIZE) {373    if (thread_id < WARP_SIZE) {374      for (int table_idx_start = WARP_SIZE; table_idx_start < hashtable_capacity; table_idx_start = table_idx_start + WARP_SIZE) {375        table_buffer[table_idx_start + thread_id] += table_buffer[table_idx_start - 1];376      }377    }378  }379 380  copy_data<int>(table_buffer, &count_sort_table[batch_idx__hash_f_idx * hashtable_capacity], hashtable_capacity, num_threads, thread_id);381 382}383 384 385__global__ void count_sort_step3_cuda_kernel(386  int *key_mask,          // [batch_size, num_key]387  int *key_hash_code,     // [batch_size, num_key, num_hash_f]388  int *count_sort_table,  // [batch_size, num_hash_f, hashtable_capacity]389  int *key_sorted_idxes,  // [batch_size, num_hash_f, num_key]390  int batch_size,391  int num_hash_f,392  int hashtable_capacity,393  int num_key394) {395 396  int batch_idx = blockIdx.y;397  int key_idx = blockIdx.x * blockDim.y + threadIdx.y;398  int hash_f_idx = threadIdx.x;399 400  int batch_idx__key_idx = batch_idx * num_key + key_idx;401  if (key_mask[batch_idx__key_idx] == 0) {402    return;403  }404 405  int batch_idx__hash_f_idx = batch_idx * num_hash_f + hash_f_idx;406 407  int hash_code = key_hash_code[batch_idx__key_idx * num_hash_f + hash_f_idx];408  int sort_idx = atomicAdd(&count_sort_table[batch_idx__hash_f_idx * hashtable_capacity + hash_code], 1);409  key_sorted_idxes[batch_idx__hash_f_idx * num_key + sort_idx] = key_idx;410 411}412 413__global__ void extract_query_info_cuda_kernel(414  int *query_mask,       // [batch_size, num_query]415  int *query_hash_code,  // [batch_size, num_query, num_hash_f]416  int *count_sort_table, // [batch_size, num_hash_f, hashtable_capacity]417  int *query_info,       // [batch_size, num_query, 2, num_hash_f]418  int batch_size,419  int num_hash_f,420  int hashtable_capacity,421  int num_query422) {423 424  int batch_idx = blockIdx.y;425  int query_idx = blockIdx.x * blockDim.y + threadIdx.y;426  int hash_f_idx = threadIdx.x;427 428  int batch_idx__query_idx = batch_idx * num_query + query_idx;429  if (query_mask[batch_idx__query_idx] == 0) {430    return;431  }432 433  int hash_code = query_hash_code[batch_idx__query_idx * num_hash_f + hash_f_idx];434  int batch_idx__hash_f_idx__hash_code = (batch_idx * num_hash_f + hash_f_idx) * hashtable_capacity + hash_code;435 436  int key_offset = select(hash_code == 0, 0, count_sort_table[batch_idx__hash_f_idx__hash_code - 1]);437  int key_count = count_sort_table[batch_idx__hash_f_idx__hash_code] - key_offset;438 439  query_info[batch_idx__query_idx * 2 * num_hash_f + hash_f_idx] = key_offset;440  query_info[(batch_idx__query_idx * 2 + 1) * num_hash_f + hash_f_idx] = key_count;441 442}443 444__global__ void lsh_weighted_cumulation_ver2_step2_cuda_kernel(445  int *query_mask,         // [batch_size, num_query]446  int *query_info,         // [batch_size, num_query, 2, num_hash_f]447  int *key_sorted_idxes,   // [batch_size, num_hash_f, num_key]448  float *query_weight,     // [batch_size, num_query, weight_dim]449  float *key_weight,       // [batch_size, num_key, weight_dim]450  float *value,            // [batch_size, num_key, value_dim]451  float *cumulation_value, // [batch_size, num_query, value_dim]452  int batch_size,453  int num_hash_f,454  int num_query,455  int num_key,456  int value_dim,457  int weight_dim458) {459 460  int batch_idx = blockIdx.z;461  int hash_f_idx = blockIdx.y;462  int query_idx = blockIdx.x;463 464  int num_threads = blockDim.y * blockDim.x;465  int thread_id = threadIdx.y * blockDim.x + threadIdx.x;466 467  int num_warps = blockDim.y;468  int warp_idx = threadIdx.y;469  int warp_thread_idx = threadIdx.x;470 471  int batch_idx__query_idx = batch_idx * num_query + query_idx;472  if (query_mask[batch_idx__query_idx] == 0) {473    return;474  }475 476  int key_offset = query_info[batch_idx__query_idx * 2 * num_hash_f + hash_f_idx];477  int key_count = query_info[(batch_idx__query_idx * 2 + 1) * num_hash_f + hash_f_idx];478 479  if (key_count == 0) {480    return;481  }482 483  extern __shared__ float buffer[];484 485  if (key_count == 1) {486    if (warp_idx == 0) {487      int key_idx = key_sorted_idxes[(batch_idx * num_hash_f + hash_f_idx) * num_key + key_offset];488      int batch_idx__key_idx = batch_idx * num_key + key_idx;489      float weight = 0;490      for (int weight_offset = 0; weight_offset < weight_dim; weight_offset = weight_offset + WARP_SIZE) {491        int weight_dim_idx = weight_offset + warp_thread_idx;492        float val = query_weight[batch_idx__query_idx * weight_dim + weight_dim_idx] * key_weight[batch_idx__key_idx * weight_dim + weight_dim_idx];493        #pragma unroll494        for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) {495          val += __shfl_xor_sync(FULL_MASK, val, offset);496        }497        weight = weight + val;498      }499      weight = weight / float(num_hash_f);500      for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) {501        int value_dim_idx = value_offset + warp_thread_idx;502        float val = value[batch_idx__key_idx * value_dim + value_dim_idx];503        atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val);504      }505    }506  } else {507    float *weight_buffer = buffer;508    int *key_idxes_buffer = (int*)&buffer[weight_dim];509 510    copy_data_nonblocking<float>(&query_weight[batch_idx__query_idx * weight_dim], weight_buffer, weight_dim, num_threads, thread_id);511 512    while (key_count > 0) {513      int work_size = min(WARP_SIZE, key_count);514      copy_data_nonblocking<int>(&key_sorted_idxes[(batch_idx * num_hash_f + hash_f_idx) * num_key + key_offset], key_idxes_buffer, work_size, num_threads, thread_id);515      __syncthreads();516      for (int work_offset = 0; work_offset < WARP_SIZE; work_offset = work_offset + num_warps) {517        int work_idx = work_offset + warp_idx;518        if (work_idx < key_count) {519          int key_idx = key_idxes_buffer[work_idx];520          int batch_idx__key_idx = batch_idx * num_key + key_idx;521          float weight = 0;522          for (int weight_offset = 0; weight_offset < weight_dim; weight_offset = weight_offset + WARP_SIZE) {523            int weight_dim_idx = weight_offset + warp_thread_idx;524            float val = weight_buffer[weight_dim_idx] * key_weight[batch_idx__key_idx * weight_dim + weight_dim_idx];525            #pragma unroll526            for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) {527              val += __shfl_xor_sync(FULL_MASK, val, offset);528            }529            weight = weight + val;530          }531          weight = weight / float(num_hash_f);532          for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) {533            int value_dim_idx = value_offset + warp_thread_idx;534            float val = value[batch_idx__key_idx * value_dim + value_dim_idx];535            atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val);536          }537        }538      }539      key_count = key_count - work_size;540      key_offset = key_offset + work_size;541    }542  }543 544}545 546__global__ void lsh_weighted_cumulation_ver3_step2_cuda_kernel(547  int *query_sorted_idxes,   // [batch_size, num_hash_f, num_query]548  int *key_mask,             // [batch_size, num_key]549  int *key_info,             // [batch_size, num_key, 2, num_hash_f]550  float *query_weight,       // [batch_size, num_query, weight_dim]551  float *key_weight,         // [batch_size, num_key, weight_dim]552  float *value,              // [batch_size, num_key, value_dim]553  float *cumulation_value,   // [batch_size, num_query, value_dim]554  int batch_size,555  int num_hash_f,556  int num_query,557  int num_key,558  int value_dim,559  int weight_dim560) {561 562  int batch_idx = blockIdx.z;563  int hash_f_idx = blockIdx.y;564  int key_idx = blockIdx.x;565 566  int num_threads = blockDim.y * blockDim.x;567  int thread_id = threadIdx.y * blockDim.x + threadIdx.x;568 569  int num_warps = blockDim.y;570  int warp_idx = threadIdx.y;571  int warp_thread_idx = threadIdx.x;572 573  int batch_idx__key_idx = batch_idx * num_key + key_idx;574  if (key_mask[batch_idx__key_idx] == 0) {575    return;576  }577 578  int query_offset = key_info[batch_idx__key_idx * 2 * num_hash_f + hash_f_idx];579  int query_count = key_info[(batch_idx__key_idx * 2 + 1) * num_hash_f + hash_f_idx];580 581  if (query_count == 0) {582    return;583  }584 585  extern __shared__ float buffer[];586 587  if (query_count == 1) {588    if (warp_idx == 0) {589      int query_idx = query_sorted_idxes[(batch_idx * num_hash_f + hash_f_idx) * num_query + query_offset];590      int batch_idx__query_idx = batch_idx * num_query + query_idx;591      float weight = 0;592      for (int weight_offset = 0; weight_offset < weight_dim; weight_offset = weight_offset + WARP_SIZE) {593        int weight_dim_idx = weight_offset + warp_thread_idx;594        float val = key_weight[batch_idx__key_idx * weight_dim + weight_dim_idx] * query_weight[batch_idx__query_idx * weight_dim + weight_dim_idx];595        #pragma unroll596        for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) {597          val += __shfl_xor_sync(FULL_MASK, val, offset);598        }599        weight = weight + val;600      }601      weight = weight / float(num_hash_f);602      for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) {603        int value_dim_idx = value_offset + warp_thread_idx;604        float val = value[batch_idx__key_idx * value_dim + value_dim_idx];605        atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val);606      }607    }608  } else {609    float *weight_buffer = buffer;610    float *value_buffer = &buffer[weight_dim];611    int *query_idxes_buffer = (int*)&buffer[weight_dim + value_dim];612 613    copy_data_nonblocking<float>(&key_weight[batch_idx__key_idx * weight_dim], weight_buffer, weight_dim, num_threads, thread_id);614    copy_data_nonblocking<float>(&value[batch_idx__key_idx * value_dim], value_buffer, value_dim, num_threads, thread_id);615 616    while (query_count > 0) {617      int work_size = min(WARP_SIZE, query_count);618      copy_data_nonblocking<int>(&query_sorted_idxes[(batch_idx * num_hash_f + hash_f_idx) * num_query + query_offset], query_idxes_buffer, work_size, num_threads, thread_id);619      __syncthreads();620      for (int work_offset = 0; work_offset < WARP_SIZE; work_offset = work_offset + num_warps) {621        int work_idx = work_offset + warp_idx;622        if (work_idx < query_count) {623          int query_idx = query_idxes_buffer[work_idx];624          int batch_idx__query_idx = batch_idx * num_query + query_idx;625          float weight = 0;626          for (int weight_offset = 0; weight_offset < weight_dim; weight_offset = weight_offset + WARP_SIZE) {627            int weight_dim_idx = weight_offset + warp_thread_idx;628            float val = weight_buffer[weight_dim_idx] * query_weight[batch_idx__query_idx * weight_dim + weight_dim_idx];629            #pragma unroll630            for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) {631              val += __shfl_xor_sync(FULL_MASK, val, offset);632            }633            weight = weight + val;634          }635          weight = weight / float(num_hash_f);636          for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) {637            int value_dim_idx = value_offset + warp_thread_idx;638            float val = value_buffer[value_dim_idx];639            atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val);640          }641        }642      }643      query_count = query_count - work_size;644      query_offset = query_offset + work_size;645    }646  }647 648}649 650__global__ void lsh_weighted_cumulation_ver4_step2_cuda_kernel(651  int *query_sorted_idxes,   // [batch_size, num_hash_f, num_query]652  int *key_mask,             // [batch_size, num_key]653  int *key_info,             // [batch_size, num_key, 2, num_hash_f]654  float *query_weight,       // [batch_size, num_query, weight_dim]655  float *key_weight,         // [batch_size, num_key, weight_dim]656  float *value,              // [batch_size, num_key, value_dim]657  float *cumulation_value,   // [batch_size, num_query, value_dim]658  int batch_size,659  int num_hash_f,660  int num_query,661  int num_key,662  int value_dim,663  int weight_dim664) {665 666  int batch_idx = blockIdx.y;667  int key_idx = blockIdx.x;668 669  int num_threads = blockDim.y * blockDim.x;670  int thread_id = threadIdx.y * blockDim.x + threadIdx.x;671 672  int num_warps = blockDim.y;673  int warp_idx = threadIdx.y;674  int warp_thread_idx = threadIdx.x;675 676  int batch_idx__key_idx = batch_idx * num_key + key_idx;677  if (key_mask[batch_idx__key_idx] == 0) {678    return;679  }680 681  extern __shared__ float buffer[];682  float *weight_buffer = buffer;683  float *value_buffer = &buffer[weight_dim];684  int *key_info_buffer = (int*)&buffer[weight_dim + value_dim];685 686  copy_data_nonblocking<float>(&key_weight[batch_idx__key_idx * weight_dim], weight_buffer, weight_dim, num_threads, thread_id);687  copy_data_nonblocking<float>(&value[batch_idx__key_idx * value_dim], value_buffer, value_dim, num_threads, thread_id);688  copy_data_nonblocking<int>(&key_info[batch_idx__key_idx * 2 * num_hash_f], key_info_buffer, 2 * num_hash_f, num_threads, thread_id);689 690  int *query_offset_buffer = key_info_buffer;691  int *query_count_buffer = &key_info_buffer[num_hash_f];692 693  const int hashtable_size = 1024 + OPTIMAL_THREADS_PER_BLOCK;694  __shared__ int hashtable_query[hashtable_size];695  __shared__ int hashtable_count[hashtable_size];696  __shared__ int inserted_query[hashtable_size];697  __shared__ int query_counter[1];698 699  int hash_f_idx_base = 0;700 701  while (true) {702 703    init_buffer_nonblocking<int>(EMPTY_VALUE, hashtable_query, hashtable_size, num_threads, thread_id);704    init_buffer_nonblocking<int>(0, hashtable_count, hashtable_size, num_threads, thread_id);705    init_buffer_nonblocking<int>(EMPTY_VALUE, inserted_query, hashtable_size, num_threads, thread_id);706    init_buffer_nonblocking<int>(0, query_counter, 1, num_threads, thread_id);707    __syncthreads();708 709    while (hash_f_idx_base < num_hash_f) {710 711      int hash_f_idx = hash_f_idx_base + warp_idx;712      int batch_idx__hash_f_idx = batch_idx * num_hash_f + hash_f_idx;713 714      int stop_flag = 0;715 716      int query_offset = query_offset_buffer[hash_f_idx];717      int query_count = query_count_buffer[hash_f_idx];718 719      while (query_count > 0) {720 721        int work_size = min(query_count, WARP_SIZE);722 723        // try inserting query to set and check whether the query is new724        int found_new_query = 0;725        int query_idx = -1;726        if (warp_thread_idx < work_size) {727          query_idx = query_sorted_idxes[batch_idx__hash_f_idx * num_query + query_offset + warp_thread_idx];728          int slot = set_insert<int>(hashtable_query, hashtable_size, query_idx);729          if (slot >= 0) {730            found_new_query = atomicAdd(&hashtable_count[slot], 1) == 0;731          }732        }733 734        // compute cumulative offset735        int position_offset = found_new_query;736        int next_position_offset = 0;737        #pragma unroll738        for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) {739          next_position_offset = __shfl_up_sync(FULL_MASK, position_offset, offset);740          if (thread_id % WARP_SIZE >= offset) {741            position_offset = position_offset + next_position_offset;742          }743        }744 745        // get the inserted query list end index746        int inserted_query_base = 0;747        if (thread_id % WARP_SIZE == WARP_SIZE - 1) {748          inserted_query_base = atomicAdd(query_counter, position_offset);749        }750        inserted_query_base = __shfl_sync(FULL_MASK, inserted_query_base, WARP_SIZE - 1);751 752        // insert new queries to list753        int insert_idx = inserted_query_base + position_offset - 1;754        if (found_new_query) {755          inserted_query[insert_idx] = query_idx;756        }757 758        // remove inserted queries from list759        query_offset_buffer[hash_f_idx] += work_size;760        query_count_buffer[hash_f_idx] -= work_size;761        query_offset += work_size;762        query_count -= work_size;763 764        // if list is almost full, stop inserting765        if (inserted_query_base + OPTIMAL_THREADS_PER_BLOCK > hashtable_size) {766          stop_flag = 1;767          break;768        }769 770      }771 772      if (stop_flag) {773        break;774      }775 776      hash_f_idx_base = hash_f_idx_base + num_warps;777 778    }779 780    __syncthreads();781 782    int num_distinct_query = query_counter[0];783 784    if (num_distinct_query > 0) {785      for (int idx_base = 0; idx_base < num_distinct_query; idx_base = idx_base + num_warps) {786        int idx = idx_base + warp_idx;787        if (idx < num_distinct_query) {788          int query_idx = inserted_query[idx];789          int batch_idx__query_idx = batch_idx * num_query + query_idx;790 791          int slot = set_lookup<int>(hashtable_query, hashtable_size, query_idx);792          int duplicate_count = hashtable_count[slot];793 794          float weight = 0;795          for (int weight_idx_base = 0; weight_idx_base < weight_dim; weight_idx_base = weight_idx_base + WARP_SIZE) {796            int weight_dim_idx = weight_idx_base + warp_thread_idx;797            float val = weight_buffer[weight_dim_idx] * query_weight[batch_idx__query_idx * weight_dim + weight_dim_idx];798            #pragma unroll799            for (int offset = 1; offset < WARP_SIZE; offset = offset << 1) {800              val += __shfl_xor_sync(FULL_MASK, val, offset);801            }802            weight = weight + val;803          }804 805          weight = (float)duplicate_count * weight / float(num_hash_f);806 807          for (int value_idx_base = 0; value_idx_base < value_dim; value_idx_base = value_idx_base + WARP_SIZE) {808            int value_dim_idx = value_idx_base + warp_thread_idx;809            float val = value_buffer[value_dim_idx];810            atomicAdd(&cumulation_value[batch_idx__query_idx * value_dim + value_dim_idx], weight * val);811          }812        }813      }814    } else {815 816      // all computation is completed if num_distinct_query == 0817      break;818 819    }820 821    __syncthreads();822 823  }824 825}826