Aluode/PerceptionLabPortable

1#include <torch/extension.h>2#include <ATen/ATen.h>3#include "fast_lsh_cumulation.h"4#include "common_cuda.h"5#include <vector>6 7std::vector<at::Tensor> fast_hash(8  at::Tensor query_mask,9  at::Tensor query_vector,10  at::Tensor key_mask,11  at::Tensor key_vector,12  int num_hash_f,13  int hash_code_len,14  bool use_cuda,15  int version16) {17  return fast_hash_ver1_kernel(18    query_mask,19    query_vector,20    key_mask,21    key_vector,22    num_hash_f,23    hash_code_len,24    use_cuda25  );26}27 28at::Tensor lsh_cumulation(29  at::Tensor query_mask,         // [batch_size, num_query]30  at::Tensor query_hash_code,    // [batch_size, num_query, num_hash_f]31  at::Tensor key_mask,           // [batch_size, num_key]32  at::Tensor key_hash_code,      // [batch_size, num_key, num_hash_f]33  at::Tensor value,              // [batch_size, num_key, value_dim]34  int hashtable_capacity,35  bool use_cuda,36  int version37) {38  return lsh_cumulation_ver1_kernel(39    query_mask,40    query_hash_code,41    key_mask,42    key_hash_code,43    value,44    hashtable_capacity,45    use_cuda46  );47}48 49at::Tensor lsh_weighted_cumulation(50  at::Tensor query_mask,         // [batch_size, num_query]51  at::Tensor query_hash_code,    // [batch_size, num_query, num_hash_f]52  at::Tensor query_weight,       // [batch_size, num_query, weight_dim]53  at::Tensor key_mask,           // [batch_size, num_key]54  at::Tensor key_hash_code,      // [batch_size, num_key, num_hash_f]55  at::Tensor key_weight,         // [batch_size, num_key, weight_dim]56  at::Tensor value,              // [batch_size, num_key, value_dim]57  int hashtable_capacity,58  bool use_cuda,59  int version60) {61  if (version == 1) {62    return lsh_weighted_cumulation_ver1_kernel(63      query_mask,64      query_hash_code,65      query_weight,66      key_mask,67      key_hash_code,68      key_weight,69      value,70      hashtable_capacity,71      use_cuda72    );73  } else if (version == 2) {74    return lsh_weighted_cumulation_ver2_kernel(75      query_mask,76      query_hash_code,77      query_weight,78      key_mask,79      key_hash_code,80      key_weight,81      value,82      hashtable_capacity,83      use_cuda84    );85  } else if (version == 3) {86    return lsh_weighted_cumulation_ver3_kernel(87      query_mask,88      query_hash_code,89      query_weight,90      key_mask,91      key_hash_code,92      key_weight,93      value,94      hashtable_capacity,95      use_cuda96    );97  } else if (version == 4) {98    return lsh_weighted_cumulation_ver4_kernel(99      query_mask,100      query_hash_code,101      query_weight,102      key_mask,103      key_hash_code,104      key_weight,105      value,106      hashtable_capacity,107      use_cuda108    );109  } else {110    return lsh_weighted_cumulation_ver3_kernel(111      query_mask,112      query_hash_code,113      query_weight,114      key_mask,115      key_hash_code,116      key_weight,117      value,118      hashtable_capacity,119      use_cuda120    );121  }122}123 124PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {125  m.def("fast_hash", &fast_hash, "Fast Hash (CUDA)");126  m.def("lsh_cumulation", &lsh_cumulation, "LSH Cumulation (CUDA)");127  m.def("lsh_weighted_cumulation", &lsh_weighted_cumulation, "LSH Weighted Cumulation (CUDA)");128}129