Aluode/PerceptionLabPortable

1// File from https://github.com/mlpen/YOSO/blob/main/encoders/backbones/efficient_attentions/yoso/yoso_v1/cuda/fast_lsh_cumulation.cu2 3#include <torch/extension.h>4#include <ATen/ATen.h>5#include "fast_lsh_cumulation.h"6#include "fast_lsh_cumulation_cuda.h"7#include "common_cuda.h"8#include "common.h"9#include <vector>10//////////////////////////////////////////////////////////////////////////////////////////////////11//////////////////////////////////////////////////////////////////////////////////////////////////12 13std::vector<at::Tensor> fast_hash_ver1_kernel(14  at::Tensor query_mask,15  at::Tensor query_vector,16  at::Tensor key_mask,17  at::Tensor key_vector,18  int num_hash_f,19  int hash_code_len,20  bool use_cuda21) {22 23  int batch_size = query_vector.size(0);24  int num_query = query_vector.size(1);25  int num_key = key_vector.size(1);26  int vector_dim = query_vector.size(2);27 28  int num_hash_per_part = vector_dim / hash_code_len;29  int num_part = max(1, ceil_divide(num_hash_f, num_hash_per_part));30 31  at::Tensor Dmat = 2 * at::randint(0, 2, {batch_size, 3, num_part, vector_dim}, query_mask.options()) - 1;32  at::Tensor query_hash_code = at::zeros({batch_size, num_query, num_hash_f}, query_mask.options());33  at::Tensor key_hash_code = at::zeros({batch_size, num_key, num_hash_f}, key_mask.options());34 35  int *query_mask_ptr = query_mask.data_ptr<int>();36  float *query_vector_ptr = query_vector.data_ptr<float>();37  int *key_mask_ptr = key_mask.data_ptr<int>();38  float *key_vector_ptr = key_vector.data_ptr<float>();39 40  int *Dmat_ptr = Dmat.data_ptr<int>();41 42  int *query_hash_code_ptr = query_hash_code.data_ptr<int>();43  int *key_hash_code_ptr = key_hash_code.data_ptr<int>();44 45  if (use_cuda) {46    {47      dim3 threads(vector_dim);48      dim3 blocks(num_part, num_query, batch_size);49      int shared_mem = vector_dim * sizeof(float);50      fast_hash_ver1_cuda_kernel<<<blocks, threads, shared_mem>>>(51        query_mask_ptr,52        query_vector_ptr,53        Dmat_ptr,54        query_hash_code_ptr,55        batch_size,56        num_query,57        vector_dim,58        num_part,59        num_hash_f,60        hash_code_len61      );62    }63    {64      dim3 threads(vector_dim);65      dim3 blocks(num_part, num_key, batch_size);66      int shared_mem = vector_dim * sizeof(float);67      fast_hash_ver1_cuda_kernel<<<blocks, threads, shared_mem>>>(68        key_mask_ptr,69        key_vector_ptr,70        Dmat_ptr,71        key_hash_code_ptr,72        batch_size,73        num_key,74        vector_dim,75        num_part,76        num_hash_f,77        hash_code_len78      );79    }80  }81 82  return {query_hash_code, key_hash_code};83 84}85 86at::Tensor lsh_cumulation_ver1_kernel(87  at::Tensor query_mask,88  at::Tensor query_hash_code,89  at::Tensor key_mask,90  at::Tensor key_hash_code,91  at::Tensor value,92  int hashtable_capacity,93  bool use_cuda94) {95 96  int batch_size = query_hash_code.size(0);97  int num_hash_f = query_hash_code.size(2);98 99  int num_query = query_hash_code.size(1);100  int num_key = key_hash_code.size(1);101  int value_dim = value.size(2);102 103  at::Tensor hashtable_value = at::empty({batch_size, num_hash_f, hashtable_capacity, WARP_SIZE}, value.options());104  at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options());105 106  if (use_cuda) {107    int threads_x = WARP_SIZE;108    int threads_y = OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE;109    int block_x_step1 = num_key / threads_y;110    int block_x_step2 = num_query / threads_y;111    int block_y = batch_size;112 113    dim3 threads(threads_x, threads_y);114    dim3 blocks_step1(block_x_step1, block_y);115    dim3 blocks_step2(block_x_step2, block_y);116 117    int *query_mask_ptr = query_mask.data_ptr<int>();118    int *query_hash_code_ptr = query_hash_code.data_ptr<int>();119    int *key_mask_ptr = key_mask.data_ptr<int>();120    int *key_hash_code_ptr = key_hash_code.data_ptr<int>();121    float *value_ptr = value.data_ptr<float>();122    float *hashtable_value_ptr = hashtable_value.data_ptr<float>();123    float *cumulation_value_ptr = cumulation_value.data_ptr<float>();124 125    for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) {126 127      cudaMemset(hashtable_value_ptr, 0, (batch_size * num_hash_f * hashtable_capacity * WARP_SIZE) * sizeof(float));128 129      lsh_cumulation_ver1_step1_cuda_kernel<<<blocks_step1, threads>>>(130        key_mask_ptr,131        key_hash_code_ptr,132        value_ptr,133        hashtable_value_ptr,134        batch_size,135        num_hash_f,136        hashtable_capacity,137        num_key,138        value_dim,139        value_offset140      );141 142      lsh_cumulation_ver1_step2_cuda_kernel<<<blocks_step2, threads>>>(143        query_mask_ptr,144        query_hash_code_ptr,145        hashtable_value_ptr,146        cumulation_value_ptr,147        batch_size,148        num_hash_f,149        hashtable_capacity,150        num_query,151        value_dim,152        value_offset153      );154    }155 156  }157 158  return cumulation_value;159 160}161 162at::Tensor lsh_weighted_cumulation_ver1_kernel(163  at::Tensor query_mask,164  at::Tensor query_hash_code,165  at::Tensor query_weight,166  at::Tensor key_mask,167  at::Tensor key_hash_code,168  at::Tensor key_weight,169  at::Tensor value,170  int hashtable_capacity,171  bool use_cuda172) {173 174  int batch_size = query_hash_code.size(0);175  int num_hash_f = query_hash_code.size(2);176 177  int num_query = query_hash_code.size(1);178  int num_key = key_hash_code.size(1);179  int value_dim = value.size(2);180  int weight_dim = query_weight.size(2);181 182  at::Tensor hashtable_value = at::zeros({batch_size, num_hash_f, hashtable_capacity, WARP_SIZE}, value.options());183  at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options());184 185  if (use_cuda) {186    int threads_x = WARP_SIZE;187    int threads_y = OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE;188    int block_x_step1 = num_key / threads_y;189    int block_x_step2 = num_query / threads_y;190    int block_y = batch_size;191 192    dim3 threads(threads_x, threads_y);193    dim3 blocks_step1(block_x_step1, block_y);194    dim3 blocks_step2(block_x_step2, block_y);195 196    int *query_mask_ptr = query_mask.data_ptr<int>();197    int *query_hash_code_ptr = query_hash_code.data_ptr<int>();198    float *query_weight_ptr = query_weight.data_ptr<float>();199    int *key_mask_ptr = key_mask.data_ptr<int>();200    int *key_hash_code_ptr = key_hash_code.data_ptr<int>();201    float *key_weight_ptr = key_weight.data_ptr<float>();202    float *value_ptr = value.data_ptr<float>();203    float *hashtable_value_ptr = hashtable_value.data_ptr<float>();204    float *cumulation_value_ptr = cumulation_value.data_ptr<float>();205 206    for (int value_offset = 0; value_offset < value_dim; value_offset = value_offset + WARP_SIZE) {207      for (int weight_idx = 0; weight_idx < weight_dim; weight_idx++) {208 209        cudaMemset(hashtable_value_ptr, 0, (batch_size * num_hash_f * hashtable_capacity * WARP_SIZE) * sizeof(float));210 211        lsh_weighted_cumulation_ver1_step1_cuda_kernel<<<blocks_step1, threads>>>(212          key_mask_ptr,213          key_hash_code_ptr,214          key_weight_ptr,215          value_ptr,216          hashtable_value_ptr,217          batch_size,218          num_hash_f,219          hashtable_capacity,220          num_key,221          value_dim,222          weight_dim,223          value_offset,224          weight_idx225        );226 227        lsh_weighted_cumulation_ver1_step2_cuda_kernel<<<blocks_step2, threads>>>(228          query_mask_ptr,229          query_hash_code_ptr,230          query_weight_ptr,231          hashtable_value_ptr,232          cumulation_value_ptr,233          batch_size,234          num_hash_f,235          hashtable_capacity,236          num_query,237          value_dim,238          weight_dim,239          value_offset,240          weight_idx241        );242      }243    }244 245  }246 247  return cumulation_value;248 249}250 251at::Tensor lsh_weighted_cumulation_ver2_kernel(252  at::Tensor query_mask,253  at::Tensor query_hash_code,254  at::Tensor query_weight,255  at::Tensor key_mask,256  at::Tensor key_hash_code,257  at::Tensor key_weight,258  at::Tensor value,259  int hashtable_capacity,260  bool use_cuda261) {262 263  int batch_size = query_hash_code.size(0);264  int num_hash_f = query_hash_code.size(2);265 266  int num_query = query_hash_code.size(1);267  int num_key = key_hash_code.size(1);268  int value_dim = value.size(2);269  int weight_dim = query_weight.size(2);270 271  at::Tensor count_sort_table = at::zeros({batch_size, num_hash_f, hashtable_capacity}, query_hash_code.options());272  at::Tensor key_sorted_idxes = at::zeros({batch_size, num_hash_f, num_key}, query_hash_code.options());273  at::Tensor query_info = at::zeros({batch_size, num_query, 2, num_hash_f}, query_hash_code.options());274  at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options());275 276  if (use_cuda) {277 278    int *query_mask_ptr = query_mask.data_ptr<int>();279    int *query_hash_code_ptr = query_hash_code.data_ptr<int>();280    float *query_weight_ptr = query_weight.data_ptr<float>();281    int *key_mask_ptr = key_mask.data_ptr<int>();282    int *key_hash_code_ptr = key_hash_code.data_ptr<int>();283    float *key_weight_ptr = key_weight.data_ptr<float>();284    float *value_ptr = value.data_ptr<float>();285 286    int *count_sort_table_ptr = count_sort_table.data_ptr<int>();287    int *key_sorted_idxes_ptr = key_sorted_idxes.data_ptr<int>();288    int *query_info_ptr = query_info.data_ptr<int>();289 290    float *cumulation_value_ptr = cumulation_value.data_ptr<float>();291 292    {293      dim3 threads_step13(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f));294      dim3 blocks_step13(num_key / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size);295      dim3 threads_step2(min(hashtable_capacity, OPTIMAL_THREADS_PER_BLOCK));296      dim3 blocks_step2(num_hash_f, batch_size);297      int shared_mem = hashtable_capacity * sizeof(float);298      count_sort_step1_cuda_kernel<<<blocks_step13, threads_step13>>>(299        key_mask_ptr,300        key_hash_code_ptr,301        count_sort_table_ptr,302        batch_size,303        num_hash_f,304        hashtable_capacity,305        num_key306      );307      count_sort_step2_cuda_kernel<<<blocks_step2, threads_step2, shared_mem>>>(308        count_sort_table_ptr,309        batch_size,310        num_hash_f,311        hashtable_capacity312      );313      count_sort_step3_cuda_kernel<<<blocks_step13, threads_step13>>>(314        key_mask_ptr,315        key_hash_code_ptr,316        count_sort_table_ptr,317        key_sorted_idxes_ptr,318        batch_size,319        num_hash_f,320        hashtable_capacity,321        num_key322      );323    }324    {325      dim3 threads(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f));326      dim3 blocks(num_query / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size);327      extract_query_info_cuda_kernel<<<blocks, threads>>>(328        query_mask_ptr,329        query_hash_code_ptr,330        count_sort_table_ptr,331        query_info_ptr,332        batch_size,333        num_hash_f,334        hashtable_capacity,335        num_query336      );337    }338    {339      dim3 threads(WARP_SIZE, OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE);340      dim3 blocks(num_query, num_hash_f, batch_size);341      int shared_mem = (weight_dim + WARP_SIZE) * sizeof(float);342      lsh_weighted_cumulation_ver2_step2_cuda_kernel<<<blocks, threads, shared_mem>>>(343        query_mask_ptr,344        query_info_ptr,345        key_sorted_idxes_ptr,346        query_weight_ptr,347        key_weight_ptr,348        value_ptr,349        cumulation_value_ptr,350        batch_size,351        num_hash_f,352        num_query,353        num_key,354        value_dim,355        weight_dim356      );357    }358  }359 360  return cumulation_value;361 362}363 364at::Tensor lsh_weighted_cumulation_ver3_kernel(365  at::Tensor query_mask,366  at::Tensor query_hash_code,367  at::Tensor query_weight,368  at::Tensor key_mask,369  at::Tensor key_hash_code,370  at::Tensor key_weight,371  at::Tensor value,372  int hashtable_capacity,373  bool use_cuda374) {375 376  int batch_size = query_hash_code.size(0);377  int num_hash_f = query_hash_code.size(2);378 379  int num_query = query_hash_code.size(1);380  int num_key = key_hash_code.size(1);381  int value_dim = value.size(2);382  int weight_dim = query_weight.size(2);383 384  at::Tensor count_sort_table = at::zeros({batch_size, num_hash_f, hashtable_capacity}, query_hash_code.options());385  at::Tensor query_sorted_idxes = at::zeros({batch_size, num_hash_f, num_query}, query_hash_code.options());386  at::Tensor key_info = at::zeros({batch_size, num_key, 2, num_hash_f}, query_hash_code.options());387  at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options());388 389  if (use_cuda) {390 391    int *query_mask_ptr = query_mask.data_ptr<int>();392    int *query_hash_code_ptr = query_hash_code.data_ptr<int>();393    float *query_weight_ptr = query_weight.data_ptr<float>();394    int *key_mask_ptr = key_mask.data_ptr<int>();395    int *key_hash_code_ptr = key_hash_code.data_ptr<int>();396    float *key_weight_ptr = key_weight.data_ptr<float>();397    float *value_ptr = value.data_ptr<float>();398 399    int *count_sort_table_ptr = count_sort_table.data_ptr<int>();400    int *query_sorted_idxes_ptr = query_sorted_idxes.data_ptr<int>();401    int *key_info_ptr = key_info.data_ptr<int>();402 403    float *cumulation_value_ptr = cumulation_value.data_ptr<float>();404 405    {406      dim3 threads_step13(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f));407      dim3 blocks_step13(num_query / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size);408      dim3 threads_step2(min(hashtable_capacity, OPTIMAL_THREADS_PER_BLOCK));409      dim3 blocks_step2(num_hash_f, batch_size);410      int shared_mem = hashtable_capacity * sizeof(float);411      count_sort_step1_cuda_kernel<<<blocks_step13, threads_step13>>>(412        query_mask_ptr,413        query_hash_code_ptr,414        count_sort_table_ptr,415        batch_size,416        num_hash_f,417        hashtable_capacity,418        num_query419      );420      count_sort_step2_cuda_kernel<<<blocks_step2, threads_step2, shared_mem>>>(421        count_sort_table_ptr,422        batch_size,423        num_hash_f,424        hashtable_capacity425      );426      count_sort_step3_cuda_kernel<<<blocks_step13, threads_step13>>>(427        query_mask_ptr,428        query_hash_code_ptr,429        count_sort_table_ptr,430        query_sorted_idxes_ptr,431        batch_size,432        num_hash_f,433        hashtable_capacity,434        num_query435      );436    }437    {438      dim3 threads(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f));439      dim3 blocks(num_key / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size);440      extract_query_info_cuda_kernel<<<blocks, threads>>>(441        key_mask_ptr,442        key_hash_code_ptr,443        count_sort_table_ptr,444        key_info_ptr,445        batch_size,446        num_hash_f,447        hashtable_capacity,448        num_key449      );450    }451    {452      dim3 threads(WARP_SIZE, OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE);453      dim3 blocks(num_key, num_hash_f, batch_size);454      int shared_mem = (weight_dim + value_dim + WARP_SIZE) * sizeof(float);455      lsh_weighted_cumulation_ver3_step2_cuda_kernel<<<blocks, threads, shared_mem>>>(456        query_sorted_idxes_ptr,457        key_mask_ptr,458        key_info_ptr,459        query_weight_ptr,460        key_weight_ptr,461        value_ptr,462        cumulation_value_ptr,463        batch_size,464        num_hash_f,465        num_query,466        num_key,467        value_dim,468        weight_dim469      );470    }471  }472 473  return cumulation_value;474 475}476 477at::Tensor lsh_weighted_cumulation_ver4_kernel(478  at::Tensor query_mask,479  at::Tensor query_hash_code,480  at::Tensor query_weight,481  at::Tensor key_mask,482  at::Tensor key_hash_code,483  at::Tensor key_weight,484  at::Tensor value,485  int hashtable_capacity,486  bool use_cuda487) {488 489  int batch_size = query_hash_code.size(0);490  int num_hash_f = query_hash_code.size(2);491 492  int num_query = query_hash_code.size(1);493  int num_key = key_hash_code.size(1);494  int value_dim = value.size(2);495  int weight_dim = query_weight.size(2);496 497  at::Tensor count_sort_table = at::zeros({batch_size, num_hash_f, hashtable_capacity}, query_hash_code.options());498  at::Tensor query_sorted_idxes = at::zeros({batch_size, num_hash_f, num_query}, query_hash_code.options());499  at::Tensor key_info = at::zeros({batch_size, num_key, 2, num_hash_f}, query_hash_code.options());500  at::Tensor cumulation_value = at::zeros({batch_size, num_query, value_dim}, value.options());501 502  if (use_cuda) {503 504    int *query_mask_ptr = query_mask.data_ptr<int>();505    int *query_hash_code_ptr = query_hash_code.data_ptr<int>();506    float *query_weight_ptr = query_weight.data_ptr<float>();507    int *key_mask_ptr = key_mask.data_ptr<int>();508    int *key_hash_code_ptr = key_hash_code.data_ptr<int>();509    float *key_weight_ptr = key_weight.data_ptr<float>();510    float *value_ptr = value.data_ptr<float>();511 512    int *count_sort_table_ptr = count_sort_table.data_ptr<int>();513    int *query_sorted_idxes_ptr = query_sorted_idxes.data_ptr<int>();514    int *key_info_ptr = key_info.data_ptr<int>();515 516    float *cumulation_value_ptr = cumulation_value.data_ptr<float>();517 518    {519      dim3 threads_step13(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f));520      dim3 blocks_step13(num_query / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size);521      dim3 threads_step2(min(hashtable_capacity, OPTIMAL_THREADS_PER_BLOCK));522      dim3 blocks_step2(num_hash_f, batch_size);523      int shared_mem = hashtable_capacity * sizeof(float);524      count_sort_step1_cuda_kernel<<<blocks_step13, threads_step13>>>(525        query_mask_ptr,526        query_hash_code_ptr,527        count_sort_table_ptr,528        batch_size,529        num_hash_f,530        hashtable_capacity,531        num_query532      );533      count_sort_step2_cuda_kernel<<<blocks_step2, threads_step2, shared_mem>>>(534        count_sort_table_ptr,535        batch_size,536        num_hash_f,537        hashtable_capacity538      );539      count_sort_step3_cuda_kernel<<<blocks_step13, threads_step13>>>(540        query_mask_ptr,541        query_hash_code_ptr,542        count_sort_table_ptr,543        query_sorted_idxes_ptr,544        batch_size,545        num_hash_f,546        hashtable_capacity,547        num_query548      );549    }550    {551      dim3 threads(num_hash_f, max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f));552      dim3 blocks(num_key / max(1, OPTIMAL_THREADS_PER_BLOCK / num_hash_f), batch_size);553      extract_query_info_cuda_kernel<<<blocks, threads>>>(554        key_mask_ptr,555        key_hash_code_ptr,556        count_sort_table_ptr,557        key_info_ptr,558        batch_size,559        num_hash_f,560        hashtable_capacity,561        num_key562      );563    }564    {565      dim3 threads(WARP_SIZE, OPTIMAL_THREADS_PER_BLOCK / WARP_SIZE);566      dim3 blocks(num_key, batch_size);567      int shared_mem = (weight_dim + value_dim + 2 * num_hash_f) * sizeof(float);568      lsh_weighted_cumulation_ver4_step2_cuda_kernel<<<blocks, threads, shared_mem>>>(569        query_sorted_idxes_ptr,570        key_mask_ptr,571        key_info_ptr,572        query_weight_ptr,573        key_weight_ptr,574        value_ptr,575        cumulation_value_ptr,576        batch_size,577        num_hash_f,578        num_query,579        num_key,580        value_dim,581        weight_dim582      );583    }584  }585 586  return cumulation_value;587 588}589