Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
PQCodeLoad.cuh
1 
2 /**
3  * Copyright (c) 2015-present, Facebook, Inc.
4  * All rights reserved.
5  *
6  * This source code is licensed under the CC-by-NC license found in the
7  * LICENSE file in the root directory of this source tree.
8  */
9 
10 // Copyright 2004-present Facebook. All Rights Reserved.
11 
12 #pragma once
13 
14 #include "../utils/PtxUtils.cuh"
15 
16 namespace faiss { namespace gpu {
17 
18 ///
19 /// This file contains loader functions for PQ codes of various byte
20 /// length.
21 ///
22 
23 // Type-specific wrappers around the PTX bfe.* instruction, for
24 // quantization code extraction
25 inline __device__ unsigned int getByte(unsigned char v,
26  int pos,
27  int width) {
28  return v;
29 }
30 
31 inline __device__ unsigned int getByte(unsigned short v,
32  int pos,
33  int width) {
34  return getBitfield((unsigned int) v, pos, width);
35 }
36 
37 inline __device__ unsigned int getByte(unsigned int v,
38  int pos,
39  int width) {
40  return getBitfield(v, pos, width);
41 }
42 
43 inline __device__ unsigned int getByte(unsigned long v,
44  int pos,
45  int width) {
46  return getBitfield(v, pos, width);
47 }
48 
49 template <int NumSubQuantizers>
50 struct LoadCode32 {};
51 
52 template<>
53 struct LoadCode32<1> {
54  static inline __device__ void load(unsigned int code32[1],
55  unsigned char* p,
56  int offset) {
57  p += offset * 1;
58  asm("ld.global.cs.u8 {%0}, [%1];" :
59  "=r"(code32[0]) : "l"(p));
60  }
61 };
62 
63 template<>
64 struct LoadCode32<2> {
65  static inline __device__ void load(unsigned int code32[1],
66  unsigned char* p,
67  int offset) {
68  p += offset * 2;
69  asm("ld.global.cs.u16 {%0}, [%1];" :
70  "=r"(code32[0]) : "l"(p));
71  }
72 };
73 
74 template<>
75 struct LoadCode32<3> {
76  static inline __device__ void load(unsigned int code32[1],
77  unsigned char* p,
78  int offset) {
79  p += offset * 3;
80  unsigned int a;
81  unsigned int b;
82  unsigned int c;
83 
84  // FIXME: this is a non-coalesced, unaligned, non-vectorized load
85  // unfortunately need to reorganize memory layout by warp
86  asm("ld.global.cs.u8 {%0}, [%1 + 0];" :
87  "=r"(a) : "l"(p));
88  asm("ld.global.cs.u8 {%0}, [%1 + 1];" :
89  "=r"(b) : "l"(p));
90  asm("ld.global.cs.u8 {%0}, [%1 + 2];" :
91  "=r"(c) : "l"(p));
92 
93  // FIXME: this is also slow, since we have to recover the
94  // individual bytes loaded
95  code32[0] = (c << 16) | (b << 8) | a;
96  }
97 };
98 
99 template<>
100 struct LoadCode32<4> {
101  static inline __device__ void load(unsigned int code32[1],
102  unsigned char* p,
103  int offset) {
104  p += offset * 4;
105  asm("ld.global.cs.u32 {%0}, [%1];" :
106  "=r"(code32[0]) : "l"(p));
107  }
108 };
109 
110 template<>
111 struct LoadCode32<8> {
112  static inline __device__ void load(unsigned int code32[2],
113  unsigned char* p,
114  int offset) {
115  p += offset * 8;
116  asm("ld.global.cs.v2.u32 {%0, %1}, [%2];" :
117  "=r"(code32[0]), "=r"(code32[1]) : "l"(p));
118  }
119 };
120 
121 template<>
122 struct LoadCode32<12> {
123  static inline __device__ void load(unsigned int code32[3],
124  unsigned char* p,
125  int offset) {
126  p += offset * 12;
127  // FIXME: this is a non-coalesced, unaligned, non-vectorized load
128  // unfortunately need to reorganize memory layout by warp
129  asm("ld.global.cs.nc.u32 {%0}, [%1 + 0];" :
130  "=r"(code32[0]) : "l"(p));
131  asm("ld.global.cs.nc.u32 {%0}, [%1 + 4];" :
132  "=r"(code32[1]) : "l"(p));
133  asm("ld.global.cs.nc.u32 {%0}, [%1 + 8];" :
134  "=r"(code32[2]) : "l"(p));
135  }
136 };
137 
138 template<>
139 struct LoadCode32<16> {
140  static inline __device__ void load(unsigned int code32[4],
141  unsigned char* p,
142  int offset) {
143  p += offset * 16;
144  asm("ld.global.cs.v4.u32 {%0, %1, %2, %3}, [%4];" :
145  "=r"(code32[0]), "=r"(code32[1]),
146  "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
147  }
148 };
149 
150 template<>
151 struct LoadCode32<20> {
152  static inline __device__ void load(unsigned int code32[5],
153  unsigned char* p,
154  int offset) {
155  p += offset * 20;
156  // FIXME: this is a non-coalesced, unaligned, non-vectorized load
157  // unfortunately need to reorganize memory layout by warp
158  asm("ld.global.cs.nc.u32 {%0}, [%1 + 0];" :
159  "=r"(code32[0]) : "l"(p));
160  asm("ld.global.cs.nc.u32 {%0}, [%1 + 4];" :
161  "=r"(code32[1]) : "l"(p));
162  asm("ld.global.cs.nc.u32 {%0}, [%1 + 8];" :
163  "=r"(code32[2]) : "l"(p));
164  asm("ld.global.cs.nc.u32 {%0}, [%1 + 12];" :
165  "=r"(code32[3]) : "l"(p));
166  asm("ld.global.cs.nc.u32 {%0}, [%1 + 16];" :
167  "=r"(code32[4]) : "l"(p));
168  }
169 };
170 
171 template<>
172 struct LoadCode32<24> {
173  static inline __device__ void load(unsigned int code32[6],
174  unsigned char* p,
175  int offset) {
176  p += offset * 24;
177  // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
178  // unfortunately need to reorganize memory layout by warp
179  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 0];" :
180  "=r"(code32[0]), "=r"(code32[1]) : "l"(p));
181  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 8];" :
182  "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
183  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 16];" :
184  "=r"(code32[4]), "=r"(code32[5]) : "l"(p));
185  }
186 };
187 
188 template<>
189 struct LoadCode32<28> {
190  static inline __device__ void load(unsigned int code32[7],
191  unsigned char* p,
192  int offset) {
193  p += offset * 28;
194  // FIXME: this is a non-coalesced, unaligned, non-vectorized load
195  // unfortunately need to reorganize memory layout by warp
196  asm("ld.global.cs.nc.u32 {%0}, [%1 + 0];" :
197  "=r"(code32[0]) : "l"(p));
198  asm("ld.global.cs.nc.u32 {%0}, [%1 + 4];" :
199  "=r"(code32[1]) : "l"(p));
200  asm("ld.global.cs.nc.u32 {%0}, [%1 + 8];" :
201  "=r"(code32[2]) : "l"(p));
202  asm("ld.global.cs.nc.u32 {%0}, [%1 + 12];" :
203  "=r"(code32[3]) : "l"(p));
204  asm("ld.global.cs.nc.u32 {%0}, [%1 + 16];" :
205  "=r"(code32[4]) : "l"(p));
206  asm("ld.global.cs.nc.u32 {%0}, [%1 + 20];" :
207  "=r"(code32[5]) : "l"(p));
208  asm("ld.global.cs.nc.u32 {%0}, [%1 + 24];" :
209  "=r"(code32[6]) : "l"(p));
210  }
211 };
212 
213 template<>
214 struct LoadCode32<32> {
215  static inline __device__ void load(unsigned int code32[8],
216  unsigned char* p,
217  int offset) {
218  p += offset * 32;
219  // FIXME: this is a non-coalesced load
220  // unfortunately need to reorganize memory layout by warp
221  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4];" :
222  "=r"(code32[0]), "=r"(code32[1]),
223  "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
224  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 16];" :
225  "=r"(code32[4]), "=r"(code32[5]),
226  "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
227  }
228 };
229 
230 template<>
231 struct LoadCode32<40> {
232  static inline __device__ void load(unsigned int code32[10],
233  unsigned char* p,
234  int offset) {
235  p += offset * 40;
236  // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
237  // unfortunately need to reorganize memory layout by warp
238  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 0];" :
239  "=r"(code32[0]), "=r"(code32[1]) : "l"(p));
240  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 8];" :
241  "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
242  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 16];" :
243  "=r"(code32[4]), "=r"(code32[5]) : "l"(p));
244  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 24];" :
245  "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
246  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 32];" :
247  "=r"(code32[8]), "=r"(code32[9]) : "l"(p));
248  }
249 };
250 
251 template<>
252 struct LoadCode32<48> {
253  static inline __device__ void load(unsigned int code32[12],
254  unsigned char* p,
255  int offset) {
256  p += offset * 48;
257  // FIXME: this is a non-coalesced load
258  // unfortunately need to reorganize memory layout by warp
259  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4];" :
260  "=r"(code32[0]), "=r"(code32[1]),
261  "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
262  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 16];" :
263  "=r"(code32[4]), "=r"(code32[5]),
264  "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
265  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 32];" :
266  "=r"(code32[8]), "=r"(code32[9]),
267  "=r"(code32[10]), "=r"(code32[11]) : "l"(p));
268  }
269 };
270 
271 template<>
272 struct LoadCode32<56> {
273  static inline __device__ void load(unsigned int code32[14],
274  unsigned char* p,
275  int offset) {
276  p += offset * 56;
277  // FIXME: this is a non-coalesced, unaligned, 2-vectorized load
278  // unfortunately need to reorganize memory layout by warp
279  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 0];" :
280  "=r"(code32[0]), "=r"(code32[1]) : "l"(p));
281  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 8];" :
282  "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
283  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 16];" :
284  "=r"(code32[4]), "=r"(code32[5]) : "l"(p));
285  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 24];" :
286  "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
287  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 32];" :
288  "=r"(code32[8]), "=r"(code32[9]) : "l"(p));
289  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 40];" :
290  "=r"(code32[10]), "=r"(code32[11]) : "l"(p));
291  asm("ld.global.cs.nc.v2.u32 {%0, %1}, [%2 + 48];" :
292  "=r"(code32[12]), "=r"(code32[13]) : "l"(p));
293  }
294 };
295 
296 template<>
297 struct LoadCode32<64> {
298  static inline __device__ void load(unsigned int code32[16],
299  unsigned char* p,
300  int offset) {
301  p += offset * 64;
302  // FIXME: this is a non-coalesced load
303  // unfortunately need to reorganize memory layout by warp
304  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4];" :
305  "=r"(code32[0]), "=r"(code32[1]),
306  "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
307  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 16];" :
308  "=r"(code32[4]), "=r"(code32[5]),
309  "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
310  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 32];" :
311  "=r"(code32[8]), "=r"(code32[9]),
312  "=r"(code32[10]), "=r"(code32[11]) : "l"(p));
313  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 48];" :
314  "=r"(code32[12]), "=r"(code32[13]),
315  "=r"(code32[14]), "=r"(code32[15]) : "l"(p));
316  }
317 };
318 
319 template<>
320 struct LoadCode32<96> {
321  static inline __device__ void load(unsigned int code32[24],
322  unsigned char* p,
323  int offset) {
324  p += offset * 96;
325  // FIXME: this is a non-coalesced load
326  // unfortunately need to reorganize memory layout by warp
327  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4];" :
328  "=r"(code32[0]), "=r"(code32[1]),
329  "=r"(code32[2]), "=r"(code32[3]) : "l"(p));
330  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 16];" :
331  "=r"(code32[4]), "=r"(code32[5]),
332  "=r"(code32[6]), "=r"(code32[7]) : "l"(p));
333  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 32];" :
334  "=r"(code32[8]), "=r"(code32[9]),
335  "=r"(code32[10]), "=r"(code32[11]) : "l"(p));
336  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 48];" :
337  "=r"(code32[12]), "=r"(code32[13]),
338  "=r"(code32[14]), "=r"(code32[15]) : "l"(p));
339  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 64];" :
340  "=r"(code32[16]), "=r"(code32[17]),
341  "=r"(code32[18]), "=r"(code32[19]) : "l"(p));
342  asm("ld.global.cs.nc.v4.u32 {%0, %1, %2, %3}, [%4 + 80];" :
343  "=r"(code32[20]), "=r"(code32[21]),
344  "=r"(code32[22]), "=r"(code32[23]) : "l"(p));
345  }
346 };
347 
348 
349 } } // namespace
Generated by   doxygen 1.8.5