27M/PreFLMR_ViT-G

1#include <pthread.h>2#include <torch/extension.h>3 4#include <algorithm>5#include <numeric>6 7typedef struct {8    int tid;9    int nthreads;10 11    int ndocs;12    int ndoc_vectors;13    int nquery_vectors;14 15    int64_t* lengths;16    float* scores;17    int64_t* offsets;18 19    float* max_scores;20} max_args_t;21 22void* max(void* args) {23    max_args_t* max_args = (max_args_t*)args;24 25    int ndocs_per_thread =26        std::ceil(((float)max_args->ndocs) / max_args->nthreads);27    int start = max_args->tid * ndocs_per_thread;28    int end = std::min((max_args->tid + 1) * ndocs_per_thread, max_args->ndocs);29 30    auto max_scores_offset =31        max_args->max_scores + (start * max_args->nquery_vectors);32    auto scores_offset =33        max_args->scores + (max_args->offsets[start] * max_args->nquery_vectors);34 35    for (int i = start; i < end; i++) {36        for (int j = 0; j < max_args->lengths[i]; j++) {37            std::transform(max_scores_offset,38                           max_scores_offset + max_args->nquery_vectors,39                           scores_offset, max_scores_offset,40                           [](float a, float b) { return std::max(a, b); });41            scores_offset += max_args->nquery_vectors;42        }43        max_scores_offset += max_args->nquery_vectors;44    }45 46    return NULL;47}48 49torch::Tensor segmented_maxsim(const torch::Tensor scores,50                               const torch::Tensor lengths) {51    auto lengths_a = lengths.data_ptr<int64_t>();52    auto scores_a = scores.data_ptr<float>();53    auto ndocs = lengths.size(0);54    auto ndoc_vectors = scores.size(0);55    auto nquery_vectors = scores.size(1);56    auto nthreads = at::get_num_threads();57 58    torch::Tensor max_scores =59        torch::zeros({ndocs, nquery_vectors}, scores.options());60 61    int64_t offsets[ndocs + 1];62    offsets[0] = 0;63    std::partial_sum(lengths_a, lengths_a + ndocs, offsets + 1);64 65    pthread_t threads[nthreads];66    max_args_t args[nthreads];67 68    for (int i = 0; i < nthreads; i++) {69        args[i].tid = i;70        args[i].nthreads = nthreads;71 72        args[i].ndocs = ndocs;73        args[i].ndoc_vectors = ndoc_vectors;74        args[i].nquery_vectors = nquery_vectors;75 76        args[i].lengths = lengths_a;77        args[i].scores = scores_a;78        args[i].offsets = offsets;79 80        args[i].max_scores = max_scores.data_ptr<float>();81 82        int rc = pthread_create(&threads[i], NULL, max, (void*)&args[i]);83        if (rc) {84            fprintf(stderr, "Unable to create thread %d: %d\n", i, rc);85        }86    }87 88    for (int i = 0; i < nthreads; i++) {89        pthread_join(threads[i], NULL);90    }91 92    return max_scores.sum(1);93}94 95PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {96    m.def("segmented_maxsim_cpp", &segmented_maxsim, "Segmented MaxSim");97}98