Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
BlockSelectHalf.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 // Copyright 2004-present Facebook. All Rights Reserved.
10 #include "blockselect/BlockSelectImpl.cuh"
11 
12 namespace faiss { namespace gpu {
13 
14 #ifdef FAISS_USE_FLOAT16
15 
16 // warp Q to thread Q:
17 // 1, 1
18 // 32, 2
19 // 64, 3
20 // 128, 3
21 // 256, 4
22 // 512, 8
23 // 1024, 8
24 
25 BLOCK_SELECT_DECL(half, true, 1);
26 BLOCK_SELECT_DECL(half, true, 32);
27 BLOCK_SELECT_DECL(half, true, 64);
28 BLOCK_SELECT_DECL(half, true, 128);
29 BLOCK_SELECT_DECL(half, true, 256);
30 BLOCK_SELECT_DECL(half, true, 512);
31 BLOCK_SELECT_DECL(half, true, 1024);
32 
33 BLOCK_SELECT_DECL(half, false, 1);
34 BLOCK_SELECT_DECL(half, false, 32);
35 BLOCK_SELECT_DECL(half, false, 64);
36 BLOCK_SELECT_DECL(half, false, 128);
37 BLOCK_SELECT_DECL(half, false, 256);
38 BLOCK_SELECT_DECL(half, false, 512);
39 BLOCK_SELECT_DECL(half, false, 1024);
40 
41 void runBlockSelect(Tensor<half, 2, true>& in,
42  Tensor<half, 2, true>& outK,
43  Tensor<int, 2, true>& outV,
44  bool dir, int k, cudaStream_t stream) {
45  FAISS_ASSERT(k <= 1024);
46 
47  if (dir) {
48  if (k == 1) {
49  BLOCK_SELECT_CALL(half, true, 1);
50  } else if (k <= 32) {
51  BLOCK_SELECT_CALL(half, true, 32);
52  } else if (k <= 64) {
53  BLOCK_SELECT_CALL(half, true, 64);
54  } else if (k <= 128) {
55  BLOCK_SELECT_CALL(half, true, 128);
56  } else if (k <= 256) {
57  BLOCK_SELECT_CALL(half, true, 256);
58  } else if (k <= 512) {
59  BLOCK_SELECT_CALL(half, true, 512);
60  } else if (k <= 1024) {
61  BLOCK_SELECT_CALL(half, true, 1024);
62  }
63  } else {
64  if (k == 1) {
65  BLOCK_SELECT_CALL(half, false, 1);
66  } else if (k <= 32) {
67  BLOCK_SELECT_CALL(half, false, 32);
68  } else if (k <= 64) {
69  BLOCK_SELECT_CALL(half, false, 64);
70  } else if (k <= 128) {
71  BLOCK_SELECT_CALL(half, false, 128);
72  } else if (k <= 256) {
73  BLOCK_SELECT_CALL(half, false, 256);
74  } else if (k <= 512) {
75  BLOCK_SELECT_CALL(half, false, 512);
76  } else if (k <= 1024) {
77  BLOCK_SELECT_CALL(half, false, 1024);
78  }
79  }
80 }
81 
82 void runBlockSelectPair(Tensor<half, 2, true>& inK,
83  Tensor<int, 2, true>& inV,
84  Tensor<half, 2, true>& outK,
85  Tensor<int, 2, true>& outV,
86  bool dir, int k, cudaStream_t stream) {
87  FAISS_ASSERT(k <= 1024);
88 
89  if (dir) {
90  if (k == 1) {
91  BLOCK_SELECT_PAIR_CALL(half, true, 1);
92  } else if (k <= 32) {
93  BLOCK_SELECT_PAIR_CALL(half, true, 32);
94  } else if (k <= 64) {
95  BLOCK_SELECT_PAIR_CALL(half, true, 64);
96  } else if (k <= 128) {
97  BLOCK_SELECT_PAIR_CALL(half, true, 128);
98  } else if (k <= 256) {
99  BLOCK_SELECT_PAIR_CALL(half, true, 256);
100  } else if (k <= 512) {
101  BLOCK_SELECT_PAIR_CALL(half, true, 512);
102  } else if (k <= 1024) {
103  BLOCK_SELECT_PAIR_CALL(half, true, 1024);
104  }
105  } else {
106  if (k == 1) {
107  BLOCK_SELECT_PAIR_CALL(half, false, 1);
108  } else if (k <= 32) {
109  BLOCK_SELECT_PAIR_CALL(half, false, 32);
110  } else if (k <= 64) {
111  BLOCK_SELECT_PAIR_CALL(half, false, 64);
112  } else if (k <= 128) {
113  BLOCK_SELECT_PAIR_CALL(half, false, 128);
114  } else if (k <= 256) {
115  BLOCK_SELECT_PAIR_CALL(half, false, 256);
116  } else if (k <= 512) {
117  BLOCK_SELECT_PAIR_CALL(half, false, 512);
118  } else if (k <= 1024) {
119  BLOCK_SELECT_PAIR_CALL(half, false, 1024);
120  }
121  }
122 }
123 
124 #endif
125 
126 } } // namespace
Generated by   doxygen 1.8.5