Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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