|
#include <pthread.h> |
|
#include <torch/extension.h> |
|
|
|
#include <algorithm> |
|
#include <numeric> |
|
|
|
typedef struct { |
|
int tid; |
|
int nthreads; |
|
|
|
int ndocs; |
|
int ndoc_vectors; |
|
int nquery_vectors; |
|
|
|
int64_t* lengths; |
|
float* scores; |
|
int64_t* offsets; |
|
|
|
float* max_scores; |
|
} max_args_t; |
|
|
|
void* max(void* args) { |
|
max_args_t* max_args = (max_args_t*)args; |
|
|
|
int ndocs_per_thread = |
|
std::ceil(((float)max_args->ndocs) / max_args->nthreads); |
|
int start = max_args->tid * ndocs_per_thread; |
|
int end = std::min((max_args->tid + 1) * ndocs_per_thread, max_args->ndocs); |
|
|
|
auto max_scores_offset = |
|
max_args->max_scores + (start * max_args->nquery_vectors); |
|
auto scores_offset = |
|
max_args->scores + (max_args->offsets[start] * max_args->nquery_vectors); |
|
|
|
for (int i = start; i < end; i++) { |
|
for (int j = 0; j < max_args->lengths[i]; j++) { |
|
std::transform(max_scores_offset, |
|
max_scores_offset + max_args->nquery_vectors, |
|
scores_offset, max_scores_offset, |
|
[](float a, float b) { return std::max(a, b); }); |
|
scores_offset += max_args->nquery_vectors; |
|
} |
|
max_scores_offset += max_args->nquery_vectors; |
|
} |
|
|
|
return NULL; |
|
} |
|
|
|
torch::Tensor segmented_maxsim(const torch::Tensor scores, |
|
const torch::Tensor lengths) { |
|
auto lengths_a = lengths.data_ptr<int64_t>(); |
|
auto scores_a = scores.data_ptr<float>(); |
|
auto ndocs = lengths.size(0); |
|
auto ndoc_vectors = scores.size(0); |
|
auto nquery_vectors = scores.size(1); |
|
auto nthreads = at::get_num_threads(); |
|
|
|
torch::Tensor max_scores = |
|
torch::zeros({ndocs, nquery_vectors}, scores.options()); |
|
|
|
int64_t offsets[ndocs + 1]; |
|
offsets[0] = 0; |
|
std::partial_sum(lengths_a, lengths_a + ndocs, offsets + 1); |
|
|
|
pthread_t threads[nthreads]; |
|
max_args_t args[nthreads]; |
|
|
|
for (int i = 0; i < nthreads; i++) { |
|
args[i].tid = i; |
|
args[i].nthreads = nthreads; |
|
|
|
args[i].ndocs = ndocs; |
|
args[i].ndoc_vectors = ndoc_vectors; |
|
args[i].nquery_vectors = nquery_vectors; |
|
|
|
args[i].lengths = lengths_a; |
|
args[i].scores = scores_a; |
|
args[i].offsets = offsets; |
|
|
|
args[i].max_scores = max_scores.data_ptr<float>(); |
|
|
|
int rc = pthread_create(&threads[i], NULL, max, (void*)&args[i]); |
|
if (rc) { |
|
fprintf(stderr, "Unable to create thread %d: %d\n", i, rc); |
|
} |
|
} |
|
|
|
for (int i = 0; i < nthreads; i++) { |
|
pthread_join(threads[i], NULL); |
|
} |
|
|
|
return max_scores.sum(1); |
|
} |
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { |
|
m.def("segmented_maxsim_cpp", &segmented_maxsim, "Segmented MaxSim"); |
|
} |
|
|
