Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
Heap.cpp
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 /* Copyright 2004-present Facebook. All Rights Reserved. */
11 /* Function for soft heap */
12
13 #include "Heap.h"
14
15
16 namespace faiss {
17
18
19
20
21
22
23
24
25
26 template <typename C>
27 void HeapArray<C>::heapify ()
28 {
29 #pragma omp parallel for
30 for (size_t j = 0; j < nh; j++)
31 heap_heapify<C> (k, val + j * k, ids + j * k);
32 }
33
34 template <typename C>
35 void HeapArray<C>::reorder ()
36 {
37 #pragma omp parallel for
38 for (size_t j = 0; j < nh; j++)
39 heap_reorder<C> (k, val + j * k, ids + j * k);
40 }
41
42 template <typename C>
43 void HeapArray<C>::addn (size_t nj, const T *vin, TI j0,
44 size_t i0, long ni)
45 {
46 if (ni == -1) ni = nh;
47 assert (i0 >= 0 && i0 + ni <= nh);
48 #pragma omp parallel for
49 for (size_t i = i0; i < i0 + ni; i++) {
50 T * __restrict simi = get_val(i);
51 TI * __restrict idxi = get_ids (i);
52 const T *ip_line = vin + (i - i0) * nj;
53
54 for (size_t j = 0; j < nj; j++) {
55 T ip = ip_line [j];
56 if (C::cmp(simi[0], ip)) {
57 heap_pop<C> (k, simi, idxi);
58 heap_push<C> (k, simi, idxi, ip, j + j0);
59 }
60 }
61 }
62 }
63
64 template <typename C>
65 void HeapArray<C>::addn_with_ids (
66 size_t nj, const T *vin, const TI *id_in,
67 long id_stride, size_t i0, long ni)
68 {
69 if (id_in == nullptr) {
70 addn (nj, vin, 0, i0, ni);
71 return;
72 }
73 if (ni == -1) ni = nh;
74 assert (i0 >= 0 && i0 + ni <= nh);
75 #pragma omp parallel for
76 for (size_t i = i0; i < i0 + ni; i++) {
77 T * __restrict simi = get_val(i);
78 TI * __restrict idxi = get_ids (i);
79 const T *ip_line = vin + (i - i0) * nj;
80 const TI *id_line = id_in + (i - i0) * id_stride;
81
82 for (size_t j = 0; j < nj; j++) {
83 T ip = ip_line [j];
84 if (C::cmp(simi[0], ip)) {
85 heap_pop<C> (k, simi, idxi);
86 heap_push<C> (k, simi, idxi, ip, id_line [j]);
87 }
88 }
89 }
90 }
91
92 template <typename C>
93 void HeapArray<C>::per_line_extrema (
94 T * out_val,
95 TI * out_ids) const
96 {
97 #pragma omp parallel for
98 for (size_t j = 0; j < nh; j++) {
99 long imin = -1;
100 typename C::T xval = C::Crev::neutral ();
101 const typename C::T * x_ = val + j * k;
102 for (size_t i = 0; i < k; i++)
103 if (C::cmp (x_[i], xval)) {
104 xval = x_[i];
105 imin = i;
106 }
107 if (out_val)
108 out_val[j] = xval;
109
110 if (out_ids) {
111 if (ids && imin != -1)
112 out_ids[j] = ids [j * k + imin];
113 else
114 out_ids[j] = imin;
115 }
116 }
117 }
118
119
120
121
122 // explicit instanciations
123
124 template class HeapArray<CMin <float, long> >;
125 template class HeapArray<CMax <float, long> >;
126 template class HeapArray<CMin <int, long> >;
127 template class HeapArray<CMax <int, long> >;
128
129
130
131
132
133 } // END namespace fasis
faiss::HeapArray::reorder
void reorder()
reorder all the heaps
Definition: Heap.cpp:35
faiss::HeapArray::per_line_extrema
void per_line_extrema(T *vals_out, TI *idx_out) const
Definition: Heap.cpp:93
faiss::HeapArray
Definition: Heap.h:351
faiss::HeapArray::addn_with_ids
void addn_with_ids(size_t nj, const T *vin, const TI *id_in=nullptr, long id_stride=0, size_t i0=0, long ni=-1)
Definition: Heap.cpp:65
faiss::HeapArray::heapify
void heapify()
prepare all the heaps before adding
Definition: Heap.cpp:27
faiss::HeapArray::addn
void addn(size_t nj, const T *vin, TI j0=0, size_t i0=0, long ni=-1)
Definition: Heap.cpp:43