Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
Pair.cuh
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
9 #pragma once
10
11 #include <cuda.h>
12 #include "MathOperators.cuh"
13 #include "WarpShuffles.cuh"
14
15 namespace faiss { namespace gpu {
16
17 /// A simple pair type for CUDA device usage
18 template <typename K, typename V>
19 struct Pair {
20 constexpr __device__ inline Pair() {
21 }
22
23 constexpr __device__ inline Pair(K key, V value)
24 : k(key), v(value) {
25 }
26
27 __device__ inline bool
28 operator==(const Pair<K, V>& rhs) const {
29 return Math<K>::eq(k, rhs.k) && Math<V>::eq(v, rhs.v);
30 }
31
32 __device__ inline bool
33 operator!=(const Pair<K, V>& rhs) const {
34 return !operator==(rhs);
35 }
36
37 __device__ inline bool
38 operator<(const Pair<K, V>& rhs) const {
39 return Math<K>::lt(k, rhs.k) ||
40 (Math<K>::eq(k, rhs.k) && Math<V>::lt(v, rhs.v));
41 }
42
43 __device__ inline bool
44 operator>(const Pair<K, V>& rhs) const {
45 return Math<K>::gt(k, rhs.k) ||
46 (Math<K>::eq(k, rhs.k) && Math<V>::gt(v, rhs.v));
47 }
48
49 K k;
50 V v;
51 };
52
53 template <typename T, typename U>
54 inline __device__ Pair<T, U> shfl_up(const Pair<T, U>& pair,
55 unsigned int delta,
56 int width = kWarpSize) {
57 return Pair<T, U>(shfl_up(pair.k, delta, width),
58 shfl_up(pair.v, delta, width));
59 }
60
61 template <typename T, typename U>
62 inline __device__ Pair<T, U> shfl_xor(const Pair<T, U>& pair,
63 int laneMask,
64 int width = kWarpSize) {
65 return Pair<T, U>(shfl_xor(pair.k, laneMask, width),
66 shfl_xor(pair.v, laneMask, width));
67 }
68
69 } } // namespace
faiss::gpu::Math
Definition: MathOperators.cuh:22
faiss::gpu::Pair
A simple pair type for CUDA device usage.
Definition: Pair.cuh:19