Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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