Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
WarpSelectHalf.cu
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 #include "warpselect/WarpSelectImpl.cuh"
10
11 namespace faiss { namespace gpu {
12
13 #ifdef FAISS_USE_FLOAT16
14
15 // warp Q to thread Q:
16 // 1, 1
17 // 32, 2
18 // 64, 3
19 // 128, 3
20 // 256, 4
21 // 512, 8
22 // 1024, 8
23
24 WARP_SELECT_DECL(half, true, 1);
25 WARP_SELECT_DECL(half, true, 32);
26 WARP_SELECT_DECL(half, true, 64);
27 WARP_SELECT_DECL(half, true, 128);
28 WARP_SELECT_DECL(half, true, 256);
29 WARP_SELECT_DECL(half, true, 512);
30 WARP_SELECT_DECL(half, true, 1024);
31
32 WARP_SELECT_DECL(half, false, 1);
33 WARP_SELECT_DECL(half, false, 32);
34 WARP_SELECT_DECL(half, false, 64);
35 WARP_SELECT_DECL(half, false, 128);
36 WARP_SELECT_DECL(half, false, 256);
37 WARP_SELECT_DECL(half, false, 512);
38 WARP_SELECT_DECL(half, false, 1024);
39
40 void runWarpSelect(Tensor<half, 2, true>& in,
41 Tensor<half, 2, true>& outK,
42 Tensor<int, 2, true>& outV,
43 bool dir, int k, cudaStream_t stream) {
44 FAISS_ASSERT(k <= 1024);
45
46 if (dir) {
47 if (k == 1) {
48 WARP_SELECT_CALL(half, true, 1);
49 } else if (k <= 32) {
50 WARP_SELECT_CALL(half, true, 32);
51 } else if (k <= 64) {
52 WARP_SELECT_CALL(half, true, 64);
53 } else if (k <= 128) {
54 WARP_SELECT_CALL(half, true, 128);
55 } else if (k <= 256) {
56 WARP_SELECT_CALL(half, true, 256);
57 } else if (k <= 512) {
58 WARP_SELECT_CALL(half, true, 512);
59 } else if (k <= 1024) {
60 WARP_SELECT_CALL(half, true, 1024);
61 }
62 } else {
63 if (k == 1) {
64 WARP_SELECT_CALL(half, false, 1);
65 } else if (k <= 32) {
66 WARP_SELECT_CALL(half, false, 32);
67 } else if (k <= 64) {
68 WARP_SELECT_CALL(half, false, 64);
69 } else if (k <= 128) {
70 WARP_SELECT_CALL(half, false, 128);
71 } else if (k <= 256) {
72 WARP_SELECT_CALL(half, false, 256);
73 } else if (k <= 512) {
74 WARP_SELECT_CALL(half, false, 512);
75 } else if (k <= 1024) {
76 WARP_SELECT_CALL(half, false, 1024);
77 }
78 }
79 }
80
81 #endif
82
83 } } // namespace