Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
BinaryDistance.cu
1 /**
2  * Copyright (c) Facebook, Inc. and its affiliates.
3  *
4  * This source code is licensed under the MIT license found in the
5  * LICENSE file in the root directory of this source tree.
6  */
7 
8 #include "../utils/DeviceTensor.cuh"
9 #include "../utils/DeviceDefs.cuh"
10 #include "../utils/DeviceUtils.h"
11 #include "../utils/Select.cuh"
12 
13 namespace faiss { namespace gpu {
14 
15 // Number of warps that the kernel is instantiated with
16 constexpr int kWarps = 8;
17 constexpr int kLanes = kWarpSize;
18 
19 constexpr int kMaxDistance = std::numeric_limits<int>::max();
20 
21 // Performs a binary matrix multiplication, returning the lowest k results in
22 // `vecs` for each `query` in terms of Hamming distance (a fused kernel)
23 // Each warp calculates distance for a single query
24 template <int NumWarpQ,
25  int NumThreadQ,
26  typename BinaryType>
27 __launch_bounds__(kWarps * kLanes)
28 __global__ void binaryDistanceAnySize(const Tensor<BinaryType, 2, true> vecs,
29  const Tensor<BinaryType, 2, true> query,
30  Tensor<int, 2, true> outK,
31  Tensor<int, 2, true> outV,
32  int k) {
33  // A matrix tile (query, k)
34  __shared__ BinaryType queryTile[kWarps][kLanes + 1]; // avoid bank conflict
35 
36  // B matrix tile (vec, k)
37  __shared__ BinaryType vecTile[kLanes][kLanes + 1]; // avoid bank conflict
38 
39  WarpSelect<int, int, false, Comparator<int>,
40  NumWarpQ, NumThreadQ, kWarps * kLanes>
41  heap(kMaxDistance, -1, k);
42 
43  int warpId = threadIdx.y;
44  int laneId = threadIdx.x;
45 
46  // Each warp handles a single query
47  int warpQuery = blockIdx.x * kWarps + warpId;
48  bool queryInBounds = warpQuery < query.getSize(0);
49 
50  // Each warp loops through the entire chunk of vectors
51  for (int blockVec = 0; blockVec < vecs.getSize(0); blockVec += kLanes) {
52  int threadDistance = 0;
53 
54  // Reduction dimension
55  for (int blockK = 0; blockK < vecs.getSize(1); blockK += kLanes) {
56  int laneK = blockK + laneId;
57  bool kInBounds = laneK < vecs.getSize(1);
58 
59  queryTile[warpId][laneId] = queryInBounds && kInBounds ?
60  query[warpQuery][laneK] : 0;
61 
62  // kWarps warps are responsible for loading 32 vecs
63 #pragma unroll
64  for (int i = 0; i < kLanes / kWarps; ++i) {
65  int warpVec = i * kWarps + warpId;
66  int vec = blockVec + warpVec;
67  bool vecInBounds = vec < vecs.getSize(0);
68 
69  vecTile[warpVec][laneId] = vecInBounds && kInBounds ?
70  vecs[vec][laneK] : 0;
71  }
72 
73  __syncthreads();
74 
75  // Compare distances
76 #pragma unroll
77  for (int i = 0; i < kLanes; ++i) {
78  threadDistance += __popc(queryTile[warpId][i] ^ vecTile[laneId][i]);
79  }
80 
81  __syncthreads();
82  }
83 
84  // Lanes within a warp are different vec results against the same query
85  // Only submit distances which represent real (query, vec) pairs
86  bool valInBounds = queryInBounds && (blockVec + laneId < vecs.getSize(0));
87  threadDistance = valInBounds ? threadDistance : kMaxDistance;
88  int id = valInBounds ? blockVec + laneId : -1;
89 
90  heap.add(threadDistance, id);
91  }
92 
93  heap.reduce();
94 
95  if (warpQuery < query.getSize(0)) {
96  heap.writeOut(outK[warpQuery].data(),
97  outV[warpQuery].data(),
98  k);
99  }
100 }
101 
102 // Version of the kernel that avoids a loop over the reduction dimension, and
103 // thus avoids reloading the query vectors
104 template <int NumWarpQ,
105  int NumThreadQ,
106  typename BinaryType,
107  int ReductionLimit = kLanes>
108 __global__ void
109 __launch_bounds__(kWarps * kLanes)
110 binaryDistanceLimitSize(const Tensor<BinaryType, 2, true> vecs,
111  const Tensor<BinaryType, 2, true> query,
112  Tensor<int, 2, true> outK,
113  Tensor<int, 2, true> outV,
114  int k) {
115  // A matrix tile (query, k)
116  __shared__ BinaryType queryTile[kWarps][kLanes + 1]; // avoid bank conflict
117 
118  // B matrix tile (vec, k)
119  __shared__ BinaryType vecTile[kLanes][kLanes + 1]; // avoid bank conflict
120 
121  WarpSelect<int, int, false, Comparator<int>,
122  NumWarpQ, NumThreadQ, kWarps * kLanes>
123  heap(kMaxDistance, -1, k);
124 
125  int warpId = threadIdx.y;
126  int laneId = threadIdx.x;
127 
128  // Each warp handles a single query
129  int laneK = laneId;
130  int warpQuery = blockIdx.x * kWarps + warpId;
131  bool kInBounds = laneK < vecs.getSize(1);
132  bool queryInBounds = warpQuery < query.getSize(0);
133 
134 
135  queryTile[warpId][laneId] = queryInBounds && kInBounds ?
136  query[warpQuery][laneK] : 0;
137 
138  // Each warp loops through the entire chunk of vectors
139  for (int blockVec = 0; blockVec < vecs.getSize(0); blockVec += kLanes) {
140  int threadDistance = 0;
141 
142  // kWarps warps are responsible for loading 32 vecs
143 #pragma unroll
144  for (int i = 0; i < kLanes / kWarps; ++i) {
145  int warpVec = i * kWarps + warpId;
146  int vec = blockVec + warpVec;
147  bool vecInBounds = vec < vecs.getSize(0);
148 
149  vecTile[warpVec][laneId] = vecInBounds && kInBounds ?
150  vecs[vec][laneK] : 0;
151  }
152 
153  __syncthreads();
154 
155  // Compare distances
156 #pragma unroll
157  for (int i = 0; i < ReductionLimit; ++i) {
158  threadDistance += __popc(queryTile[warpId][i] ^ vecTile[laneId][i]);
159  }
160 
161  __syncthreads();
162 
163  // Lanes within a warp are different vec results against the same query
164  // Only submit distances which represent real (query, vec) pairs
165  bool valInBounds = queryInBounds && (blockVec + laneId < vecs.getSize(0));
166  threadDistance = valInBounds ? threadDistance : kMaxDistance;
167  int id = valInBounds ? blockVec + laneId : -1;
168 
169  heap.add(threadDistance, id);
170  }
171 
172  heap.reduce();
173 
174  if (warpQuery < query.getSize(0)) {
175  heap.writeOut(outK[warpQuery].data(),
176  outV[warpQuery].data(),
177  k);
178  }
179 }
180 
181 template <typename BinaryType>
182 void runBinaryDistanceAnySize(Tensor<BinaryType, 2, true>& vecs,
183  Tensor<BinaryType, 2, true>& query,
184  Tensor<int, 2, true>& outK,
185  Tensor<int, 2, true>& outV,
186  int k, cudaStream_t stream) {
187  dim3 grid(utils::divUp(query.getSize(0), kWarps));
188  dim3 block(kLanes, kWarps);
189 
190  if (k == 1) {
191  binaryDistanceAnySize<1, 1, BinaryType>
192  <<<grid, block, 0, stream>>>(
193  vecs, query, outK, outV, k);
194  } else if (k <= 32) {
195  binaryDistanceAnySize<32, 2, BinaryType>
196  <<<grid, block, 0, stream>>>(
197  vecs, query, outK, outV, k);
198  } else if (k <= 64) {
199  binaryDistanceAnySize<64, 3, BinaryType>
200  <<<grid, block, 0, stream>>>(
201  vecs, query, outK, outV, k);
202  } else if (k <= 128) {
203  binaryDistanceAnySize<128, 3, BinaryType>
204  <<<grid, block, 0, stream>>>(
205  vecs, query, outK, outV, k);
206  } else if (k <= 256) {
207  binaryDistanceAnySize<256, 4, BinaryType>
208  <<<grid, block, 0, stream>>>(
209  vecs, query, outK, outV, k);
210  } else if (k <= 512) {
211  binaryDistanceAnySize<512, 8, BinaryType>
212  <<<grid, block, 0, stream>>>(
213  vecs, query, outK, outV, k);
214  } else if (k <= 1024) {
215  binaryDistanceAnySize<1024, 8, BinaryType>
216  <<<grid, block, 0, stream>>>(
217  vecs, query, outK, outV, k);
218  }
219 #if GPU_MAX_SELECTION_K >= 2048
220  else if (k <= 2048) {
221  binaryDistanceAnySize<2048, 8, BinaryType>
222  <<<grid, block, 0, stream>>>(
223  vecs, query, outK, outV, k);
224  }
225 #endif
226 }
227 
228 template <typename BinaryType, int ReductionLimit>
229 void runBinaryDistanceLimitSize(Tensor<BinaryType, 2, true>& vecs,
230  Tensor<BinaryType, 2, true>& query,
231  Tensor<int, 2, true>& outK,
232  Tensor<int, 2, true>& outV,
233  int k, cudaStream_t stream) {
234  dim3 grid(utils::divUp(query.getSize(0), kWarps));
235  dim3 block(kLanes, kWarps);
236 
237  if (k == 1) {
238  binaryDistanceLimitSize<1, 1, BinaryType, ReductionLimit>
239  <<<grid, block, 0, stream>>>(
240  vecs, query, outK, outV, k);
241  } else if (k <= 32) {
242  binaryDistanceLimitSize<32, 2, BinaryType, ReductionLimit>
243  <<<grid, block, 0, stream>>>(
244  vecs, query, outK, outV, k);
245  } else if (k <= 64) {
246  binaryDistanceLimitSize<64, 3, BinaryType, ReductionLimit>
247  <<<grid, block, 0, stream>>>(
248  vecs, query, outK, outV, k);
249  } else if (k <= 128) {
250  binaryDistanceLimitSize<128, 3, BinaryType, ReductionLimit>
251  <<<grid, block, 0, stream>>>(
252  vecs, query, outK, outV, k);
253  } else if (k <= 256) {
254  binaryDistanceLimitSize<256, 4, BinaryType, ReductionLimit>
255  <<<grid, block, 0, stream>>>(
256  vecs, query, outK, outV, k);
257  } else if (k <= 512) {
258  binaryDistanceLimitSize<512, 8, BinaryType, ReductionLimit>
259  <<<grid, block, 0, stream>>>(
260  vecs, query, outK, outV, k);
261  } else if (k <= 1024) {
262  binaryDistanceLimitSize<1024, 8, BinaryType, ReductionLimit>
263  <<<grid, block, 0, stream>>>(
264  vecs, query, outK, outV, k);
265  }
266 #if GPU_MAX_SELECTION_K >= 2048
267  else if (k <= 2048) {
268  binaryDistanceLimitSize<2048, 8, BinaryType, ReductionLimit>
269  <<<grid, block, 0, stream>>>(
270  vecs, query, outK, outV, k);
271  }
272 #endif
273 }
274 
275 void runBinaryDistance(Tensor<unsigned char, 2, true>& vecs,
276  Tensor<unsigned char, 2, true>& query,
277  Tensor<int, 2, true>& outK,
278  Tensor<int, 2, true>& outV,
279  int k, cudaStream_t stream) {
280  FAISS_ASSERT(k <= GPU_MAX_SELECTION_K);
281  FAISS_ASSERT(vecs.getSize(1) == query.getSize(1));
282 
283  FAISS_ASSERT(outK.getSize(1) == k);
284  FAISS_ASSERT(outV.getSize(1) == k);
285 
286  // For the optimized uint32 kernel, we handle 32 * 8 = 256 max dims
287  constexpr int kReductionLimit32 = 8;
288 
289  // For the optimized uint8 kernel, we handle 8 * 16 = 128 max dims
290  constexpr int kReductionLimit8 = 16;
291 
292  // All other cases (large or small) go through the general kernel
293 
294  if (vecs.getSize(1) % sizeof(unsigned int) == 0 &&
295  (vecs.getSize(1) / sizeof(unsigned int)) <= kReductionLimit32) {
296  auto vecs32 = vecs.castResize<unsigned int>();
297  auto query32 = query.castResize<unsigned int>();
298 
299  // Optimize for vectors with dimensions a multiple of 32 that are less than
300  // 32 * kReductionLimit (256) dimensions in size
301  runBinaryDistanceLimitSize<unsigned int, kReductionLimit32>(
302  vecs32, query32, outK, outV, k, stream);
303 
304  } else if (vecs.getSize(1) <= kReductionLimit8) {
305  // Optimize for vectors with dimensions a multiple of 32 that are less than
306  // 32 * kReductionLimit (256) dimensions in size
307  runBinaryDistanceLimitSize<unsigned char, kReductionLimit8>(
308  vecs, query, outK, outV, k, stream);
309  } else {
310  // Arbitrary size kernel
311  runBinaryDistanceAnySize<unsigned char>(
312  vecs, query, outK, outV, k, stream);
313  }
314 }
315 
316 } } // namespace
Generated by   doxygen 1.8.5