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